每次一加载模型GPU的显存就疯狂增加,看了卷积神经网络中参数量的计算与模型占用显存的计算原理才知道原因。之前还一直怀疑自己写的有bug。感谢https://www.jianshu.com/p/b8d48c99a47c的分享。
下面是我自己结合文章做的一些总结。
占用显存的计算
也就是说我们的模型占用的显存在不同阶段的计算方式是不同的,下面以Vgg16为例。
1.load model
(1)计算方法
(2)计算示例
如上图所示
-
计算 memory for model params
-
number for model params
以input为224*224的ImageNet(架构使用VGG16)为例,经过计算可得Vgg16的参数量为138357544byte/1024/1024=138M -
memory for params
138M*4=528MB(乘以4是32位的float需要占用4个字节, 取决于params的数据类型) -
计算memory for layer outputs
-
number for layer outputs
由上图可知计算得15237608byte约15.2M -
Memory for layer ouputs
15.2M*4=58.12MB/image(乘以4是32位的float需要占用4个字节,取决于image的数据类型))
2.When Training
当模型在训练的时候除了memory for paramaters 、memory for layer outputs之外。因为我们使用了Adam还要加上memory for param gradients,memory for moment等。
此外,,在训练的时候,因为有forward 和 backward,所以参数量要乘以2。
-
计算memory for model -------- 528*4MB=1.54GB
-
memory for params
138357544byte/1024/1024*4=528MB -
memory for SGD
528MB -
memory for momentum
528MB -
memory for others
528MB -
计算memory for layer
-
bach_size=128 (memory for layer outputs and error)
128X58.12X2=14.53GB -
total memory
1.54GB+14.53GB=16.07GB