首页 > 分享 > PyTorch单机多卡训练（无废话）

PyTorch单机多卡训练（无废话）

萌宠菠菠乐园
2024-10-17 03:53

目前大家基本都在使用DistributedDataParallel（简称DDP）用来训练，该方法主要用于分布式训练，但也可以用在单机多卡。

第一步：初始化分布式环境，主要用来帮助进程间通信

torch.distributed.init_process_group(backend='nccl') 1

第二步：负责创建 args.local_rank 变量，并接受 torch.distributed.launch 注入的值

归根到底是创建一个变量，来接收torch.distributed.launch 注入。
目前代码中常见的两种方式：

local_rank = int(os.environ["LOCAL_RANK"]) # 这种是从自定义config文件中获取LOCAL_RANK 1

另外一种是parser：

parser = argparse.ArgumentParser() parser.add_argument("--local_rank", type=int, default=-1) args = parser.parse_args() 123

第三步：每个进程根据自己的local_rank设置应该使用的GPU

torch.cuda.set_device(args.local_rank) torch.manual_seed(hps.train.seed) # 非必要，尽可能固定种子 12

第四步：分布式数据和模型

# 分布式数据 train_sampler = DistributedSampler(train_dataset) train_loader = torch.utils.data.DataLoader(train_dataset, sampler=train_sampler, batch_size=batch_size) # 分布式模型 model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], output_device=args.local_rank) 12345

第五步：运行

nproc_per_node表示每个node有多少个进程，每个GPU对应一个进程。
nnodes表示使用几个节点，一个节点对应一台电脑。

python -m torch.distributed.launch --nproc_per_node=2 --nnodes=1 train.py 1

解释

torch.distributed.launch参数解析（终端运行命令的参数）：

--> python -m torch.distributed.launch --help usage: launch.py [-h] [--nnodes NNODES] [--node_rank NODE_RANK] [--nproc_per_node NPROC_PER_NODE] [--master_addr MASTER_ADDR] [--master_port MASTER_PORT] [--use_env] [-m] [--no_python] [--logdir LOGDIR] training_script ... 123456 nnodes：节点的数量，通常一个节点对应一个主机，方便记忆，直接表述为主机node_rank：节点的序号，从0开始nproc_per_node：一个节点中显卡的数量master_addr：master节点的ip地址，也就是0号主机的IP地址，该参数是为了让其他节点知道0号节点的位，来将自己训练的参数传送过去处理master_port：master节点的port号，在不同的节点上master_addr和master_port的设置是一样的，用来进行通信 torch.ditributed.launch相关环境变量解析（代码中os.environ中的参数）： WORLD_SIZE：os.environ[“WORLD_SIZE”]所有进程的数量LOCAL_RANK：os.environ[“LOCAL_RANK”]每张显卡在自己主机中的序号，从0开始RANK：os.environ[“RANK”]进程的序号，一般是1个gpu对应一个进程
在这里插入图片描述 一般WORLD_SIZE和RANK和LOCAL_RANK在一起声明

rank = dist.get_rank() local_rank = int(os.environ["LOCAL_RANK"]) n_gpus = dist.get_world_size() 123 多机和单机的一些概念问题 rank
多机多卡：代表某一台机器
单机多卡：代表某一块GPUworld_size
多机多卡：代表有几台机器
单机多卡：代表有几块GPUlocal_rank
多机多卡：代表某一块GPU的编号
单机多卡：代表某一块GPU的编号

参考文献

1.https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel
2. https://github.com/jia-zhuang/pytorch-multi-gpu-training
3. https://blog.csdn.net/weixin_44966641/article/details/121872773
4. https://blog.csdn.net/magic_ll/article/details/122359490