julia有 pytorch包吗_PyTorch 有哪些坑/bug？

说一个 distributed 的坑。

一般如果用 DistributedDataParallel (分布式并行)的时候，每个进程单独跑在一个 GPU 上，多个卡的显存占用用该是均匀的，比如像这样的：

其实一般来说，在 Distributed 模式下，相当于你的代码分别在多个 GPU 上独立的运行，代码都是设备无关的。比如你写 t = torch.zeros(100, 100).cuda()，在4个进程上运行的程序会分别在4个 GPUs 上初始化 t。所以显存的占用会是均匀的。

但是有的时候你会发现另外几个进程会在0卡上占一部分显存，导致0卡显存出现瓶颈，可能会导致cuda-out-of-memory 错误。比如这样的：

我发现我的代码中有一处很隐蔽的地方会导致这种情况发生：当你用

checkpoint = torch.load("checkpoint.pth")

model.load_state_dict(checkpoint["state_dict"])

这样load一个 pretrained model 的时候，torch.load() 会默认把load进来的数据放到0卡上，这样4个进程全部会在0卡占用一部分显存。

解决的方法也很简单，就是把load进来的数据map到cpu上：

checkpoint = torch.load("checkpoint.pth", map_location=torch.device('cpu'))

model.load_state_dict(checkpoint["state_dict"])

广告一下最近看的两篇有意思的论文:朱见深：Learnable Group Convolutions:可以学习的分组卷积​zhuanlan.zhihu.com