目录
一、问题描述
二、解决方案
三、测试
一、问题描述
最近跑一些pytorch代码的时候遇到很多都是采用pytorch的分布式torch.distributed来训练的,相比于传统的nn.DataParallel,使用分布式的训练方式可以显著提升GPU使用率, 从而加快训练速度。
一般常见的pytorch分布式训练命令如下:
$ export CUDA_VISIBLE_DEVICES=0,1
$ python -m torch.distributed.launch --nproc_per_node=2 tools/train.py --model bisenetv2
上述第一条命令用来设置当前环境变量中GPU数量。第二条命令为正式的启动脚本,该启动脚本与一般的使用类似python demo.py这种形式不同,整个python -m torch.distributed.launch为启动命令,实际的启动文件为:/home/dl/anaconda3/lib/python3.7/site-packages/torch/distributed/launch.py。也就是说该命令会自动的去调用torch包中的分布式启动文件launch.py来执行。后面的--nproc_per_node=2 tools/train.py --model bisenetv2全部为执行参数。
一般情况下,正式训练时会使用上述命令直接在终端中输入进行训练,但是,我们经常需要调试代码,一边运行一边查看中间变量,从而了解整个代码的运行机理。我个人喜欢用开源的VS Code,但是在VS Code中如何调试呢?谷歌了一下,有这种想法的人很多,但大多数都踩到了不同的坑里面,即使有解决方案也是用于pycharm的,为此本文专门进行了多次尝试,终于解决了这个问题,在此记录,以方便之后遇到这个问题的读者。
二、解决方案
在VS Code中想要调试Python脚本很简单,只需要创建一个launch.json文件即可。如果没有launch.json文件,只需要单机下图中“python:当前文件”旁的齿轮按钮即可创建一个launch.json文件。
下面是最关键的地方,用于为debug设置配置参数,具体如下:
{
// Use IntelliSense to learn about possible attributes.
// Hover to view descriptions of existing attributes.
// For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
"version": "0.2.0",
"configurations": [
{
"name": "Python: 当前文件",
"type": "python",
"request": "launch",
"program": "/home/dl/anaconda3/lib/python3.7/site-packages/torch/distributed/launch.py",//可执行文件路径
"console": "integratedTerminal",
"args": [
"--nproc_per_node=1",
"tools/train.py",
"--model",
"bisenetv1",
],
"env": {"CUDA_VISIBLE_DEVICES":"0"},
}
]
}
其中,program参数用于设置使用torch分布式包中的launch.py文件来作为启动脚本,具体路径请参照具体的torch安装路径来修改。args用于设置每个参数。env用于设置环境变量。具体debug时,建议只用1个GPU来进行调试,所以nproc_per_node设置为1,CUDA_VISIBLE_DEVICES设置为0。
到这里如果直接按F5调式运行,可以勉强运行起来,如下所示:
但是会发现一直卡在这里,然后在终端中输入命令nvidia-smi,效果如下:
我的GPU没有一个正常运行起来的,每一个使用率全部是0%,但是这个时候程序也没有挂,也可以进行单步调试,通过单步调试发现问题出在这个地方:
## train loop
for it, (im, lb) in enumerate(dl):
im = im.cuda()
lb = lb.cuda()
也就是在循环取数据的时候卡死在这一步了。然后在google继续查原因,终于发现问题了,在设置dataloader的时候,一般都会这么写:
dl = DataLoader(
ds,
batch_size=batchsize,
shuffle=shuffle,
drop_last=drop_last,
num_workers=4,
pin_memory=True,
)
return dl
这里注意,有一个num_workers参数,如果是命令行运行,这样设置没有任何问题,但是在VS Code调试模式下这样设置就有问题了,google有人给出的解释是说“我们想通过这个参数来启动多个进程来加载数据,但是VS Code调试模型并不能成功的开启新进程”,到这里问题就变得显然了。只需要将num_workers设置为0即可。如下所示:
dl = DataLoader(
ds,
batch_size=batchsize,
shuffle=shuffle,
drop_last=drop_last,
num_workers=0,
pin_memory=True,
)
三、测试
按照步骤二进行相关修改,然后按F5启动调试,此时在终端中输入命令
nvidia-smi
查看GPU使用情况,如下图所示:
这个时候可以看到我的4块GPU中的第1块已经高负荷运行起来了,而且我也可以在VS Code中顺利的进行单步debug了。