I have a tensorflow model where each tensor in a batch has a different size. Would it be possible to get the correct gradients if I concatenate all the losses and run the optimizer on them like in this example:
import tensorflow as tf
v1 = tf.range(9,dtype=tf.float32)
v2 = tf.range(6,dtype=tf.float32)
v1 = tf.reshape(v1,[3,3])
v2 = tf.reshape(v2,[2,3])
gt1 = tf.constant([2,5,4])
gt2 = tf.constant([1,5])
with tf.variable_scope("var"):
w = tf.get_variable('w', [3,7], dtype=tf.float32)
r1 = v1 @ w
r2 = v2 @ w
loss1 = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=gt1, logits=r1)
loss2 = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=gt2, logits=r2)
loss = tf.concat([loss1, loss2],axis=0)
optimizer = tf.train.AdamOptimizer().minimize(loss)
with tf.Session() as sess:
# print the output of ta_final_result
sess.run(tf.global_variables_initializer())
print(sess.run(w))
print(sess.run(optimizer))
print(sess.run(w))
This is exactly equivalent to summing loss down to a scalar before passing it to minimize. In fact a reduce_sum will be implicitly added to the graph; you can try passing non-scalars to tf.gradients and see what happens:
import tensorflow as tf
session = tf.InteractiveSession()
v = tf.get_variable("v", shape=[])
session.run(v.assign(2.))
grad = tf.gradients([v ** 2., v ** 3.], [v])
session.run(grad)
[16.0]
Which is 2*2 + 3*2^2.