TensorFlow如何实现张量收缩
2018-10-12 16:19 更新
tf.einsum
einsum(
equation,
*inputs
)
定义在:tensorflow/python/ops/special_math_ops.py.
参考指南:数学>减少张量的计算
任意维度张量之间的广义收缩.
这个函数返回一个张量,其元素其元素是由等式定义的,这是由爱因斯坦求和公式所启发的速写形式定义的.作为示例,考虑将两个矩阵 A 和 B 相乘以形成矩阵C.C的元素由下式给出:
C[i,k] = sum_j A[i,j] * B[j,k]
相应的等式是:
ij,jk->ik
一般来说, 方程是从较熟悉的元素方程得到:
- 删除变量名称、括号和逗号;
- 用 "*" 替换 ",";
- 删除总和标志;
- 将输出移到右侧,并将 "=" 替换为 "->>".
许多常见操作可以用这种方式来表示.例如:
# 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]
此函数的行为类似于 numpy.einsum,但不支持:
- 椭圆(下标像:ij...,jk...->ik...)
- 一个轴在单个输入上出现多次的下标(例如 ijj、k->> ik).
- 在多个输入之间求和的下标(例如 ij、ij、jk->> ik).
ARGS:
- equation:一种描述收缩的 str,与 numpy.einsum 格式相同.
- * inputs:合同的输入(每个张量),其形状应与 equation 一致.
返回:
返回收缩的张量,形状由 equation 决定.
注意:
- ValueError:如果 equation 格式不正确,equation 隐含的输入数与 len(inputs) 不匹配,一个轴出现在输出下标中,但不显示在任何输入中,输入的维数与其下标中的索引数不同,或者输入形状沿特定轴线不一致.