python - tensorflow 中的张量收缩

Question

我有一个形状为 (?,4) 的张量 weights 和一个形状为 (?,4, 1024)。

我想根据相应的权重对每行embeddings中的4个张量进行加权平均来收缩张量，最后产生一个张量形状 (?,1024) 的输出。

我该怎么做？我尝试使用 output = tf.tensordot(weights, embeddings, axes = [1,1]) 但它产生了形状为 (?,?,1024) 的张量.

Answer 1

你可以这样做:

import tensorflow as tf

weights = tf.placeholder(tf.float32, [None, 4])
embeddings = tf.placeholder(tf.float32, [None, 4, 1024])
output = tf.einsum('ij,ijk->ik', weights, embeddings)

可以用矩阵乘积来表达同样的东西，不知道在性能上会不会有什么不同:

output = tf.squeeze(tf.expand_dims(weights, 1) @ embeddings, 1)

您也可以只进行乘法和减法，尽管原则上这会因为有一个中间张量而导致性能变差。

output = tf.reduce_sum(tf.expand_dims(weights, 2) * embeddings, axis=1)