tf.train.write_graph使用所有内存并导致MemoryError

时间:2018-05-16 17:00:43

标签: python tensorflow

我试图将简单的MNIST模型图以文本格式导出到原始缓冲区文件。

实际上,我正在使用tf.keras模块,但后来我决定切换,因为很难在keras顺序模型中重命名/为输入和输出张量分配名称。

这是我的代码:

import tensorflow as tf
import numpy as np

(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()

y_train = tf.keras.utils.to_categorical(y_train, 10)
y_test = tf.keras.utils.to_categorical(y_test, 10)

x_test = np.reshape(x_test, (-1, 784))/255
x_train = np.reshape(x_train, (-1, 784))/255

X = tf.placeholder(dtype=tf.float32, shape=[None, 784])
Y = tf.placeholder(dtype=tf.float32, shape=[None, 10])

h1 = tf.get_variable('h1',shape=[784, 16])
b1 = tf.zeros(shape=[16])

h2 = tf.get_variable('h2', shape=[16,10])
b2 = tf.zeros(shape=[10])

l1 = tf.nn.relu(tf.matmul(X, h1) + b1)
logits = tf.matmul(l1, h2) + b2

# l1 = tf.layers.dense(inputs=X, units=16, activation=tf.nn.relu)
# logits = tf.layers.dense(inputs=l1, units=10, activation=tf.nn.softmax)

loss = tf.losses.softmax_cross_entropy(onehot_labels=Y, logits=logits)
optimizer = tf.train.AdamOptimizer(learning_rate=0.01).minimize(loss)
accuracy = tf.reduce_mean(tf.cast(tf.equal(tf.argmax(logits, -1), tf.argmax(Y, -1)), tf.float32))

data_pipeline = tf.data.Dataset.from_tensor_slices((x_train, y_train))
iter = data_pipeline.shuffle(1000).repeat().batch(1024).make_one_shot_iterator()
next_item = iter.get_next()

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    for i in range(5):
        for i2 in range(1000):
            x, y = sess.run(next_item)
            if i2 % 100 == 0:
                lss = sess.run(loss, feed_dict={X: x, Y: y})
                print('Loss at {} => {} is {}'.format(i, i2, lss))
                sess.run(optimizer, feed_dict={X: x, Y: y})
    print('Final Accuracy : ', sess.run(accuracy, feed_dict={X: x_train, Y: y_train}))
    # tf.train.write_graph(sess.graph, './', 'temp.pbtxt')
    # This line here is causing the Memory Error

系统信息和图书馆信息

  • 操作系统:Ubuntu 18.04 LTS
  • 总内存:8.0 GB
  • 交换:4 GB
  • Tensorflow 1.8
  • 二进制安装

更多信息

正如我前面提到的,当我使用keras时,我很容易导出keras顺序,并且能够使用我从model.pbtxt获得的会话来编写backend。模型的大小和配置与此代码相同。

使用tensorflow 1.4,即使使用4 GB RAM,我也能轻松导出和保存这些模型。现在我有8个,为什么会抛出MemoryError?对不起确切的日志,它只显示一次,现在每次运行此代码时,我的系统崩溃,因为它耗尽所有内存

当我运行代码时,我收到一个警告:

  

2018-05-16 21:42:08.844531:W tensorflow / core / framework / allocator.cc:101] 376320000的分配超过系统内存的10%。

问题

如何解决此问题及其原因?

修改

我删除了正常numpy修改的tf.data.Dataset,一切正常。我还是想知道,为什么不使用tf.data.Dataset修复这个?这是否会为图表添加额外的节点?在大数据和大型机器上进行培训时,我们应该只使用这些管道吗?

1 个答案:

答案 0 :(得分:0)

尝试将批量大小设置为较低的值,它应该可以解决问题。

例如

iter = data_pipeline.shuffle(1000).repeat().batch(256).make_one_shot_iterator()