torch.no_grad 降低显存占用

python pytorch 计算机视觉

这个函数一般在测试的时候使用，用来降低显存占用

• 在训练的时候我们需要学习 weight 和 bias 的梯度，但是在测试的时候我们不需要获取他们的梯度，因此可以节省大量的显存
• 在训练时我们 crop 成了一个一个块，例如 size 为 256，在测试时如果将整幅图片都送入网络，有时会 out of memory，因此需要 no_grad 语句

使用很简单，在测试时将测试代码放到他下面即可

with torch.no_grad():   # 可以显著的降低显存
	for i, img in enumerate(tqdm(test_loader), 0):
    	img = img.cuda()
    	output = model(img)

本文由 Yonghui Wang 创作，采用 知识共享署名4.0 国际许可协议进行许可
本站文章除注明转载/出处外，均为本站原创或翻译，转载前请务必署名
最后编辑时间为: Dec 19, 2024 12:13 pm