【Pytorch】模型转GPU计算

需要转换的对象

• 模型
• 损失函数
• 数据（特征数据、标签）

1. .cuda()

'''
.cuda(0) 指定第0块GPU
'''
if torch.cuda.is_available():
    model = model.cuda()
    loss_fn = loss_fn.cuda()
    X = X.cuda()
    y = y.cuda()
    

2. .to(device)

# 第二种方法
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
device0 = torch.device('cuda:0')

model.to(device)
loss_fn.to(device)
X = X.to(device)
y = y.to(device)

3. 多GPU并行计算

3.1 单进程多GPU训练(DP)模式 torch.nn.DataParallel

• 并行的多卡都是由一个进程进行控制，在进行梯度的传播时，是在主GPU上进行的。

1. 将模型布置到多个指定GPU上
   model = torch.nn.DataParallel(model,device_ids=device_list)
2. 指定模型布置的主GPU
   main_dev = 0
model.to(main_dev)
3. 优化器放到主GPU
   optimizer = optim.Adam(model.parameters(), lr=0.001)
optimizer.to(device)
4. 数据放到主GPU
   data.to(main_dev)

device_list = [0, 1, 2, 3]
main_dev = device_list[0]  #主GPU，也就是分发任务和结果回收的GPU，也是梯度传播更新的GPU
model = torch.nn.DataParallel(model,device_ids=device_list)
model.to(main_dev)
 
for data in train_dataloaders: 
   model.train(True)
   inputs, labels = data
   inputs = Variable(inputs.to(main_dev))  #将数据放到主要GPU
   labels = Variable(labels.to(main_dev)) 

4. 限定GPU可用

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1,2"

# 从可用GPU中搜索第0和1的GPU
Model = nn.DataParallel(Model, device_ids=[0,1])