Keras实现PyTorch的Bilinear双线性变换3D输入维度适配问题

核心问题原因 #

你当前的维度报错根源是3D时序输入（带batch、时序步长两个前置维度）没有和双线性权重的多维度做对齐，原生的tf.matmul默认只处理最后两维的矩阵乘法，直接传入3D输入+权重会触发错误的广播逻辑。
你给出的张量参数对应的目标输出应为 [batch_size=72, seq_len=10, out_features=24]，正好匹配后续LSTM层要求的[batch, timesteps, feature]输入格式。

实现方案 #

方案1：使用einsum实现（最简洁易读，无维度转置负担） #

tf.einsum可以直接指定多维度的运算逻辑，完美适配双线性变换的维度匹配需求，计算逻辑和PyTorch的torch.nn.Bilinear完全对齐：

def call(self, x1, x2):
    # x1形状: [batch, seq_len, in1_features=6]
    # x2形状: [batch, seq_len, in2_features=6]
    # self.w形状: [out_features=24, in1_features=6, in2_features=6]
    # self.b形状: [out_features=24]
    bilinear_out = tf.einsum('b i j, o j k, b i k -> b i o', x1, self.w, x2)
    # 偏置广播叠加
    return bilinear_out + self.b[tf.newaxis, tf.newaxis, :]

维度对应说明：

• b：批量维度，i：时序步长维度，o：输出特征维度
• j：x1的特征维度，k：x2的特征维度
• 运算逻辑完全匹配双线性公式 $y = x_1^T A x_2 + b$，且自动处理批量、时序维度的并行计算。

方案2：使用matmul拆分实现（不使用einsum的场景） #

如果需要用原生矩阵乘法实现，按如下维度调整步骤即可：

def call(self, x1, x2):
    # 第一步：x1和权重相乘，输出[batch, seq_len, out_features, in2_features]
    x1_expand = x1[:, :, tf.newaxis, :] # [72,10,1,6]
    out1 = tf.matmul(x1_expand, self.w) # [72,10,24,6]
    # 第二步：和x2相乘，输出[batch, seq_len, out_features, 1]
    x2_expand = x2[:, :, tf.newaxis, :, tf.newaxis] # [72,10,1,6,1]
    out2 = tf.matmul(out1[..., tf.newaxis, :], x2_expand) # [72,10,24,1]
    # 压缩多余维度，叠加偏置
    out = tf.squeeze(out2, axis=[-1, -2]) + self.b
    return out

验证说明 #

两种方案输出的形状均为[72, 10, 24]，无需额外调整维度即可直接输入到后续LSTM层。

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