【深度学习】使用硬件加速模型训练速度

藓类少女2024-11-14 18:44

一、 单机器单GPU

特点
  • 配置简单：无需多 GPU 或分布式环境的复杂配置，适合资源有限的场景。
  • 适合小规模模型：对于计算量不大的模型（如中小型 CNN、RNN），单 GPU 可以处理大多数常见任务。
优势
  • 简单易用：适合深度学习初学者以及较小规模的实验。
  • 无需并行配置：仅需设置一个 GPU，代码和数据同步操作也更简单。
实现方法
  • 常用框架 ：可以直接在代码中指定 GPU 使用，如 TensorFlow 的 with tf.device('/GPU:0') 或 PyTorch 的 model.to('cuda')。
  • 适用场景：适合开发和调试中小规模模型，以及在本地进行初步实验。

PyTorch 示例：

python 复制代码 
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
from torch.cuda.amp import autocast, GradScaler

# 检查是否有 GPU 可用
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 示例模型（简单的多层感知机）
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc1 = nn.Linear(10, 50)
        self.fc2 = nn.Linear(50, 1)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 创建模型并将其加载到 GPU
model = SimpleModel().to(device)

# 创建数据集和数据加载器
X = torch.randn(1000, 10)
y = torch.randn(1000, 1)
dataset = TensorDataset(X, y)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4, pin_memory=True)

# 定义损失函数和优化器
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练循环
epochs = 5
for epoch in range(epochs):
	# 使用 PyTorch 的 AMP 自动混合精度训练
	scaler = GradScaler()
	for inputs, targets in dataloader:
	    inputs, targets = inputs.to(device), targets.to(device)
	    with autocast():
	        outputs = model(inputs)
	        loss = criterion(outputs, targets)
	    scaler.scale(loss).backward()
	    scaler.step(optimizer)
	    scaler.update()
    
    print(f"Epoch {epoch + 1}/{epochs}, Loss: {loss.item()}")
示例说明

  • dataloader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4, pin_memory=True)

    当 pin_memory=True 时，DataLoader 将把数据加载到页锁定内存中，使得从 CPU 到 GPU 的传输更加高效。

    并发流程：

    1. 批次准备 ：4 个 CPU 进程在 DataLoader 中并发从数据集加载数据，并执行数据增强、预处理，将数据固定在内存中。
    2. 数据传输：CPU 将批次数据通过 PCIe 通道传输到 GPU 的显存。
    3. GPU 执行计算：GPU 在显存中运行当前批次的前向传播和反向传播。
    4. 并发准备下一个批次 ：在 GPU 计算当前批次的过程中，CPU 上的 DataLoader 进程已经开始准备下一个批次数据。
    5. 重复以上过程：CPU 和 GPU 以此方式并行处理每个批次，避免 GPU 等待数据传输，提升整体训练速度。

  • amp（包含 autocast 和 GradScaler）：提供完整的混合精度训练方案，包含 autocast 的精度选择和 GradScaler 的梯度缩放，能保证训练的数值稳定性。

二、 单机器多GPU

在单台机器上配置多个 GPU，可以通过数据并行、模型并行和管道并行等策略实现多 GPU 训练。

1. 数据并行（Data Parallelism）

数据并行是最常用的多 GPU 训练策略，将数据划分到不同的 GPU 上，并在每个 GPU 上独立地计算前向传播和反向传播。每个 GPU 得到的梯度会在同步步骤中合并，更新全局模型参数。

  • 工作流程 ：
    1. 将批次数据划分到不同的 GPU 上。
    2. 每个 GPU 独立计算前向和反向传播，生成局部梯度。
    3. 汇总所有 GPU 的梯度，更新全局模型参数。
  • 适用场景：适合数据量大、模型规模适中的场景，是目前最常用的并行训练方式。
代码示例

PyTorch提供了多种方式来实现多GPU训练，常用的方法是使用 torch.nn.DataParallel 或 torch.nn.parallel.DistributedDataParallel。

    1. 使用 torch.nn.DataParallel

