pytorch单机多卡训练_数据并行DataParallel

1.单机多卡概述

单卡多级的模型训练,即并行训练,可分为数据并行和模型并行两种.

数据并行是指,多张 GPUs 使用相同的模型副本,但采用不同 batch 的数据进行训练.

模型并行是指,多张 GPUs 使用同一 batch 的数据,分别训练模型的不同部分.

2.DataParallel源码

2.1 需要传入的参数

module(Module):被并行运算的模型

device_ids=None: CUDA devices

output_device=None:输出设备位置

2.2 forward

复制代码
检查设备是否合理;
如果合理遍历模型参数和其缓存区, 检查参数和缓冲区的设备是否与src_device_obj相同如果不同抛RuntimeError。
                源码说模型和参数必须放到device_ids[0]
                将输入数据根据设备数量分发到不同的设备上。

3.DataParallel案例

通过 PyTorch 使用 GPU 非常简单。您可以将模型放在 GPU 上

复制代码
device = torch.device("cuda:0")
model.to(device)

然后,您可以将所有张量复制到 GPU:

复制代码
mytensor = my_tensor.to(device)

--注意:mytensor.to(device)等所有tensor操作都是copy数据然后重载,不改变原tensor

。但是,Pytorch 默认情况下仅使用一个 GPU。通过使用以下命令使模型并行运行,您可以轻松地在多个 GPU 上运行操作 DataParallel :

复制代码
model = nn.DataParallel(model)

导入Pytorch模块并导入参数

复制代码
import torch
from torch.utils.data import Dataset,DataLoader
import torch.nn as nn


# Parameters and DataLoaders
input_size = 5
output_size = 2

batch_size = 30
data_size = 100

device

复制代码
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

设计一个dummy数据集

复制代码
class RandomDataset(Dataset):
	def __init__(self, size, length):
    	self.len = length
  	  	self.data = torch.randn(length, size)

	def __getitem__(self, index):
	    return self.data[index]
	
	def __len__(self):
 		return self.len
    
rand_loader = DataLoader(dataset=RandomDataset(input_size, data_size),
                         batch_size=batch_size, shuffle=True)

Our model

复制代码
class Model(nn.Module):
	def __init__(self, input_size, output_size):
    	super(Model, self).__init__()
    	self.fc = nn.Linear(input_size, output_size)

	def forward(self, input):
    	output = self.fc(input)
    	print("\tIn Model: input size", input.size(),
          "output size", output.size())

    	return output

创建模型和DataParallel

复制代码
model = Model(input_size,output_size)       #初始化参数对应,不着急进行设备关联
if torch.cuda.device_count() >1:            #判别一下是否多GPU
    print("可以进行数据并行训练")
    model = nn.DataParallel(model)          #是的话可以进行初始化操作

model = model.to(device)                    #pytorch一般而已都是重新赋值操作而非修改源数据

数据并行加载数据

复制代码
for data in rand_loader:
    input = data.to(device)
    output = model(input)
    print("\tIn Model: input size", input.size())

运行模型

复制代码
for data in rand_loader:
    input = data.to(device)
    output = model(input)
    print("Outside: input size", input.size(),
          "output_size", output.size())

/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/nn/modules/linear.py:116: UserWarning:

Attempting to run cuBLAS, but there was no current CUDA context! Attempting to set the primary context... (Triggered internally at .../aten/src/ATen/cuda/CublasHandlePool.cpp:135.)

复制代码
    In Model: input size torch.Size([8, 5]) output size torch.Size([8, 2])
    In Model: input size torch.Size([8, 5]) output size torch.Size([8, 2])
    In Model: input size torch.Size([8, 5]) output size torch.Size([8, 2])
    In Model: input size torch.Size([6, 5]) output size torch.Size([6, 2])

Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])

In Model: input size torch.Size([8, 5]) output size torch.Size([8, 2])

In Model: input size torch.Size([8, 5]) output size torch.Size([8, 2])

In Model: input size torch.Size([8, 5]) output size torch.Size([8, 2])

In Model: input size torch.Size([6, 5]) output size torch.Size([6, 2])

Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])

In Model: input size torch.Size([8, 5]) output size torch.Size([8, 2])

In Model: input size torch.Size([8, 5]) output size torch.Size([8, 2])

In Model: input size torch.Size([6, 5]) output size torch.Size([6, 2])

In Model: input size torch.Size([8, 5]) output size torch.Size([8, 2])

Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])

In Model: input size torch.Size([3, 5]) output size torch.Size([3, 2])

In Model: input size torch.Size([3, 5]) output size torch.Size([3, 2])

In Model: input size torch.Size([3, 5]) output size torch.Size([3, 2])

In Model: input size torch.Size([1, 5]) output size torch.Size([1, 2])

Outside: input size torch.Size([10, 5]) output_size torch.Size([10, 2])

如果您没有 GPU 或只有一个 GPU,则当我们批量处理 30 个输入和 30 个输出时,模型将按预期获得 30 个输入和 30 个输出。但如果你有多个 GPU,那么你可以获得这样的结 果。

Let's use 2 GPUs!

In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])

In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])

Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])

In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])

In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])

Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])

In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])

In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])

Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])

In Model: input size torch.Size([5, 5]) output size torch.Size([5, 2])

In Model: input size torch.Size([5, 5]) output size torch.Size([5, 2])

Outside: input size torch.Size([10, 5]) output_size torch.Size([10, 2])

Let's use 3 GPUs!

In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])

In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])

In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])

Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])

In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])

In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])

In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])

Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])

In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])

In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])

In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])

Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])

In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])

In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])

In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])

Outside: input size torch.Size([10, 5]) output_size torch.Size([10, 2])

Let's use 8 GPUs!

In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])

In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])

In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])

In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])

In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])

In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])

In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])

In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])

Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])

In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])

In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])

In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])

In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])

In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])

In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])

In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])

In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])

Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])

In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])

In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])

In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])

In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])

In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])

In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])

In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])

In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])

Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])

In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])

In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])

In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])

In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])

In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])

Outside: input size torch.Size([10, 5]) output_size torch.Size([10, 2])

相关推荐
杭州默安科技6 分钟前
大模型AI Agent的工作原理与安全挑战
人工智能·安全
bug404_7 分钟前
jetson orin nano super AI模型部署之路(三)stable diffusion部署
人工智能·stable diffusion
Qiming_v8 分钟前
如何使用stable diffusion 3获得最佳效果
人工智能·stable diffusion
liruiqiang0526 分钟前
循环神经网络 - 通用近似定理 & 图灵完备
人工智能·rnn·深度学习·神经网络·机器学习
Panesle42 分钟前
广告推荐算法:COSMO算法与A9算法的对比
人工智能·算法·机器学习·推荐算法·广告推荐
hunteritself1 小时前
DeepSeek重磅升级,豆包深度思考,ChatGPT原生生图,谷歌Gemini 2.5 Pro!| AI Weekly 3.24-3.30
人工智能·深度学习·chatgpt·开源·语音识别·deepseek
不要不开心了1 小时前
Scala内容
开发语言·pytorch·flask·scala·dash
RadNIkMan1 小时前
Python学习(二)操作列表
网络·python·学习
Panesle1 小时前
transformer架构与其它架构对比
人工智能·深度学习·transformer
半盏茶香1 小时前
启幕数据结构算法雅航新章,穿梭C++梦幻领域的探索之旅——堆的应用之堆排、Top-K问题
java·开发语言·数据结构·c++·python·算法·链表