解决RuntimeError: running_mean should contain 36864 elements not 4096

一、问题描述

在模型微调最后2层时，BatchNorm1d如题报错。

def get_model():
    model_pre = models.resnet50(pretrained=True) # 获取预训练模型

    # 冻结预训练模型中所有参数
    for param in model_pre.parameters():
        param.requires_grad = False

    # 替换ResNet最后的两层网络，返回一个新的模型（迁移学习）
    model_pre.avgpool = AdaptiveConcatPool2d() # 池化层替换
    model_pre.fc = nn.Sequential(
            nn.Flatten(), # 所有维度拉平
            # nn.BatchNorm1d(4096), # 正则化处理
            nn.BatchNorm1d(36864), # 正则化处理
            nn.Dropout(0.5), # 丢掉神经元
            # nn.Linear(4096, 512), # 线性层处理
            nn.Linear(36864, 512), # 线性层处理
            nn.ReLU(), # 激活函数
            nn.BatchNorm1d(512), # 正则化处理
            nn.Dropout(p=0.5), # 丢掉神经元
            nn.Linear(512, 2), # 线性层
            nn.LogSoftmax(dim=1) # 损失函数
    )
    return model_pre

torch.nn.Conv2d(in_channels, out_channels, kernel_size, 
stride=1, padding=0, dilation=1, groups=1, bias=True, padding_mode='zeros')

二、解决方法

一般在卷积层Conv2d后添加正则化BNBatchNormal，使得数据在relu激活前不会因为数据过大而导致网络不稳定，而我在代码中BatchNorm2d的输入通道数与前一层Conv2d的输出通道数不一致，导致报这个错，两者修改一直即可（这里修改为36864即可）。