如何将Scenic项目中Flax格式预训练Checkpoints转换为模型？

我尝试使用Google Research Scenic项目中MBT的预训练模型，该模型为Flax Checkpoint格式，希望将其转换为可直接使用的模型，请问具体该如何操作？

1. 环境准备与Checkpoint加载 #

首先确保安装必要依赖：

pip install scenic flax jax tensorflow

然后通过Scenic的模型定义加载Checkpoint参数：

from scenic.projects.mbt import mbt_model
from flax.training import checkpoints

# 匹配Checkpoint对应的模型配置（根据你使用的MBT型号调整，比如mbt_small/mbt_large）
config = mbt_model.get_mbt_config()
model = mbt_model.MBTModel(config)

# 加载Checkpoint（直接指向截图中的checkpoint目录即可）
ckpt_dir = "你的Checkpoint目录路径"
params = checkpoints.restore_checkpoint(ckpt_dir, target=None)

2. 两种转换/使用方式 #

方式一：保留Flax格式，封装推理接口

如果继续使用Flax/JAX生态，直接封装推理函数即可：

import jax.numpy as jnp

def run_inference(inputs):
    # inputs需符合MBT的输入要求：比如图像张量（形状[batch, H, W, 3]）、文本token张量等
    logits = model.apply({'params': params}, inputs, train=False)
    return logits

# 示例调用（根据实际输入模态调整）
sample_image = jnp.ones((1, 224, 224, 3))
output = run_inference(sample_image)

方式二：转换为通用的TensorFlow SavedModel格式

如果需要跨框架使用，可通过jax2tf转换为TensorFlow兼容模型：

from jax.experimental import jax2tf
import tensorflow as tf

# 封装为TensorFlow模块
class MBTModule(tf.Module):
    def __init__(self, flax_model, params):
        self.flax_model = flax_model
        self.params = params
    
    @tf.function(input_signature=[tf.TensorSpec(shape=(None, 224, 224, 3), dtype=tf.float32)])
    def predict(self, inputs):
        return jax2tf.convert(self.flax_model.apply, enable_xla=False)(
            {'params': self.params}, inputs, train=False
        )

# 转换并保存
tf_model = MBTModule(model, params)
tf.saved_model.save(tf_model, "转换后的模型保存路径")

# 加载使用示例
loaded_model = tf.saved_model.load("转换后的模型保存路径")
tf_output = loaded_model.predict(tf.ones((1, 224, 224, 3)))

关键注意事项 #

• 模型配置必须与Checkpoint严格匹配：比如模型尺寸、输入模态（单模态/多模态）、输入分辨率等，否则会出现参数不匹配错误
• 截图中的checkpoint文件是Flax的索引文件，加载时直接传入上级目录即可
• 若遇到依赖版本冲突，建议参考Scenic官方文档的环境要求配置虚拟环境

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