Tensorflow 批量标量乘法

Question

我想创建一个接受 n 个输入的张量流层，n-1 是数据点，最后一个是长度为 n-1 的权重向量。然后，将（n-1 中的 i）个数据点中的每个数据点乘以存储在权重向量的索引 i 中的值。然后，加权数据点的结果被累积并作为单个数据点返回。我遇到的问题是在 TensorFlow 中批量高效地完成此操作。由于额外的批量维度，我最终需要将 n-1 个形状的张量 (None 9, 256, 256, 1) 乘以 n-1 个形状的张量 (None, )。

为了实现这一目标，我尝试使用函数

tf.math.multiply(tensor, weights)

、

tf.math.scalar_mul(weights1, tensor1)

和

tf.linalg.matvec(tensor1, weights1)

代码如下：

class LinearCombination(tf.keras.layers.Layer):
    def __init__(self, **kwargs):
        super(LinearCombination, self).__init__(**kwargs)

    def build(self, input_shape):
        # Ensure that the input shape matches the expected shape
        print(input_shape)

        num_tensors = len(input_shape[0])
        num_weights = input_shape[1][1]
        assert num_tensors == num_weights, f"Number of tensors{num_tensors} and number of weights{num_weights} must match."

        super(LinearCombination, self).build(input_shape)  # Be sure to call this at the end

    def call(self, inputs):
        # Multiply each tensor by its corresponding weight and sum them up
        print("\n\n\n\nlayer called")
        tensors = inputs[0]
        tensor1 = tensors[0]
        tensor2 = tensors[1]
        tensor3 = tensors[2]

        print(f"\ntensors: {tensors}")
        print(f"t1{tensor1}")
        print(f"t2{tensor2}")
        print(f"t3{tensor3}\n\n")

        weights = inputs[1]
        weights1 = weights[:, 0]
        #weights1 = tf.expand_dims(tf.expand_dims(tf.expand_dims(weights1, axis=-1), axis=-1), axis=-1)

        weights2 = weights[:, 1]
        #weights2 = tf.expand_dims(tf.expand_dims(tf.expand_dims(weights2, axis=-1), axis=-1), axis=-1)

        weights3 = weights[:, 2]
        #weights3 = tf.expand_dims(tf.expand_dims(tf.expand_dims(weights3, axis=-1), axis=-1), axis=-1)


        print(f"\n\nweights: {weights}")
        print(f"w1{weights1}")
        print(f"w2{weights2}")
        print(f"w3{weights3}\n\n")



        #tensor1 = tf.math.multiply(tensor1, weights1)
        #tensor2 = tf.math.multiply(tensor2, weights2)
        #tensor3 = tf.math.multiply(tensor3, weights3)
        #tensor1 = tf.math.scalar_mul(weights1, tensor1)
        #tensor2 = tf.math.scalar_mul(weights2, tensor2)
        #tensor3 = tf.math.scalar_mul(weights3, tensor3)
        #tensor1 = tf.linalg.matvec(tensor1, weights1)
        #tensor2 = tf.linalg.matvec(tensor2, weights2)
        #tensor3 = tf.linalg.matvec(tensor3, weights3)
        print(f"\nresults:")
        print(f"r1{tensor1}")
        print(f"r2{tensor2}")
        print(f"r3{tensor3}\n\n")


        out = tf.math.add(tensor1, tensor2)
        out = tf.math.add(tensor3, out)
        print(f"final output: {out}")

        return out

    def compute_output_shape(self, input_shape):
        return input_shape[0][1]  # Output shape matches the shape of each input tensor

所有乘法注释掉的张量保持正确的形状，调用函数中的打印语句如下：

tensors: [<tf.Tensor 'Placeholder:0' shape=(None, 9, 256, 256, 1) dtype=float32>, <tf.Tensor 'Placeholder_1:0' shape=(None, 9, 256, 256, 1) dtype=float32>, <tf.Tensor 'Placeholder_2:0' shape=(None, 9, 256, 256, 1) dtype=float32>]
t1Tensor("Placeholder:0", shape=(None, 9, 256, 256, 1), dtype=float32)
t2Tensor("Placeholder_1:0", shape=(None, 9, 256, 256, 1), dtype=float32)
t3Tensor("Placeholder_2:0", shape=(None, 9, 256, 256, 1), dtype=float32)




weights: Tensor("Placeholder_3:0", shape=(None, 3), dtype=float32)
w1Tensor("linear_combination/strided_slice:0", shape=(None,), dtype=float32)
w2Tensor("linear_combination/strided_slice_1:0", shape=(None,), dtype=float32)
w3Tensor("linear_combination/strided_slice_2:0", shape=(None,), dtype=float32)



results:
r1Tensor("Placeholder:0", shape=(None, 9, 256, 256, 1), dtype=float32)
r2Tensor("Placeholder_1:0", shape=(None, 9, 256, 256, 1), dtype=float32)
r3Tensor("Placeholder_2:0", shape=(None, 9, 256, 256, 1), dtype=float32)


final output: Tensor("linear_combination/Add_1:0", shape=(None, 9, 256, 256, 1), dtype=float32)

如果我们不硬编码我们有三个张量这一事实，那就太好了。

Answer 1

可以使用

tf.einsum

来实现这样的操作。这是更新后的

call()

函数：

    def call(self, inputs):
        # Multiply each tensor by its corresponding weight and sum them up
        print("\n\n\n\nlayer called")
        tensors = inputs[0]  # List of M tensor with shape (BATCH, 9, 256, 256, 1)
        weights = inputs[1]  # Tensor with shape (BATCH, M)
        
        # [(BATCH, 1, 9, 256, 256, 1)] ==concat=> [(BATCH, M, 9, 256, 256, 1)]
        big_tensor = tf.concat([tf.expand_dims(t, axis=1) for t in tensors], axis=1)
        out = tf.einsum('nmabcd,nm->nabcd', weights, big_tensor)
        
        print(f"final output: {out}")

        return out