迷人的bug--torch.load

利用Google Colab跑了50代的EDSR超分神经网络，然后把网络模型下载到win10上做测试，结果，一直出错，卡了好久

结果百度到这一文章：Pytorch load深度模型时报错：RuntimeError: cuda runtime error (10) : invalid device ordinal

 

引起这种报错的原因是因为pytorch在save模型的时候会把显卡的信息也保存，当重新load的时候，发现不是同一一块显卡就报错invalid device ordinal

知道之后，我赶紧跑的colab上做测试，结果又报错

训练默认使用了GPU，而我测试任然使用的是CPU的，所以出现以上错误

 

下面是可以使用的GPU下的测试：

from __future__ import print_function
import argparse
import torch
import torch.backends.cudnn as cudnn
from PIL import Image
from torchvision.transforms import ToTensor

import numpy as np

# ===========================================================
# Argument settings
# ===========================================================
parser = argparse.ArgumentParser(description='PyTorch Super Res Example')
parser.add_argument('--input', type=str, required=False, default='/content/drive/My Drive/app/EDSR_medical/path/rgb512.tif', help='input image to use')
parser.add_argument('--model', type=str, default='/content/drive/My Drive/app/EDSR_medical/path/EDSR_model_50.pth', help='model file to use')
parser.add_argument('--output_filename', type=str, default='/content/drive/My Drive/app/EDSR_medical/path/out_EDSR_50.tif', help='where to save the output image')
args = parser.parse_args()
print(args)


# ===========================================================
# input image setting
# ===========================================================
GPU_IN_USE = torch.cuda.is_available()
img = Image.open(args.input)
red_c, green_c, blue_c = img.split()

red_c.save('/content/drive/My Drive/app/EDSR_medical/path/red_c.tif')
green_c.save('/content/drive/My Drive/app/EDSR_medical/path/green_c.tif')
all_black = Image.open('/content/drive/My Drive/app/EDSR_medical/path/all_black_1024.tif')
# ===========================================================
# model import & setting
# ===========================================================
device = torch.device('cuda' if GPU_IN_USE else 'cpu')
model = torch.load(args.model)#, map_location=lambda storage, loc: storage
model = model.to(device)
data = (ToTensor()(red_c)).view(1, -1, y.size[1], y.size[0])
data = data.to(device)

if GPU_IN_USE:
    cudnn.benchmark = True


# ===========================================================
# output and save image
# ===========================================================
out = model(data)
out = out.cpu()
out_img_r = out.data[0].numpy()
out_img_r *= 255.0
out_img_r = out_img_r.clip(0, 255)
out_img_r = Image.fromarray(np.uint8(out_img_r[0]), mode='L')

out_img_r.save('/content/drive/My Drive/app/EDSR_medical/path/test_r_1024_edsr_50.tif')

out_img_g = all_black
out_img_b = all_black
out_img = Image.merge('RGB', [out_img_r, out_img_g, out_img_b]).convert('RGB')
#out_img = out_img_r
out_img.save(args.output_filename)
print('output image saved to ', args.output_filename)