Pytorch-DistributedDataParallel（DDP）进行多 GPU 训练

使用 ResNet18 作为模型，并使用 CIFAR-10 数据集进行训练:

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms, models
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel

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

    # 初始化进程组
    dist.init_process_group("nccl", rank=rank, world_size=world_size)

def cleanup():
    dist.destroy_process_group()

def train(rank, world_size):
    setup(rank, world_size)

    # 假设你有一个简单的模型，例如在 ImageNet 上预训练的 ResNet18
    model = models.resnet18(pretrained=True)
    # 替换最后的全连接层以适应你的任务
    model.fc = nn.Linear(512, num_classes)  # 替换 num_classes 为你的分类数

    # 将模型移到 GPU 上
    model = model.to(rank)

    # 使用 DistributedDataParallel 将模型包装起来
    model = DistributedDataParallel(model, device_ids=[rank])

    # 定义损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.001)

    # 加载数据集（以 CIFAR-10 为例）
    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
    train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)

    # 使用 DistributedSampler 对数据集进行分布式采样
    train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset, num_replicas=world_size, rank=rank)
    train_loader = DataLoader(train_dataset, batch_size=64, shuffle=False, sampler=train_sampler)

    # 训练模型
    num_epochs = 5
    for epoch in range(num_epochs):
        for images, labels in train_loader:
            images, labels = images.to(rank), labels.to(rank)

            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

        print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item()}')

    cleanup()

if __name__ == "__main__":
    # 启动多个进程进行分布式训练
    os.environ["CUDA_VISIBLE_DEVICES"]='0,1,2,3'
    world_size = 4  # 假设有4个GPU
    mp.spawn(train, args=(world_size,), nprocs=world_size, join=True)

但是我们往往是运行两个程序，使用多gpu训练，所以端口设置这里需要自动化一些，避开已有端口，防止端口冲突 ：

def setup(rank, world_size):    
    os.environ['MASTER_ADDR'] = 'localhost' #'127.0.0.1'
    # os.environ['MASTER_PORT'] = '12345'
    # 固定端口号
    base_port = 12345
    # 检查端口是否被占用，如果被占用，自动选择另一个端口
    current_port = base_port
    while True:
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        result = sock.connect_ex(('localhost', current_port))
        sock.close()
        if result == 0:
            current_port += 1
        else:
            os.environ['MASTER_PORT'] = str(current_port)
            break
    # 初始化进程组
    dist.init_process_group("nccl", rank=rank, world_size=world_size)