前言
关于 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 分布式训练的代码了。
附一张图:
需要引起注意的是,如果强制终止 debug,经常会造成某个卡的显存不能正常释放,此时可以在 nvidia-smi
中查看没有终止的进程,手动将其 kill 掉。
错误做法
下图为错误的 configuration,这样设置可以正常运行代码,但是 debug 时会出现 ValueError: Unexpected option: --local_rank=1 。具体原因未知。