Keras中自定义Bicubic Lambda层训练正常但模型无法保存求助

Keras 2.3.1版本保存HDF5格式模型时，需要将所有层的配置序列化为JSON格式。你封装的Lambda层内部的bic函数捕获了外部的静态shp、scale等不可直接JSON序列化的变量，且没有将自定义参数注册到层配置中，因此保存时抛出序列化异常。

方案1：仅保存权重（最省事，适合调试场景） #

避开序列化整个模型结构的需求，只保存模型权重：

• 保存时将generator.save(model_save_dir + 'gen_model%d.h5' % e)替换为：

generator.save_weights(model_save_dir + 'gen_weights%d.h5' % e)

• 加载时先初始化完整的generator结构，再调用加载权重方法即可：

generator = build_generator() # 你的生成器初始化函数
generator.load_weights(model_save_dir + 'gen_weights%d.h5' % e)

方案2：修改Lambda层实现支持序列化 #

调整Lambda层的参数传递方式，通过官方提供的arguments参数传入自定义配置，避免捕获外部不可序列化变量：

from keras.layers import Lambda
import tensorflow as tf

def bicubic(x, scale=2):
    # 动态计算输出尺寸，不依赖外部传入的静态shape
    h = tf.shape(x)[1] * scale
    w = tf.shape(x)[2] * scale
    return tf.image.resize_images(x, [h, w], method=tf.image.ResizeMethod.BICUBIC)

def bicubic_lambda(scale=2):
    def outshape(input_shape):
        dims = [input_shape[0], input_shape[1] * scale, input_shape[2] * scale, input_shape[3]]
        return tuple(dims)
    # 通过arguments参数传递scale，Keras会自动将该参数加入可序列化的层配置
    return Lambda(bic, output_shape=outshape, arguments={'scale': scale}, name='bicubic_lambda')

修改后直接调用generator.save()即可正常保存，加载时需要传入custom_objects：

from keras.models import load_model
model = load_model('gen_model.h5', custom_objects={'bicubic': bicubic})

方案3：自定义上采样层（最稳定，适合正式上线场景） #

完全放弃Lambda层，自定义实现可序列化的Keras层，实现get_config方法注册配置参数：

from keras.layers import Layer
import tensorflow as tf

class BicubicUpSampling(Layer):
    def __init__(self, scale=2, **kwargs):
        self.scale = scale
        super(BicubicUpSampling, self).__init__(**kwargs)
    
    def call(self, x):
        h = tf.shape(x)[1] * self.scale
        w = tf.shape(x)[2] * self.scale
        return tf.image.resize_images(x, [h, w], method=tf.image.ResizeMethod.BICUBIC)
    
    def compute_output_shape(self, input_shape):
        return (input_shape[0], input_shape[1] * self.scale, input_shape[2] * self.scale, input_shape[3])
    
    def get_config(self):
        config = super(BicubicUpSampling, self).get_config()
        config.update({'scale': self.scale})
        return config

使用时直接用BicubicUpSampling(scale=2)替换原有Lambda层调用即可，保存和加载方式和原生层几乎一致，加载时传入自定义类即可：

model = load_model('gen_model.h5', custom_objects={'BicubicUpSampling': BicubicUpSampling})

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