带有大过滤器的 tensorflow conv2d 的内存使用情况

Question

我有一个带有一些相对较大的135 x 135 x 1 x 3卷积滤波器的 tensorflow 模型。我发现这tf.nn.conv2d对于如此大的过滤器变得无法使用 - 它试图使用超过 60GB 的内存，此时我需要杀死它。这是重现我的错误的最小脚本：

import tensorflow as tf
import numpy as np

frames, height, width, channels = 200, 321, 481, 1
filter_h, filter_w, filter_out = 5, 5, 3  # With this, output has shape (200, 317, 477, 3)
# filter_h, filter_w, filter_out = 7, 7, 3  # With this, output has shape (200, 315, 475, 3)
# filter_h, filter_w, filter_out = 135, 135, 3  # With this, output will be smaller than the above with shape (200, 187, 347, 3), but memory usage explodes

images = np.random.randn(frames, height, width, channels).astype(np.float32)

filters = tf.Variable(np.random.randn(filter_h, filter_w, channels, filter_out).astype(np.float32))
images_input = tf.placeholder(tf.float32)
conv = tf.nn.conv2d(images_input, filters, strides=[1, 1, 1, 1], padding="VALID")

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    result = sess.run(conv, feed_dict={images_input: images})

print result.shape

首先，任何人都可以解释这种行为吗？为什么内存使用会随着过滤器的大小而爆炸？（注意：我也尝试改变我的尺寸以使用单个conv3d而不是一批conv2ds，但这有同样的问题）

其次，除了将操作分解为 200 个单独的单图像卷积之外，有人可以提出其他解决方案吗？

编辑：重新阅读 上的文档后tf.nn.conv2d()，我在其工作原理的解释中注意到了这一点：

1. 将滤波器展平为形状为 的二维矩阵[filter_height * filter_width * in_channels, output_channels]。
2. 从输入张量中提取图像块以形成形状为 的虚拟张量[batch, out_height, out_width, filter_height * filter_width * in_channels]。
3. 对于每个补丁，右乘滤波器矩阵和图像补丁向量。

我最初只是简单地将其作为对过程的描述，但是如果 tensorflow 实际上是从引擎盖下的图像中提取和存储单独的过滤器大小的“补丁”，那么粗略计算表明中间计算在我的情况下，涉及需要~130GB，远远超过我可以测试的限制。这可能会回答我的第一个问题，但如果是这样，谁能解释为什么当我仍然只在 CPU 上调试时 TF 会这样做？

Answer 1

我最初只是将此作为过程的描述，但如果张量流实际上从引擎盖下的图像中提取和存储单独的过滤器大小的“补丁”，那么粗略的计算表明中间计算就我而言，所涉及的需要约 130GB，远远超出了我可以测试的限制。

正如您自己所发现的，这就是内存消耗较大的原因。Tensorflow 这样做是因为滤波器通常很小，并且计算矩阵乘法比计算卷积快得多。

谁能解释为什么当我仍然只在 CPU 上调试时 TF 会这样做？

您也可以在没有 GPU 的情况下使用张量流，因此 CPU 实现不仅仅用于调试。它们还针对速度进行了优化，矩阵乘法在 CPU 和 GPU 上都更快。

为了使大型过滤器的卷积成为可能，您必须在 C++ 中实现大型过滤器的卷积，并将其作为新的操作添加到张量流中。