自定义层必须继承tf.keras.layers.Layer,正确实现__init__、build、call及get_config方法;build中用input_shape[1:]获取维度,call中必须接收training参数并传递给子层,缺失get_config会导致保存加载失败。

自定义层必须继承 tf.keras.layers.Layer
不继承这个基类,就无法被 model.fit() 或 model.call() 正确识别,会报错 AttributeError: 'MyLayer' object has no attribute 'built'。Keras 的构建机制依赖该基类的生命周期方法(如 build()、call()),跳过它等于绕开整个层管理逻辑。
常见错误是只写一个普通类,或误继承 tf.keras.Model —— 后者用于封装完整模型,不是单层。
- 必须显式调用
super().__init__(**kwargs)在__init__中 -
build(self, input_shape)里定义权重,用self.add_weight(),不能在__init__里直接创建tf.Variable -
call(self, inputs, training=None)是前向逻辑主入口,training参数必须保留(即使不用),否则启用 dropout 或 batch norm 时出错
build() 中的 input_shape 是什么格式
它是 TensorShape 对象,形如 (None, 32, 32, 3)(batch 维为 None),不是元组也不是列表。直接解包会失败,比如 h, w, c = input_shape 报错,因为首维是 None。
正确做法是用索引访问,且跳过第 0 维(batch):
立即学习“Python免费学习笔记(深入)”;
def build(self, input_shape):
# input_shape: (None, height, width, channels)
h, w, c = input_shape[1], input_shape[2], input_shape[3]
self.kernel = self.add_weight(
shape=(h * w, c),
initializer='glorot_uniform',
trainable=True
)
注意:如果输入是 1D(如 (None, 128)),input_shape[1] 就是特征维度;若不确定维度数,可用 input_shape[-1] 安全取通道/特征数。
如何支持 training=True/False 切换行为
很多自定义层需要在训练和推理时表现不同(比如自定义 dropout、噪声注入、条件归一化)。仅靠判断 training 布尔值不够——你得把它传给底层可训练操作,否则 model.evaluate() 仍会执行训练逻辑。
- 所有子层(如
self.dropout)必须在__init__中声明,并接收training参数 - 在
call()中显式传入:self.dropout(inputs, training=training) - 不要用
tf.keras.backend.learning_phase(),它已被弃用,且在 tf.function 下不可靠
示例:带训练开关的随机掩码层
class RandomMaskLayer(tf.keras.layers.Layer):
def __init__(self, rate=0.1, **kwargs):
super().__init__(**kwargs)
self.rate = rate
<pre class="brush:php;toolbar:false;">def call(self, inputs, training=None):
if training:
noise = tf.random.uniform(tf.shape(inputs)) > self.rate
return inputs * tf.cast(noise, inputs.dtype)
return inputs
保存与加载自定义层会失败?检查 get_config()
用 model.save('my_model.h5') 或 tf.keras.models.load_model() 时,若没实现 get_config(),会报 ValueError: Unknown layer: RandomMaskLayer。Keras 序列化时只存配置,不存代码逻辑。
必须同时提供 get_config() 和类方法 from_config()(后者由父类默认实现,通常不用重写):
def get_config(self):
config = super().get_config()
config.update({'rate': self.rate})
return config
注意:所有构造参数都得进 config,包括 name、dtype 等父类参数;如果用了非标准类型(如自定义函数),要转成可序列化的形式(字符串或基本类型)。
真正容易被忽略的是:即使你只用 tf.keras.models.load_model(..., custom_objects={...}) 手动注册,get_config 缺失仍会导致从 SavedModel 格式加载失败(尤其是跨环境部署时)。


















