如何使用 PyTorch 编码器重塑多通道图像？

Question

我有一个尺寸为 [18, 512, 512] 的张量，表示 512x512 图像上最多 18 个特定对象的灰度热图。为了为我的条件 GAN 生成该图像的合适表示，我需要使用编码器将该张量重塑为 [512, 4, 4] 形状。但是，我无法理解这种转换是如何实现的，因为给定的维度对于直接线性或卷积转换来说似乎太不匹配了。

class HeatmapEncoder(torch.nn.Module):
  def __init__(self):
    # source = 18x512x512
    # target = 512x4x4
    self.encoder = torch.nn.Sequential(
        nn.Linear(),
        nn.ReLU(),
        nn.Linear()
    )

    def forward(self, x):
      pass

这里可以使用 start_dim=0 的 `nn.Flatten()'，但结果将是一个简化的张量，不能用作线性层的输入。 解码器部分现在并不是特别重要，因为我只需要热图的低维表示来调节我的 GAN，而不是重新创建这些图像。

Answer 1

您可以尝试几种不同的方法来解决您的问题，例如将其视为一个大向量，然后慢慢将其缩小到您的大小，或者置换尺寸并应用不同的操作等。因为没有完整的代码来测试它实际问题，我不能说哪个会更好，但我的第一直觉说基于卷积的降维非常适合这个问题。先写代码，再谈：

class ReduceConv(torch.nn.Module):
  def __init__(self, nin, nout, activ=nn.ReLU):
    super(ReduceConv, self).__init__()

    # source = Batch x nin x H x W
    # target = Batch x nout x (H/2) x (W/2)

    self.conv = nn.Sequential(
        nn.Conv2d(
            nin, nout,
            kernel_size=3,
            stride=1,
            padding=1),
        nn.Conv2d(
            nout, nout,
            kernel_size=3,
            stride = 2,
            padding = 1),
        nn.BatchNorm2d(nout),
        activ()
    )

  def forward(self, x):
    return self.conv(x)

class HeatmapEncoder(torch.nn.Module):
  def __init__(self):
    super(HeatmapEncoder, self).__init__()

    # source = 18x512x512
    # target = 512x4x4
    self.encoder = torch.nn.Sequential(
        ReduceConv(18, 32),       # out->  32 256 256 
        ReduceConv(32, 64),       # out->  64 128 128
        ReduceConv(64, 64),       # out->  64 64 64
        ReduceConv(64, 64),       # out->  64 32 32
        ReduceConv(64, 128),      # out-> 128 16 16
        ReduceConv(128, 256),     # out-> 256 8 8
        ReduceConv(256, 512)      # out-> 512 4 4
    )

  def forward(self, x):
    return self.encoder(x)

# 10 is batch size
inp = torch.rand(10, 18, 512, 512)
enc = HeatmapEncoder()
out = enc(inp)
print(inp.shape)      # torch.Size([10, 18, 512, 512])
print(out.shape)      # torch.Size([10, 512, 4, 4]) 

它本质上只是一堆卷积层。请注意，在每个 ReduceConv 层中，输入维度通过使用 stride=2 卷积减半。从技术上讲，您不需要 ReduceConv 层中的第一个卷积，但现在是深度学习时代，越多越好 :) 我还在每次归约后添加了 BatchNorm 以及一个激活函数。当输入到第一个卷积时，18 被视为通道。这样通道最多可建立 512 个，而宽度和高度在每次操作后减半。此编码器模型可能不是最好或最有效的模型，但它应该足以解决您的问题。