How to do batch inner product in Tensorflow?
-
16-10-2019 - |
문제
I have two tensor a:[batch_size, dim]
b:[batch_size, dim]
.
I want to do inner product for every pair in the batch, generating c:[batch_size, 1]
, where c[i,0]=a[i,:].T*b[i,:]
. How?
해결책
There is no native .dot_product
method. However, a dot product between two vectors is just element-wise multiply summed, so the following example works:
import tensorflow as tf
# Arbitrarity, we'll use placeholders and allow batch size to vary,
# but fix vector dimensions.
# You can change this as you see fit
a = tf.placeholder(tf.float32, shape=(None, 3))
b = tf.placeholder(tf.float32, shape=(None, 3))
c = tf.reduce_sum( tf.multiply( a, b ), 1, keep_dims=True )
with tf.Session() as session:
print( c.eval(
feed_dict={ a: [[1,2,3],[4,5,6]], b: [[2,3,4],[5,6,7]] }
) )
The output is:
[[ 20.]
[ 92.]]
다른 팁
Another option worth checking out is [tf.einsum][1]
- it's essentially a simplified version of Einstein Notation.
Following along with Neil and dumkar's examples:
import tensorflow as tf
a = tf.placeholder(tf.float32, shape=(None, 3))
b = tf.placeholder(tf.float32, shape=(None, 3))
c = tf.einsum('ij,ij->i', a, b)
with tf.Session() as session:
print( c.eval(
feed_dict={ a: [[1,2,3],[4,5,6]], b: [[2,3,4],[5,6,7]] }
) )
The first argument to einsum
is an equation representing the axes to be multiplied and summed over. The basic rules for an equation are:
- Input-tensors are described by a comma-separated string of dimension-labels
- Repeated labels indicate that the corresponding dimensions will be multiplied
- The output-tensor is described by another string of dimension-labels representing corresponding inputs (or products)
- Labels that are missing from the output string are summed over
In our case, ij,ij->i
means that our inputs will be 2 matrices of equal shape (i,j)
, and our output will be a vector of shape (i,)
.
Once you get the hang of it, you'll find that einsum
generalizes a huge number of other operations:
X = [[1, 2]]
Y = [[3, 4], [5, 6]]
einsum('ab->ba', X) == [[1],[2]] # transpose
einsum('ab->a', X) == [3] # sum over last dimension
einsum('ab->', X) == 3 # sum over both dimensions
einsum('ab,bc->ac', X, Y) == [[13,16]] # matrix multiply
einsum('ab,bc->abc', X, Y) == [[[3,4],[10,12]]] # multiply and broadcast
Unfortunately, einsum
takes a pretty hefty performance hit when compared to a manual multiply+reduce. Where performance is critical, I'd definitely recommend sticking with Neil's solution.
Taking the diagonal of tf.tensordot also does what you want, if you set axis to e.g.
[[1], [1]]
I have adapted Neil Slater's example:
import tensorflow as tf
# Arbitrarity, we'll use placeholders and allow batch size to vary,
# but fix vector dimensions.
# You can change this as you see fit
a = tf.placeholder(tf.float32, shape=(None, 3))
b = tf.placeholder(tf.float32, shape=(None, 3))
c = tf.diag_part(tf.tensordot( a, b, axes=[[1],[1]]))
with tf.Session() as session:
print( c.eval(
feed_dict={ a: [[1,2,3],[4,5,6]], b: [[2,3,4],[5,6,7]] }
) )
which now also gives:
[ 20. 92.]
This might be suboptimal for large matrices though (see discussion here)