PyTorch CNN不同输入尺寸加载模型报权重不匹配如何解决

问题背景 #

• 训练集、验证集输入尺寸为256，待预测测试集输入尺寸为496
• 已保存训练完成的模型权重，加载权重用于测试集预测时报错

原模型定义 #

class Net(nn.Module):
    def __init__(self, shape):
        super(Net,self).__init__()
        self.conv1 = nn.Conv1d(shape,1,1)
        self.batch1 = nn.BatchNorm1d(1)
        self.avgpl1 = nn.AvgPool1d(1, stride=1)
        self.fc1 = nn.Linear(1,3)
    
    #forward method 
    def forward(self,x):
        x = self.conv1(x)
        x = self.batch1(x)
        x = F.relu(x)
        x = self.avgpl1(x)
        x = torch.flatten(x,1)
        x = F.log_softmax(self.fc1(x))
        return x

报错信息 #

RuntimeError: Error(s) in loading state_dict for Net:
size mismatch for conv1.weight: copying a param with shape 
torch.Size([1, 256, 1]) from checkpoint, the shape in current model is torch.Size([1, 494, 1]).

报错核心：加载state_dict时conv1.weight尺寸不匹配，checkpoint存储的权重形状为torch.Size([1, 256, 1])，当前初始化模型的对应权重形状为torch.Size([1, 494, 1])。

问题根源 #

1. 模型定义时混淆了Conv1d的参数含义：nn.Conv1d第一个参数是输入通道数（特征维度），不是输入序列长度，原代码将输入序列长度作为in_channels传入，导致卷积核尺寸和输入序列长度绑定，输入长度变化时卷积核维度直接不匹配。
2. 池化层使用固定核大小的AvgPool1d，没有做维度对齐，配合错误的卷积参数定义，无法适配可变长度输入。

解决方案 #

根据实际场景二选一处理：

方案1：测试集预处理逻辑对齐训练集（优先选择） #

如果训练时输入的256是固定特征维度，测试集出现494/496的维度是预处理逻辑错误导致：

• 核对训练时的预处理流程（特征筛选、切片、变换等步骤），将测试集输入维度处理为和训练集完全一致的256
• 初始化模型时传入shape=256，即可正常加载训练好的权重做预测

方案2：修正模型结构适配任意输入长度（需要重新训练或迁移权重） #

如果确实需要适配不同长度的输入，需要修改模型结构，从根源上消除维度绑定问题：

1. 将Conv1d的in_channels固定为数据的实际特征通道数（比如单变量时序传入1，多变量时序传入变量数，不要传入序列长度）
2. 将固定核大小的平均池化替换为AdaptiveAvgPool1d(1)做全局自适应池化，不管输入序列长度是多少，池化后输出维度固定，可直接对接后续全连接层
3. 补充log_softmax的dim参数避免运行警告
   修正后的模型代码如下：

class Net(nn.Module):
    def __init__(self, in_channels):
        super(Net,self).__init__()
        self.conv1 = nn.Conv1d(in_channels, 1, 1)
        self.batch1 = nn.BatchNorm1d(1)
        self.avgpl1 = nn.AdaptiveAvgPool1d(1) # 自适应全局池化，适配任意序列长度
        self.fc1 = nn.Linear(1, 3)
    
    def forward(self,x):
        x = self.conv1(x)
        x = self.batch1(x)
        x = F.relu(x)
        x = self.avgpl1(x)
        x = torch.flatten(x, 1)
        x = F.log_softmax(self.fc1(x), dim=1)
        return x

• 如果可以重新训练：用修正后的模型在训练集上重新训练，后续不管输入序列长度是256还是496，都可以直接加载权重预测，不会出现维度不匹配问题
• 如果不能重新训练：先按训练时的shape=256初始化模型加载完整权重，再将conv1替换为适配新通道数的卷积层，把原卷积权重复制到新卷积对应位置，新增通道对应的权重做零初始化/随机初始化后，可直接用于推理。

内容的提问来源于stack exchange，提问作者user13071308