pytorch分布式训练示例小计

示例地址：https://pytorch.org/tutorials/intermediate/ddp_tutorial.html

完整代码如下：

import os
import sys
import tempfile
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim
import torch.multiprocessing as mp

from torch.nn.parallel import DistributedDataParallel as DDP

# On Windows platform, the torch.distributed package only
# supports Gloo backend, FileStore and TcpStore.
# For FileStore, set init_method parameter in init_process_group
# to a local file. Example as follow:
# init_method="file:///f:/libtmp/some_file"
# dist.init_process_group(
#    "gloo",
#    rank=rank,
#    init_method=init_method,
#    world_size=world_size)
# For TcpStore, same way as on Linux.

def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'

    # initialize the process group
    dist.init_process_group("gloo", rank=rank, world_size=world_size)
    print(f'setup rank{rank} world_size={world_size}')

def cleanup():
    dist.destroy_process_group()




class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.net1 = nn.Linear(10, 10)
        self.relu = nn.ReLU()
        self.net2 = nn.Linear(10, 5)

    def forward(self, x):
        return self.net2(self.relu(self.net1(x)))


def demo_basic(rank,world_size):
    print(f"Running basic DDP example on rank {rank} world_size {world_size}.")
    setup(rank, world_size)

    # create model and move it to GPU with id rank
    model = ToyModel().to(rank)
    ddp_model = DDP(model, device_ids=[rank])

    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

    optimizer.zero_grad()
    outputs = ddp_model(torch.randn(20, 10)) #应该是 outputs = ddp_model(torch.randn(20, 10).to(rank))
    labels = torch.randn(20, 5).to(rank)
    loss_fn(outputs, labels).backward()
    optimizer.step()

    cleanup()


def run_demo(demo_fn, world_size):
    mp.spawn(demo_fn,
             args=(world_size,),#这个参数只能这样传，rank通过start_processes函数中的for循环和nprocs传递过去了
             nprocs=world_size,
             join=True)

    # def start_processes(fn, args=(), nprocs=1, join=True, daemon=False, start_method='spawn'):
    #     _python_version_check()
    #     mp = multiprocessing.get_context(start_method)
    #     error_queues = []
    #     processes = []
    #     for i in range(nprocs):
    #         error_queue = mp.SimpleQueue()
    #         process = mp.Process(
    #             target=_wrap,
    #             args=(fn, i, args, error_queue),
    #             daemon=daemon,
    #         )
    #         process.start()
    #         error_queues.append(error_queue)
    #         processes.append(process)


# 这里demo_basic(rank,world_size)有两个参数，但
# 参考start_processes函数后，
# mp.spawn(demo_fn,
#          args=(world_size,),#这个参数只能这样传，
#          nprocs=world_size,
#          join=True)
# demo_basic(rank,world_size)中的rank通过start_processes函数中的for循环和nprocs传递过去了

if __name__ == '__main__':
    demo_fn=demo_basic
    world_size=4
    run_demo(demo_fn, world_size)

输出：

Running basic DDP example on rank 2 world_size 4.
Running basic DDP example on rank 3 world_size 4.
Running basic DDP example on rank 0 world_size 4.
Running basic DDP example on rank 1 world_size 4.
setup rank0 world_size=4
setup rank1 world_size=4
setup rank2 world_size=4
setup rank3 world_size=4