为许多小型矩阵向量乘法优化 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, but actually you can compute the same result without any tiling just using 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)

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

Optimizing Tensorflow for many small matrix-vector multiplications

matrix-multiplication

tensorflow