Pytorch模型测试时显存一直上升导致爆显存

问题描述

首先说明: 由于我的测试集很大, 因此需要对测试集进行分批次推理.

在写代码的时候发现进行训练的时候大概显存只占用了2GB左右, 而且训练过程中显存占用量也基本上是不变的. 而在测试的时候, 发现显存在每个batch数据推理后逐渐增加, 直至最后导致爆显存, 程序fail.

这里放一下我测试的代码:

	y, y_ = torch.Tensor(), torch.Tensor()
        for batch in tqdm(loader):
            x, batch_y = batch[0], batch[1]
            batch_y_ = model(x)
            y = torch.cat([y, batch_y], dim=0)
            y_ = torch.cat([y_, batch_y_], dim=0)

解决方法

遇到问题后我就进行单步调试, 然后观察显存的变化. 发现在模型推理这一步, 每一轮次显存都会增加.

batch_y_ = model(x)

这里令人费解的是, 模型推理实际上在训练和测试中都是存在的, 为什么训练的时候就不会出现这个问题呢.

最后发现其实是在训练的时候有这样一步与测试不同:

self.optimizer.zero_grad()

⭐️ 在训练时, 每一个batch后都会将模型的梯度进行一次清零. 而测试的时候我则没有加这一步, 这样的话每次模型再做推理的时候都会产生新的梯度, 并累积到显存当中.

清楚了问题, 那么解决方法也就随之而来, 在测试的时候让模型不要记录梯度就好, 因为其实也用不到:

with torch.no_grad():
	test()

总结

  • 训练的时候每一个batch结束除了梯度反向传播, 还要提前清理梯度
  • 梯度如果不清理的话, 会在显存中累积下来
  • 40
    点赞
  • 56
    收藏
    觉得还不错? 一键收藏
  • 8
    评论
PyTorch模型的大小超过单张显卡的容量,可以采用以下策略来解决: 1. 模型分割:将大型模型分割为多个较小的子模型,每个子模型最适合放入单张显卡进行计算。这要求模型的架构能够容易地进行拆分和重新组合,以保持整体性能。 2. 模型并行:将模型参数划分为多个部分,并在多个显卡上进行并行计算。这样可以将计算负载分散到多个显卡上,以增加可用的显存大小。在每个显卡上运行的子模型共享参数,并通过梯度聚合来更新模型参数。 3. 内存管理:优化显存的使用,例如减少中间变量和不必要的计算。可以使用in-place操作或Tensor流水线来最小化内存占用。此外,可以手动释放不再需要的Tensor以及临存储。 4. 混合精度计算:使用低精度的计算来减少显存占用PyTorch支持半精度浮点数(FP16)的训练和推理,可以通过缩小参数和激活数据的精度来减少显存使用量。 5. 多个显卡使用:如果单张显卡的容量无法满足需要,可以考虑使用多张显卡进行计算。PyTorch通过DataParallel和DistributedDataParallel等模块支持在多个显卡上进行并行计算,并自动处理数据切片和梯度聚合。 总之,当PyTorch模型超过单张显卡容量,我们可以采用模型分割、模型并行、内存管理、混合精度计算和多个显卡使用等策略来解决这个问题。这些策略的应用取决于具体的模型和硬件配置。

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 8
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

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

抵扣说明:

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

余额充值