Pytorch：获取最后一层的正确尺寸

Question

Pytorch 新手来了！我正在尝试微调 VGG16 模型以预测 3 个不同的类别。我的部分工作涉及将 FC 层转换为 CONV 层。但是，我的预测值不在 0 到 2（3 个类别）之间。

有人能给我指出一个关于如何计算最后一层正确尺寸的好资源吗？

以下是 VGG16 的原始 fC 层：

(classifier): Sequential(

    (0): Linear(in_features=25088, out_features=4096, bias=True)

    (1): ReLU(inplace)

    (2): Dropout(p=0.5)

    (3): Linear(in_features=4096, out_features=4096, bias=True)

    (4): ReLU(inplace)

    (5): Dropout(p=0.5)

    (6): Linear(in_features=4096, out_features=1000, bias=True)

  )

我将 FC 层转换为 CONV 的代码：

 def convert_fc_to_conv(self, fc_layers):

        # Replace first FC layer with CONV layer

        fc = fc_layers[0].state_dict()

        in_ch = 512

        out_ch = fc["weight"].size(0)

        first_conv = nn.Conv2d(512, out_ch, kernel_size=(1, 1), stride=(1, 1))

        conv_list = [first_conv]

        for idx, layer in enumerate(fc_layers[1:]):

            if isinstance(layer, nn.Linear):

                fc = layer.state_dict()

                in_ch = fc["weight"].size(1)

                out_ch = fc["weight"].size(0)

                if idx == len(fc_layers)-4:

                    in_ch = 3

                conv = nn.Conv2d(out_ch, in_ch, kernel_size=(1, 1), stride=(1, 1))

                conv_list += [conv]

            else:

                conv_list += [layer]

            gc.collect()

        avg_pool = nn.AvgPool2d(kernel_size=2, stride=1, ceil_mode=False)

        conv_list += [avg_pool, nn.Softmax()]

        top_layers = nn.Sequential(*conv_list)  

        return top_layers