python - 不同等级的 Matmul

Question

我有 3 个张量
X 形状(1, c, h, w)，假设(1, 20, 40, 50)
Fx 形状(num, w, N)，假设(1000, 50, 10)
Fy shape (num, N, h)，假设(1000, 10, 40)

我想做的是Fy * (X * Fx)(*表示matmul)
X * Fx shape (num, c, h, N)，假设(1000, 20, 40, 10)
Fy * (X * Fx) shape (num, c, N, N)，假设(1000, 20, 10, 10)

我正在使用 tf.tile 和 tf.expand_dims 来做到这一点
但我认为它使用大量内存(tile 复制数据对吗？)，而且很慢
尝试找到更快更好的方法并使用较小的内存来完成

# X: (1, c, h, w)
# Fx: (num, w, N)
# Fy: (num, N, h)

X = tf.tile(X, [tf.shape(Fx)[0], 1, 1, 1])  # (num, c, h, w)
Fx_ex = tf.expand_dims(Fx, axis=1)  # (num, 1, w, N)
Fx_ex = tf.tile(Fx_ex, [1, c, 1, 1])  # (num, c, w, N)
tmp = tf.matmul(X, Fxt_ex)  # (num, c, h, N)

Fy_ex = tf.expand_dims(Fy, axis=1)  # (num, 1, N, h)
Fy_ex = tf.tile(Fy_ex, [1, c, 1, 1])  # (num, c, N, h)
res = tf.matmul(Fy_ex, tmp) # (num, c, N, N)

Answer 1

mythical einsum 的案例，我猜:

>>> import numpy as np
>>> X = np.random.rand(1, 20, 40, 50)
>>> Fx = np.random.rand(100, 50, 10)
>>> Fy = np.random.rand(100, 10, 40)
>>> np.einsum('nMh,uchw,nwN->ncMN', Fy, X, Fx).shape
(100, 20, 10, 10)

它在 tf 中的工作方式应该与在 numpy 中几乎相同(在某些 tf 版本中不允许使用大写索引，我看到).尽管如果您以前从未见过该符号，那么这无疑超过了正则表达式的不可读性。