DataParallel 是一种简单直接的多GPU训练方法，适合小规模的多GPU训练任务。它会自动将数据分配到多个GPU上，并将每个GPU的输出合并。

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

# 定义模型
model = MyModel()

# 将模型包装成多GPU模型
if torch.cuda.device_count() > 1:
    print(f"使用 {torch.cuda.device_count()} 个 GPU 进行训练")
    model = nn.DataParallel(model)

# 将模型移动到GPU
model = model.cuda()

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters())

# 加载数据
dataset = MyDataset()
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

# 训练循环
for epoch in range(num_epochs):
    for data, target in dataloader:
        data, target = data.cuda(), target.cuda()

        # 前向传播
        output = model(data)
        loss = criterion(output, target)

        # 反向传播和优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item()}")
    1. 使用 torch.nn.parallel.DistributedDataParallel

DistributedDataParallel (DDP) 是一种更高效的分布式训练方式，推荐在多GPU和分布式训练任务中使用，尤其适合更大规模的训练。

在单机器多GPU的情况下，也可以用DDP，它会在每个GPU上启动一个进程来进行并行计算。以下是使用DDP的代码示例。

设置DDP的关键步骤
  1. 设置进程组 ：使用 torch.multiprocessing 来启动多个进程，每个进程对应一个GPU。
  2. 初始化DDP模型 ：用 DistributedDataParallel 包装模型。
  3. 数据划分 ：使用 DistributedSampler 来分配数据到不同的进程。
  4. 运行代码 ：需通过torch.multiprocessing.spawn启动多进程。或者直接使用torchrun 命令启动脚本python -m torchrun --nproc_per_node=NUM_GPUS_YOU_HAVE your_script.py --arg1 --arg2
python 复制代码 
import torch
import torch.nn as nn
import torch.multiprocessing as mp
import torch.distributed as dist
from torch.utils.data import DataLoader, DistributedSampler

# 初始化进程函数
def train(rank, world_size):
    # 初始化进程组
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)
    
    # 定义模型，并移动到相应GPU
    model = MyModel().to(rank)
    model = nn.parallel.DistributedDataParallel(model, device_ids=[rank])

    # 定义损失函数和优化器
    criterion = nn.CrossEntropyLoss().to(rank)
    optimizer = torch.optim.Adam(model.parameters())

    # 创建数据集和分布式采样器
    dataset = MyDataset()
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True)
    dataloader = DataLoader(dataset, batch_size=32, sampler=sampler)

    # 训练循环
    for epoch in range(num_epochs):
        sampler.set_epoch(epoch)  # 每个epoch设定不同的采样
        for data, target in dataloader:
            data, target = data.to(rank), target.to(rank)
            
            # 前向传播
            output = model(data)
            loss = criterion(output, target)

            # 反向传播和优化
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        if rank == 0:  # 只在主进程上打印
            print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item()}")

    # 清理进程组
    dist.destroy_process_group()

# 主函数
if __name__ == "__main__":
    world_size = torch.cuda.device_count()  # 获取可用GPU数量
    mp.spawn(train, args=(world_size,), nprocs=world_size, join=True)
解释
  • train 函数中包含了每个进程需要执行的代码，rank表示每个GPU对应的编号。
  • DistributedDataParallel 在每个GPU上创建一个独立的模型副本，且 device_ids 参数指定了每个模型副本运行的设备。
  • DistributedSampler 将数据分配给每个进程，以确保在多进程中数据不会重复。
2. 模型并行（Model Parallelism）

模型并行是指将模型的不同部分分配到不同的 GPU 上，每个 GPU 只负责一部分模型的计算。这种方法适合非常大的模型，单个 GPU 无法容纳所有参数时使用。

  • 工作流程 ：
    1. 将模型划分成多个子部分（如按层或模块）。
    2. 数据在 GPU 间传递，每个 GPU 负责一个模型部分的前向和反向传播。
  • 适用场景：适合超大规模模型（如 GPT-3），在单个 GPU 内存不足以容纳全部参数时使用。
