• pytorch 多GPU训练

    代码库地址: mnist

    目录​​​​​​​

    普通单机单卡训练流程

    DDP分布式训练

    horovod方式

    PyTorch 并行训练极简 Demo

    1. 程序开始时执行dist.init_process_group('nccl')，结束时执行dist.destroy_process_group()。
    2. 用torchrun --nproc_per_node=GPU_COUNT main.py运行脚本。
    3. 进程初始化后用rank = dist.get_rank()获取当前的GPU ID，把模型和数据都放到这个GPU上。
    4. 封装一下模型
    ddp_model = DistributedDataParallel(model, device_ids=[device_id])

    5. 封装一下DataLoader

    1.     dataset = MyDataset()
    2.     sampler = DistributedSampler(dataset)
    3.     dataloader = DataLoader(dataset, batch_size=2, sampler=sampler)

    6. 训练时打乱数据。sampler.set_epoch(epoch)

    7. 保存只在单卡上进行。

    1. if rank == 0:
    2.     torch.save(ddp_model.state_dict(), ckpt_path)
    3. dist.barrier()

    8. 读取数据时注意map_location，也要注意参数名里的module。

    1. map_location = {'cuda:0': f'cuda:{device_id}'}
    2. state_dict = torch.load(ckpt_path, map_location=map_location)
    3. ddp_model.load_state_dict(state_dict)

     

    普通单机单卡训练流程

    以mnist为例，主要包括数据加载、模型构建、优化器和迭代训练等部分 

    1. import argparse
    2. import torch
    3. import torch.nn as nn
    4. import torchvision
    5. import torchvision.transforms as transforms
    6. from datetime import datetime
    7. from tqdm import tqdm
    8.  
    9. class ConvNet(nn.Module):
    10.     def __init__(self, num_classes=10):
    11.         super(ConvNet, self).__init__()
    12.         self.layer1 = nn.Sequential(
    13.             nn.Conv2d(1, 16, kernel_size=5, stride=1, padding=2),
    14.             nn.BatchNorm2d(16),
    15.             nn.ReLU(),
    16.             nn.MaxPool2d(kernel_size=2, stride=2))
    17.         self.layer2 = nn.Sequential(
    18.             nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2),
    19.             nn.BatchNorm2d(32),
    20.             nn.ReLU(),
    21.             nn.MaxPool2d(kernel_size=2, stride=2))
    22.         self.fc = nn.Linear(7*7*32, num_classes)
    23.  
    24.     def forward(self, x):
    25.         out = self.layer1(x)
    26.         out = self.layer2(out)
    27.         out = out.reshape(out.size(0), -1)
    28.         out = self.fc(out)
    29.         return out
    30.  
    31. def train(gpu, args):
    32.     torch.manual_seed(0)
    33.     model = ConvNet()
    34.     torch.cuda.set_device(gpu)
    35.     model.cuda(gpu)
    36.     batch_size = 100
    37.     # define loss function (criterion) and optimizer
    38.     criterion = nn.CrossEntropyLoss().cuda(gpu)
    39.     optimizer = torch.optim.SGD(model.parameters(), 1e-4)
    40.     # Data loading code
    41.     train_dataset = torchvision.datasets.MNIST(root='./data', train=True, transform=transforms.ToTensor(), download=True)
    42.     train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size,shuffle=True, num_workers=0,pin_memory=True)
    43.  
    44.     start = datetime.now()
    45.     for epoch in range(args.epochs):
    46.         if gpu == 0:
    47.             print("Epoch: {}/{}".format(epoch+1, args.epochs))
    48.         pbar = tqdm(train_loader)
    49.         for images, labels in pbar:
    50.             images = images.cuda(non_blocking=True)
    51.             labels = labels.cuda(non_blocking=True)
    52.             # Forward pass
    53.             outputs = model(images)
    54.             loss = criterion(outputs, labels)
    55.  
    56.             # Backward and optimize
    57.             optimizer.zero_grad()
    58.             loss.backward()
    59.             optimizer.step()
    60.             if gpu == 0:
    61.                 msg = 'Loss: {:.4f}'.format(loss.item())
    62.                 pbar.set_description(msg)
    63.     
    64.     if gpu == 0:
    65.         print("Training complete in: " + str(datetime.now() - start))
    66.  
    67. def main():
    68.     parser = argparse.ArgumentParser()
    69.     parser.add_argument('-n', '--nodes', default=1, type=int, metavar='N')
    70.     parser.add_argument('-g', '--gpus', default=2, type=int, help='number of gpus per node')
    71.     parser.add_argument('-nr', '--nr', default=0, type=int, help='ranking within the nodes')
    72.     parser.add_argument('--epochs', default=10, type=int, metavar='N', help='number of total epochs to run')
    73.     args = parser.parse_args()
    74.     train(0, args)
    75.  
    76. if __name__ == '__main__':
    77.     main()

     在2080Ti上训练2个epoch耗时1分12秒.

    DDP分布式训练

    需要的改动 0.引入必须的库

    1. import os
    2. import torch.multiprocessing as mp
    3. import torch.distributed as dist

     1.修改main函数

    1. def main():
    2.     parser = argparse.ArgumentParser()
    3.     ...
    4.     args = parser.parse_args()
    5.     args.world_size = args.gpus * args.nodes 
    6.     if args.world_size > 1:
    7.         os.environ['MASTER_ADDR'] = '127.0.0.1'                 #
    8.         os.environ['MASTER_PORT'] = '8889'                      #
    9.         mp.spawn(train, nprocs=args.gpus, args=(args,))         #
    10.     else:
    11.         train(0, args)

    2.初始化通信库

    1. def train(gpu, args):
    2.     if args.world_size > 1:
    3.         rank = args.nr * args.gpus + gpu
    4.         dist.init_process_group(backend='nccl', init_method='env://', world_size=args.world_size, rank=rank)

    3.送给每个node的数据需要打乱，有请DistributedSampler 

    1.     if args.world_size > 1:
    2.         model = nn.parallel.DistributedDataParallel(model, device_ids=[gpu])
    3.         train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset,num_replicas=args.world_size,rank=rank)
    4.         shuffle = False
    5.     else:
    6.         train_sampler = None
    7.         shuffle = True
    8.     train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size,shuffle=shuffle, num_workers=0,pin_memory=True,sampler=train_sampler)

    这里我们首先计算出当前进程序号：rank = args.nr * args.gpus + gpu，然后就是通过dist.init_process_group初始化分布式环境，其中backend参数指定通信后端，包括mpi, gloo, nccl，这里选择nccl，这是Nvidia提供的官方多卡通信框架，相对比较高效。mpi也是高性能计算常用的通信协议，不过你需要自己安装MPI实现框架，比如OpenMPI。gloo倒是内置通信后端，但是不够高效。init_method指的是如何初始化，以完成刚开始的进程同步；这里我们设置的是env://，指的是环境变量初始化方式，需要在环境变量中配置4个参数：MASTER_PORT，MASTER_ADDR，WORLD_SIZE，RANK，前面两个参数我们已经配置，后面两个参数也可以通过dist.init_process_group函数中world_size和rank参数配置。其它的初始化方式还包括共享文件系统（https://pytorch.org/docs/stable/distributed.html#shared-file-system-initialization）以及TCP（https://pytorch.org/docs/stable/distributed.html#tcp-initialization），比如采用TCP作为初始化方法init_method='tcp://10.1.1.20:23456'，其实也是要提供master的IP地址和端口。注意这个调用是阻塞的，必须等待所有进程来同步，如果任何一个进程出错，就会失败。

    完整代码, 搜索add可快速直达修改的地方

    1. import argparse
    2. import torch
    3. import torch.nn as nn
    4. import torchvision
    5. import torchvision.transforms as transforms
    6. from datetime import datetime
    7. from tqdm import tqdm
    8.  
    9. # add 0
    10. import os
    11. import torch.multiprocessing as mp
    12. import torch.distributed as dist
    13.  
    14. class ConvNet(nn.Module):
    15.     def __init__(self, num_classes=10):
    16.         super(ConvNet, self).__init__()
    17.         self.layer1 = nn.Sequential(
    18.             nn.Conv2d(1, 16, kernel_size=5, stride=1, padding=2),
    19.             nn.BatchNorm2d(16),
    20.             nn.ReLU(),
    21.             nn.MaxPool2d(kernel_size=2, stride=2))
    22.         self.layer2 = nn.Sequential(
    23.             nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2),
    24.             nn.BatchNorm2d(32),
    25.             nn.ReLU(),
    26.             nn.MaxPool2d(kernel_size=2, stride=2))
    27.         self.fc = nn.Linear(7*7*32, num_classes)
    28.  
    29.     def forward(self, x):
    30.         out = self.layer1(x)
    31.         out = self.layer2(out)
    32.         out = out.reshape(out.size(0), -1)
    33.         out = self.fc(out)
    34.         return out
    35.  
    36. def train(gpu, args):
    37.     # add 2
    38.     if args.world_size > 1:
    39.         rank = args.nr * args.gpus + gpu
    40.         dist.init_process_group(backend='nccl', init_method='env://', world_size=args.world_size, rank=rank)
    41.     torch.manual_seed(0)
    42.     model = ConvNet()
    43.     torch.cuda.set_device(gpu)
    44.     model.cuda(gpu)
    45.     batch_size = 100
    46.     # define loss function (criterion) and optimizer
    47.     criterion = nn.CrossEntropyLoss().cuda(gpu)
    48.     optimizer = torch.optim.SGD(model.parameters(), 1e-4)
    49.     # Data loading code
    50.     train_dataset = torchvision.datasets.MNIST(root='./data', train=True, transform=transforms.ToTensor(), download=True)
    51.     # add 3
    52.     if args.world_size > 1:
    53.         model = nn.parallel.DistributedDataParallel(model, device_ids=[gpu])
    54.         train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset,num_replicas=args.world_size,rank=rank)
    55.         shuffle = False
    56.     else:
    57.         train_sampler = None
    58.         shuffle = True
    59.     train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size,shuffle=shuffle, num_workers=0,pin_memory=True,sampler=train_sampler)
    60.  
    61.     start = datetime.now()
    62.     for epoch in range(args.epochs):
    63.         if gpu == 0:
    64.             print("Epoch: {}/{}".format(epoch+1, args.epochs))
    65.         pbar = tqdm(train_loader)
    66.         for i, (images, labels) in enumerate(pbar):
    67.             images = images.cuda(non_blocking=True)
    68.             labels = labels.cuda(non_blocking=True)
    69.             # Forward pass
    70.             outputs = model(images)
    71.             loss = criterion(outputs, labels)
    72.  
    73.             # Backward and optimize
    74.             optimizer.zero_grad()
    75.             loss.backward()
    76.             optimizer.step()
    77.             if gpu == 0:
    78.                 msg = 'Loss: {:.4f}'.format(loss.item())
    79.                 pbar.set_description(msg)
    80.     
    81.     if gpu == 0:
    82.         print("Training complete in: " + str(datetime.now() - start))
    83.  
    84. def main():
    85.     parser = argparse.ArgumentParser()
    86.     parser.add_argument('-n', '--nodes', default=1, type=int, metavar='N')
    87.     parser.add_argument('-g', '--gpus', default=2, type=int, help='number of gpus per node')
    88.     parser.add_argument('-nr', '--nr', default=0, type=int, help='ranking within the nodes')
    89.     parser.add_argument('--epochs', default=10, type=int, metavar='N', help='number of total epochs to run')
    90.     args = parser.parse_args()
    91.     # add 1
    92.     args.world_size = args.gpus * args.nodes 
    93.     if args.world_size > 1:
    94.         os.environ['MASTER_ADDR'] = '127.0.0.1'              #
    95.         os.environ['MASTER_PORT'] = '8889'                      #
    96.         mp.spawn(train, nprocs=args.gpus, args=(args,))         #
    97.     else:
    98.         train(0, args)
    99.  
    100. if __name__ == '__main__':
    101.     main()

     耗时缩小到了48秒，还是很显著的.

    horovod方式

    Horovod是一个专注于分布式训练的深度学习框架，通过Horovod可以为Tensorflow、Keras、Pytorch和MXNet提供分布式训练的能力。使用Horovod进行分布式训练主要能给我们带来两个好处：

    • 易用性：仅需要通过几行代码的修改，就可以将单机的Tensorflow、Keras、Pytorch或MXNet的代码更改为同时支持单机单卡、单机多卡、多机多卡的训练代码
    • 高性能：高性能主要体现在扩展性上，一般的训练框架随着worker的增加（尤其是在node增加）训练性能会逐渐下降（主要是由于node间网络通信性能比较低引起）；Horovod在128机4 P100 25Gbit/s的RoCE网络条件下，Inception V3 和 ResNet-101均可以获得90%的扩展效率。

    它的安装非常简单，pip install horovod即可

    通过Horovod编写分布式训练代码，一般为6个步骤：

    • 添加hvd.init()来初始化Horovod；
    • 为每个worker分配GPU，一般一个worker process对应一个GPU，对应关系通过rank id来映射。例如pytorch中为torch.cuda.set_device(hvd.local_rank())
    • 随着world_size的变化，batch_size也在变化，因此我们也要随着world_size的变化来调整lr，一般为原有的lr 值乘以world_size；
    • 将原有深度学习框架的optimizer通过horovod中的hvd.DistributedOptimizer进行封装；
    • rank 0 将初始的variable广播给所有worker: hvd.broadcast_parameters(model.state_dict(), root_rank=0)
    • 仅在worker 0上进行checkpoint的save

    完整代码，可搜索add一键直达修改的地方

    1. import argparse
    2. import torch
    3. import torch.nn as nn
    4. import torchvision
    5. import torchvision.transforms as transforms
    6. from datetime import datetime
    7. from tqdm import tqdm
    8. # add 0
    9. import horovod.torch as hvd
    10.  
    11. class ConvNet(nn.Module):
    12.     def __init__(self, num_classes=10):
    13.         super(ConvNet, self).__init__()
    14.         self.layer1 = nn.Sequential(
    15.             nn.Conv2d(1, 16, kernel_size=5, stride=1, padding=2),
    16.             nn.BatchNorm2d(16),
    17.             nn.ReLU(),
    18.             nn.MaxPool2d(kernel_size=2, stride=2))
    19.         self.layer2 = nn.Sequential(
    20.             nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2),
    21.             nn.BatchNorm2d(32),
    22.             nn.ReLU(),
    23.             nn.MaxPool2d(kernel_size=2, stride=2))
    24.         self.fc = nn.Linear(7*7*32, num_classes)
    25.  
    26.     def forward(self, x):
    27.         out = self.layer1(x)
    28.         out = self.layer2(out)
    29.         out = out.reshape(out.size(0), -1)
    30.         out = self.fc(out)
    31.         return out
    32.  
    33. def train(gpu, args):
    34.     torch.manual_seed(0)
    35.     model = ConvNet()
    36.     torch.cuda.set_device(gpu)
    37.     model.cuda(gpu)
    38.     batch_size = 100
    39.     # define loss function (criterion) and optimizer
    40.     criterion = nn.CrossEntropyLoss().cuda(gpu)
    41.     optimizer = torch.optim.SGD(model.parameters(), 1e-4)
    42.     # Data loading code
    43.     train_dataset = torchvision.datasets.MNIST(root='./data', train=True, transform=transforms.ToTensor(), download=True)
    44.     # add 2
    45.     if hvd.size() > 1:
    46.         train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset, num_replicas=hvd.size(), rank=hvd.rank())
    47.         shuffle = False
    48.         optimizer = hvd.DistributedOptimizer(optimizer, named_parameters=model.named_parameters(), op=hvd.Average)
    49.         hvd.broadcast_parameters(model.state_dict(), root_rank=0)
    50.         hvd.broadcast_optimizer_state(optimizer, root_rank=0)
    51.     else:
    52.         train_sampler = None
    53.         shuffle = True
    54.     train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size,shuffle=shuffle, num_workers=0, sampler=train_sampler, pin_memory=True)
    55.  
    56.     start = datetime.now()
    57.     for epoch in range(args.epochs):
    58.         if gpu == 0:
    59.             print("Epoch: {}/{}".format(epoch+1, args.epochs))
    60.         if gpu == 0:
    61.             pbar = tqdm(train_loader)
    62.         else:
    63.             pbar = train_loader
    64.         for images, labels in pbar:
    65.             images = images.cuda(non_blocking=True)
    66.             labels = labels.cuda(non_blocking=True)
    67.             # Forward pass
    68.             outputs = model(images)
    69.             loss = criterion(outputs, labels)
    70.  
    71.             # Backward and optimize
    72.             optimizer.zero_grad()
    73.             loss.backward()
    74.             optimizer.step()
    75.             if gpu == 0:
    76.                 msg = 'Loss: {:.4f}'.format(loss.item())
    77.                 pbar.set_description(msg)
    78.     
    79.     if gpu == 0:
    80.         print("Training complete in: " + str(datetime.now() - start))
    81.  
    82. def main():
    83.     parser = argparse.ArgumentParser()
    84.     parser.add_argument('-n', '--nodes', default=1, type=int, metavar='N')
    85.     parser.add_argument('-g', '--gpus', default=2, type=int, help='number of gpus per node')
    86.     parser.add_argument('-nr', '--nr', default=0, type=int, help='ranking within the nodes')
    87.     parser.add_argument('--epochs', default=10, type=int, metavar='N', help='number of total epochs to run')
    88.     args = parser.parse_args()
    89.     # add 1
    90.     hvd.init()
    91.     train(hvd.local_rank(), args)
    92.  
    93. if __name__ == '__main__':
    94.     main()

    启动命令:

    1. # Horovod 1 单机训练
    2. #horovodrun -np 1 -H localhost:1 python train_horovod.py
    3. # Horovod 2 双卡训练
    4. horovodrun -np 2 -H localhost:2 python train_horovod.py
    5. # Horovod 4 四卡训练
    6. #horovodrun -np 4 -H localhost:4 python train_horovod.py

    这种方式下双卡仅需41秒 

    虽然单卡的速度没有变快，但是处理任务的worker多了，吞吐就上去了，4卡仅需要20秒 


    分布式训练之旅   PyTorch分布式训练简明教程

    PyTorch分布式训练基础--DDP使用

    主要变动的位置包括:
    1. 启动的方式引入了一个多进程机制； 
    2. 引入了几个环境变量； 
    3. DataLoader多了一个sampler参数； 
    4. 网络被一个DistributedDataParallel(net)又包裹了一层； 
    5. ckpt的保存方式发生了变化。

    Pytorch - 分布式通信原语（附源码） 

    Pytorch - 多机多卡极简实现（附源码）

    Pytorch - DDP实现分析

    Pytorch - 使用Horovod分布式训练

    Pytorch - Horovod分布式训练源码分析 

  • 相关阅读:
    Python 快速入门
    Ubuntu20.04 Mysql基本操作知识
    数据结构笔记（王道考研） 第六章：图
    我的Quick Latex For Obsidian-Setting（持续更新）
    K8s安全一
    设计一个实用好看的餐边柜或者酒柜需要知道这5点。福州中宅装饰，福州装修
    ZTree自定义不可展开节点的折叠图标
    1024程序节特辑：一文读懂小程序支付流程
    Linux/Ubuntu/Debian 常用用户管理命令
    matlab计算图像信噪比SNR
  • 原文地址：https://blog.csdn.net/minstyrain/article/details/127731963