Pytorch gpu加速方法

原文：

https://www.zhihu.com/question/274635237

• relu 用 inplace=True
• 用 eval() 和 with torch.no_grad():
• 每个 batch 后认真的把所有参数从 GPU 拿出来后删除
• 虽然很多回答建议用, 但我建议不要用 torch.cuda.empty_cache() , 这只是释放 GPU 缓存而使得 nvidia-smi 能看得见 pytorch 自动释放的内存而已. 99% 的用户不需要使用这个命令. 并有用户反应每次用反而会减慢 1~2s.[1]
• 注意: 当每张 GPU 里面的 batch_size 太小(<8)时用 batch_norm 会导致训练不稳定, 除非你用以下所说的 APEX 来实现多 GPU sync_bn
• torch.backends.cudnn.deterministic = True 用不用对 GPU 内存占用和效率都没有什么太大的影响. 建议开着.
• 不要用 .cpu() 来取 GPU 里面出来的图片. 这样做的话训练时长可能翻倍.

实现: 研究 pytorch 官方架构就会发现大部分 forward pass 都是 `x = self.conv(x)` 的形式, 很少 introduce new variable. 所以: (1) 把不需要的变量都由 `x` 代替; (2) 变量用完后用 `del` 删除.

例子

def forward(self, x):
    conv2 = self.conv2(self.conv1(x)) #1/4
    del x
    conv3 = self.conv3(conv2) #1/8
    conv4 = self.conv4(conv3) #1/16
    conv5 = self.conv5(conv4) #1/32

    center_64 = self.center_conv1x1(self.center_global_pool(conv5))

    d5 = self.decoder5(self.center(conv5), conv5)
    del conv5
    d4 = self.decoder4(d5, conv4)
    del conv4
    d3 = self.decoder3(d4, conv3)
    del conv3
    d2 = self.decoder2(d3, conv2)
    del conv2

如果你按照上面的方法把 pin_memory 开启了的话, 请数据放入 GPU 的时候把 non_blocking 开启. 这样如果你只把数据放入 GPU 而不把数据从 GPU 拿出来再做计算的话就会加快很多 (据用户报告可加速 50%). 就算你把 GPU 中数据拿出来 (ie. 用了 .cpu() 命令, 最坏的结果也是与 non_blocking=False 相当:

"""Sync Point"""
image = image.cuda(non_blocking=True)
labels = labels.cuda(non_blocking=True).float()

"""Async Point"""
prediction = net(image)