如何为未知批量大小的输入创建可训练线性层？（Keras/TensorFlow）

我之前也碰到过一模一样的需求——要给形状为[?, m, n]的3D输入做类似x*w + b的线性变换，还得严格保留输入的原始形状。你遇到的维度不匹配问题确实很典型，毕竟可训练参数不能依赖动态的批量大小，直接用2D的[m,n]参数和3D输入相乘又会报错。

最终解决方案代码 #

import tensorflow as tf
from tensorflow.keras.layers import Layer
from tensorflow.keras.initializers import RandomNormal

class LinearLayer(Layer):
    def __init__(self, **kwargs):
        super(LinearLayer, self).__init__(**kwargs)
    
    def build(self, input_shape):
        # 只针对空间维度[m,n]定义可训练参数，避开动态的批量维度
        self.w = self.add_weight(
            name='w', 
            shape=(input_shape[1], input_shape[2]), 
            initializer=RandomNormal(), 
            trainable=True
        )
        self.b = self.add_weight(
            name='b', 
            shape=(input_shape[1], input_shape[2]), 
            initializer=RandomNormal(), 
            trainable=True
        )
        super(LinearLayer, self).build(input_shape)
    
    def call(self, x, **kwargs):
        # 获取当前输入的动态批量大小（构建时可能是None）
        batch_size = tf.shape(x)[0]
        # 把w从[m,n]扩展为[1,m,n]，再复制到批量维度
        w_expand = tf.expand_dims(self.w, axis=0)
        w_tile = tf.tile(w_expand, multiples=[batch_size, 1, 1])
        # 逐元素相乘后加b，b会自动广播到[batch_size,m,n]
        return tf.multiply(x, w_tile) + self.b
    
    def compute_output_shape(self, input_shape):
        # 严格保持输出形状和输入一致
        return input_shape

关键细节解释 #

• 参数定义：在build方法里，我们用input_shape[1]和input_shape[2]来指定w和b的形状，完全避开了动态的批量维度（input_shape[0]是None），这样参数是固定大小的可训练张量，不会随批量变化。
• 维度匹配：在call方法中，先用tf.shape(x)[0]获取当前输入的实际批量大小（因为构建时批量是未知的，不能用input_shape[0]），然后通过tf.expand_dims给w增加一个批量维度，再用tf.tile把这个维度复制到和输入一致的批量数，这样就能和输入x做逐元素乘法了。
• 自动广播：b的形状是[m,n]，TensorFlow会自动把它广播到[batch_size,m,n]，所以不需要额外做tile操作，节省了计算资源。

这个层的本质是对每个空间位置(i,j)（对应m行n列的每个元素）做独立的线性变换x[:,i,j] * w[i,j] + b[i,j]，批量维度只是重复这个变换，参数在所有样本间共享。

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