PyTorch指定单GPU和多GPU训练及保存-加载模型(含CPU)的总结
| 更新：2020.10.25 | [email protected]
前言：本博客基本涵盖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.py - python代码：
# 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