3. 管道并行（Pipeline Parallelism）

管道并行是一种结合数据并行和模型并行的方式，将模型分成多个阶段，数据分批次流经每个阶段的不同 GPU，从而在多个 GPU 之间形成"流水线"。

  • 工作流程 ：
    1. 将模型按层次或功能模块划分成多个阶段，每个阶段在不同的 GPU 上。
    2. 数据逐批通过每个阶段的 GPU，同时保持并行计算。
  • 适用场景：适合具有分层结构的模型，尤其在超大批次的训练任务中提升效率。

三、多机器多GPU

在分布式训练中，模型可以跨越多台机器（节点）进行训练，每个节点可以有一个或多个 GPU。分布式训练的核心在于如何在各节点之间有效地通信和同步梯度。

1. 分布式数据并行（Distributed Data Parallelism）

分布式数据并行是分布式训练中最常用的模式，它与单机器多GPU 的数据并行类似，但在不同机器的 GPU 间进行并行计算。每台机器上会有一个完整的模型副本，训练数据在所有机器间分布。

  • 梯度同步：常用的梯度同步方法是全归约（All-Reduce），它在每个节点中同步并合并梯度更新。
  • 实现框架 ：
    • PyTorch Distributed Data Parallel (DDP)：使用 PyTorch 自带的 DDP 模块，可实现多节点的分布式训练，适合在多 GPU 集群中。
    • TensorFlow MirroredStrategy ：TensorFlow 的 tf.distribute.MirroredStrategy 支持多 GPU、单机多 GPU 和多机多 GPU 的训练。
    • Horovod：Uber 开发的分布式训练框架，支持 TensorFlow、Keras 和 PyTorch，使用环形 All-Reduce 方式进行高效的梯度同步。
  • 适用场景：适合数据量大且可分割的大批量训练任务。
DDP示例
python 复制代码 
import os
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler
from torchvision import datasets, transforms

# 定义模型
class MyModel(torch.nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.fc = torch.nn.Linear(784, 10)
    
    def forward(self, x):
        return self.fc(x.view(-1, 784))

# 训练函数
def train():
    # 从环境变量中获取 rank 和 world_size
    rank = int(os.environ["RANK"])
    world_size = int(os.environ["WORLD_SIZE"])

    # 初始化分布式进程组
    dist.init_process_group(
        backend="nccl",
        init_method="env://",
        rank=rank,
        world_size=world_size
    )

    # 获取当前节点的 GPU 数量
    local_rank = rank % torch.cuda.device_count()
    torch.cuda.set_device(local_rank)

    # 创建模型并移动到对应的 GPU 上
    model = MyModel().to(local_rank)
    model = DDP(model, device_ids=[local_rank])

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

    # 数据加载器，使用 DistributedSampler 分配数据
    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
    dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True)
    dataloader = DataLoader(dataset, batch_size=64, sampler=sampler)

    # 训练循环
    for epoch in range(10):
        sampler.set_epoch(epoch)
        for batch_idx, (data, target) in enumerate(dataloader):
            data, target = data.to(local_rank), target.to(local_rank)

            # 前向传播
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)

            # 反向传播和优化
            loss.backward()
            optimizer.step()

            if batch_idx % 100 == 0 and rank == 0:  # 仅主进程打印日志
                print(f"Epoch [{epoch+1}/10], Step [{batch_idx}/{len(dataloader)}], Loss: {loss.item():.4f}")

    # 清理进程组
    dist.destroy_process_group()

if __name__ == "__main__":
    train()

在使用 torchrun 启动分布式训练时，torchrun 会自动计算和设置 RANK 和 WORLD_SIZE 环境变量，这样每个进程都能知道其在分布式系统中的全局位置。这是 torchrun 的一个重要特性，帮助简化分布式训练的设置流程。

解释

