PyTorch指定单GPU和多GPU训练及保存-加载模型(含CPU)的总结

| 更新：2020.10.25 | fjy2035@foxmail.com

---

前言：本博客基本涵盖single-gpu和multi-gpu的使用，及训练模型的保存和加载。更复杂功能，修改后亦可得到。 查看gpu使用情况和哪些用户在使用gpu：（watch -n [time] nvidia-smi）和（gpustat -cpu） https://github.com/wookayin/gpustat https://pypi.org/project/gpustat/ 关闭服务器 GPU 占用线程：kill -9 PID

注意：Train/Test过程中 inputs 和 labels，以及待训练 model 均加载到GPU中。对小模型来说，多GPU并行运算反而耗时，大模型bath_size远大于GPU数(或加宽加深Hidden-layers)，GPU优势才能体现。增大bath_size，导致预测准确率降低，可增大epoch。

因为pytorch是在第0块gpu上初始化，占用一定空间的显存，所以使用不当会遇到out of memory的问题。以下探讨涵盖single-GPU和Multi-GPU在训练前指定GPU、保存和加载训练模型、GPU和CPU互加载模型三个过程。

1. PyTorch使用指定GPU - 单GPU

直接使用代码 model.cuda(), PyTorch默认从0开始的单GPU: model = Model() if torch.cuda.is_available(): model = model.cuda()

有两种方法可直接指定单GPU:

在终端shell：CUDA_VISIBLE_DEVICES=1 python main.py，表示只有第1块gpu可见，其他gpu不可用。第1块gpu编号已变成第0块，如果依然使用cuda:1会报invalid device ordinal；以下同效。 python代码（2选1）： os.environ["CUDA_VISIBLE_DEVICES"] = "1" # 官方推荐使用 "CUDA_VISIBLE_DEVICES" model = Model() if torch.cuda.is_available(): model = model.cuda() #使用第一个GPU images = images.cuda() labels = labels.cuda() or # 直接定义设备device，并指定起始位置GPU："cuda:0"。或"cuda:1"作为起始位置，编号为0 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # 单GPU运行，且多GPU时可指定起始位置/编号 net = self.model.to(device) # 等效于self.model.cuda() images = self.images.to(device) labels = self.labels.to(device)

Note，“cuda:0"或"cuda"都代表起始device_id为0，系统默认从0开始。可根据需要修改起始位置，如“cuda:1”等效"cuda:0"或"cuda”。

# 任取一个，torch版本不同会有差别 torch.cuda.device(id) # id 是GPU编号 or torch.cuda.set_device(id) or torch.device('cuda') 单GPU中保存训练模型（2选1） state = {'model': self.model.state_dict(), 'epoch': ite} torch.save(state, self.model.name()) or # 直接保存 torch.save(self.model.state_dict(), 'Mymodel.pth') # 当前目录 测试，单GPU/CPU中加载 single-gpu 训练模型（3选1）详解参考第3部分：[GPU和CPU互加载模型参数] (3. PyTorch使用指定GPU训练 - 其他问题详解（含CPU）) checkpoint = torch.load(self.model.name()) self.model.load_state_dict(checkpoint['model']) or # 直接加载 self.model.load_state_dict(torch.load('Mymodel.pth')) or # load gpu or cpu if torch.cuda.is_available(): # gpu self.model.load_state_dict(torch.load('Mymodel.pth')) else: # cpu 官方推荐CPU的加载方式 checkpoint = torch.load(self.model.name(),map_location=lambda storage, loc: storage) self.model.load_state_dict(checkpoint['model'])

2. PyTorch使用指定GPU - 多GPU（DataParallel）

仍有两种方法可直接指定多GPU:

