反向传播表现出二次存储器消耗

Question

我在TensorFlow遇到了一个奇怪的问题。我设置了一个非常简单的分类问题，四个输入变量，一个二进制输出变量，一个权重和偏置层，输出通过一个sigmoid到0或1。

问题是，内存消耗是训练数据记录数量的二次方！只有5,000条记录，它已经是900兆字节;在10000，它可以达到几千兆字节。由于我想最终使用至少数万条记录，这是一个问题。

它特别发生在反向传播步骤中;当我只是尝试评估成本函数时，内存消耗与记录数量呈线性关系，如预期的那样。

代码如下。我究竟做错了什么？

import numpy as np
import os
import psutil
import tensorflow as tf
process = psutil.Process(os.getpid())
sess = tf.InteractiveSession()

# Parameters
learning_rate = 0.01
random_seed = 1
tf.set_random_seed(random_seed)

# Data
data = np.loadtxt('train.csv', delimiter=',', dtype=np.float32)
train_X = data[:, :-1]
train_Y = data[:, -1]

rows = np.shape(train_X)[0]
cols = np.shape(train_X)[1]

# Inputs and outputs
X = tf.placeholder(np.float32, shape=(rows, cols))
Y = tf.placeholder(np.float32, shape=rows,)

# Weights
W = tf.Variable(tf.random_normal((cols, 1)))
b = tf.Variable(tf.random_normal(()))

# Model
p = tf.nn.sigmoid(tf.matmul(X, W) + b)
cost = tf.reduce_sum((p-Y)**2/rows)
optimizer = tf.train.GradientDescentOptimizer(learning_rate).minimize(cost)
tf.global_variables_initializer().run()

# Just one optimizer step is enough to demonstrate the problem
optimizer.run({X: train_X, Y: train_Y})

# Memory consumption is quadratic in number of rows
print('{0:,} bytes'.format(process.memory_info().peak_wset))

Answer 1

事实证明这又是形状问题。使用matmul我在那里的方式，生成形状的输出（n，1）。在期望形状（n，）的上下文中使用它，静默地生成二次爆炸。

解决方案是挤压。具体来说，tf.squeeze(tf.matmul(X, W))。

Answer 2

因为backprop需要额外的内存来跟踪每个操作的梯度（尽管我无法弄清楚它最终是二次的），因此内存消耗会像这样爆炸。

解决方案：小批量生产

在培训模型方面，这通常是goto方法。将您的训练数据分成几个小批量，每个小批量包含固定数量的样本（这很少超过200个样本），一次一个小批量地将其提供给优化器。因此，如果你的batch_size=64然后加入优化器的train_X和train_Y将分别为(64, 4)和(64,)形状。

我会尝试这样的事情

batch_size = 64
for i in range(rows):
    batch_X = train_X[i*batch_size : (i + 1)*batch_size]
    batch_Y = train_Y[i*batch_size : (i + 1)*batch_size]

    optimizer.run({X: batch_X, Y:batch_Y})