Pytorch 多块GPU的使用详解

yipeiwu_com5年前Python基础

注:本文针对单个服务器上多块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的使用详解就是小编分享给大家的全部内容了,希望能给大家一个参考,也希望大家多多支持【听图阁-专注于Python设计】。

相关文章

Django model反向关联名称的方法

如果你在 ForeignKey或 ManyToManyField字段上使用 related_name属性,你必须总是为该字段指定一个唯一的反向名称。但在抽象基类上这样做就会引发一个很严重...

python的tkinter布局之简单的聊天窗口实现方法

python的tkinter布局之简单的聊天窗口实现方法

本文实例展示了一个python的tkinter布局的简单聊天窗口。分享给大家供大家参考之用。具体方法如下: 该实例展示的是一个简单的聊天窗口,可以实现下方输入聊天内容,点击发送,可以增加...

Python设计模式之备忘录模式原理与用法详解

Python设计模式之备忘录模式原理与用法详解

本文实例讲述了Python设计模式之备忘录模式原理与用法。分享给大家供大家参考,具体如下: 备忘录模式(Memento Pattern):不破坏封装性的前提下捕获一个对象的内部状态,并在...

Flask模拟实现CSRF攻击的方法

Flask模拟实现CSRF攻击的方法

CSRF CSRF全拼为Cross Site Request Forgery,译为跨站请求伪造。 CSRF指攻击者盗用了你的身份,以你的名义发送恶意请求。 包括:以你名义发送邮件,发消息...

python使用sorted函数对列表进行排序的方法

本文实例讲述了python使用sorted函数对列表进行排序的方法。分享给大家供大家参考。具体如下: python提供了sorted函数用于对列表进行排序,并且可以按照正序或者倒序进行排...