pytorch 单机多gpu运行

Posted antoniosu

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了pytorch 单机多gpu运行相关的知识,希望对你有一定的参考价值。

官方原理图

技术图片

前向传播过程:将数据按照batch维度分发到各个GPU上(平均分配),而后将模型拷贝到GPU,各GPU并行前向传播,将各个输出(o1、02、03、04)汇总到总的GPU。

后向传播过程:在总GPU上并行计算得到损失,并得到初始梯度;将各梯度分发到各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()

以上是关于pytorch 单机多gpu运行的主要内容,如果未能解决你的问题,请参考以下文章

pytorch单机多卡训练

Pytorch使用分布式训练,单机多卡

Pytorch使用分布式训练,单机多卡

使用 pytorch 闪电进行多 GPU 训练时出错

docker容器内运行pytorch多gpu报错 RuntimeError: NCCL Error 2: unhandled system error

pytorch multi-gpu train