torchrun 使用以下命令行参数来推断和设置这些环境变量：

  • --nnodes：指定分布式系统中的总节点（机器）数量。
  • --nproc_per_node：指定每台机器上启动的进程数量（通常等于 GPU 数量）。
  • --node_rank ：指定当前机器的节点编号（从 0 开始），用于确定当前机器的 RANK 起始位置。
  • --master_addr 和 --master_port：指定主节点的 IP 地址和端口，确保所有节点可以连接到主节点进行通信。

torchrun 根据这些参数计算并设置 WORLD_SIZE 和 RANK，每个进程启动时就能自动获得这些环境变量。

  • WORLD_SIZE：全局进程数，即分布式系统中的总进程数。
  • RANK ：每个进程的全局唯一编号，用来标识当前进程在整个分布式系统中的位置。
    • 每台机器上的 RANK 是按顺序递增的，从主节点开始。
    • 例如，如果有两台机器，每台机器启动 4 个进程，那么 RANK 在主节点（node_rank=0）上为 0-3，在第二台机器（node_rank=1）上为 4-7。

假设有两台机器，每台机器上分别有 4 个 GPU，配置如下：

    1. 主节点（机器 1）启动命令：
bash 复制代码 
torchrun --nnodes=2 --nproc_per_node=4 --node_rank=0 --master_addr="<主节点IP地址>" --master_port="12355" script.py

  • WORLD_SIZE 自动设置为 2 * 4 = 8。

  • RANK 自动为 0-3，对应机器 1 上的 4 个进程。

    1. 第二台机器（机器 2）启动命令：
bash 复制代码 
torchrun --nnodes=2 --nproc_per_node=4 --node_rank=1 --master_addr="<主节点IP地址>" --master_port="12355" script.py
  • WORLD_SIZE 仍为 8。
  • RANK 自动为 4-7，对应机器 2 上的 4 个进程。
注意

torchrun 不支持不同的 nproc_per_node，我们需要手动启动每个机器的进程，确保每个进程的 RANK 和 WORLD_SIZE 环境变量设置正确。

假设我们有两台机器，且每台机器上 GPU 数量不同：机器 1有 4 个 GPU，机器 2有 2 个 GPU。

主节点（机器 1）启动命令：

python 复制代码 
# 启动第一个进程
RANK=0 torchrun --nproc_per_node=1 --nnodes=2 --node_rank=0 --master_addr="<主节点IP地址>" --master_port="12355" --world_size=6 script.py &

# 启动第二个进程
RANK=1 torchrun --nproc_per_node=1 --nnodes=2 --node_rank=0 --master_addr="<主节点IP地址>" --master_port="12355" --world_size=6 script.py &

# 启动第三个进程
RANK=2 torchrun --nproc_per_node=1 --nnodes=2 --node_rank=0 --master_addr="<主节点IP地址>" --master_port="12355" --world_size=6 script.py &

# 启动第四个进程
RANK=3 torchrun --nproc_per_node=1 --nnodes=2 --node_rank=0 --master_addr="<主节点IP地址>" --master_port="12355" --world_size=6 script.py &

第二台机器（机器 2）启动命令：

python 复制代码 
export WORLD_SIZE=6

# 启动第一个进程
RANK=4 torchrun --nproc_per_node=1 --nnodes=2 --node_rank=1 --master_addr="<主节点IP地址>" --master_port="12355" --world_size=6 script.py &

# 启动第二个进程
RANK=5 torchrun --nproc_per_node=1 --nnodes=2 --node_rank=1 --master_addr="<主节点IP地址>" --master_port="12355" --world_size=6 script.py &
2. 分布式模型并行（Distributed Model Parallelism）

分布式模型并行将模型拆分到不同节点的 GPU 上，通常用于单节点 GPU 无法容纳完整模型的超大规模模型（如 GPT-3）。

  • 参数分配：模型参数按功能模块或层次划分到不同的 GPU。
  • 前向与反向传播：每个 GPU 只计算自己负责的部分，然后将输出传递给下一层的 GPU。
  • 适用场景：适用于超大模型的训练（如 Transformer、大规模语言模型等），尤其在内存受限时。
3. 混合并行（Hybrid Parallelism）

