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