pytorchgpu测试_Pytorch中利用GPU训练,CPU测试

要先利用GPU训练,CPU测试,那么在模型训练时候,是能保存模型的参数而不能保存整个模型,可见保存、加载、测试

这里主要讲一点重要的,即在pytorch 1.6的版本中训练模型保存时,不能直接使用

torch.save(state_r, model_out_r_path)

否则,在CPU测试时,由于版本的不兼容会导致出错,正确使用如下:

torch.save(state_g, model_out_g_path, _use_new_zipfile_serialization=False)

而后,在模型加载时候,要限定map_location = cpu的加载:

checkpoint_r = torch.load(opt.model_r,map_location='cpu')

model_rt.load_state_dict(checkpoint_r['model'])

model_rt.eval()

optimizer_rt.load_state_dict(checkpoint_r['optimizer'])

epochs_rt = checkpoint_r['epoch']

GPU下训练的模型即可方便的在CPU环境中测试了

若模型已经训练保存,但是有没有使用_use_new_zipfile_serialization=False来进行约束,那么,可以在pytorch 1.6中直接加载模型,然后再次使用torch.save进行保存为非zip格式:

#在torch 1.6版本中重新加载一下网络参数

model = MyModel().cuda() # 先预加载模型

model.load_state_dict(torch.load(pre_model)) #加载模型参数,pre_model为直接没有设置_use_new_zipfile_serialization=False而模型的参数模型

#重新保存网络参数

torch.save(model.state_dict(), pre_model,_use_new_zipfile_serialization=False)

另外,网络模型的参数一般较大,在CPU中测试时,可能会导致内存爆满而使得程序无法运行,此时,可以将模型的输入进行修改。

以下,是我将input的图像设置为512*512,导致了16G内存爆满

之后,我将input image裁剪成128 * 128进行CPU测试,便可大功告成。

具体的CPU测试代码如下,可供参考与理解:

from __future__ import print_function

import argparse

import torch

from PIL import Image

from torchvision.transforms import Compose, ToTensor, CenterCrop

import torchvision

from model import LapSRN_r, LapSRN_g

import torch.optim as optim

# Argument settingsss

parser = argparse.ArgumentParser(description='PyTorch LapSRN')

parser.add_argument('--input', type=str, required=False, default='/content/drive/My Drive/TestImg/val/44-512pix-speed7-ave1.tif', help='input image to use')

parser.add_argument('--model_r', type=str, default='LapSRN_r_epoch_10.pth', help='model file to use')

parser.add_argument('--model_g', type=str, default='LapSRN_g_epoch_10.pth', help='model file to use')

parser.add_argument('--outputHR2', type=str, default='./73_LapSRN_R_epochs100_HR2.tif', help='where to save the output image')

parser.add_argument('--outputHR4', type=str, default='./73_LapSRN_R_epochs100_HR4.tif', help='where to save the output image')

parser.add_argument('--cuda', action='store_true', help='use cuda')

opt = parser.parse_args()

print(opt)

model_rt = LapSRN_r()

model_gt = LapSRN_g()

optimizer_rt = optim.Adagrad(model_rt.parameters(), lr=1e-3, weight_decay=1e-5)

optimizer_gt = optim.Adagrad(model_gt.parameters(), lr=1e-3, weight_decay=1e-5)

checkpoint_r = torch.load(opt.model_r,map_location='cpu')

model_rt.load_state_dict(checkpoint_r['model'])

model_rt.eval()

optimizer_rt.load_state_dict(checkpoint_r['optimizer'])

epochs_rt = checkpoint_r['epoch']

checkpoint_g = torch.load(opt.model_g,map_location='cpu')

model_gt.load_state_dict(checkpoint_g['model'])

model_rt.eval()

optimizer_gt.load_state_dict(checkpoint_g['optimizer'])

epochs_gt = checkpoint_g['epoch']

transform = Compose(

[

ToTensor(),

])

img = Image.open(opt.input).convert('RGB')

r, g, _ = img.split()

r = transform(r)

r = r.unsqueeze(0)

g = transform(g)

g = g.unsqueeze(0)

HR_2_r, HR_4_r = model_rt(r)

HR_2_g, HR_4_g = model_gt(g)

black_2 = torch.zeros(1, 128 * 2, 128 * 2).unsqueeze(0)

HR_2 = torch.cat((HR_2_r.squeeze(0), HR_2_g.squeeze(0), black_2.squeeze(0))).unsqueeze(0)

black_4 = torch.zeros(1, 128 * 4, 128 * 4).unsqueeze(0)

HR_4 = torch.cat((HR_4_r.squeeze(0), HR_4_g.squeeze(0), black_4.squeeze(0))).unsqueeze(0)

torchvision.utils.save_image(HR_2, opt.outputHR2, padding=0)

torchvision.utils.save_image(HR_4, opt.outputHR4, padding=0)

print("saved !")

你可能感兴趣的:(pytorchgpu测试)