一种DDP写法的记录（欢迎交流指正）

"""

从torch.distributed.launch的源码，可以看出launch实际上主要完成的工作：

1.参数定义与传递。解析环境变量，并将变量传递到子进程中。

2.起多进程。调用subprocess.Popen启动多进程。

用launch方式需要注意的位置：

需要添加一个解析local_rank的参数

parser.add_argument("--local_rank", type=int)

dist初始化的方式int_method取env

dist.init_process_group("gloo", init_method='env://')

DDP的设备都需要指定

local_rank net=torch.nn.parallel.DistributedDataParallel(net, device_ids=[args.local_rank], output_device=args.local_rank)

"""

# 下面为实际例子

import torch
import torchvision
# 改动1
import torch.utils.data.distributed
import argparse
# 改动2
import torch.distributed as dist
from torchvision import transforms
parser = argparse.ArgumentParser()
# 改动3
parser.add_argument("--local_rank", type=int)
args = parser.parse_args()
# 改动4
torch.cuda.set_device(args.local_rank)

def main():
    # 改动5
    dist.init_process_group("nccl", init_method='env://')
    trans = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
    data_set = torchvision.datasets.MNIST('~/DATA/',
                                          train=True,
                                          transform=trans,
                                          target_transform=None,
                                          download=True)
    train_sampler = torch.utils.data.distributed.DistributedSampler(data_set)
    data_loader_train = torch.utils.data.DataLoader(dataset=data_set,
                                                    batch_size=256,
                                                    sampler=train_sampler,
                                                    num_workers=16,
                                                    pin_memory=True)

    net = torchvision.models.resnet101(num_classes=10)
    net.conv1 = torch.nn.Conv2d(1, 64, (7, 7), (2, 2), (3, 3), bias=False)
    net = net.cuda()

    # DDP输出方式修改
    # 改动6
    net = torch.nn.parallel.DistributedDataParallel(net, device_ids=[args.local_rank],
                                                    output_device=args.local_rank)
    criterion = torch.nn.CrossEntropyLoss()
    opt = torch.optim.Adam(net.parameters(), lr=0.001)
    for epoch in range(1):
        for i, data in enumerate(data_loader_train):
            images, labels = data
            # 改动7
            images = images.to(args.local_rank, non_blocking=True)
            labels = labels.to(args.local_rank, non_blocking=True)
            
            opt.zero_grad()
            outputs = net(images)
            loss = criterion(outputs, labels)
            loss.backward()
            opt.step()
            if i % 10 == 0:
                print("loss: {}".format(loss.item()))



if __name__ == '__main__':
    main()

 """

假设一共有两台机器（节点1和节点2），每个节点上有8张卡， 节点1的IP地址为192.168.1.1 占用的端口12355（端口可以更换），启动的方式如下：

>>> #节点1

>>>python -m torch.distributed.launch --nproc_per_node=8 --nnodes=2 --node_rank=0 --master_addr="192.168.1.1" --master_port=12355 MNIST.py

>>> #节点2

>>>python -m torch.distributed.launch --nproc_per_node=8 --nnodes=2 --node_rank=1 --master_addr="192.168.1.1" --master_port=12355 MNIST.py

如果只要启动一台机器，只需要将nnodes数量改为1，就是单机多卡的另一种方式。设置如下： python -m torch.distributed.launch --nproc_per_node=8 --nnodes=1 --node_rank=0 --master_addr=localhost --master_port=12355 MNIST.py

"""