迷人的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(, -, y.size[], y.size[])

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[].numpy()

out_img_r *= 255.0

out_img_r = out_img_r.clip(, )

out_img_r = Image.fromarray(np.uint8(out_img_r[]), 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)

巴特西

迷人的bug--torch.load

最新文章

热门文章