在终端shell：CUDA_VISIBLE_DEVICES=0,1,3 python main.pypython代码： # gpu_ids = [0, 1, 3] # 或 os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,3" # os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(map(str, [0, 1, 3])) os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,3" # CUDA_VISIBLE_DEVICES 表当前可被python程序检测到的显卡 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # 多GPU时可指定起始位置/编号 # 若不加if项，也不报错，但训练可能会变成单GPU if torch.cuda.device_count() > 1: # 查看当前电脑可用的gpu数量，或 if len(gpu_ids) > 1: print("Let's use", torch.cuda.device_count(), "GPUs!") # self.model = torch.nn.DataParallel(self.model, device_ids=gpu_ids) self.model = torch.nn.DataParallel(self.model) # 声明所有设备 net = self.model.to(device) # 从指定起始位置开始，将模型放到gpu或cpu上 images = self.images.to(device) # 模型和训练数据都放在主设备 labels = self.labels.to(device)

Note：使用多GPU训练，单用 model = torch.nn.DataParallel(model)，默认所有存在的显卡都会被使用。

多GPU中保存训练模型（3选1） if isinstance(self.model,torch.nn.DataParallel): # 判断是否并行 self.model = self.model.module state = {'model': self.model.state_dict(), 'epoch': ite} torch.save(state, self.model.name()) # No-module or if isinstance(self.model, torch.nn.DataParallel): torch.save(self.model.module.stat_dict, 'Mymodel') # No-module else: torch.save(self.model.stat_dict, 'Mymodel') # No-module or # 直接保存 torch.save(self.model.state_dict(), 'Mymodel.pth') # is-module 测试，单GPU/多GPU/CPU加载 multi-gpu 训练模型：（3选1）详解参考第3部分：[GPU和CPU互加载模型参数] (3. PyTorch使用指定GPU训练 - 其他问题详解（含CPU）) # ################## 方法 1: add net = torch.nn.DataParallel(net) # 加上module net.load_state_dict(torch.load("model/cnn_train.pth")) # 加上module，再加载model # ################## 方法 2: remove (2选1) net.load_state_dict({k.replace('module.', ''): v for k, v in torch.load("model/cnn_train.pth").items()}) or from collections import OrderedDict state_dict = torch.load("model/cnn_train.pth") # 当前路径 model 文件下 new_state_dict = OrderedDict() # create new OrderedDict that does not contain `module.` for k, v in state_dict.items(): # remove `module.` name = k[7:] # 或 name = k.replace('module.', '') new_state_dict[name] = v net.load_state_dict(new_state_dict)

3. PyTorch使用指定GPU训练 - 其他问题详解（含CPU）

DataParallel：torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0) （1）DataParallel 实现在module级别上的数据并行使用，返回新模型，即将model在每个GPU分别保存一份。 （2）DataParallel 将输入tensor自动划分并分配到多GPU上的多个模型，即每个GPU计算tensor的一部分，所以输入batch_size应大于设备量GPU。 （3）DataParallel 在每个model完成计算后，收集与合并结果然后可返回到某一个GPU集中处理。 Note：多GPU训练使用DataParallel对网络进行封装，因此在原网络结构中添加了一层module。 module：多GPU并行处理的模型 device_ids：GPU编号（默认全部GPU） output_device：输出位置（默认device_ids[0]或cuda:0) dim：tensors被分散的维度，默认0 gpu_ids = [3, 4, 6, 7] # 或os.environ["CUDA_VISIBLE_DEVICES"] = "3,4,6,7" device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # 多GPU时可指定起始位置/编号 # 若不加if项，也不报错，但训练可能会变成单GPU if torch.cuda.device_count() > 1: # 查看当前电脑可用的gpu数量，或 if len(gpu_ids) > 1: print("Let's use", torch.cuda.device_count(), "GPUs!") self.model = torch.nn.DataParallel(self.model, device_ids=gpu_ids) # 声明所有可用设备 net = self.model.to(device) # 模型放在主设备 images = self.images.to(device) # 训练数据放在主设备 labels = self.labels.to(device) 训练过程中，若用model的子模块： model = Net() # 在单GPU中 out = model.fc(input) model = Net() # 在DataParallel中，调用并行网络中定义的网络层 model = torch.nn.DataParallel(model) out = model.module.fc(input) 测试过程中，GPU和CPU互加载模型参数：参考博客 [gpu和cpu互加载模型参数] (https://blog.csdn.net/bc521bc/article/details/85623515) # 假设只保存了模型的参数(model.state_dict())到文件名为modelparameters.pth, model = Net() # cpu -> cpu or gpu -> gpu: checkpoint = torch.load('modelparameters.pth') model.load_state_dict(checkpoint) # cpu -> gpu 1 torch.load('modelparameters.pth', map_location=lambda storage, loc: storage.cuda(1)) # gpu 1 -> gpu 0 torch.load('modelparameters.pth', map_location={'cuda:1':'cuda:0'}) # gpu -> cpu torch.load('modelparameters.pth', map_location=lambda storage, loc: storage) # 特殊情况 torch.load(opt.model,map_location='cpu')

