TensorFlow技术问题：已实现get_config()的自定义TF Layer仍无法进行JSON序列化

问题背景 #

我为了在模型内部嵌入数据缩放逻辑，自定义了FeatureMinMaxScaler()和FeatureMinMaxDescaler()两个Keras子类化层，用来替代外部的MinMaxScaler操作——毕竟要把多个独立模型链接成新模型，外部缩放的方式行不通。我也实现了get_config()方法，但保存模型时却遇到了这个错误：

TypeError: ('Not JSON Serializable:', <tf.Tensor: shape=(13,), dtype=float64, numpy= array([ 88.9762, 100. , 27.74 , 1. , 0.871 , 8.725 , 100. , 10.7103, 24. , 711. , 22. , 396.9 , 37.97 ])>)

错误成因 #

Keras保存模型时，会把层的配置序列化成JSON格式，而JSON只能处理Python原生数据类型（比如列表、数字、字符串），你在get_config()里直接返回了Tensor对象，Tensor是TensorFlow的专属对象，没办法被JSON序列化，这就是报错的根源。

解决步骤 #

要解决这个问题，我们需要在get_config()里把Tensor转换成可序列化的Python类型（比如列表），同时在__init__方法里，把序列化后的数据再还原成Tensor。具体要做这两步：

1. 在get_config()中，将每个Tensor属性转换为numpy数组的列表形式（用.numpy().tolist()）；
2. 在__init__方法中，判断传入的参数是否是列表，如果是就先转成numpy数组，再转换成Tensor。

修改后的完整代码 #

import numpy as np
import tensorflow as tf
import tensorflow.keras as krs
from sklearn.preprocessing import MinMaxScaler

(x,y),(xtest,ytest) = tf.keras.datasets.boston_housing.load_data()
scaler = MinMaxScaler()
scaler2 = MinMaxScaler()
scaler.fit(x)
transformed = scaler.transform(x)

class FeatureMinMaxScaler( krs.layers.Layer ):
    def __init__(self, max_data, min_data, limits = (0.,1.) ):
        super(FeatureMinMaxScaler, self).__init__(trainable=False)
        # 处理传入的参数：如果是列表，先转成numpy数组再转Tensor
        if isinstance(max_data, list):
            max_data = np.array(max_data)
        if isinstance(min_data, list):
            min_data = np.array(min_data)
        if isinstance(limits, list):
            limits = np.array(limits)
            
        self.max_data = tf.convert_to_tensor(max_data,dtype='float64')
        self.min_data = tf.convert_to_tensor(min_data,dtype='float64')
        self.limits = tf.convert_to_tensor(np.array(limits),dtype='float64')
        self.range = tf.constant(limits[1]-limits[0],dtype='float64')
        self.inv_denominator = tf.divide(tf.constant(1.,dtype='float64'), tf.subtract(self.max_data,self.min_data) )
        self.inv_denominator = tf.where(tf.math.is_inf(self.inv_denominator),tf.constant(0.,dtype='float64'),self.inv_denominator)
        self.scaling_multiplier = tf.multiply(self.inv_denominator,self.range)
        self.range_min = tf.constant(limits[0],dtype='float64')
        self.num_outputs = max_data.size
        
    def build(self, input_shape):
        pass
    
    def call(self, input):
        return tf.cast(tf.add( tf.multiply( tf.subtract(tf.cast(input,dtype='float64'), self.min_data), self.scaling_multiplier ), self.range_min ),dtype='float32')
    
    def get_config(self):
        # 将Tensor转换为可序列化的列表
        config = super().get_config()
        config.update({
            'max_data': self.max_data.numpy().tolist(),
            'min_data': self.min_data.numpy().tolist(),
            'limits': self.limits.numpy().tolist(),
            'range': self.range.numpy().tolist(),
            'inv_denominator': self.inv_denominator.numpy().tolist(),
            'scaling_multiplier': self.scaling_multiplier.numpy().tolist(),
            'range_min': self.range_min.numpy().tolist(),
            'num_outputs': self.num_outputs
        })
        return config

class FeatureMinMaxDescaler( krs.layers.Layer ):
    def __init__(self, max_data, min_data, limits = (0.,1.) ):
        super(FeatureMinMaxDescaler, self).__init__(trainable=False)
        # 处理传入的参数
        if isinstance(max_data, list):
            max_data = np.array(max_data)
        if isinstance(min_data, list):
            min_data = np.array(min_data)
        if isinstance(limits, list):
            limits = np.array(limits)
            
        self.max_data = tf.convert_to_tensor(max_data,dtype='float64')
        self.min_data = tf.convert_to_tensor(min_data,dtype='float64')
        self.limits = tf.convert_to_tensor(np.array(limits),dtype='float64')
        self.range = tf.constant(limits[1]-limits[0],dtype='float64')
        self.inv_denominator = tf.divide(tf.constant(1.,dtype='float64'), self.range)
        self.scaling_multiplier = tf.multiply(self.inv_denominator,tf.subtract(self.max_data,self.min_data))
        self.range_min = tf.constant(limits[0],dtype='float64')
        self.num_outputs = max_data.size
        
    def build(self, input_shape):
        pass
    
    def call(self, input):
        return tf.cast(tf.add( tf.multiply( tf.subtract(tf.cast(input,dtype='float64'), self.range_min), self.scaling_multiplier ), self.min_data ),dtype='float32')
    
    def get_config(self):
        config = super().get_config()
        config.update({
            'max_data': self.max_data.numpy().tolist(),
            'min_data': self.min_data.numpy().tolist(),
            'limits': self.limits.numpy().tolist(),
            'range': self.range.numpy().tolist(),
            'inv_denominator': self.inv_denominator.numpy().tolist(),
            'scaling_multiplier': self.scaling_multiplier.numpy().tolist(),
            'range_min': self.range_min.numpy().tolist(),
            'num_outputs': self.num_outputs
        })
        return config

# 测试模型保存
inputs = krs.Input(shape=(x.shape[1],))
outputs = FeatureMinMaxScaler(scaler.data_max_, scaler.data_min_)(inputs)
model = krs.models.Model(inputs = inputs, outputs = outputs )
model.compile(optimizer='adam')
krs.models.save_model(model,"./serialization_test.mdl")
print("模型保存成功！")

额外优化建议 #

其实range、inv_denominator、scaling_multiplier这些属性是可以通过max_data、min_data、limits重新计算出来的，完全不需要在get_config()里保存它们——这样能简化配置，减少序列化的数据量。比如可以把这些派生属性的计算移到build()方法或者直接在call()里计算，或者在__init__里只保存原始参数，派生属性每次重新计算。这样你的get_config()只需要存max_data、min_data、limits和num_outputs就够了，代码会更简洁。

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