torch.distributed.launch命令介绍
我们在训练分布式时候,会使用到 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 ...
torch.ditributed.launch参数解析(终端运行命令的参数):
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对应一个进程
单机多卡 的分布式
python -m torch.distributed.launch --nproc_per_node 4 --master_port='29500' train.py
多机多卡 的分布式
在0号机器上调用
python -m torch.distributed.launch --nproc_per_node 4 --nnodes 2 --node_rank 0 --master_addr='172.18.39.122' --master_port='29500' train.py
在1号机器上调用
python -m torch.distributed.launch --nproc_per_node 4 --nnodes 2 --node_rank 1 --master_addr='172.18.39.122' --master_port='29500' train.py
注意:
命令中的【–master_addr=‘172.18.39.122’】指的是0号机器的IP,在0号机器上运行的命令中【node_rank】必须为0
只有当【nnodes】个机器全部运行,代码才会进行分布式的训练操作,否则一直处于等待状态
单机多卡 训练
只需要说明 想要使用GPU的[编号]、[数量]即可。由于不需要不同机器之间的通信,就少了其余4个参数的设定
export CUDA_VISIBLE_DEVICES=0,1
python -m torch.distributed.launch --nproc_per_node=2 train.py
单机单卡训练
当工程提供的是分布式训练代码,但我们只想用单张显卡运行。
机器上只有一张显卡:
python -m torch.distributed.launch train.py
机器上有多张显卡:
export CUDA_VISIBLE_DEVICES=1
python -m torch.distributed.launch train.py