python - 对 3D 张量行进行广播操作 (Tensorflow)

Question

我有一个维度为 [BATCH_SIZE, 128] 的张量 p，以及维度为 [528, 128] 的张量 A、B。我想构建一个大小为 [BATCH_SIZE, 528] 的新张量 p'，其中 j 列定义为:

tf.reduce_prod(self.A[j,:] * p + self.B[j,:], axis=1)

我目前使用 for 循环对此进行了强力实现，但速度非常慢。有什么方法可以使用广播或其他方式来加快速度吗？

ps = []
for j in xrange(self.A.shape[0]):
    a = self.A[j,:]
    b = self.B[j,:]
    ps.append(tf.reduce_prod(a * p + b, axis=1))
p = tf.stack(ps, axis=1)

Answer 1

试试这个:

result = tf.reduce_prod(self.A * p[:, tf.newaxis] + self.B, axis=2)