4. 完整代码示意

# coding: utf-8 # coding: GBK import torch import torchvision import torchvision.transforms as transforms import numpy as np import torch.nn as nn import torch.nn.functional as F import torch.optim as optim import matplotlib.pyplot as plt from torch.autograd import Variable from torch.backends import cudnn # 若使用服务器多卡训练 import os from collections import OrderedDict # 指定对程序可见的GPU编号 # 表示只有第0,1,3块GPU可见，其他GPU不可用，并且第1块GPU默认编号就是第0块 os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,3' # torch.cuda.current_device() # torch.cuda.initialized = True # 定义数据转换transformer transform = transforms.Compose( [transforms.ToTensor(), # (H,W,C)转换为(C,H,W) 并且值为[0, 1.] # transforms.Resize((32, 32)), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))] # 归一化 ) # 加载数据 train_set = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform) train_loader = torch.utils.data.DataLoader(train_set, batch_size=10, shuffle=True, num_workers=0) test_set = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform) test_loader = torch.utils.data.DataLoader(test_set, batch_size=10, shuffle=False, num_workers=0) classes = ['plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck'] # ############################################################ 定义网络 简单的CNN class CNN(nn.Module): def __init__(self): super(CNN, self).__init__() self.conv1 = nn.Conv2d(3, 6, 5) self.conv2 = nn.Conv2d(6, 16, 5) self.pool = nn.MaxPool2d(2, 2) self.fc1 = nn.Linear(16 * 5 * 5, 120) self.fc2 = nn.Linear(120, 84) self.fc3 = nn.Linear(84, 10) def forward(self, x): h1 = self.pool(F.relu(self.conv1(x))) h2 = self.pool(F.relu(self.conv2(h1))) h2 = h2.view(-1, 16 * 5 * 5) h3 = self.fc1(h2) h4 = self.fc2(h3) h5 = self.fc3(h4) return h5 # 实例化模型 net = CNN() # 使用（多）GPU训练 # 定义device，“cuda:0” 只代表起始的device_id为 0 device = torch.device('cuda:2' if torch.cuda.is_available() else 'cpu') print("GPU or CPU is available: ", device) if torch.cuda.device_count() > 1: # multi-gpu print('Lets use', torch.cuda.device_count(), 'GPUs!') net = nn.DataParallel(net) net.to(device) # 定义损失函数(loss function)和优化器(optimizer) criterion = nn.CrossEntropyLoss() # classification criterion and regression criterion optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9) # ############################################################ 训练 net.train() # 在训练时启用BN层和Dropout层，对模型进行更改 for epoch in range(1): # 循环遍历数据集的次数 running_loss = 0. # enumerate 将一个可遍历的数据对象(如列表、元组或字符串)组合为一个索引序列，并列出数据和数据下标，常用于for for i, data in enumerate(train_loader, 0): images, labels = data # get the inputs # 需要将输入网络的数据复制到GPU images = images.to(device) labels = labels.to(device) optimizer.zero_grad() # 清空过往梯度缓存区域 # 经典四步 outs = net(images) loss = criterion(outs, labels) # forward，前向传播 loss.backward() # backward，后向传播，计算当前梯度 optimizer.step() # optimize，根据梯度更新网络参数 # 打印loss running_loss += loss.item() if i % 2000 == 1999: # print every 2000 mini-batches print('[epoch %d, iter %d] loss : %.3f' % (epoch + 1, i + 1, running_loss / 2000)) running_loss = 0. print('Finish Training!') torch.save(net.state_dict(), 'model/cnn_train.pth') # multi-gpu has module，single-gpu or cpu has No-module print('Finish save the model!') # ############################################################ 测试 # os.environ['CUDA_VISIBLE_DEVICES'] = '0' device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') # train，device：(cuda：0-n)；test，device: (cuda:0) print('Test is running:', device) # ################## load 'model.pth' # def load_gpu_cpu_pth(self, net, path): # single-gpu, multi-gpu and cpu, auto - load and convert types in 'model.pth' if 'gpu' if torch.cuda.is_available() else 'cpu' == 'gpu': state_dict = torch.load("model/cnn_train.pth") else: state_dict = torch.load("model/cnn_train.pth", map_location=lambda storage, loc: storage) new_state_dict = OrderedDict() # create new OrderedDict that does not contain `module.` if isinstance(net, torch.nn.DataParallel): # 判断模型net是否并行 print('\nThe source model is isinstance in test') if list(state_dict.keys())[0][:6] == 'module': print("The loaded model always contains 'module'") # 直接加载 -- Model is module net.load_state_dict(state_dict) else: print("The loaded model is adding 'module'...") # Method 1: add -- Model is No-module net = torch.nn.DataParallel(net) # add module net.load_state_dict(state_dict) # then load model print("Finish loading 'model.pth'\n") else: print('\nThe source model is not isinstance in test') if list(state_dict.keys())[0][:6] == 'module': print("The loaded model is removing 'module'") # Method 2: remove (2选1) # net.load_state_dict({k.replace('module.', ''): v for k, v in torch.load("model/cnn_train.pth").items()}) for k, v in state_dict.items(): # remove `module.` name = k[7:] # 或 name = k.replace('module.', '') new_state_dict[name] = v net.load_state_dict(new_state_dict) else: print("The loaded model always contains 'module'") # 直接加载 -- Model is No-module net.load_state_dict(state_dict) print("Finish loading 'model.pth'\n") # ################## test net.to(device) net.eval() # 在评测时不启用BN层和Dropout层，冻结后这两个操作不会对模型进行更改 correct_test = 0 total_test = 0 for epoch in range(1): # range(start, stop[, step])，默认从0开始，range(0)是空集 for data in test_loader: images_test, labels_test = data # 需要将测试网络的数据复制到GPU images_test = images_test.to(device) labels_test = labels_test.to(device) # 评估预测 # 虽然使用net.eval()，但在验证阶段有时报错out of memory，可能是梯度不回传，造成梯度累加。故取消验证阶段的loss。 with torch.no_grad(): outs_test = net(images_test) _, predict = torch.max(outs_test.data, 1) total_test += labels_test.size(0) correct_test += (predict == labels_test).sum().item() print('Accuracy of the network on the 10000 test images: %d %%' % ( 100 * correct_test / total_test)) print('Finish Testing!')

5. 拓展其他博客

[1] CPU加载GPU训练model和GPU加载CPU训练model： https://www.ptorch.com/news/74.html [2] 单机多卡并行训练、多机多GPU训练和DistributedDataParallel解决显存使用不平衡： https://blog.csdn.net/weixin_47196664/article/details/106542016?utm_medium=distribute.wap_relevant.none-task-blog-title-2