pytorch分布式数据训练结合学习率周期及混合精度

文章目录

正如标题所写,我们正常的普通训练都是单机单卡或单机多卡。而往往一个高精度的模型需要训练时间很长,所以DDP分布式数据并行和混合精度可以加速模型训练。混精可以增大batch size.

如下提供示例代码,经过官网查阅验证的。原始代码由百度文心一言提供。

问题:pytorch 分布式数据并行与学习率周期和混合精度训练的示例代码

该问题分别问了(当前2023年10月16日)百度文心一言,google bard,科大讯飞星火,chatgpt4,bing。目前看来只有文心给出最接近正确答案。当前版本 pythorch 2.0.1。关于混精 torch.cuda.amp.autocast(args...)`` is equivalent to ``torch.autocast("cuda", args...),混精可以参考:

1、SPAWN方式

在使用这些助手(都是免费的)时,发现对于不同问题,各自的优劣不同,所以可以都试一下。

在PyTorch中,分布式数据并行(Distributed Data Parallel,DDP)可以在多个GPU上并行训练模型,提高训练速度。学习率周期(Learning Rate Schedules)用于调整训练过程中学习率的变化,而混合精度训练则可以提高训练速度并减少内存使用。

以下是一个使用PyTorch的分布式数据并行、学习率周期和混合精度训练的示例代码:

python 复制代码
import os
import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.optim.lr_scheduler import ExponentialLR, CosineAnnealingLR
from torch.cuda.amp import GradScaler, autocast



def setup(rank, world_size):
    """
    Args:
        rank: Unique identifier of each process
        world_size: Total number of processes
    """
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'

    # initialize the process group
    dist.init_process_group(backend="nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)

def cleanup():
    dist.destroy_process_group()
# 定义模型
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 10)

    def forward(self, x):
        return self.fc(x)

# 训练函数
def train(rank, world_size):
    print(f"Running basic DDP example on rank {rank}.")
    gpu=rank
    setup(rank=rank,world_size=world_size,)
    torch.manual_seed(0)
    model = SimpleModel().cuda(gpu)
    ddp_model = DDP(model, device_ids=[gpu])

    # 使用混合精度训练
    scaler = GradScaler()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.01)

    # 使用学习率周期
    scheduler1 = ExponentialLR(optimizer, gamma=0.9)
    scheduler2 = CosineAnnealingLR(optimizer, T_max=50, eta_min=0)

    # 模拟数据
    data = torch.randn(32, 10).to(gpu)
    target = torch.randn(32, 10).to(gpu)

    for epoch in range(100):
        optimizer.zero_grad()
        with autocast():  # 启用混合精度训练
            output = ddp_model(data)
            loss = nn.MSELoss()(output, target)
        scaler.scale(loss).backward()  # 使用scaler进行梯度缩放
        scaler.step(optimizer)  # 更新参数
        scaler.update()  # 更新GradScaler状态
        scheduler1.step()  # 使用ExponentialLR调整学习率
        scheduler2.step()  # 使用CosineAnnealingLR调整学习率
        print(f"Rank {rank}, Epoch {epoch}, Loss {loss.item()}")
    cleanup()
if __name__ == "__main__":
    world_size = torch.cuda.device_count()  # 获取可用GPU的数量
    mp.spawn(train, args=(world_size,), nprocs=world_size, join=True)  # 在每个GPU上运行train函数

这个示例代码在每个GPU上并行训练一个简单的线性模型。每个进程(即每个GPU)都有自己的模型副本,并且每个进程都独立计算梯度。然后,所有进程都会聚集他们的梯度并平均,然后用于一次总体参数更新。这个过程会根据学习率周期来调整每个epoch后的学习率

本部分参考官方的:https://pytorch.org/tutorials/beginner/ddp_series_multigpu.html?highlight=torch multiprocessing 是写单GPU和多GPU的区别。

2、torchrun 方式

首先是写一个ddp.py,内容如下:

python 复制代码
import os
import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.optim.lr_scheduler import ExponentialLR, CosineAnnealingLR
from torch.cuda.amp import GradScaler, autocast

    
# 定义模型
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 10)

    def forward(self, x):
        return self.fc(x)

# 训练函数
def train():
    dist.init_process_group("nccl")
    rank = dist.get_rank()
    print(f"Start running basic DDP example on rank {rank}.")
    gpu = rank %  torch.cuda.device_count()
    torch.manual_seed(0)
    model = SimpleModel().to(gpu)
    ddp_model = DDP(model, device_ids=[gpu])

    # 使用混合精度训练
    scaler = GradScaler()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.01)

    # 使用学习率周期
    scheduler1 = ExponentialLR(optimizer, gamma=0.9)
    scheduler2 = CosineAnnealingLR(optimizer, T_max=50, eta_min=0)

    # 模拟数据
    data = torch.randn(32, 10).to(gpu)
    target = torch.randn(32, 10).to(gpu)

    for epoch in range(100):
        optimizer.zero_grad()
        with autocast():  # 启用混合精度训练
            output = ddp_model(data)
            loss = nn.MSELoss()(output, target)
        scaler.scale(loss).backward()  # 使用scaler进行梯度缩放
        scaler.step(optimizer)  # 更新参数
        scaler.update()  # 更新GradScaler状态
        scheduler1.step()  # 使用ExponentialLR调整学习率
        scheduler2.step()  # 使用CosineAnnealingLR调整学习率
        print(f"Rank {rank}, Epoch {epoch}, Loss {loss.item()}")
    dist.destroy_process_group()
if __name__ == "__main__":
    train()

单机多卡,执行:

shell 复制代码
torchrun --nproc_per_node=4 --standalone ddp.py

如果是多机多卡:

shell 复制代码
torchrun --nnodes=2 --nproc_per_node=8 --rdzv_id=100 --rdzv_backend=c10d --rdzv_endpoint=$MASTER_ADDR:29400 elastic_ddp.py

本部分参考:

https://pytorch.org/tutorials/intermediate/ddp_tutorial.html#save-and-load-checkpoints

相关推荐
zhangfeng11331 小时前
pytorch 的交叉熵函数,多分类,二分类
人工智能·pytorch·分类
Seeklike1 小时前
11.22 深度学习-pytorch自动微分
人工智能·pytorch·深度学习
YRr YRr2 小时前
如何使用 PyTorch 实现图像分类数据集的加载和处理
pytorch·深度学习·分类
z千鑫17 小时前
【人工智能】PyTorch、TensorFlow 和 Keras 全面解析与对比:深度学习框架的终极指南
人工智能·pytorch·深度学习·aigc·tensorflow·keras·codemoss
学不会lostfound18 小时前
三、计算机视觉_05MTCNN人脸检测
pytorch·深度学习·计算机视觉·mtcnn·p-net·r-net·o-net
Mr.谢尔比19 小时前
李宏毅机器学习课程知识点摘要(1-5集)
人工智能·pytorch·深度学习·神经网络·算法·机器学习·计算机视觉
做程序员的第一天20 小时前
在PyTorch中,钩子(hook)是什么?在神经网络中扮演什么角色?
pytorch·python·深度学习
Nerinic21 小时前
PyTorch基础2
pytorch·python
曼城周杰伦1 天前
自然语言处理:第六十二章 KAG 超越GraphRAG的图谱框架
人工智能·pytorch·神经网络·自然语言处理·chatgpt·nlp·gpt-3
Joyner20181 天前
pytorch训练的双卡,一个显卡占有20GB,另一个卡占有8GB,怎么均衡?
人工智能·pytorch·python