PyTorch训练与测试模型分别在GPU和CPU上的灵活转换

由于抢不到实验室的gpu资源，决定把毕设需要的训练部署在远程cpu服务器上（gpu服务器太贵搞不起），为了测试模型在gpu和cpu之间来回转换时是否会影响模型的性能，特别做了以下实验。

- 首先明确测试的四种情况：

Device_Train	Device_Test
GPU	GPU
GPU	CPU
CPU	GPU
CPU	CPU

1. 把数据部署在GPU/CPU上：

   （1）GPU：

   device = torch.device('cuda:0'if torch.cuda.is_available() else 'cpu')
   

   (2) CPU:

   device = torch.device('cpu') #默认即在cpu上
   

   (3) 把网络和数据集部署到对应的设备上：

   net = Network().to(device)
   images,labels = images.to(device),labels.to(device)
   

2. 保存模型

   torch.save(net.state_dict(),path)
   

   一般用PyTorch训练出来的模型保存的文件名以.pt或者.pth结尾。

   如果用到了多个GPU并行处理(torch.nn.DataParallel)，则保存模型时用以下方式保存：

   torch.save(net.module.state_dict(),path)
   

3. 载入模型

   net = Network()
   net.load_state_dict(torch.load(path,map_location=device))
   net.to(device)
   

4. 训练和测试的区别

（1）训练时，对网络进行设置：

net.train()

开启BN和DropOut

（2）测试时，对网络进行设置：

net.eval()

关闭BN和DropOut的功能，固定参数，避免在测试过程中修改网络的参数。

- 实验结果

​ 本实验的数据集用的是FashinMNIST，训练代码详见本文：https://blog.csdn.net/weixin_41848012/article/details/104776337

Device_Train	Device_Test	Accuracy	Time
GPU	GPU	0.8551	0.8207
GPU	CPU	0.8551	0.9764
CPU	GPU	0.8719	0.8231
CPU	CPU	0.8719	0.9388

​ 从表格里面可以看出，训练好的模型用不同的设备加载时，不影响性能，只影响程序运行的时间。GPU和CPU之间可以实现数据部署的自由转换。