PyTorch加载保存的ResNet模型后测试精度骤降原因排查

问题背景 #

训练ResNet50模型并将权重保存为myresnet.pt，训练阶段验证精度达95%，独立测试集精度为93%。但在新建空白脚本中加载该权重文件后，测试集精度骤降至20.6%。

训练代码（文件1） #

## This is file 1 ##
model = resnet50()
model.to(device)
optimizer = Adam(model.parameters(), eps=1e-08, lr = 0.001, weight_decay=1e-4, betas=(0.9, 0.999))
criterion = nn.CrossEntropyLoss()
scheduler = lr_scheduler.MultiplicativeLR(optimizer, lr_lambda=lmbda)

model.train()
train_model(model, criterion, optimizer, scheduler, num_epochs=num_epochs)
torch.save(model.state_dict(), 'myresnet.pt')

model.eval()
loss, acc, y_pred, y_true = test_model(model, criterion)

权重加载代码（文件2） #

## This is file 2 ##
model = models.resnet50()
state_dict = torch.load('myresnet.pt')
model.load_state_dict(state_dict)
model.eval()
model.to(device)
loss, acc, y_pred, y_true = test_model(model, criterion)

已尝试操作 #

• 多次运行排除文件损坏可能
• 将加载代码追加到原训练脚本后，测试精度恢复正常

相关DataLoader代码 #

batch_size  = 4
image_size = [32, 32]
random_seed = int(time.time()//1000)
random.seed(random_seed)

def random_ratio_3d(): return [randrange(0, 100)/100, randrange(0, 100)/100, randrange(0, 100)/100]
tmp_mean, tmp_std = random_ratio_3d(), random_ratio_3d()

#data_train_path = 'data/train/'
data_test_path = 'data/test/'
#train_dataset = ImageFolder(data_train_path, Compose([Resize(image_size), ToTensor(), Normalize(mean=tmp_mean, std=tmp_std)]))
test_dataset = ImageFolder(data_test_path, Compose([Resize(image_size), ToTensor(), Normalize(mean=tmp_mean, std=tmp_std)]))

#train_idx, valid_idx = train_test_split(list(range(len(train_dataset))), test_size=0.2, random_state=random_seed)
datasets = {}
#datasets['train'] = Subset(train_dataset, train_idx)
#datasets['valid'] = Subset(train_dataset, valid_idx)
datasets['test']  = test_dataset

dataloaders, batch_num = {}, {}
num_workers = 6 # half of cpu core number
#dataloaders['train'] = DataLoader(datasets['train'], batch_size=batch_size, shuffle=True, num_workers=num_workers)
#dataloaders['valid'] = DataLoader(datasets['valid'],batch_size=batch_size, shuffle=True, num_workers=num_workers)
dataloaders['test']  = DataLoader(datasets['test'], batch_size=batch_size, shuffle=True, num_workers=num_workers)
#batch_num['train'], batch_num['valid'], batch_num['test'] = len(dataloaders['train']), len(dataloaders['valid']), len(dataloaders['test'])
batch_num['test'] = len(dataloaders['test'])

问题原因分析 #

核心原因是测试数据的归一化参数与训练时不匹配：

1. DataLoader代码中，tmp_mean和tmp_std通过random_ratio_3d()随机生成，且随机种子random_seed基于当前时间戳生成。训练脚本和新加载脚本的运行时间不同，导致生成的均值、标准差完全不同。
2. 训练时模型是基于训练脚本生成的那组均值/标准差归一化后的图像学习的，新脚本测试时用了另一组随机参数处理测试图像，数据分布完全偏离模型训练时的分布，因此精度暴跌。
3. 把加载代码追加到原训练脚本后，复用了训练时生成的tmp_mean和tmp_std，归一化参数匹配，精度自然恢复正常。

另外还有两个次要问题：

• 测试集DataLoader设置了shuffle=True，测试阶段无需打乱数据，建议改为shuffle=False
• 加载模型时，建议先执行model.to(device)再加载权重，避免设备不匹配的潜在问题（但这不是本次精度暴跌的原因）

解决方法：

• 训练时将使用的tmp_mean和tmp_std保存为文件（比如用numpy.save），测试时加载该文件，确保归一化参数完全一致
• 固定随机种子，避免每次运行生成不同的归一化参数（但更可靠的方式是直接保存训练时的均值/标准差）

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