Pytorch 节省内存、显存的一些技巧

本文分享了在PyTorch中节省内存和显存的多种技巧,包括使用inplace操作、降低计算精度、合理利用eval和requires_grad=False、以及适时调用torch.cuda.empty_cache()等方法,特别适用于处理大型模型和高精度图像处理任务。

摘要生成于 C知道 ,由 DeepSeek-R1 满血版支持, 前往体验 >

Pytorch中有一些节省内存、显存的技巧,我结合自己的经验以及如下网址所载内容进行描述:
技巧

inplace 操作

比如在relu或者LeakyRelu里面使用inplace,可以减少对内存的消耗;这种操作根据我个人的经验是比较有效的,尤其是在一些ResNet结构单元使用比较多的模型上;具体可以参考下文进行尝试:

  1. 比如ResNet 和 DenseNet 可以将 batchnorm 和relu打包成inplace,在bp时再重新计算。使用到了pytorch新的checkpoint特性,有以下两个代码。由于需要重新计算bn后的结果,所以会慢一些。
    model1
    model2
    在这里插入图片描述

删除loss

每次循环结束时 删除 loss,可以节约很少显存,但聊胜于无。可见如下issue
issue
这个方案我没有使用过,不太好评价

降低计算的精度,比如float32 变为float16

这个是一个大杀器,基本上有减少内存尤其是显存消耗的模型,都可以这么做,我个人的经验是在CNN图像处理领域,这么做不会带来显著的模型质量的下降,这里给出一个英伟达的apex,根据参考文档作者的建议,这个是比较有效的工具。
apex

分清楚eval 和 requires_grad = False

对于不需要bp的forward,如validation 请使用 torch.no_grad , 注意model.eval() 不等于 torch.no_grad() 请看如下讨论
eval vs. no grad
关于这种操作的效果,我没有使用过,因此难以评论。

使用torch.cuda.empty_cache()

在确定的地方释放显存
Can we release GPU cache
这种方式我使用过,效果很棒

总结

我个人的使用经验是,如果你什么都不清楚的话,建议使用降低精度的方法试一试。

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值