### PyTorch单机多卡训练示例 对于希望在PyTorch框架内利用单机上的多个GPU来加速模型训练的需求，下面提供了一个完整的代码实例以及一些注意事项。 #### 数据集准备与加载 为了确保数据能够被不同进程中的各个线程正确读取，在创建`DataLoader`对象时需注意参数配置。特别是当采用分布式模式时，应禁用shuffle选项以防止错误发生[^1]： ```python from torchvision import datasets, transforms import torch.utils.data.distributed as dist_data transform=transforms.Compose([ transforms.ToTensor(), ]) dataset = datasets.MNIST('./data', train=True, download=True, transform=transform) train_sampler = dist_data.DistributedSampler(dataset) train_loader = torch.utils.data.DataLoader( dataset, batch_size=batch_size, sampler=train_sampler, shuffle=False) # 关闭shuffle以防报错 ``` #### 初始化环境变量并启动子进程 通过命令行工具`torch.distributed.launch`可以轻松地管理这些细节。具体来说，可以通过指定`nproc_per_node`参数告知程序要使用的GPU数量，并执行相应的Python脚本文件[^2]: ```bash python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE your_training_script.py ``` #### 设置Distributed Data Parallel (DDP) 为了让每个进程中都拥有独立的优化器副本并且只更新本地权重，同时又能同步梯度信息给其他节点，需要使用到`DistributedDataParallel`类来进行封装处理: ```python import os import torch.nn as nn import torch.optim as optim import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP def setup(rank, world_size): os.environ['MASTER_ADDR'] = 'localhost' os.environ['MASTER_PORT'] = '12355' # initialize the process group dist.init_process_group("nccl", rank=rank, world_size=world_size) model = Model().to(device) ddp_model = DDP(model, device_ids=[device]) optimizer = optim.SGD(ddp_model.parameters(), lr=0.001) for epoch in range(num_epochs): ... setup(local_rank, num_gpus_used) ``` 上述代码展示了如何在一个典型的深度学习项目里集成单机多卡功能，从而充分利用硬件资源提升效率。