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

时间:2017-12-27 11:55:24

标签: tensorflow neural-network backpropagation

我遇到了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))

2 个答案:

答案 0 :(得分:2)

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

解决方案是挤压。具体而言,tf.squeeze(tf.matmul(X, W))

答案 1 :(得分:1)

因为backprop需要额外的内存来跟踪每个操作的梯度(尽管我无法弄清楚它是如何最终呈二次方式),但内存消耗却是这样才有意义。

解决方案:小批量

这通常是培训模型的goto方法。将您的训练数据分成几个小批量,每个小批量包含固定数量的样本(这很少超过200个样本),一次一个小批量地将其提供给优化器。因此,如果您的batch_size=64,那么投放到优化程序的train_Xtrain_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})