pycharm pytorch_Pycharm 下如何 debug torch.distributed

前言

关于 torch.distributed 的用法，请参考：

Todd：PyTorch Parallel Training（单机多卡并行、混合精度、同步BN训练指南文档）

纵横：当代研究生应当掌握的并行训练方法（单机多卡）

https://github.com/tczhangzhi/pytorch-distributed

里面的内容大同小异，看过一个就够了，这里不作赘述。

上述文章给出的代码调用方式，都是在命令行输入如下命令：

CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.launch --nproc_per_node=4 main.py

我只有两个卡，所以相应的命令就是：

CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 main.py

但是这样运行代码 debug 起来很不方便，能不能用 pycharm 运行代码呢？谷歌了一下，有这种想法的人很多，但大多数都踩到了不同的坑里面，甚至有人认为 pycharm 不支持分布式训练。但本人经过尝试，终于找到了在 pycharm 下 debug torch.distributed 的方法，在此记录，以方便之后遇到这个问题的同学。

对我帮助最大的是这个链接中的内容：how to run 'python-m' command in pycharm

方法

方法说起来很简单，只需要在 Pycharm 的 Configuration 中作一些设置即可。

• 打开 Run -> Edit Configurations...
• Script path 不再是你自己代码的路径，而是 launch.py 文件的保存路径，例如我的是：

/home/xn/miniconda3/envs/deeplearning/lib/python3.6/site-packages/torch/distributed/launch.py

• 设置 Parameters：

--nproc_per_node=2 main.py

• 在 Environment variables 中添加 CUDA_VISIBLE_DEVICES=0,1 。
• 剩下的 Python interpreter 和 Working directory 就按照通常情况来设置即可。

经过这些步骤，就可以在 Pycharm 中 debug 分布式训练的代码了。

附一张图：

我的主函数为 main_dist.py 故作相应修改

需要引起注意的是，如果强制终止 debug，经常会造成某个卡的显存不能正常释放，此时可以在 nvidia-smi 中查看没有终止的进程，手动将其 kill 掉。

错误做法

下图为错误的 configuration，这样设置可以正常运行代码，但是 debug 时会出现 ValueError: Unexpected option: --local_rank=1 。具体原因未知。