具有 2 个输入的 Pytorch UNet

Question

我正在创建一个 UNet，它根据建筑物地图的 100x100x1 灰度图像生成 100x100x1 小区服务覆盖图。但是，我还想添加基站位置的输入，这将是一个具有 x 位置、y 位置和高度三个通道的张量。我已经有一个数据集，其中基站位置与用于训练模型的地面实况图像相关联，但是我需要构建一个可以组合两个输入的模型。建筑物地图的形状为(batch size, 1, 100, 100)，基站位置的形状为(batch size, 3)。我将在下面粘贴模型。在哪里以及如何有效地组合输入？我正在使用来自 gitHub 的模型，其中“Down”是 2 个转换层和一个 maxpool，“Up”则完全相反。

数据示例（不用担心 pyplot 轴它们不正确）：

基站信息为24（x）、23（y）、23（高度）

UNet（当前仅配置 1 个输入）

class UNet(nn.Module):
    def __init__(self, n_channels, n_classes, bilinear=False):
        super(UNet, self).__init__()
        self.n_channels = n_channels
        self.n_classes = n_classes
        self.bilinear = bilinear

        self.inc = (DoubleConv(n_channels, 64))
        self.down1 = (Down(64, 128))
        self.down2 = (Down(128, 256))
        self.down3 = (Down(256, 512))
        factor = 2 if bilinear else 1
        self.down4 = (Down(512, 1024 // factor))

        self.up1 = (Up(1024, 512 // factor, bilinear))
        self.up2 = (Up(512, 256 // factor, bilinear))
        self.up3 = (Up(256, 128 // factor, bilinear))
        self.up4 = (Up(128, 64, bilinear))
        self.outc = (OutConv(64, n_classes))

    def forward(self, building_map, base_station_info):
        x1 = self.inc(building_map)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)
        x5 = self.down4(x4)
        x = self.up1(x5, x4)
        x = self.up2(x, x3)
        x = self.up3(x, x2)
        x = self.up4(x, x1)
        logits = self.outc(x)
        return logits