①首先要明白训练模型时显存主要用在如下四个地方:(1、模型参数;2、梯度;3、优化器;4、激活值)
以两层NN为举例(第一次NN是4*3,第二层NN是3*2假设batch_size=3;Adam优化器)
参数量:
模型参数显存占用:92字节
梯度显存占用:92字节
优化器显存占用:184字节
激活值显存占用:60字节
总结:
1.梯度和模型结构的参数量是一样。
2.Adam优化器占用显存是模型结构的2倍。
3.激活值就是每层输出值,它略小于模型结构参数。