如何在Keras中实现批量矩阵乘法自定义层？

核心问题分析 #

你的代码主要问题是维度未对齐：原einsum表达式未考虑输入的batch维度，且输入与权重的维度映射关系需要明确对应。以下是修正后的完整实现：

完整自定义层代码 #

import tensorflow as tf

class BatchMatMulLayer(tf.keras.layers.Layer):
    def __init__(self, num_parallel, output_units, **kwargs):
        self.num_parallel = num_parallel
        self.output_units = output_units
        super().__init__(**kwargs)

    def build(self, input_shape):
        # 从输入形状中获取特征维度
        input_units = input_shape[-1]
        # 权重形状：(并行批次数量, 输入特征数, 输出特征数)
        self.dense_weights = self.add_weight(
            shape=(self.num_parallel, input_units, self.output_units),
            initializer='random_normal',
            trainable=True,
            name='dense_weights'
        )
        # 偏置形状：(并行批次数量, 1, 输出特征数)，利用广播机制适配输出维度
        self.dense_biases = self.add_weight(
            shape=(self.num_parallel, 1, self.output_units),
            initializer='zeros',
            trainable=True,
            name='dense_biases'
        )
        super().build(input_shape)

    def call(self, inputs):
        # 使用einsum完成批量矩阵乘法：
        # inputs维度：(batch_size, 并行批次数量, 输入特征数) -> 标记为b,c,i
        # weights维度：(并行批次数量, 输入特征数, 输出特征数) -> 标记为c,i,o
        # 输出维度：(batch_size, 并行批次数量, 输出特征数) -> 标记为b,c,o
        outputs = tf.einsum('bci,cio->bco', inputs, self.dense_weights)
        # 偏置自动广播到输出维度
        outputs += self.dense_biases
        return outputs

    def compute_output_shape(self, input_shape):
        # 返回输出形状，适配Keras的shape推断
        return (input_shape[0], self.num_parallel, self.output_units)

关键修正点 #

• 补充batch维度：einsum表达式加入b维度（对应样本批次），解决维度不匹配问题。
• 明确维度映射：确保输入、权重、输出的维度一一对应，避免张量形状冲突。
• 保留Keras层规范：继承tf.keras.layers.Layer并正确实现build、call、compute_output_shape方法，保证层能正常参与模型训练。

测试示例 #

# 初始化层：3个并行批次，输出5维特征
layer = BatchMatMulLayer(num_parallel=3, output_units=5)
# 构造输入：2个样本，每个样本含3个并行批次，每个批次4维特征
inputs = tf.random.normal((2, 3, 4))
outputs = layer(inputs)
print(outputs.shape)  # 预期输出：(2, 3, 5)

内容的提问来源于stack exchange，提问作者K V