*# Dockerfile# Start FROM Nvidia PyTorch image https://ngc.nvidia.com/catalog/containers/nvidia:pytorch
***FROM** nvcr.io/nvidia/pytorch:20.03-py3

*## main.py文件
***import** torch

*# 构造模型
*model **=** nn**.**Linear(10, 10)**.**to(local_rank)

*# 前向传播
*outputs **=** model(torch**.**randn(20, 10)**.**to(rank))
labels **=** torch**.**randn(20, 10)**.**to(rank)
loss_fn **=** nn**.**MSELoss()
loss_fn(outputs, labels)**.**backward()
*# 后向传播
*optimizer **=** optim**.**SGD(model**.**parameters(), lr**=**0.001)
optimizer**.**step()

*## Bash运行
*python main**.**py

*## main.py文件
***import** torch
*# 新增：
***import** torch.distributed **as** dist

*# 新增：从外面得到local_rank参数***import** argparse
parser **=** argparse**.**ArgumentParser()
parser**.**add_argument("--local_rank", default**=-**1)
FLAGS **=** parser**.**parse_args()
local_rank **=** FLAGS**.**local_rank

*# 新增：DDP backend初始化*torch**.**cuda**.**set_device(local_rank)
dist**.**init_process_group(backend**=**'nccl')  *# nccl是GPU设备上最快、最推荐的后端# 构造模型*device **=** torch**.**device("cuda", local_rank)
model **=** nn**.**Linear(10, 10)**.**to(device)
*# 新增：构造DDP model*model **=** DDP(model, device_ids**=**[local_rank], output_device**=**local_rank)

*# 前向传播*outputs **=** model(torch**.**randn(20, 10)**.**to(rank))
labels **=** torch**.**randn(20, 10)**.**to(rank)
loss_fn **=** nn**.**MSELoss()
loss_fn(outputs, labels)**.**backward()
*# 后向传播*optimizer **=** optim**.**SGD(model**.**parameters(), lr**=**0.001)
optimizer**.**step()


*## Bash运行# 改变：使用torch.distributed.launch启动DDP模式，#   其会给main.py一个local_rank的参数。这就是之前需要"新增:从外面得到local_rank参数"的原因*python **-**m torch**.**distributed**.**launch **--**nproc_per_node 4 main**.**py

model **=** torch**.**nn**.**DataParallel(model)

model = DDP(model, device_ids=[local_rank], output_device=local_rank)

*## main.py文件
***import** torch
**import** argparse

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

*# 新增2：从外面得到local_rank参数，在调用DDP的时候，其会自动给出这个参数，后面还会介绍。所以不用考虑太多，照着抄就是了。
*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')

*# 新增4：定义并把模型放置到单独的GPU上，需要在调用`model=DDP(model)`前做哦。
#       如果要加载模型，也必须在这里做哦。
*device **=** torch**.**device("cuda", local_rank)
model **=** nn**.**Linear(10, 10)**.**to(device)
*# 可能的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()

*## Bash运行**# 假设我们只在一台机器上运行，可用卡数是8*
python -m torch.distributed.launch --nproc_per_node 8 main.py

*# 假设我们只用4,5,6,7号卡
*CUDA_VISIBLE_DEVICES**=**"4,5,6,7" 
python -m torch.distributed.launch --nproc_per_node 4 main.py
*# 假如我们还有另外一个实验要跑，也就是同时跑两个不同实验。
#    这时，为避免master_port冲突，我们需要指定一个新的。这里我随便敲了一个。
*CUDA_VISIBLE_DEVICES**=**"4,5,6,7" python -m torch.distributed.launch --nproc_per_node 4 \
    --master_port 53453 main.py

**def** **demo_fn**(rank, world_size):
    dist**.**init_process_group("nccl", rank**=**rank, world_size**=**world_size)
    *# lots of code.***...****def** **run_demo**(demo_fn, world_size):
    mp**.**spawn(demo_fn,
             args**=**(world_size,),
             nprocs**=**world_size,
             join**=**True)