TensorFlow批量标量乘法：自定义线性组合层优化实现

需要实现一个TensorFlow层，接收n个输入：其中n-1个是形状为(None,9,256,256,1)的数据张量，最后一个是长度为n-1的权重向量（形状(None, n-1)）。功能是将第i个数据张量与权重向量的第i个元素相乘，再将所有加权结果累加返回单个数据张量。

当前实现存在两个问题：

• 代码硬编码了3个张量的逻辑，无法适配任意数量的输入张量
• 批量场景下的元素乘法效率待优化，尝试过tf.math.multiply、tf.math.scalar_mul、tf.linalg.matvec等函数但未找到通用高效的实现方式

现有代码如下：

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)

优化后的通用实现利用TensorFlow的广播机制和批量操作，避免硬编码，同时提升效率。核心思路是将所有数据张量堆叠到一个新维度，将权重向量扩展为匹配的维度后进行逐元素相乘，最后沿堆叠维度求和得到结果。

优化后的代码：

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

    def build(self, input_shape):
        num_tensors = len(input_shape[0])
        num_weights = input_shape[1][1]
        assert num_tensors == num_weights, f"张量数量{num_tensors}必须与权重数量{num_weights}匹配"
        super().build(input_shape)

    def call(self, inputs):
        tensors = inputs[0]
        weights = inputs[1]
        
        # 1. 将所有数据张量堆叠到新维度（维度1），形状变为(None, num_tensors, 9, 256, 256, 1)
        stacked_tensors = tf.stack(tensors, axis=1)
        
        # 2. 扩展权重维度以匹配堆叠后的张量形状，形状变为(None, num_tensors, 1, 1, 1, 1)
        expanded_weights = tf.expand_dims(tf.expand_dims(tf.expand_dims(tf.expand_dims(weights, axis=-1), axis=-1), axis=-1), axis=-1)
        
        # 3. 逐元素相乘（广播机制自动匹配维度），再沿堆叠维度求和
        weighted_sum = tf.reduce_sum(stacked_tensors * expanded_weights, axis=1)
        
        return weighted_sum

    def compute_output_shape(self, input_shape):
        # 输出形状与单个输入张量一致
        return input_shape[0][0]

关键优化点说明： #

• 通用适配：通过tf.stack将任意数量的输入张量堆叠，无需硬编码每个张量的索引
• 高效广播：扩展权重维度后，利用TensorFlow的广播机制完成批量元素相乘，避免循环操作，提升计算效率
• 简洁求和：使用tf.reduce_sum直接沿堆叠维度累加结果，替代多次手动加法

内容的提问来源于stack exchange，提问作者Max Conway