如何将使用自定义标准化函数的Keras TextVectorization层的配置保存到pickle文件中，并重新加载它？

Question

我拥有一个使用了自定义标准化函数的 Keras TextVectorization 层。

def custom_standardization(input_string, preserve=['[', ']'], add=['¿']):
    strip_chars = string.punctuation
    for item in add:
        strip_chars += item
    for item in preserve:
        strip_chars = strip_chars.replace(item, '')
    lowercase = tf.strings.lower(input_string)
    output = tf.strings.regex_replace(lowercase, f'[{re.escape(strip_chars)}]', '')
    return output

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

target_vectorization.adapt(train_spanish_texts)

我希望保存这个适应后的配置，以便推理模型能够使用。一种方法是如此处所述，将weights和config分开保存为pickle文件并重新加载。

但是，当调用target_vectorization.get_config()时，返回的配置中包含了自定义标准化函数的引用，

{'name': 'text_vectorization_5',
 'trainable': True,
 ...
 'standardize': ,
 ...
 'vocabulary_size': 15000}

这会被保存进pickle文件中。尝试使用以下代码加载此配置时，

keras.layers.TextVectorization.from_config(
    pickle.load(open('ckpts/spanish_vectorization.pkl', 'rb'))['config']
)

会遇到TypeError: Could not parse config: 错误，因为保存的文件中并未包含关于这个自定义标准化函数的任何信息。

在这种情况下，如何妥善保存TextVectorization层的权重和配置，以便推理模型能够使用呢？

siride · Answer

问题
该问题与序列化自定义标准化可调用对象有关。请参阅这里的文档：(tf.keras.layers.TextVectorization)。

解决方案
根据文档，您应使用装饰器将该层注册为 Keras 可序列化对象，具体做法是创建一个带有以下装饰器的包装类 (tf.keras.saving.register_keras_serializable)。

我已经使用您的自定义函数进行了最小可行的测试示例，该示例在 Python 3.9.12 和 keras/tensorflow 2.15 版本上工作正常：

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

def custom_standardization(input_string, preserve=['[', ']'], add=['¿']):
    strip_chars = string.punctuation
    for item in add:
        strip_chars += item
    for item in preserve:
        strip_chars = strip_chars.replace(item, '')
    lowercase = tf.strings.lower(input_string)
    output = tf.strings.regex_replace(lowercase, f'[{re.escape(strip_chars)}]', '')
    return output

@keras.utils.register_keras_serializable(package='custom_layers', name='TextVectorizer')
class TextVectorizer(keras.layers.Layer):
    def __init__(self, custom_standardization, **kwargs):
        super().__init__(**kwargs)
        self.custom_standardization = custom_standardization
        self.vectorizer = tf.keras.layers.TextVectorization(
            standardize=self.custom_standardization,
            max_tokens=1000,  # 根据数据集调整此参数
            output_mode='int'
        )

    def call(self, inputs):
        return self.vectorizer(inputs)

    def get_config(self):
        base_config = super().get_config()
        return {**base_config}

# 示例用法
text_vectorizer = TextVectorizer(custom_standardization)

# 使 TextVectorization 层适应训练数据
train_data = tf.constant(['Hello [World]!', 'Another [example].'])
text_vectorizer.vectorizer.adapt(train_data)

# 构建层以初始化 TextVectorization 层
text_vectorizer.build(input_shape=(None,))

# 创建一个模型来包含 TextVectorization 层
model = tf.keras.Sequential([text_vectorizer])
model.build(input_shape=())
# 保存模型的权重
model.save_weights('text_vectorizer_weights.tf')

# 将权重加载到新实例的 TextVectorization 中
loaded_text_vectorizer = TextVectorizer(custom_standardization)
loaded_text_vectorizer.build(input_shape=(None,))

# 创建一个模型来包含加载的 TextVectorization 层
loaded_model = tf.keras.Sequential([loaded_text_vectorizer])

# 再次使 TextVectorization 层适应相同的训练数据（实际应用中可能不需要此步骤）
loaded_text_vectorizer.vectorizer.adapt(train_data)

