从全连接层反推输入尺寸大小

全连接层：输出 512维

batch_size = 128
nn.Linear(128*8*8, 512),

输入尺寸(3通道图片）：32x32=1024

batch_size = 128
channel = 3
torch.Size([128, 3, 32, 32])

完整的网络结构：

import torch.nn as nn

class _Encoder(nn.Module):
    def __init__(self, layers):
        super(_Encoder, self).__init__()
        self.layers = nn.Sequential(*layers)

    def forward(self, x):
        x = self.layers(x)
        x = x.view(x.size(0), -1)
        return x


class _Decoder(nn.Module):
    def __init__(self, output_size):
        super(_Decoder, self).__init__()
        self.layers = nn.Sequential(
            nn.Linear(128*8*8, 512),
            nn.BatchNorm1d(512),
            nn.ReLU(),
            nn.Linear(512, output_size)
        )

    def forward(self, x):
        x = self.layers(x)
        return x


def Model(num_classes, num_channels):
    layers = [
        nn.Conv2d(num_channels, 32, kernel_size=3, padding=1),
        nn.BatchNorm2d(32),
        nn.ReLU(),
        nn.Conv2d(32, 64, kernel_size=3, padding=1, stride=2),
        nn.BatchNorm2d(64),
        nn.ReLU(),
        nn.Conv2d(64, 64, kernel_size=3, padding=1),
        nn.BatchNorm2d(64),
        nn.ReLU(),
        nn.Conv2d(64, 128, kernel_size=3, padding=1, stride=2),
        nn.BatchNorm2d(128),
        nn.ReLU(),
        nn.Conv2d(128, 128, kernel_size=3, padding=1),
        nn.BatchNorm2d(128),
        nn.ReLU(),
    ]

    encoders = [_Encoder(layers=layers) for _ in num_classes]
    print(num_classes, encoders) # num_classes: [2, 2, 2, 2, 2, 2, 2, 2, 2, 2]
    return [_Model(output_size=cls, encoder=encoder) for cls, encoder in zip(num_classes, encoders)]