如何在Tensorflow圖形中保留計算值（在GPU上）？

Question

我們如何確保計算的值不會被複制回CPU/python內存，但仍然可用於下一步的計算？

下面的代碼顯然沒有做到這一點：

import tensorflow as tf 

a = tf.Variable(tf.constant(1.),name="a") 
b = tf.Variable(tf.constant(2.),name="b") 
result = a + b 
stored = result 

with tf.Session() as s: 
    val = s.run([result,stored],{a:1.,b:2.}) 
    print(val) # 3 
    val=s.run([result],{a:4.,b:5.}) 
    print(val) # 9 
    print(stored.eval()) # 3 NOPE: 

錯誤：試圖使用未初始化值_recv_b_0

Answer 1

答案是通過將其存儲使用the assign operation對存儲在tf.Variable值：

工作代碼：

import tensorflow as tf 
with tf.Session() as s: 
    a = tf.Variable(tf.constant(1.),name="a") 
    b = tf.Variable(tf.constant(2.),name="b") 
    result = a + b 
    stored = tf.Variable(tf.constant(0.),name="stored_sum") 
    assign_op=stored.assign(result) 
    val,_ = s.run([result,assign_op],{a:1.,b:2.}) 
    print(val) # 3 
    val=s.run(result,{a:4.,b:5.}) 
    print(val[0]) # 9 
    print(stored.eval()) # ok, still 3