注：本文针对单个服务器上多块GPU的使用，不是多服务器多GPU的使用。
 
在一些实验中，由于Batch_size的限制或者希望提高训练速度等原因，我们需要使用多块GPU。本文针对Pytorch中多块GPU的使用进行说明。
 
1. 设置需要使用的GPU编号
 
import os
 
os.environ["CUDA_VISIBLE_DEVICES"] = "0,4"
 
ids = [0,1]
 
比如我们需要使用第0和第4块GPU，只用上述三行代码即可。
 
其中第二行指程序只能看到第1块和第4块GPU；
 
第三行的0即为第二行中编号为0的GPU；1即为编号为4的GPU。
 
2.更改网络，可以理解为将网络放入GPU
 
class CNN(nn.Module):
 
def __init__(self):
 
super(CNN,self).__init__()
 
self.conv1 = nn.Sequential(
 
......
 
)
 
......
 
self.out = nn.Linear(Liner_input,2)
 
......
 
def forward(self,x):
 
x = self.conv1(x)
 
......
 
output = self.out(x)
 
return output,x
 
cnn = CNN()
 
# 更改，.cuda()表示将本存储到CPU的网络及其参数存储到GPU！
 
cnn.cuda()
 
3. 更改输出数据(如向量/矩阵/张量):
 
for epoch in range(EPOCH):
 
epoch_loss = 0.
 
for i, data in enumerate(train_loader2):
 
image = data['image'] # data是字典，我们需要改的是其中的image
 
#############更改！！！##################
 
image = Variable(image).float().cuda()
 
############################################
 
label = inputs['label']
 
#############更改！！！##################
 
label = Variable(label).type(torch.LongTensor).cuda()
 
############################################
 
label = label.resize(BATCH_SIZE)
 
output = cnn(image)[0]
 
loss = loss_func(output, label) # cross entropy loss
 
optimizer.zero_grad() # clear gradients for this training step
 
loss.backward() # backpropagation, compute gradients
 
optimizer.step()
 
... ...
 
4. 更改其他CPU与GPU冲突的地方
 
有些函数必要在GPU上完成，例如将Tensor转换为Numpy，就要使用data.cpu().numpy()，其中data是GPU上的Tensor。
 
若直接使用data.numpy()则会报错。除此之外，plot等也需要在CPU中完成。如果不是很清楚哪里要改的话可以先不改，等到程序报错了，再哪里错了改哪里，效率会更高。例如：
 
... ...
 
#################################################
 
pred_y = torch.max(test_train_output, 1)[1].data.cpu().numpy()
 
accuracy = float((pred_y == label.cpu().numpy()).astype(int).sum()) / float(len(label.cpu().numpy()))
 
假如不加.cpu()便会报错，此时再改即可。
 
5. 更改前向传播函数，从而使用多块GPU
 
以VGG为例：
 
class VGG(nn.Module):
 
def __init__(self, features, num_classes=2, init_weights=True):
 
super(VGG, self).__init__()
 
... ...
 
def forward(self, x):
 
#x = self.features(x)
 
#################Multi GPUS#############################
 
x = nn.parallel.data_parallel(self.features,x,ids)
 
x = x.view(x.size(0), -1)
 
# x = self.classifier(x)
 
x = nn.parallel.data_parallel(self.classifier,x,ids)
 
return x
 
... ...
 
然后就可以看运行结果啦，nvidia-smi查看GPU使用情况：
 

 
可以看到0和4都被使用啦
 
以上这篇Pytorch 多块GPU的使用详解就是小编分享给大家的全部内容了，希望能给大家一个参考，也希望大家多多支持我们。
 
时间： 2019-12-29

查看GPU状态和信息
 
nvidia-smi
 
在这里插入图片描述
 
定时刷新状态，-n（秒），每秒刷新一次
 
watch -n 1 nvidia-smi
 
指定可见GPU
 
终端跑程序指定GPU
 
CUDA_VISIBLE_DEVICES=0 python XXX.py
 
或者在python文件中加入
 
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
 
多块GPU可选
 
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3"
 
pytorch使用单GPU
 
model = model.cuda()
 
pytorch使用多GPU
 
model = torch.nn.DataParallel(model, device_ids=[0,1,2,3]).cuda()
 
使用了gpu0,1,2,3 共4块同时跑
 
其他
 
训练后保存的模型参数，再次导入模型时，出现报错：
 
Missing key(s) in state_dict: ...
 
解决办法：
 在导入参数前加上
 import torch.nn as nn
 from torch.backends import cudnn
 
model = torch.nn.DataParallel(model)
cudnn.benchmark = True

代码
 
# 模型，训练模式
net = Net(opt)
net.train()


# 损失函数，优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), lr, (b1, b2))


# 设定gpu
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3"

# 如果epoch不为0，就加载已经训练过的ckpt
if epoch != 0:
    ckpt = os.path.join(opt.out_path, 'net.pth')
    
    if torch.cuda.is_available():
        device = torch.device("cuda:0")
        net.to(device)
        criterion.to(device)

        # 多块gpu只转移模型
        if torch.cuda.device_count() > 1:
            net = nn.DataParallel(net, device_ids=[0,1,2,3])
            net.to(device)

        net.load_state_dict(torch.load(ckpt))

# 否则不加载
else:
    if torch.cuda.is_available():
        device = torch.device("cuda:0")
        net.to(device)
        criterion.to(device)
        
        if torch.cuda.device_count() > 1:
            net = nn.DataParallel(net, device_ids=[0,1,2,3])
            net.to(device)

 
模型转移
 
主要记载在GPU保存的模型读取到CPU，以及CPU保存的模型读取到GPU，主要的不同在于是否有module，无非就是加上和去掉的问题
 
CPU->GPU
 
# cpu保存的模型在gpu加载的方法
from collections import OrderedDict
new_state_dict = OrderedDict()
for key, value in torch.load(ckpt).items():
    name = 'module.' + key # 加上module
    new_state_dict[name] = value
    
net.load_state_dict(new_state_dict)

 
GPU->CPU
 
from collections import OrderedDict
state_dict_new = OrderedDict()
for key, value in torch.load(ckpt).items():
    name = key[7:] # 去掉module
    state_dict_new[name] = value
    
net.load_state_dict(state_dict_new)

 
或者还可以写得更简单：
 
net.load_state_dict({k.replace('module.', ''): v for k, v in torch.load(model_path, map_location=torch.device('cpu')).items()})
 
微调网络
 
官方方法
 
# 冻结除fc以外的层，再次写一遍fc层，默认求梯度
for param in net.parameters():
	param.requires_grad = False
	net.fc = nn.Linear(512, 100)

 
新方法
 
# 更加灵活决定那些层是不求梯度的
for param in list(net.parameters())[1:49]:
    param.requires_grad = False

首先batch size要大于现有的gpu的数量。
 
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")#cuda:0代表起始的； #device_id为0,如果直接是cuda,同样默认是从0开始，可以根据实际需要修改起始位置，如cuda:1
 
if torch.cuda.device_count() > 1: #查看当前电脑的可用的gpu的数量，若gpu数量>1,就多gpu训练
	model = torch.nn.DataParallel(model)#多gpu训练,自动选择gpu
# 
model.to(device)