6-4 关于pytorch 模型GPU转CPU加载

1

raws = Variable(raws)
# raws = raws.cuda()

device = torch.device('cpu')
#device = torch.device('cuda')
model.load_state_dict(torch.load(model_path, map_location=device))

2 训练时使用了DataParallel：

https://blog.csdn.net/c654528593/article/details/81539441

3其他：

参考地块分割云端cpu版本的注意点：

  （1）#加载模型
        NAME_MODEL = 'dikuai_mclinknet34_epoch_best_model2150'  # 有效
        model = Unet_res34(n_classes)
        model.eval()
        model = torch.nn.DataParallel(model)#训练时batch_size=2,现在test加载为1,所以报错：
        # （size mismatch for b7.bias: copying a param with shape torch.Size([1]) from checkpoint, the shape in current model is torch.Size([2]).）,加入这一步就可以了.
        #这里会自动检测使用GPU或CPU即使模型加载时使用cpu(如果电脑有GPU，还是会使用GPU)

device = torch.device('cpu')
        # device = torch.device('cuda')
        model.load_state_dict(torch.load('weights/' + NAME_MODEL + '.pth', map_location=device),False)#从另外一台机器得到训练模型进行加载时，设置为False，否则会报错

"""
 （2）tta部分：

#训练时batch_size = 2, 所以导致输出多了一个维度（2）需要去掉,

mask2 = mask2[0, 0]

mask22=mask22.transpose(1, 0, 2)

mask22=mask22[0]

mask3 = mask2[0] + mask22

mask3_copy = mask2[0] + mask22

  (3) 需注意训练时
/home/comin/Downloads/DeepGlobe-Road-Extraction-Challenge-master/framework.py中屏蔽#self.net = torch.nn.DataParallel(self.net, device_ids=range(torch.cuda.device_count()))
训练时batch_size=2或者batch_size=1
"""

  (4) n_classes=1 #注意此处计算loss时需要为2,不计算时需要为1（与训练时参数一致）
***ps:(2)和(3)不需要关注，是错误的

最后

以上就是年轻店员为你收集整理的6-4 关于pytorch 模型GPU转CPU加载的全部内容，希望文章能够帮你解决6-4 关于pytorch 模型GPU转CPU加载所遇到的程序开发问题。

如果觉得靠谱客网站的内容还不错，欢迎将靠谱客网站推荐给程序员好友。