pytorch 单机多gpu运行

import os
import torch
import torch.nn as nn
import Encoder
os.environ[‘CUDA_VISIBLE_DEVICES‘]=‘0,1,2‘# 这里的0 就是主gpu,  1、2的模型和数据由主gpu分发
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

encoder=Encoder()
encoder = encoder.to(device) # 这里将模型复制到gpu ,默认是cuda(‘0‘)，即转到第一个GPU 2
decoder_optimizer = torch.optim.Adam(params=encoder.parameters(),lr=encoder_lr, betas = (0.8,0.999)) #定义优化器
criterion = nn.CrossEntropyLoss()
if torch.cuda.device_count() > 1:
    encoder = torch.nn.DataParallel(encoder)  # 前提是model已经在cuda上了

# 前向传播时数据也要放到GPU中,即复制到主gpu里
for batch_idx, (data, label) in train_data:
    data=data.to(device)
    label=label.to(device)
    prediction = encoder(data)
    # 这里的prediction 预测结果是由两个gpu合并过的，并行计算只存在在前向传播里
    # 前向传播每个gpu计算量为 batch_size/torch.cuda.device_count() ,等前向传播完了将结果合到主gpu
    
    loss = criterion(prediction, label)  # 计算loss
    optimizer.zero_grad()
    loss.backward()