笔者在知乎上发现了一个叫“Pytorch有什么节省内存(显存)的小技巧??”的问题,很有意思,下面与大家分享一些大佬们的见解,希望对你的研究有帮助。

知乎高质量回答

作者:郑哲东

在不修改网络结构的情况下, 有如下操作:同意

@Jiaming

,  尽可能使用inplace操作, 比如relu 可以使用 inplace=True

进一步,比如ResNet 和 DenseNet 可以将 batchnorm 和relu打包成inplace,在bp时再重新计算。使用到了pytorch新的checkpoint特性,有以下两个代码。由于需要重新计算bn后的结果,所以会慢一些。gpleiss/efficient_densenet_pytorch

mapillary/inplace_abn


3. 每次循环结束时 删除 loss,可以节约很少显存,但聊胜于无。可见如下issue

Tensor to Variable and memory freeing best practices

4. 使用float16精度混合计算。我用过

@NVIDIA英伟达

apex,很好用,可以节约将近50%的显存,但是要小心一些不安全的操作如 mean和sum,溢出fp16。

NVIDIA/apex


5. 对于不需要bp的forward,如validation 请使用 torch.no_grad ,  注意model.eval() 不等于 torch.no_grad() 请看如下讨论。

'model.eval()' vs 'with torch.no_grad()'

6. torch.cuda.empty_cache() 这是del的进阶版,使用nvidia-smi 会发现显存有明显的变化。但是训练时最大的显存占用似乎没变。大家可以试试。

How can we release GPU memory cache?

另外,会影响精度的骚操作还有:

把一个batchsize=64分为两个32的batch,两次forward以后,backward一次。但会影响 batchnorm等和batchsize相关的层。

相关链接:

老外写的提高pytorch效率的方法,包含data prefetch等

Optimizing PyTorch training code

作者:Gary

一般呢,神经网络显存的占用可以简单分为这三部分:网络模型自身参数占用的显存。

模型计算时(包括forward/backward/optimizer)所产生的中间变量或参数也有占用显存。

编程框架自身一些额外的开销。

依据个人一些小经验,改变网络结构和不改变其结构的节省显存的方法有:减小Batch-size(这哪门子算trick,哈哈,- -!)

出自https://oldpan.me/archives/how-to-use-memory-pytorch,牺牲计算速度减少显存用量,将计算过程分为两半,先计算一半模型的结果,保存中间结果再计算后面一半的模型。如下# 输入

input = torch.rand(1, 10)

# 假设我们有一个非常深的网络

layers = [nn.Linear(10, 10) for _ in range(1000)]

model = nn.Sequential(*layers)

output = model(input)

### 可进行如下更改

# 首先设置输入的input=>requires_grad=True

# 如果不设置可能会导致得到的gradient为0

input = torch.rand(1, 10, requires_grad=True)

layers = [nn.Linear(10, 10) for _ in range(1000)]

# 定义要计算的层函数,可以看到我们定义了两个

# 一个计算前500个层,另一个计算后500个层

def run_first_half(*args):

x = args[0]

for layer in layers[:500]:

x = layer(x)

return x

def run_second_half(*args):

x = args[0]

for layer in layers[500:-1]:

x = layer(x)

return x

# 我们引入新加的checkpoint

from torch.utils.checkpoint import checkpoint

x = checkpoint(run_first_half, input)

x = checkpoint(run_second_half, x)

# 最后一层单独调出来执行

x = layers[-1](x)

x.sum.backward()  # 这样就可以了使用pooling,减小特征图的size。

减少全连接层的使用。

relu(inplace=true),inplace_abn

使用半精度float16。

optimizer的变换使用,理论上,sgd

Depthwise Convolution。

暂时想到这些,最后贴一张模型大小和准确率的图,忘记是哪篇paper了,侵删


作者:Jiaming

针对显存OOM:

0.利用nvidia-smi -lms检查是是否有显存泄露情况,即显存消耗越来越高

1.尽量使用inplace操作/flag,如nn.ReLU(inplace=True)

2.减少batch size/input size

3.使用0.4.0的新feature: checkpoint

作者:不疯魔不像话

之前看到pytorch1.0新更新的教程,模型实在太大的话,在多gpu机器上可以将模型模型拆成两部分,在两张卡上运行。

初始化模型的时候,两部分模型放在两张卡上。第一部分结算结束后,将中间结果移到第二张卡上,计算结束后在发给输出卡。

如何还想并行的话,也可以用pytorch新加的低一级的并行元语实现。

所以,如果之前是8张卡一起跑8个模型,现在就变成8张卡一起跑四个模型。

链接:PyTorch官方多gpu例子


官方给的例子是将一部分模型放在了cpu上,一部分在gpu上。当我们的模型过大时,完全可以仿照这种方法,利用torch.nn和数据的to()方法,将模型放在两张卡上,中间结果利用to()方法进行传递。

ps:这种方法肯定会影响效率,但当你的实在没办法的时候,我觉得这还是可以接受的。

— 完 —编辑:忆臻