PyTorch AutoEncoder - декодированный размер вывода не совпадает с входным

Question

Я создаю пользовательский автоэнкодер для обучения на наборе данных. Моя модель выглядит следующим образом:

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

        self.encoder = nn.Sequential(
        nn.Conv2d(in_channels = 3, out_channels = 32, kernel_size=3,stride=1),
        nn.ReLU(inplace=True),
        nn.Conv2d(in_channels = 32, out_channels = 64, kernel_size=3,stride=1),
        nn.ReLU(inplace=True),
        nn.Conv2d(in_channels = 64, out_channels = 128, kernel_size=3,stride=1),
        nn.ReLU(inplace=True),
        nn.Conv2d(in_channels=128,out_channels=256,kernel_size=5,stride=2),
        nn.ReLU(inplace=True),
        nn.Conv2d(in_channels=256,out_channels=512,kernel_size=5,stride=2),
        nn.ReLU(inplace=True),
        nn.Conv2d(in_channels=512,out_channels=1024,kernel_size=5,stride=2),
        nn.ReLU(inplace=True)
        )

        self.decoder = nn.Sequential(
        nn.ConvTranspose2d(in_channels=1024,out_channels=512,kernel_size=5,stride=2),
        nn.ReLU(inplace=True),
        nn.ConvTranspose2d(in_channels=512,out_channels=256,kernel_size=5,stride=2),
        nn.ReLU(inplace=True),
        nn.ConvTranspose2d(in_channels=256,out_channels=128,kernel_size=5,stride=2),
        nn.ReLU(inplace=True),
        nn.ConvTranspose2d(in_channels=128,out_channels=64,kernel_size=3,stride=1),
        nn.ReLU(inplace=True),
        nn.ConvTranspose2d(in_channels=64,out_channels=32,kernel_size=3,stride=1),
        nn.ReLU(inplace=True),
        nn.ConvTranspose2d(in_channels=32,out_channels=3,kernel_size=3,stride=1),
        nn.ReLU(inplace=True)
        )


    def forward(self,x):
        x = self.encoder(x)
        print(x.shape)
        x = self.decoder(x)
        return x



def unit_test():
    num_minibatch = 16
    img = torch.randn(num_minibatch, 3, 512, 640).cuda(0)
    model = AutoEncoder().cuda()
    model = nn.DataParallel(model)
    output = model(img)
    print(output.shape)

if __name__ == '__main__':
    unit_test()

Как видите, мой входной размер равен (3, 512, 640), но мой выходной сигнал после прохождения его через декодер равен (3, 507, 635). Я что-то упускаю при добавлении слоев Conv2D Transpose?

Буду признателен за любую помощь. Спасибо

Answer 1

Несоответствие вызвано различными формами вывода слоя ConvTranspose2d. Вы можете добавить output_padding из 1 к первому и третьему слою транспонирования свертки, чтобы решить эту проблему.

т.е. nn.ConvTranspose2d(in_channels=1024,out_channels=512,kernel_size=5,stride=2, output_padding=1) и nn.ConvTranspose2d(in_channels=256,out_channels=128,kernel_size=5,stride=2, output_padding=1)

Согласно документации :

Когда шаг> 1, Conv2d отображает несколько входных фигур в одну и ту же выходную форму. output_padding предоставляется для устранения этой неоднозначности путем эффективного увеличения расчетной выходной формы на одной стороне.

---

Формы слоев декодера перед добавлением output_padding:

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
   ConvTranspose2d-1        [-1, 512, 123, 155]      13,107,712
              ReLU-2        [-1, 512, 123, 155]               0
   ConvTranspose2d-3        [-1, 256, 249, 313]       3,277,056
              ReLU-4        [-1, 256, 249, 313]               0
   ConvTranspose2d-5        [-1, 128, 501, 629]         819,328
              ReLU-6        [-1, 128, 501, 629]               0
   ConvTranspose2d-7         [-1, 64, 503, 631]          73,792
              ReLU-8         [-1, 64, 503, 631]               0
   ConvTranspose2d-9         [-1, 32, 505, 633]          18,464
             ReLU-10         [-1, 32, 505, 633]               0
  ConvTranspose2d-11          [-1, 3, 507, 635]             867
             ReLU-12          [-1, 3, 507, 635]               0

После добавления отступов:

================================================================
   ConvTranspose2d-1        [-1, 512, 124, 156]      13,107,712
              ReLU-2        [-1, 512, 124, 156]               0
   ConvTranspose2d-3        [-1, 256, 251, 315]       3,277,056
              ReLU-4        [-1, 256, 251, 315]               0
   ConvTranspose2d-5        [-1, 128, 506, 634]         819,328
              ReLU-6        [-1, 128, 506, 634]               0
   ConvTranspose2d-7         [-1, 64, 508, 636]          73,792
              ReLU-8         [-1, 64, 508, 636]               0
   ConvTranspose2d-9         [-1, 32, 510, 638]          18,464
             ReLU-10         [-1, 32, 510, 638]               0
  ConvTranspose2d-11          [-1, 3, 512, 640]             867
             ReLU-12          [-1, 3, 512, 640]               0