Inplace ABNSync 与 pytorch GPU多卡并行的一点坑

说在前头

众所周知，torch.nn.DataParallel(variable, device_ids=None) 是 pytorch进行多卡并行的函数。每次只要将变量用这个函数转换，即可转化成多卡并行的变量。Inplace ABN 是可以代替 BatchNorm 的一个新函数，效果比 pytorch 自带的 BatchNorm 好，而且可以节省显存。由于笔者之前没用过 pytorch 的并行模型，所以使用这个 repo 的功能的时候遇到很多坑。

1、torch.cuda()

这个函数将变量转化成可以在显卡上跑的模式。其实有个参数，指定了GPU 的类型。如果不加参数，默认指定为 GPU 0，也就是这个变量转化到 GPU 0 上。这其实就有点问题了，因为多卡并行的时候其实是一个主模块和很多子模块。主模块负责进行 getter/scatter 操作，负责同步。如果你使用了 GPU 1，2，3 进行并行，而 torch.cuda()中没加参数，就会默认用 GPU 0 当主模块。这样就会引发问题，因为 GPU 0 实际没有参与运算。所以可以使用 torch.cuda(device_ids[0])。

2、ninja 的问题

使用 Inplace ABN 需要使用 ninja，而 ninja 是个类似 gcc 的编译器。这个通过很多渠道可以安装，常用的比如 conda install ninja。然后安装后，Inplace ABN 调用它的时候可能显示找不到 ninja 这个命令。这时候将 ninja 可执行文件移动到 /usr/sbin 下即可。

3、libcudart.so.9.1 找不到

可以把 libcudart.so.9.1 所在的文件夹路径放到 ~/.bashrc 中，比如

export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/cuda/lib64"

4、os.environ[“CUDA_VISIBLE_DEVICES”] 设置无效

未解，在 pytorch 中似乎这样并不能指定使用的 GPU ID。

5、 Inplace ABNSync 使用中的编译相关问题

subprocess.CalledProcessError: Command '['ninja', '-v']' returned non-zero exit status 1.
ImportError: /tmp/torch_extensions/inplace_abn/inplace_abn.so: undefined symbol: _ZN2at5ErrorC1ENS_14SourceLocationESs
可能是 pytorch 版本的问题。参见 github 的 issue
issue32
issue71

6、Inplace ABNSync 同步时卡住不动

Inplace ABNSync 默认使用 GPU 0 来同步。同时类中有个参数可以设置 GPU 的 ID。设置 ID 之后可解决。

7、总结

善用 github 的搜索功能。很多已经关闭的 issue 可能会有很大的启发性。一定程度来讲，在 github 搜索比 google 更容易解决问题。