PyTorch中的loss.backward()和optimizer.step()区别

一句话解释

loss.backward()作用是将当前GPU储存gradient的计算图反传;optimizer.step()是利用计算图中储存的gradient更新网络参数。

关于累加loss

一句话解释,n个backward使用一个step

详解

假如想要使用batch size(bs)为100训练模型,但是老黄卖的大memeory的卡太贵了,只能fit进batch size为10的数据,那怎么办呢?

在PyTorch中可以这样搞:使用10的bs对网络进行forward10次,每次都调用loss.backward(),但是每10次forward才调用一次optimizer.step(),这样就实现了batch size为100的效果。其原理就是step和bachward的区别。

小坑预警

pytorch中,网络每次forward都会生成一次计算图用于存储gradient,因此必须每次forward之后调用loss.bachward()反传以及清空当前计算图,否则多次forward不反传会在GPU中累加计算图,导致GPU内存不够。

Ref