注：本文針對單個伺服器上多塊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,ids)
    return x
... ...

然後就可以看執行結果啦，nvidia-smi檢視GPU使用情況：

可以看到0和4都被使用啦

以上這篇Pytorch 多塊GPU的使用詳解就是小編分享給大家的全部內容了，希望能給大家一個參考，也希望大家多多支援我們。