Reputation: 5233
In Tensorflow, I have a float tensor T with shape [batch_size, 3]. For example, T[0] = [4, 4, 3]
.
I want to turn that into a size 5 one hot in order to yield entries from an embedding dictionary. In the above case, that would look like
T[0] = [[0, 0, 0, 0, 1], [0, 0, 0, 0, 1], [0, 0, 0, 1, 0]].
If I can get it into this format, then I can multiply it by the embedding dictionary. However, this is in the middle of the graph and I need the gradients to flow through it. Is there a clever way to use stop_gradient a la How Can I Define Only the Gradient for a Tensorflow Subgraph? to make this work? I'm coming up short.
Upvotes: 1
Views: 69
Reputation: 5233
I was able to solve this in the following way:
expanded = tf.expand_dims(inputs, 2)
embedding_input = tf.cast(tf.one_hot(tf.to_int32(inputs), 5), inputs.dtype)
embedding_input = tf.stop_gradient(embedding_input - expanded) + expanded
Upvotes: 1