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

藓类少女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 模型。
  • 大规模深度神经网络：适合需要更大数据处理能力的深度神经网络。
相关推荐
Guofu_Liao2 分钟前
大语言模型中Softmax函数的计算过程及其参数描述
人工智能·语言模型·自然语言处理
非自律懒癌患者2 分钟前
Transformer中的Self-Attention机制如何自然地适应于目标检测任务
人工智能·算法·目标检测
IT闫7 分钟前
使用微信小程序调用飞桨PaddleX平台自行训练的模型——微信小程序用训练的牡丹花模型Demo测试
人工智能·paddlepaddle
Jurio.23 分钟前
Conda 管理项目环境
人工智能·python·深度学习·conda·virtualenv·pip
曼城周杰伦35 分钟前
自然语言处理:第六十二章 KAG 超越GraphRAG的图谱框架
人工智能·pytorch·神经网络·自然语言处理·chatgpt·nlp·gpt-3
Donvink38 分钟前
多模态大语言模型——《动手学大模型》实践教程第六章
人工智能·深度学习·语言模型·自然语言处理·llama
Joyner20181 小时前
pytorch训练的双卡，一个显卡占有20GB，另一个卡占有8GB，怎么均衡？
人工智能·pytorch·python
我爱学Python!1 小时前
解决复杂查询难题：如何通过 Self-querying Prompting 提高 RAG 系统效率？
人工智能·程序人生·自然语言处理·大模型·llm·大语言模型·rag
AI视觉网奇1 小时前
pytorch3d linux安装
linux·人工智能·pytorch
OBOO鸥柏1 小时前
OBOO鸥柏28.6寸液晶广告屏：创新技术引领智能显示新时代
人工智能·科技·大屏端·广告一体机
热门推荐
01(欧拉)openEuler系统添加网卡文件配置流程、(欧拉)openEuler系统手动配置ipv6地址流程、(欧拉)openEuler系统网络管理说明02PyTorch机器学习实现液态神经网络03【HarmonyOS】HUAWEI DevEco Studio 下载地址汇总04玄机平台应急响应—webshell查杀05Coze扣子平台完整体验和实践（附国内和国际版对比）06RAG 实践- Ollama+RagFlow 部署本地知识库0704 - matlab m_map地学绘图工具基础函数 - 设置网格08Unity中PICO实现 隔空取物 和 接触抓取物体09Ubuntu 20.04使用Livox mid 360 测试 FAST_LIO10【目标跟踪】相机运动补偿