CNN内存消耗

数据挖掘 机器学习 神经网络 卷积神经网络 显卡 参数估计
2021-10-12 04:20:16

我希望能够估计所提出的模型是否足够小,可以在具有给定内存量的 GPU 上进行训练

如果我有一个像这样的简单 CNN 架构:

  • Input: 50x50x3
  • C1:32 个 3x3 内核,带填充(我猜实际上它们实际上是 3x3x3 给定输入深度?)
  • P1:2x2,步幅为 2
  • C2:64 个 3x3 内核,带填充
  • P2:2x2,步幅为 2
  • FC: 500 个神经元
  • Output:softmax 10 类
  • 小批量大小为 64

假设 32bit 浮点值,如何计算训练时网络每一层的内存消耗?那么训练这样一个模型所需的总内存呢?

3个回答

我假设C1,C2等表示卷积层,而P1表示P2池化层,FC表示全连接层。

我们可以像这样计算前向传递所需的内存:

一张图片

如果您正在使用 float32 值,那么按照上面@Alexandru Burlacu 提供的链接,您将拥有:

Input: 50x50x3 = 7,500 = 7.5K

C1: 50x50x32 = 80,000 = 80K

P1: 25x25x32 = 20,000 = 20K

C2:25x25x64 = 40,000 = 40K

P2: 12x12x64 = 9,216 = 9.2K <- 这是一个问题(我的近似值在这里是一个非常随意的猜测)。与其使用 50、25、'12.5'​​,不如使用 32 的倍数更有意义。我听说从记忆的角度来看,使用 32 的倍数也更有效。据我所知,这是一个坏主意的原因是 2x2 池不能正确划分空间。如果我错了,请随时纠正我。

FC: 1x500 = 500 = 0.5K

Output:1 x 10 = 10 = 0.01K(几乎没有)

总内存:7.5K + 80K + 20K + 40K + 0.5K = 157.2K * 4 字节 = 628.8 KB

这是一张图片。

小批量

如果您使用 64 的 minibatch 大小,那么您一次将其中 64 个读入内存并一起执行操作,像这样按比例放大所有内容:

Input: 64x50x50x3 = 480,000 = 480K = 0.48M

C1: 64x50x50x32 = 5,120,000 = 5.12M

P1: 64x25x25x32 = 1,280,000 = 1.28M

C2: 64x25x25x64 = 2,560,000 = 2.56M

P2:64x12x12x64 = 589,824 = 590K = 0.59M

FC: 64x500 = 32,000 = 32K = 0.032M

Output:1x10x64 = 640 = 0.64K = 0.00064M(我们不在乎,这很小)

总内存:10M x 4 bytes ~ 40MB(我说的是近似值,因为网站上也有近似值)

编辑:我误读了网站,对不起。

根据该网站,向后传递需要大约三倍,因为需要存储:

  • 每个神经元的激活和相关梯度——它们大小相等;

  • 与参数大小相同的权重(参数)的梯度;

  • 动量的价值,如果你正在使用它;

  • 某种杂记(这部分我没看懂)

也许这个链接会给你一个关于如何计算任意神经网络的内存使用的解释。链接中的波纹管解释了 VGGNet 模型的内存使用情况。单击此处并向下滚动一点))

在训练一个 convNet 时,所需的总内存包括以下内容:

  • 参数记忆
  • 中间层输出的内存
  • 每个参数的梯度记忆
  • 如果您使用 Momentum、RMSprop、Adams 等优化器,则需要额外的内存
  • 用于实现的杂项内存

一个很好的粗略近似值是参数数量 x 3 x 4(如果您使用 32 位浮点数)字节

好吧,现在这就是您计算参数数量的方式:

  • 卷积层:(内核宽度 x 内核高度)x 通道数 x 深度 + 深度(仅当存在偏差时才添加深度)
  • FC层:输入的数量*输出的数量+输出(输出加上偏置的数量)
  • 最大池层:无参数

现在只需将所有参数的数量相加并使用我提到的公式。