【深度学习】单机多卡 | DataParallel将计算任务在多个 GPU 上并行执行，可以在多个 GPU 上分摊工作负载，从而加快训练速度

是Yu欸2024-07-02 16:55

【深度学习】单机多卡 | DataParallel将计算任务在多个 GPU 上并行执行，可以在多个 GPU 上分摊工作负载，从而加快训练速度


🌈你好呀！我是 是Yu欸 🌌 2024每日百字篆刻时光，感谢你的陪伴与支持 ~ 🚀 欢迎一起踏上探险之旅，挖掘无限可能，共同成长！

写在最前面

希望在单机多卡的模式下运行我的模型代码，加快训练速度。

请教吕博：如何更改代码？

其中，提到模型先用DP方式运行

DP是什么？又被学到了一个知识点。

在深度学习和分布式计算领域，DP 通常指的是 DataParallel 。DataParallel 是一种将计算任务在多个 GPU 上并行执行的方法。它在单机多卡环境中非常有用，可以在多个 GPU 上分摊工作负载，从而加快训练速度。

DataParallel (DP) 简介

torch.nn.DataParallel 是 PyTorch 中的一个工具，可以让模型在多个 GPU 上并行运行。它通过将输入批次拆分成多个子批次，每个子批次发送到不同的 GPU 上，并行执行前向传播和反向传播，然后将每个 GPU 上的梯度聚合到主 GPU 上进行参数更新。

使用 DataParallel 的场景

  • 单机多卡训练 : 当你有一台机器配备了多块 GPU，并希望利用所有的 GPU 资源来加速模型训练时，DataParallel 是一个简单而有效的解决方案。
  • 简化代码 : 相比于更复杂的分布式训练方案，DataParallel 提供了一种较为简化的方式来实现多 GPU 并行训练，通常只需要对模型进行简单包装。

使用 DataParallel 的基本步骤

  1. 定义模型: 创建你的神经网络模型。
  2. 包装模型 : 使用 torch.nn.DataParallel 包装你的模型。
  3. 将模型和数据迁移到 GPU : 使用 .to(device) 将模型和输入数据迁移到合适的设备上。
  4. 训练模型: 按照常规方式训练模型。

代码部分

train.py

仅展示相关部分

python 复制代码 
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3"
# import部分省略

def evaluate(model, device, dataloader):
    model.eval()
    total_loss, total_step = 0.0, 0.0

    # 使用with torch.no_grad()来禁用梯度计算
    with torch.no_grad():
        # 对dataloader中的每个batch进行遍历
        for step, batch in enumerate(dataloader):
            # 将batch中的数据移动到指定设备上
            batch = tuple(t.to(device) for t in batch)
            input_ids, attention_mask, decoder_input_ids, decoder_attention_mask, labels = batch

            # 通过模型进行前向传播，并获取输出结果
            outputs = model(input_ids, attention_mask=attention_mask, decoder_input_ids=decoder_input_ids,
                            decoder_attention_mask=decoder_attention_mask, labels=labels)

            # 获取模型输出结果中的loss值
            loss = outputs['loss']
            # 累加总损失和总步数
            total_loss += loss.item()
            total_step += 1

    # 返回总损失和None
    return total_loss / total_step, None


# 设置随机数种子、日志路径
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 初始化tokenizer、添加特殊的tokens
# define dataloader 定义数据加载器
# 每批的个数4/梯度累计个数4
batch_size = int(args.batch_size / args.gradient_accumulation_steps)
# from processing.dataset import BartDataset
# return input_ids, attention_mask, decoder_input_ids, decoder_attention_mask, labels
trian_dataset = BartDataset(tokenizer, args, mode='train')
train_dataloader = DataLoader(
    dataset=trian_dataset,
    batch_size=batch_size,
    shuffle=True,
    collate_fn=trian_dataset.collate_fn,
    num_workers=20 # 优化数据加载
)

eval_dataset = BartDataset(tokenizer, args, mode='test')
eval_dataloader = DataLoader(
    dataset=eval_dataset,
    batch_size=batch_size,
    shuffle=False,
    collate_fn=eval_dataset.collate_fn,
    num_workers=20 # 优化数据加载
)

# define model 实例化模型

# 检查GPU数量并设置DataParallel
if torch.cuda.device_count() > 1:
    print(f"Using {torch.cuda.device_count()} GPUs")
    net = nn.DataParallel(model)
