解决网络训练正常、测试显存爆炸的问题(model.eval()和torch.no_grad()的区别)

问题分析:

网络训练时正常,但测试中途却出现显存爆炸时。训练网络基本快占满了一张卡,这样测试时是有可能会显存爆炸。因为可能由于某张图测试的时候需要更大的显存,图像输入偏大等不确定因素。但测试的时候可以以更少的显存进行测试的(eg:例如之前测试时11g显存会爆炸,使用后仅用了2g多)。那就是使用with torch.no_grad()语句。测试代码修改大体如下所示:

model.eval()
for i, image in enumerate(testLoader):
    # ...... (省略代码)
	with torch.no_grad():  # 主要将网络部分置于此即可
    	output = model(image)  
    # ......

附model.eval()和torch.no_grad()的区别

  • model.eval()时不启用 Batch Normalization 和 Dropout。如果模型中有BN层(Batch Normalization)和Dropout,在测试时添加model.eval()。model.eval()是保证BN层能够用全部训练数据的均值和方差,即测试过程中要保证BN层的均值和方差不变。对于Dropout,model.eval()是利用到了所有网络连接,即不进行随机舍弃神经元。
  • with torch.no_grad()则主要是用于停止autograd模块的工作,以起到加速和节省显存的作用。它的作用是将该with语句包裹起来的部分停止梯度的更新,从而节省了GPU算力和显存,但是并不会影响dropout和BN层的行为。因此,测试的时候加上此语句也不会影响测试精度的,只是停止了梯度更新而已。

参考:

Pytorch:model.train()和model.eval()用法和区别,以及model.eval()和torch.no_grad()的区别

  • 13
    点赞
  • 15
    收藏
    觉得还不错? 一键收藏
  • 2
    评论

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

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

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

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

抵扣说明:

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

余额充值