显示二次内存消耗的反向传播

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

问题是,内存消耗是训练数据记录数的二次方!只有5000条记录,已经是900兆字节了;在10000字节时,它会运行到几GB。由于我希望最终使用至少数万条记录,这是一个问题。

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

代码如下。我做错了什么?

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))