else:
    print("Using single GPU or CPU")
    net = model

# define criterion 定义损失函数

# define optimizer优化器
# 参考：https://blog.csdn.net/hottie_xiaomiao/article/details/124392847
# 打印每一次迭代元素的名字和param
param_optimizer = list(model.named_parameters())
no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
# 参数组：每组参数可以指定自己的优化器参数，即可使用不同的优化策略
optimizer_grouped_parameters = [
    {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01},
    {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
]
optimizer = torch.optim.AdamW(optimizer_grouped_parameters, lr=args.bart_lr)
# 计算总步数
total_steps = int(len(trian_dataset) * args.epochs / args.gradient_accumulation_steps)
# 初始化学习率调整器
scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=int(args.warmup_steps * total_steps),
                                            num_training_steps=total_steps)

# Begin training
# logger部分省略

# 定义一个初始的最好评估损失，即正无穷大
best_eval_loss = float('inf')
# 定义当前的步骤和当前损失
current_step, current_loss = 0, 0
# 定义全局步骤数
global_step = 0

# 对于每个 epoch 进行循环
for epoch in range(args.epochs):
    # 将模型设置为训练模式
    model.train()
    # 对训练数据集进行循环
    for step, batch in enumerate(train_dataloader):
        # 将batch中的每一个tensor都移动到指定的设备上（如GPU）
        batch = tuple(t.to(device) for t in batch)
        # 从batch中获取输入，注意这里的命名方式要和模型中forward函数中的输入命名相同
        input_ids, attention_mask, decoder_input_ids, decoder_attention_mask, labels = batch
        # 将输入传递给模型进行前向计算
        outputs = model(input_ids, attention_mask=attention_mask, decoder_input_ids=decoder_input_ids,
                        decoder_attention_mask=decoder_attention_mask, labels=labels)
        # TODO:1
        # print(outputs)

        # 获取模型的损失
        loss = outputs['loss']

        # 记录当前损失和步骤数，用于计算平均损失
        current_loss += loss.item()
        current_step += 1

        # 如果使用了梯度累积，则将损失除以累积步骤数
        if args.gradient_accumulation_steps > 1:
            loss = loss / args.gradient_accumulation_steps

        # 反向传播计算梯度
        loss.backward()
        # 将梯度进行裁剪，以防止梯度爆炸
        clip_grad_norm_(model.parameters(), args.max_clip_norm)

        # 如果达到了梯度累积的步骤数，则进行一次优化更新
        if (step + 1) % args.gradient_accumulation_steps == 0:
            optimizer.step()
            scheduler.step()
            optimizer.zero_grad()
            global_step += 1

        # 如果当前步骤是一个log间隔的倍数，则打印日志信息，清空当前步骤和当前损失

    # 在训练完一个 epoch 后，对模型在验证集上进行评估
    eval_loss, _ = evaluate(model, device, eval_dataloader)
    logger.info("Eval loss: {:.6f}, the best loss: {:.6f}".format(eval_loss, best_eval_loss))
    # 如果当前的评估损失比之前的最好评估损失更小，则更新最好评估损失
    if eval_loss < best_eval_loss:
        best_eval_loss = eval_loss

    # 创建一个输出目录，用于存储模型的输出
    # 在日志中输出检查点保存路径，将模型、tokenizer、args的设置保存到检查点路径中

简单的代码示例

以下是使用 DataParallel 在多 GPU 上运行模型的一个简单示例：

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

# 设置环境变量，指定使用的GPU
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,6,7"

