在pycharm、clion中运行/调试torch分布式训练

记录一下使用JetBrains IDE工具中分布式调试

在很多深度学习研究中，都会进行分布式训练，其一般会使用shell脚本启动训练；但在调试时，用shell脚本就显的不是太方便。而使用IDE调试则会方便很多，可以加快撸代码和bug分析解决的速度。所以两者结合，怎么用IDE进行分布式调试，是大多数AI工程师的一个小需求。其主要难点在于不知道shell脚本中的命令如何配置到IDE中。下面本博文以pytorch框架和Clion为例，来进行配置演示，此处IDE换成PyCharm也基本是一样的配置。

开始

首先，找到项目的Readme文件或shell运行脚本，查看提供的分布式训练命令。例如是

#!/bin/bash
export CUDA_VISIBLE_DEVICES="0，1，2，3"
export OMP_NUM_THREADS=2
...
python -m torch.distributed.launch --nproc_per_node=4 tools/train.py --cfg xxx.yaml

如上，其中export行是设置环境相关命令；python -m torch.distributed.launch 表示调用torch.distributed.launch 这个.py文件进行分布式训练；–nproc_per_node=4 说明创建节点数为4，这个值通常与训练使用的GPU数量一致。tools/train.py --cfg xxx.yaml 是真正的训练文件，后面的–cfg xxx.yaml 是train.py 使用时需要给出的执行参数名称和值。

设置

经过上面的分析，可以知道其实是先运行的pytorch中的launch.py脚本，其调用了项目中train.py进行启动运行。所以可以将对应项分别填到IDE中即可。例如，点击clion的Edit Configurations…

出现Run/Debug Configurations界面，如图所示：

1.将${your pytorch path}/distributed/launch.py路径填入Script path中。当然也可以像下方参考文章中的建立软连接。
2.将启动需要的参数填入Parameters中。
3.将环境设置变量填入Environment variables中。

经过以上配置，就可以实现分布式调试啦。
注：是运行launch.py文件，不是运行train.py文件。

参考

[1] 如何在pycharm中运行/调试torch分布式训练