加载MLP模型state_dict时参数尺寸不匹配问题求助

核心原因 #

你加载的checkpoint权重是用不同压缩率（compression_rate）训练的MLP模型生成的，导致中间层维度和当前模型不一致：

• checkpoint中fc1输出维度为627（对应compression_rate≈0.8，784×0.8≈627）
• 当前模型fc1输出维度为78（对应compression_rate≈0.1，784×0.1≈78）

具体解决方案 #

方案1：匹配训练时的compression_rate #

确保当前加载模型的compression_rate和训练checkpoint时使用的完全一致，这样模型结构会和训练时对齐，参数形状自然匹配：

# 替换为训练checkpoint时实际使用的compression_rate值
compression_rate = 0.8  
model = MLP()
model.load_state_dict(torch.load("your_checkpoint.pth"))

方案2：手动修改state_dict适配当前模型 #

如果必须使用当前的小维度模型，可从checkpoint权重中截取对应维度的参数，调整后加载（注：此方法会丢弃部分预训练参数，可能影响模型性能）：

import torch

# 加载原始checkpoint
checkpoint = torch.load("your_checkpoint.pth")
# 初始化当前模型
compression_rate = 0.1
model = MLP()

# 调整fc1参数：截取前78行/元素适配当前输出维度
checkpoint['fc1.weight'] = checkpoint['fc1.weight'][:78, :]
checkpoint['fc1.bias'] = checkpoint['fc1.bias'][:78]
# 调整fc2参数：截取前78列适配当前输入维度
checkpoint['fc2.weight'] = checkpoint['fc2.weight'][:, :78]

# 加载调整后的权重
model.load_state_dict(checkpoint)

方案3：排查是否加载了错误模型的权重 #

确认你加载的checkpoint是否来自MLP模型，而非MLP_MNIST——两者结构完全不同，误加载必然导致形状不匹配。若属于这种情况，换用对应模型初始化即可：

# 若checkpoint是MLP_MNIST的权重，初始化对应模型
model = MLP_MNIST()
model.load_state_dict(torch.load("mlp_mnist_checkpoint.pth"))

预防措施 #

训练时将压缩率和权重一起保存，避免后续加载时结构不匹配：

torch.save({
    'model_state_dict': model.state_dict(),
    'compression_rate': compression_rate
}, "checkpoint_with_rate.pth")

加载时先读取压缩率，再初始化对应结构的模型：

checkpoint = torch.load("checkpoint_with_rate.pth")
compression_rate = checkpoint['compression_rate']
model = MLP()
model.load_state_dict(checkpoint['model_state_dict'])

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