🧠 记录并行训练并且定制一份自己特色的代码

import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torch.utils.data.distributed import DistributedSampler
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel

def setup():
    dist.init_process_group('nccl')

def clearnup():
    dist.destroy_process_group()

class ToyModel(nn.Module):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.layer = nn.Linear(1, 1)
    
    def forward(self, x):
        return self.layer(x)

class ToyDataset(Dataset):
    def __init__(self):
        super().__init__()
        self.data = torch.tensor([1, 2, 3, 4], dtype=torch.float32)
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, index):
        return self.data[index, index+1]

ckpt_path = ''

def main():
    setup()
    rank = dist.get_rank()
    pid = os.getpid()
    print(f'current pid is {pid}')
    print(f'current rank is {rank}')
    deviced_id = rank % torch.cuda.device_count()

    dataset = ToyDataset()
    sampler = DistributedSampler(dataset=dataset)
    dataloader = DataLoader(dataset, batch_size=2, sampler=sampler)

    model = ToyModel().to(deviced_id)
    model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model) # 设置多个gpu的BN同步
    ddp_model = DistributedDataParallel(model, device_ids=[deviced_id])

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

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

    map_location = {'cuda:0': f'cuda:{deviced_id}'}
    state_dict = torch.load(ckpt_path, map_location=map_location)
    print(f'rank {rank}: {state_dict}')
    ddp_model.load_state_dict(state_dict)

    for epoch in range(2):
        sampler.set_epoch(epoch) 
        """这句莫忘，否则相当于没有shuffle数据, dataloader设置方式如下，注意shuffle与sampler是冲突的，并行训练需要设置sampler，此时务必要把shuffle设为False。
        但是这里shuffle=False并不意味着数据就不会乱序了，而是乱序的方式交给sampler来控制，实质上数据仍是乱序的。
        """
        for x in dataloader:
            print(f'epoch {epoch}, rank{rank} data: {x}')
            x = x.to(deviced_id)
            y = ddp_model(x)
            
            optimizer.zero_grad()
            loss = loss_fn(x, y)
            loss.backward()
            optimizer.step()
    
    clearnup()

if __name__ == "__main__":
    main()