VS Code中如何调试pytorch分布式训练脚本torch.distributed

目录

一、问题描述

二、解决方案

三、测试

---

一、问题描述

最近跑一些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了。