python pytorch gpu_Pytorch 多块GPU的使用详解

注：本文针对单个服务器上多块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. 更改输出数据(