Pytorch在多GPU下训练模型，分别在单GPU以及多GPU加载模型的方法

我们通常都会把模型拿到服务器上去训练，然后再将模型参数下载到本地端进行测试等工作。但是通常情况下，我们在训练时会使用多块GPU，但是在测试时，有时是多GPU的操作环境，有时是单GPU的操作环境，这就需要我们以不同的方式去加载模型。

一、 多GPU训练，多GPU加载

1.1 训练

os.environ["CUDA_VISIBLE_DEVICES"] = "0,1" #可视化0,1这两块GPU
cnn = model.VGG16().train().cuda()  # 搭建模型
cnn= nn.DataParallel(cnn,device_ids=[0,1]) #将模型同时使用0,1两块GPU进行训练

2.2 保存模型

这里使用的是保存模型中的参数，而不是整个模型都保存下来

 torch.save(cnn.state_dict(), filepath) #保存模型的参数

2.3 加载模型

filepath = "/home/jerry/Pytorch_project/MSP-ReCNN/MSP_CNN/Msp_33.pkl" #模型保存的地址
img_path="/datasets/Dset_Jerry/TrainData_Class" #测试数据集地址
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1" #0,1 GPU可见
net = MDe_CNN() #创建模型结构
net = nn.DataParallel(net, device_ids=[0]) #将数据加载到GPU0上
net.load_state_dict(torch.load(filepath), strict=False) #加载模型参数，strict设置成False
net.eval().cuda()
print("Download done!\n")
test(net,img_path)

二、多GPU训练，单GPU加载

训练和保存模型和多GPU的一致，主要区别就是加载方式

2.1 加载模型

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") #可见化设备，CPU也可
net =  model.ResNet().eval() #定义模型
net = nn.DataParallel(net)  # 在加载时需要把网络也转成DataParallel的
net.to(device)  # 放到定义的设备上
net.load_state_dict(torch.load(args.modle_dir_test),strict=False) #加载模型
test(net)