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运行的主要内容,如果未能解决你的问题,请参考以下文章
docker容器内运行pytorch多gpu报错 RuntimeError: NCCL Error 2: unhandled system error