张量流可以扫描并仅保留最终结果吗?

时间:2019-06-10 22:07:57

标签: python tensorflow

以下代码使用tf.scan计算数组的总和。 张量流会分配一个6元素数组来容纳六个部分和吗? 还是张量流足够聪明,只分配一个1元素数组来保存部分和?

如何确定?

import tensorflow as tf
import numpy as np

elems = np.array([1, 2, 3, 4, 5, 6])
s = tf.scan(lambda a, x: a + x, elems)
# s == [1, 3, 6, 10, 15, 21]
final_s = s[-1]  # final_s == 21
with tf.Session() as sess:
    print(sess.run(final_s))

以下代码导致16 GB机器上的内存不足。

我想张量流永远不会尝试优化扫描。

import tensorflow as tf
import numpy as np

elems = tf.ones(int(1e9), dtype=np.float32)  # 4GB
s = tf.scan(lambda a, x: a + x, elems)
final_s = s[-1]
with tf.Session() as sess:
    print(sess.run(final_s))

0 个答案:

没有答案