Python DistributedDataParallel（DDP）训练模型

Python 训练模型

---

---

前言

最近在搞DistributedDataParallel（DDP）模型训练，踩了不少坑，总结一下，以后遇到问题，可以看一下。

---

一、DistributedDataParallel（DDP）

DistributedDataParallel（DDP）支持多机多卡分布式训练。说效率更高，所以想用一下。
重要参考

二、使用步骤

1. 简单命令

运行时，一般创建一个sh文件，例如 dist_train.sh。 里面输入代码运行命令。

#!/usr/bin/env bash

set -x
NGPUS=$1
PY_ARGS=${@:2}

while true
do
    PORT=$(( ((RANDOM<<15)|RANDOM) % 49152 + 10000 ))
    status="$(nc -z 127.0.0.1 $PORT < /dev/null &>/dev/null; echo $?)"
    if [ "${status}" != "0" ]; then
        break;
    fi
done
echo $PORT

export CUDA_VISIBLE_DEVICES=1,2,3,4,5,6,7
python -m torch.distributed.launch --nproc_per_node=${NGPUS} --rdzv_endpoint=localhost:${PORT} train.py --launcher pytorch ${PY_ARGS}

torch.distributed.launch 的参数，因为是多台机器，多张显卡。
介绍一些参数：
–nnodes 有多少台机器 #默认是1
–node_rank 当前是哪台机器 # 默认是1
–nproc_per_node 每台机器有多少进程 # 建议一个进程一张卡，因此就是卡的数量
实现方式：在每台机子上都运行一次torch.distributed.launch，每个torch.distributed.launch会启动n个进程，并给每个进程一个–local_rank=i的参数。参考
word size 是 rank的总数量。

---

单台机器多张卡情况下使用那些卡，其实就是把那些卡暴露在Linux（Win）环境中。
在Linux有两种方式

1. 在sh文件中： export CUDA_VISIBLE_DEVICES=1,2,3,4,5,6,7
2. 在代码开始加入： os.environ[‘CUDA_VISIBLE_DEVICES’]=‘0,1,2’

在Win中：

1. set CUDA_VISIBLE_DEVICES=1,2,3,4,5,6,7
2. os.environ[‘CUDA_VISIBLE_DEVICES’]=‘0,1,2’
   torch可以直接获取系统环境中的gpus数量， num_gpus = torch.cuda.device_count()

---

新增1:依赖
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
新增2:从外面得到local_rank参数，在调用DDP的时候，其会自动给出这个参数
每个进程对应一个rank，也是local_rank。多卡的时候有区分
在使用torch.distributed.launch会自动获取local_rank，返回在args.local_rank**.
新增3：DDP backend配置初始化

1. torch.cuda.set_device(args.local_rank) # 根据local_rank来设定当前使用哪块GPU
2. torch.distributed.init_process_group(backend=‘nccl’, init_method=‘env://’) # 初始化DDP，使用默认backend(nccl)就行。如果是CPU模型运行，需要选择其他后端。

新增4：初始化DDP模型
net = model.cuda() # 将模型传入到cuda中
net = torch.nn.parallel.DistributedDataParallel(net, device_ids = [args.local_rank]) #DDP模型
新增5：DDP模型的数据分发
pytorch中使用torch.utils.data.distributed.DistributedSampler实现数据的分发，将数据分发到每张卡上。
sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
trainloader = torch.utils.data.DataLoader(my_trainset, batch_size=batch_size, sampler=train_sampler) # 这里的batch_size指的是每个进程下的batch_size。也就是说，总batch_size是这里的batch_size再乘以并行数(world_size)。
新增6：模型保存

1. save模型的时候，和DP模式一样，有一个需要注意的点：保存的是model.module而不是model。
   因为model其实是DDP model，参数是被model=DDP(model)包起来的。
2. 我只需要在进程0上保存一次就行了，避免多次保存重复的东西。

if dist.get_rank() == 0:
#save code# 看下面例子

例子

## main.py文件
import torch
import argparse

# 新增1:依赖
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

# 新增2：从外面得到local_rank参数，在调用DDP的时候，其会自动给出这个参数，后面还会介绍。所以不用考虑太多，照着抄就是了。
#       argparse是python的一个系统库，用来处理命令行调用，如果不熟悉，可以稍微百度一下，很简单！
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", default=-1)
FLAGS = parser.parse_args()
local_rank = FLAGS.local_rank

# 新增3：DDP backend初始化
#   a.根据local_rank来设定当前使用哪块GPU
torch.cuda.set_device(local_rank)
#   b.初始化DDP，使用默认backend(nccl)就行。如果是CPU模型运行，需要选择其他后端。
dist.init_process_group(backend='nccl', init_method='env://')

# 新增4：定义并把模型放置到单独的GPU上，需要在调用`model=DDP(model)`前做哦。
#       如果要加载模型，也必须在这里做哦。
model = nn.Linear(10, 10).cuda()
# 可能的load模型...

# 新增5：之后才是初始化DDP模型
model = DDP(model, device_ids=[local_rank], output_device=local_rank)

my_trainset = torchvision.datasets.CIFAR10(root='./data', train=True)


# 新增1：使用DistributedSampler，DDP帮我们把细节都封装起来了。用，就完事儿！
#       sampler的原理，后面也会介绍。
train_sampler = torch.utils.data.distributed.DistributedSampler(my_trainset)
# 需要注意的是，这里的batch_size指的是每个进程下的batch_size。也就是说，总batch_size是这里的batch_size再乘以并行数(world_size)。
trainloader = torch.utils.data.DataLoader(my_trainset, batch_size=batch_size, sampler=train_sampler)

for epoch in range(num_epochs):
    # 新增2：设置sampler的epoch，DistributedSampler需要这个来维持各个进程之间的相同随机数种子
    trainloader.sampler.set_epoch(epoch)
    # 后面这部分，则与原来完全一致了。
    for data, label in trainloader:
        prediction = model(data)
        loss = loss_fn(prediction, label)
        loss.backward()
        optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)
        optimizer.step()

# 1. save模型的时候，和DP模式一样，有一个需要注意的点：保存的是model.module而不是model。
#    因为model其实是DDP model，参数是被`model=DDP(model)`包起来的。
# 2. 我只需要在进程0上保存一次就行了，避免多次保存重复的东西。
if dist.get_rank() == 0:
    model_state_dict = net.state_dict()
    state = {'model': model_state_dict, 'optimizer': optimizer.state_dict()}
    # state = {'model': model_state_dict}
    assert os.path.exists(save_path)
    model_path = os.path.join(save_path, 'ep%03d.pth' % epoch)
    torch.save(state, model_path)

## 2.读入数据

代码如下（示例）：

data = pd.read_csv(
    'https://labfile.oss.aliyuncs.com/courses/1283/adult.data.csv')
print(data.head())

该处使用的url网络请求的数据。

---

总结

提示：这里对文章进行总结：

例如：以上就是今天要讲的内容，本文仅仅简单介绍了pandas的使用，而pandas提供了大量能使我们快速便捷地处理数据的函数和方法。