### 使用 `torch.distributed.run` 启动分布式训练 为了使用 `python -m torch.distributed.run --nproc_per_node 3` 来运行 `train.py` 分布式训练脚本，可以按照以下方式构建完整的命令： ```bash python -m torch.distributed.run --nproc_per_node=3 train.py [其他参数] ``` 其中 `[其他参数]` 是指传递给 `train.py` 的任何额外参数。此命令通过指定 `--nproc_per_node=3` 参数告知系统在每个节点上启动三个进程来进行分布式训练[^1]。 需要注意的是，在较新的 PyTorch 版本中，推荐使用的模块名已更改为 `torch.distributed.run` 而不是旧版本中的 `torch.distributed.launch`[^2]。因此建议确认所使用的PyTorch版本并相应调整命令。 对于调试目的而言，如果希望在类似 PyCharm 的 IDE 中设置断点进行调试，则可以在IDE的运行配置里加入上述提到的启动选项以及必要的环境变量配置[^3]。 ### 示例代码片段展示如何定义多GPU训练入口函数 假设 `train.py` 文件中有如下结构用于支持单机多卡训练: ```python import argparse import torch.nn as nn from torch.utils.data import DataLoader, Dataset from torchvision.models.segmentation import deeplabv3_resnet50 def main_worker(gpu, ngpus_per_node, args): # 初始化Distributed Data Parallel (DDP) rank = gpu * ngpus_per_node + args.rank model = deeplabv3_resnet50(pretrained=True).cuda() # 将模型封装成 DDP 实例 ddp_model = nn.parallel.DistributedDataParallel(model) dataset = CustomDataset() # 自定义数据集类实例化 dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=(rank==0)) for epoch in range(args.epochs): pass # 训练逻辑 if __name__ == '__main__': parser = argparse.ArgumentParser(description='Train Model') parser.add_argument('--batch-size', type=int, default=8, help='input batch size for training ') parser.add_argument('--epochs', type=int, default=10, help='number of epochs to train ') parsed_args = parser.parse_args() ngpus_per_node = torch.cuda.device_count() world_size = ngpus_per_node torch.multiprocessing.spawn(main_worker, nprocs=ngpus_per_node, args=(world_size, parsed_args)) ```