针对许多小型矩阵向量乘法优化 Tensorflow

Question

要构建胶囊网络训练脚本，我需要计算许多小的矩​​阵向量乘法。 每个权重矩阵的大小最多为 20 x 20。
权重矩阵数量超过900个。

我很好奇tf.matmul 或tf.linalg.matvec 是最好的选择。 谁能给我一个优化训练脚本的提示？

Answer 1

编辑：

查看您所指的notebook，您似乎有以下参数：

batch_size = 50
caps1_n_caps = 1152
caps1_n_dims = 8
caps2_n_caps = 10
caps2_n_dims = 16

然后你有一个张量w，形状为(caps1_n_caps, caps2_n_caps, caps2_n_dims, caps1_n_dims)（在笔记本中，它有一个初始尺寸，大小为1，我正在跳过）和另一个张量caps1_output，形状为(batch_size, caps1_n_caps, caps1_n_dims)。您需要将它们组合起来以生成具有形状(batch_size, caps1_n_caps, caps1_n_dims, caps2_n_dims) 的caps2_predicted。

在笔记本中，他们将张量平铺以便使用tf.linalg.matmul 对其进行操作，但实际上您可以仅使用tf.einsum 来计算相同的结果而无需任何平铺：

import tensorflow as tf

batch_size = 50
caps1_n_caps = 1152
caps1_n_dims = 8
caps2_n_caps = 10
caps2_n_dims = 16
w = tf.zeros((caps1_n_caps, caps2_n_caps, caps2_n_dims, caps1_n_dims), dtype=tf.float32)
caps1_output = tf.zeros((batch_size, caps1_n_caps, caps1_n_dims), dtype=tf.float32)
caps2_predicted = tf.einsum('ijkl,bil->bilk', w, caps1_output)
print(caps2_predicted.shape)
# (50, 1152, 8, 16)

---

我不确定我是否完全理解了你想要什么，但你说你想要计算类似的东西：

û_ij = W_ij × u_i

对于几个矩阵W和向量u的集合。假设您有 900 个矩阵和向量，矩阵大小为 20×20，向量大小为 20，您可以将它们表示为两个张量，ws，形状为(900, 20, 20)，和us，形状为(900, 20)。如果你这样做，你的结果us_hat，形状为(900, 20, 20)，将被简单地计算为：

us_hat = ws * tf.expand_dims(us, axis=-1)