pytorch 保存模型，加载预训练模型问题

最新推荐文章于 2024-07-26 20:52:07 发布

「已注销」

最新推荐文章于 2024-07-26 20:52:07 发布

阅读量791

点赞数

CC 4.0 BY-SA版权

分类专栏： PyTorch

本文链接：https://blog.youkuaiyun.com/wangchy29/article/details/107797310

在使用PyTorch进行模型训练时，发现加载预训练模型后，在验证集上的PSNR（峰值信噪比）结果与训练时的验证集PSNR存在显著差异。经过检查，解决方案是在测试前对网络模型调用net.eval()来切换到评估模式。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

在写pytorch代码时，遇到问题：加载预训练模型时在验证集上测试的psnr结果与训练时验证集的psnr差异特别大。

源代码：

pretrained_dict = torch.load('epochG_515.pth')
net.load_state_dict(pretrained_dict)
net = prepare(net)
valdata = Data(root=os.path.join(args.dir_data, args.data_val), args=args, train=False)
valset = DataLoader(valdata, batch_size=1, shuffle=False, num_workers=1)

val_psnr = 0
val_ssim = 0
with torch.no_grad():

    timer_test = util.timer()
    for batch, (lr, hr, filename) in enumerate(valset):
        lr, hr = prepare(lr), prepare(hr)
        sr = net(lr)
        print(sr.shape, hr.shape)
        val_psnr = val_psnr + cal_psnr(hr[0].data.cpu(), sr[0].data.cpu())
        val_ssim = val_ssim + cal_ssim(hr[0].data.cpu(), sr[0].data.cpu())
    print("Test psnr: {:.3f}".format(val_psnr / (len(valset))))
    print('Forward: {:.2f}s\n'.format(timer_test.toc()))
    print(val_ssim

最低0.47元/天解锁文章