# 加载权重到新模型中
loaded_model.load_weights('text_vectorizer_weights.tf')

# 编译模型
loaded_model.compile()

# 测试加载的层
text_input = tf.constant(['Hello [World]!'])
output = loaded_model(text_input)
print(output)

georg · Answer

解决方案是围绕 TextVectorization 对象定义一个包装器，并将自定义标准化器作为一个方法使用。此外，我们需要在保存配置到pickle文件时排除可调用对象。以下是修复后的代码：

@keras.utils.register_keras_serializable(package='custom_layers', name='TextVectorizer')
class TextVectorizer(layers.Layer):
    '''英文-西班牙语文本向量化器'''

    def __init__(self, max_tokens=None, output_mode='int', output_sequence_length=None, standardize='lower_and_strip_punctuation', vocabulary=None, config=None):
        super().__init__()
        if config:
            self.vectorization = layers.TextVectorization.from_config(config)
        else:
            self.max_tokens = max_tokens
            self.output_mode = output_mode
            self.output_sequence_length = output_sequence_length
            self.vocabulary = vocabulary
            if standardize != 'lower_and_strip_punctuation':
                self.vectorization = layers.TextVectorization(max_tokens=self.max_tokens,
                                                              output_mode=self.output_mode,
                                                              output_sequence_length=self.output_sequence_length,
                                                              vocabulary=self.vocabulary,
                                                              standardize=self.standardize)
            else:
                self.vectorization = layers.TextVectorization(max_tokens=self.max_tokens,
                                                              output_mode=self.output_mode,
                                                              output_sequence_length=self.output_sequence_length,
                                                              vocabulary=self.vocabulary)

    def standardize(self, input_string, preserve=['[', ']'], add=['¿']) -> str:
        strip_chars = string.punctuation
        for item in add:
            strip_chars += item
        for item in preserve:
            strip_chars = strip_chars.replace(item, '')
        lowercase = tf.strings.lower(input_string)
        output = tf.strings.regex_replace(lowercase, f'[{re.escape(strip_chars)}]', '')
        return output

    def __call__(self, *args, **kwargs):
        return self.vectorization(*args, **kwargs)

    def get_config(self):
        config = self.vectorization.get_config()
        return {k: v for k, v in config.items() if not callable(v)}

    @classmethod
    def from_config(cls, config):
        return cls(config=config)

    def set_weights(self, weights):
        self.vectorization.set_weights(weights)

    def adapt(self, dataset):
        self.vectorization.adapt(dataset)

    def get_vocabulary(self):
        return self.vectorization.get_vocabulary()

在训练阶段适应并保存权重：

vocab_size = 15000
sequence_length = 20

source_vectorization = TextVectorizer(max_tokens=vocab_size,
                                      output_mode='int',
                                      output_sequence_length=sequence_length)

target_vectorization = TextVectorizer(max_tokens=vocab_size,
                                      output_mode='int',
                                      output_sequence_length=sequence_length + 1,
                                      standardize='spanish')

# 假设train_pairs已定义
train_english_texts = [pair[0] for pair in train_pairs]
train_spanish_texts = [pair[1] for pair in train_pairs]

source_vectorization.adapt(train_english_texts)
target_vectorization.adapt(train_spanish_texts)

with open('ckpts/english_vectorization.pkl', 'wb') as f:
    pickle.dump({'config': source_vectorization.get_config(), 'weights': source_vectorization.get_weights()}, f)

with open('ckpts/spanish_vectorization.pkl', 'wb') as f:
    pickle.dump({'config': target_vectorization.get_config(), 'weights': target_vectorization.get_weights()}, f)

在推理阶段加载并使用它们：

with open('ckpts/english_vectorization.pkl', 'rb') as f:
    vectorization_data = pickle.load(f)
source_vectorization = TextVectorizer.from_config(vectorization_data['config'])
source_vectorization.set_weights(vectorization_data['weights'])

with open('ckpts/spanish_vectorization.pkl', 'rb') as f:
    vectorization_data = pickle.load(f)
target_vectorization = TextVectorizer.from_config(vectorization_data['config'])
target_vectorization.set_weights(vectorization_data['weights'])