TensorFlow子类化自定义训练模型的保存与TFLite转换求助

我太懂你这种头疼了！用TensorFlow子类API搭的自定义模型，还自己写训练循环（完全不用.fit()/.compile()那种），尤其是像NMT带注意力、图像字幕这种复杂模型，保存和转TFLite确实比常规模型麻烦不少。我来给你一步步捋清楚解决方案：

自定义模型保存的核心是确保模型架构、权重和推理逻辑都被正确序列化，因为没用到.compile()，我们得手动处理几个关键点：

• 先搞定自定义层的序列化：如果你的模型包含自定义层（比如注意力层），必须实现get_config()方法，不然SavedModel会丢失层的初始化参数。举个例子：

class AttentionLayer(tf.keras.layers.Layer):
    def __init__(self, units):
        super().__init__()
        self.units = units
        # 初始化层内变量（比如W1, W2, V等）

    def call(self, inputs):
        # 注意力计算逻辑
        ...

    def get_config(self):
        # 必须把自定义参数加入config
        config = super().get_config()
        config.update({"units": self.units})
        return config

• 用SavedModel保存推理逻辑：TFLite对SavedModel的支持最好，所以我们要把模型的推理流程封装成一个明确的函数（和训练时的逻辑分开，比如去掉teacher forcing、把dropout设为training=False），再保存。
  比如针对NMT模型，你可以写一个推理函数：

class NMTModel(tf.keras.Model):
    def __init__(self, encoder, decoder, vocab_size):
        super().__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.vocab_size = vocab_size

    # 自定义推理函数，封装生成逻辑
    @tf.function
    def inference(self, encoder_input, start_token):
        # 编码器输出
        encoder_output, hidden = self.encoder(encoder_input)
        decoder_input = tf.expand_dims([start_token], 0)
        result = []

        # 循环生成序列（用tf.while_loop替代Python循环，适配TFLite）
        for _ in range(max_decoder_seq_len):
            predictions, hidden, _ = self.decoder(decoder_input, hidden, encoder_output)
            predicted_id = tf.argmax(predictions[0]).numpy()
            result.append(predicted_id)
            if predicted_id == end_token:
                break
            decoder_input = tf.expand_dims([predicted_id], 0)
        return result

然后保存为SavedModel，同时指定输入签名（因为没有.compile()，TFLite需要明确输入形状）：

# 定义输入签名（根据你的模型调整形状和 dtype）
encoder_input_signature = tf.TensorSpec(shape=(None, max_encoder_seq_len), dtype=tf.int32)
start_token_signature = tf.TensorSpec(shape=(), dtype=tf.int32)

# 保存模型
tf.saved_model.save(
    your_nmt_model,
    "./nmt_saved_model",
    signatures={
        "serving_default": your_nmt_model.inference.get_concrete_function(
            encoder_input=encoder_input_signature,
            start_token=start_token_signature
        )
    }
)

• 可选：保存训练状态：如果需要恢复训练，可以用tf.train.Checkpoint保存模型和优化器的状态：

checkpoint = tf.train.Checkpoint(model=your_model, optimizer=your_optimizer)
checkpoint.save("./training_checkpoints/ckpt")
# 恢复时
checkpoint.restore(tf.train.latest_checkpoint("./training_checkpoints"))

有了SavedModel，转TFLite就顺畅多了，注意几个适配移动端的细节：

# 加载SavedModel
converter = tf.lite.TFLiteConverter.from_saved_model("./nmt_saved_model")

# 启用必要的选项：支持TF原生操作（如果模型用了TFLite内置没有的ops）
converter.target_spec.supported_ops = [
    tf.lite.OpsSet.TFLITE_BUILTINS,
    tf.lite.OpsSet.SELECT_TF_OPS
]
# 允许自定义操作（如果有自己写的TF ops）
converter.allow_custom_ops = True
# 如果模型有动态形状（比如可变序列长度），启用这个
converter.experimental_enable_resource_variables = True

# 执行转换
tflite_model = converter.convert()

# 保存TFLite模型
with open("./nmt_model.tflite", "wb") as f:
    f.write(tflite_model)

转换完一定要验证输出和原TF模型一致，避免踩坑：

# 加载TFLite模型
interpreter = tf.lite.Interpreter(model_path="./nmt_model.tflite")
interpreter.allocate_tensors()

# 获取输入输出张量信息
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# 准备测试输入
test_encoder_input = tf.random.uniform((1, max_encoder_seq_len), dtype=tf.int32, minval=0, maxval=vocab_size)
test_start_token = tf.constant(start_token, dtype=tf.int32)

# 设置TFLite输入
interpreter.set_tensor(input_details[0]['index'], test_encoder_input.numpy())
interpreter.set_tensor(input_details[1]['index'], test_start_token.numpy())

# 运行推理
interpreter.invoke()

# 获取TFLite输出
tflite_output = interpreter.get_tensor(output_details[0]['index'])

# 和原TF模型输出对比
tf_output = your_nmt_model.inference(test_encoder_input, test_start_token)
assert tf.reduce_all(tf.abs(tf.convert_to_tensor(tf_output) - tflite_output) < 1e-5), "模型输出不一致，请检查推理逻辑！"

• 推理逻辑必须用TF原生操作：比如生成序列时要用tf.while_loop而不是Python的for循环，不然TFLite无法正确转换控制流。
• 区分训练/推理模式：自定义层里如果有training参数（比如dropout、batch norm），推理时一定要设为False，不然转换后的模型输出会和原模型不一致。
• 输入形状要明确：如果移动端场景允许固定序列长度，尽量设置固定的输入形状，能减少TFLite转换的复杂度；如果必须动态长度，一定要启用动态形状支持。

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