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

在使用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 
### 如何在 PyTorch加载预训练模型 为了在 PyTorch加载预训练模型,通常有两种主要方式:一种是从官方提供的预训练权重文件中直接加载;另一种则是通过自定义路径加载本地保存模型。下面分别介绍这两种方法并提供相应的代码实例。 #### 方法一:从官方资源加载预训练模型 许多流行的神经网络架构已经在 torchvision.models 或其他库中实现了,并附带了 ImageNet 数据集上的预训练参数。只需指定 `pretrained=True` 即可轻松获取这些预训练好的模型。 ```python import torch from torchvision import models model = models.resnet18(pretrained=True) # 加载 ResNet-18 预训练模型 ``` 此段代码展示了如何利用 torchvision 库中的 resnet18 函数来创建一个带有 ImageNet 上预训练权重的 ResNet-18 模型[^1]。 #### 方法二:从本地磁盘加载保存模型 当需要恢复之前训练过的特定版本的模型时,则可以采用这种方式。这涉及到先序列化(即保存)整个模型的状态字典到硬盘上,在后续运行期间再反序列化回来。 ```python # 假设 'checkpoint.pth' 是先前保存的一个检查点文件名 checkpoint_path = "path/to/checkpoint.pth" device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') # 初始化相同结构的新模型对象 model = TheModelClass(*args, **kwargs) # 将模型移动至 GPU (如果可用的话) model.to(device) # 使用 map_location 参数确保即使是在不同设备间也能正确读取数据 state_dict = torch.load(checkpoint_path, map_location=device)[^2] # 载入状态字典更新当前模型参数 model.load_state_dict(state_dict) # 设置为评估模式 model.eval() ``` 上述代码片段说明了怎样使用 `torch.load()` 来处理可能存在于 CPU 和 GPU 不同环境之间的兼容性问题,同时也强调了设置好目标计算平台的重要性。
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值