# 定义设备
globalDevice = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 定义一个简单的CNN模型
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv = nn.Conv2d(3, 16, 3, 1)
        self.fc = nn.Linear(16 * 26 * 26, 10)

    def forward(self, x):
        x = self.conv(x)
        x = torch.relu(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

# 实例化模型
cnn = CNN().to(globalDevice)

# 检查GPU数量并设置DataParallel
if torch.cuda.device_count() > 1:
    print(f"Using {torch.cuda.device_count()} GPUs")
    net = nn.DataParallel(cnn)
else:
    print("Using single GPU or CPU")
    net = cnn

# 定义数据集和数据加载器
class SimpleDataset(Dataset):
    def __init__(self, size):
        self.size = size

    def __len__(self):
        return self.size

    def __getitem__(self, idx):
        return torch.randn(3, 28, 28), torch.tensor(1)

dataset = SimpleDataset(1000)
dataloader = DataLoader(dataset, batch_size=64, shuffle=True, num_workers=20)

# 定义优化器和损失函数
optimizer = optim.SGD(net.parameters(), lr=0.01)
criterion = nn.CrossEntropyLoss()

# 简单的训练过程
for epoch in range(args.epochs):
    for inputs, labels in dataloader:
        inputs, labels = inputs.to(globalDevice), labels.to(globalDevice)
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
    print(f"Epoch {epoch+1}, Loss: {loss.item()}")

代码解析

  1. 环境变量设置:

    python 复制代码 
    os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,6,7"

    指定要使用的 GPU。

  2. 定义设备:

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

    根据是否有可用的 GPU 设置设备。

  3. 定义模型:

    python 复制代码 
    class CNN(nn.Module):
    ...

  4. 包装模型:

    python 复制代码 
    if torch.cuda.device_count() > 1:
    net = nn.DataParallel(cnn)
else:
    net = cnn

    如果检测到多个 GPU，使用 DataParallel 包装模型。

  5. 数据加载器:

    python 复制代码 
    dataloader = DataLoader(dataset, batch_size=64, shuffle=True, num_workers=20)

    使用 num_workers 参数优化数据加载。

  6. 训练过程:

    python 复制代码 
    for epoch in range(2):
    ...

DataParallel 的局限性

  • 数据并行粒度 : DataParallel 进行的是数据并行操作，每个 GPU 处理一部分数据批次。这可能导致 GPU 利用率不均衡，尤其是在有计算负载差异的情况下。
  • 单节点限制 : DataParallel 主要用于单节点多 GPU。如果需要跨节点并行（分布式训练），应该考虑使用 torch.nn.parallel.DistributedDataParallel。

小结

DataParallel 是 PyTorch 提供的一种简单易用的多 GPU 并行方法，适合单节点多卡训练。通过这种方法，可以在多个 GPU 上分摊计算任务，提高训练速度和效率。对于更复杂的分布式计算任务，可以考虑使用 DistributedDataParallel。

欢迎大家添加好友交流。

相关推荐
2403_87573687几秒前
道品科技智慧农业中的自动气象检测站
网络·人工智能·智慧城市
荒Huang3 分钟前
Linux挖矿病毒（kswapd0进程使cpu爆满）
linux·运维·服务器
海阔天空_201313 分钟前
Python pyautogui库：自动化操作的强大工具
运维·开发语言·python·青少年编程·自动化
桥田智能16 分钟前
气爪在自动化装配线中是如何应用的？
运维·自动化
零意@21 分钟前
ubuntu切换不同版本的python
windows·python·ubuntu
学术头条24 分钟前
AI 的「phone use」竟是这样练成的，清华、智谱团队发布 AutoGLM 技术报告
人工智能·科技·深度学习·语言模型
准橙考典25 分钟前
怎么能更好的通过驾考呢？
人工智能·笔记·自动驾驶·汽车·学习方法
ai_xiaogui28 分钟前
AIStarter教程：快速学会卸载AI项目【AI项目管理平台】
人工智能·ai作画·语音识别·ai写作·ai软件
思忖小下32 分钟前
Python基础学习_01
python
孙同学要努力33 分钟前
《深度学习》——深度学习基础知识（全连接神经网络）
人工智能·深度学习·神经网络
热门推荐
01【HarmonyOS】HUAWEI DevEco Studio 下载地址汇总02(欧拉)openEuler系统添加网卡文件配置流程、(欧拉)openEuler系统手动配置ipv6地址流程、(欧拉)openEuler系统网络管理说明03组基轨迹建模 GBTM的介绍与实现（Stata 或 R）04【AIGC】重塑未来的科技巨轮05全面解析：构建基于深度学习的安全帽检测系统（UI界面+YOLO代码+数据集）06【经验分享】Ubuntu22.04安装微信(linux官方版)07基于YOLOv10深度学习的CT扫描图像肾结石智能检测系统【python源码+Pyqt5界面+数据集+训练代码】深度学习实战、目标检测08Ubuntu 20.04使用Livox mid 360 测试 FAST_LIO09RAG 实践- Ollama+RagFlow 部署本地知识库10【TC3xx芯片】TC3xx芯片电源管理系统PMS详解