混合并行结合数据并行和模型并行的优势，用于训练规模极大的模型。可以同时在不同 GPU 上进行数据并行训练，并在每个 GPU 内部按模型并行分配参数。

  • 实现框架 ：
    • Megatron-LM：NVIDIA 提供的框架，支持混合并行训练超大规模 Transformer 模型。
    • DeepSpeed：微软开发的分布式训练优化库，支持数据并行、模型并行、流水线并行以及零冗余优化。
  • 适用场景：超大规模模型的训练，适合在分布式多 GPU 集群中实现。

四、TPU（张量处理单元）

TPU 是 Google 专门为深度学习设计的硬件加速器，尤其针对矩阵运算进行了优化。

TPU 的特点
  • 高度优化的矩阵运算：TPU 专为深度学习而设计，适合卷积神经网络（CNN）和循环神经网络（RNN）等高密度计算任务。
  • 高内存带宽：TPU 拥有很高的内存带宽，可以加速训练过程中数据的加载和计算。
使用 TPU 的方法
  • Google Colab 和 Google Cloud Platform（GCP）：可以在 Google Colab 免费使用 TPU，GCP 提供的 TPU 资源适合更大规模任务。
  • TensorFlow 高效兼容：TPU 最适合 TensorFlow 框架的使用，Keras 和 TensorFlow 可以在单行代码中切换 TPU 支持。
TPU 使用场景
  • Transformer 模型：如 BERT、GPT 等 NLP 模型。
  • 大规模深度神经网络：适合需要更大数据处理能力的深度神经网络。
相关推荐
扫地的小何尚1 小时前
NVIDIA RTX 系统上使用 llama.cpp 加速 LLM
人工智能·aigc·llama·gpu·nvidia·cuda·英伟达
埃菲尔铁塔_CV算法3 小时前
深度学习神经网络创新点方向
人工智能·深度学习·神经网络
艾思科蓝-何老师【H8053】4 小时前
【ACM出版】第四届信号处理与通信技术国际学术会议（SPCT 2024）
人工智能·信号处理·论文发表·香港中文大学
weixin_452600694 小时前
《青牛科技 GC6125：驱动芯片中的璀璨之星，点亮 IPcamera 和云台控制（替代 BU24025/ROHM）》
人工智能·科技·单片机·嵌入式硬件·新能源充电桩·智能充电枪
学术搬运工4 小时前
【珠海科技学院主办，暨南大学协办 | IEEE出版 | EI检索稳定 】2024年健康大数据与智能医疗国际会议（ICHIH 2024）
大数据·图像处理·人工智能·科技·机器学习·自然语言处理
右恩4 小时前
AI大模型重塑软件开发：流程革新与未来展望
人工智能
图片转成excel表格5 小时前
WPS Office Excel 转 PDF 后图片丢失的解决方法
人工智能·科技·深度学习
ApiHug5 小时前
ApiSmart x Qwen2.5-Coder 开源旗舰编程模型媲美 GPT-4o, ApiSmart 实测！
人工智能·spring boot·spring·ai编程·apihug
哇咔咔哇咔6 小时前
【科普】简述CNN的各种模型
人工智能·神经网络·cnn
李歘歘6 小时前
万字长文解读深度学习——多模态模型CLIP、BLIP、ViLT
人工智能·深度学习
热门推荐
01玄机平台应急响应—webshell查杀02Ubuntu 20.04使用Livox mid 360 测试 FAST_LIO03【HarmonyOS】HUAWEI DevEco Studio 下载地址汇总04新版微信小程序获取用户手机号05文件或文件夹名称中有空格如何批量去除06组基轨迹建模 GBTM的介绍与实现（Stata 或 R）07【TC3xx芯片】TC3xx芯片电源管理系统PMS详解08优化手机性能，解决卡顿问题：关闭这3个微信开关，释放内存空间09基于YOLOv10深度学习的CT扫描图像肾结石智能检测系统【python源码+Pyqt5界面+数据集+训练代码】深度学习实战、目标检测10【软件工程】UML用例图介绍和实例说明