使用Keras牛津宠物分割代码出现CustomMaskWarning如何解决

警告翻译 #

/usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/utils/generic_utils.py:497: 自定义掩码层警告：自定义掩码层需要配置，并且必须重写get_config方法。加载模型时，必须将自定义掩码层传入custom_objects参数。
类别=CustomMaskWarning

触发原因 #

该警告是Keras在序列化（保存）模型时检测到代码中使用了自定义层，而自定义层未显式实现参数序列化方法导致的，不处理的话后续加载保存的模型时会直接报错。

永久解决方法 #

1. 给所有自定义层补充get_config方法 #

以代码中常用的自定义数据增强层为例，修改前后的代码如下：
修改前：

class Augment(tf.keras.layers.Layer):
    def __init__(self, seed=42):
        super().__init__()
        self.augment_inputs = tf.keras.layers.RandomFlip(mode="horizontal", seed=seed)
        self.augment_labels = tf.keras.layers.RandomFlip(mode="horizontal", seed=seed)

    def call(self, inputs, labels):
        inputs = self.augment_inputs(inputs)
        labels = self.augment_labels(labels)
        return inputs, labels

修改后：

class Augment(tf.keras.layers.Layer):
    def __init__(self, seed=42, **kwargs):
        # 补充传入父类初始化参数
        super().__init__(**kwargs)
        # 将初始化参数存为实例属性
        self.seed = seed
        self.augment_inputs = tf.keras.layers.RandomFlip(mode="horizontal", seed=seed)
        self.augment_labels = tf.keras.layers.RandomFlip(mode="horizontal", seed=seed)

    def call(self, inputs, labels):
        inputs = self.augment_inputs(inputs)
        labels = self.augment_labels(labels)
        return inputs, labels

    # 新增参数序列化方法
    def get_config(self):
        config = super().get_config()
        # 将自定义参数加入配置
        config.update({
            "seed": self.seed
        })
        return config

如果有其他自定义层，按照相同逻辑修改即可：将__init__的自定义参数都存为实例属性，在get_config中把参数写入配置字典。

2. （可选）注册自定义层简化加载流程 #

给自定义层加注册装饰器，后续加载模型时不需要额外传参数：

@tf.keras.utils.register_keras_serializable()
class Augment(tf.keras.layers.Layer):
    # 其余代码和上面修改后的版本一致

3. 加载模型时适配自定义层 #

如果没有做注册，加载模型时需要手动传入自定义对象：

model = tf.keras.models.load_model("你的模型保存路径.h5", custom_objects={"Augment": Augment})

临时解决方法 #

如果暂时不需要保存模型，只想屏蔽警告，可以在代码开头加入以下内容：

import warnings
from tensorflow.python.keras.utils.generic_utils import CustomMaskWarning

warnings.filterwarnings("ignore", category=CustomMaskWarning)

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