tf.einsum()简单使用

tf.einsum

einsum( equation, *inputs)
一般来说, 方程是从较熟悉的元素方程得到：

删除变量名称、括号和逗号；
用 “*” 替换 “，”；
删除总和标志；
将输出移到右侧，并将 “=” 替换为 “->>”。
许多常见操作可以用这种方式来表示。例如:

# Matrix multiplication
>>> einsum('ij,jk->ik', m0, m1)  # output[i,k] = sum_j m0[i,j] * m1[j, k]

# Dot product
>>> einsum('i,i->', u, v)  # output = sum_i u[i]*v[i]

# Outer product
>>> einsum('i,j->ij', u, v)  # output[i,j] = u[i]*v[j]

# Transpose
>>> einsum('ij->ji', m)  # output[j,i] = m[i,j]

# Batch matrix multiplication
>>> einsum('aij,ajk->aik', s, t)  # out[a,i,k] = sum_j s[a,i,j] * t[a, j, k]

例子

# 初始化矩阵
x = tf.constant([[1., 2., 3.],
                 [1., 2., 3.]])
y = tf.constant([[2., 3., 4.],
                 [2., 3., 4.]])
z = tf.constant([[3., 4.],
                 [3., 4.],
                 [3., 4.]])
# 外积
out = tf.multiply(x, y)
out1 = tf.einsum('ij,ij->ij', x, y)
# 点积
dot = tf.matmul(x, z)
dot1 = tf.einsum('ij,jk->ik', x, z)
# 转置
trans = tf.transpose(x, [1, 0])
trans1 = tf.einsum('ij->ji', x)
with tf.Session() as sess:
    print('外积')
    print('out\n{}'.format(sess.run(out)))
    print('out1\n{}'.format(sess.run(out1)))
    print('点积')
    print('dot\n{}'.format(sess.run(dot)))
    print('dot1\n{}'.format(sess.run(dot1)))
    print('转置')
    print('trans\n{}'.format(sess.run(trans)))
    print('trans1\n{}'.format(sess.run(trans1)))

输出

外积
out
[[ 2.  6. 12.]
 [ 2.  6. 12.]]
out1
[[ 2.  6. 12.]
 [ 2.  6. 12.]]
点积
dot
[[18. 24.]
 [18. 24.]]
dot1
[[18. 24.]
 [18. 24.]]
转置
trans
[[1. 1.]
 [2. 2.]
 [3. 3.]]
trans1
[[1. 1.]
 [2. 2.]
 [3. 3.]]

这里可以看到使用einsum是一种矩阵运算的简便方法。