def testIndexedSlicesGradientInCondInWhileLoop(self): with ops.Graph().as_default(): embedding_matrix = tf.get_variable( "embedding_matrix", [5, 5], initializer=tf.random_normal_initializer()) def Cond(it, _): return it < 5 def Body(it, cost): embedding = embedding_ops.embedding_lookup(embedding_matrix, [0]) cost = tf.cond(tf.equal(it, 3), lambda: tf.square(cost), lambda: cost + tf.reduce_sum(embedding)) return it + 1, cost _, cost = control_flow_ops.While( Cond, Body, [tf.constant(0), tf.constant(0.0)]) dynamic_grads = tf.gradients(cost, [embedding_matrix])[0] dynamic_grads = tf.segment_sum(dynamic_grads.values, dynamic_grads.indices) embedding = embedding_ops.embedding_lookup(embedding_matrix, [0]) static = tf.square( tf.reduce_sum(embedding) + tf.reduce_sum(embedding) + tf.reduce_sum(embedding)) + tf.reduce_sum(embedding) static_grads = tf.gradients(static, [embedding_matrix])[0] static_grads = tf.segment_sum(static_grads.values, static_grads.indices) with self.test_session() as sess: sess.run(tf.initialize_all_variables()) self.assertAllEqual(*sess.run([static_grads, dynamic_grads]))
def testIndexedSlicesGradient(self): with ops.Graph().as_default(): embedding_matrix = tf.get_variable( "embedding_matrix", [5, 5], initializer=tf.random_normal_initializer()) def Cond(it, _): return it < 5 def Body(it, cost): embedding = embedding_ops.embedding_lookup(embedding_matrix + 0.0, [0]) cost += tf.reduce_sum(embedding) return it + 1, cost _, cost = control_flow_ops.While( Cond, Body, [tf.constant(0), tf.constant(0.0)]) optimizer = momentum.MomentumOptimizer(0.1, 0.9) train_op = optimizer.minimize(cost) with self.test_session() as sess: sess.run(tf.initialize_all_variables()) for _ in range(10): sess.run([train_op])