如何将带自定义标准化函数的Keras TextVectorization层保存到Pickle并重新加载

方法一：直接保存整个适配后的TextVectorization层（推荐） #

由于你的自定义标准化逻辑完全基于TensorFlow原生操作（tf.strings.lower、tf.strings.regex_replace等），可以直接用TensorFlow的saved_model格式保存整个层，自动包含适配后的词汇表和标准化逻辑，无需手动处理配置与权重：

保存代码

import tensorflow as tf

# 保存适配完成的层到指定目录
tf.saved_model.save(target_vectorization, "target_vectorization_layer")

加载代码

# 加载保存的层
loaded_vectorization = tf.saved_model.load("target_vectorization_layer")

# 测试使用
test_text = "¿Hola, mundo! [test]"
print(loaded_vectorization(test_text))

加载后无需重新定义custom_standardization函数，可直接调用层处理文本。

方法二：封装为可序列化的自定义标准化层 #

如果需要使用Keras原生的save/from_config流程，可以将标准化逻辑封装成继承自keras.layers.Layer的自定义层，让参数可被序列化：

定义自定义标准化层

import string
import re
import tensorflow as tf
from tensorflow import keras

class CustomStandardization(keras.layers.Layer):
    def __init__(self, preserve=['[', ']'], add=['¿'], **kwargs):
        super().__init__(**kwargs)
        self.preserve = preserve
        self.add = add
        # 预计算需要剔除的字符集合
        strip_chars = string.punctuation
        for item in add:
            strip_chars += item
        for item in preserve:
            strip_chars = strip_chars.replace(item, '')
        self.regex_pattern = f'[{re.escape(strip_chars)}]'

    def call(self, input_string):
        lowercase = tf.strings.lower(input_string)
        output = tf.strings.regex_replace(lowercase, self.regex_pattern, '')
        return output

    def get_config(self):
        # 序列化自定义参数
        config = super().get_config()
        config.update({
            'preserve': self.preserve,
            'add': self.add
        })
        return config

使用自定义层初始化TextVectorization

target_vectorization = keras.layers.TextVectorization(
    max_tokens=vocab_size,
    output_mode='int',
    output_sequence_length=sequence_length + 1,
    standardize=CustomStandardization()
)

target_vectorization.adapt(train_spanish_texts)

保存与加载

# 保存整个层
target_vectorization.save("target_vectorization")

# 加载（需确保CustomStandardization类已定义）
loaded_vectorization = keras.layers.TextVectorization.load("target_vectorization")

问题根源说明 #

你之前遇到的错误是因为get_config()只能序列化JSON兼容的参数，自定义函数的内存引用无法被序列化。上述两种方法要么直接保存层的完整计算图（方法一），要么将函数逻辑转化为可序列化的层参数（方法二），解决了自定义逻辑的序列化问题。

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