将MoViNets模型转CoreML时遇NotImplementedError问题求助

问题原因 #

StatefulPartitionedCall是TensorFlow用于封装有状态子图或分区计算逻辑的操作，常见于包含循环层、自定义状态化组件，或导出SavedModel时使用了带状态的tf.function的场景。Core ML Tools当前版本对该操作的直接转换支持不完善，因此抛出NotImplementedError。

解决方法 #

1. 重新导出无状态版本的SavedModel

如果是预训练MoViNets模型，先加载模型并包装成无状态的推理函数，再重新导出：

import tensorflow as tf
from tensorflow.keras.applications import MoViNet

# 加载预训练模型（根据你的模型参数调整）
model = MoViNet(model_name='movinet_a0', num_classes=1000, pretrained=True)

# 定义无状态推理函数，明确输入签名
@tf.function(input_signature=[tf.TensorSpec(shape=(None, 16, 224, 224, 3), dtype=tf.float32)])
def infer(inputs):
    # 强制关闭训练模式，避免状态更新
    return model(inputs, training=False)

# 重新导出SavedModel
tf.saved_model.save(model, 'stateless_movinet', signatures={'serving_default': infer})

之后用这个新的SavedModel目录执行Core ML转换。

2. 直接转换模型推理函数而非整个SavedModel

加载SavedModel后提取其推理签名，直接转换该函数：

import coremltools as ct
import tensorflow as tf

# 加载原始SavedModel
loaded_model = tf.saved_model.load(saved_dir)
infer_signature = loaded_model.signatures['serving_default']

# 定义输入张量类型（需匹配模型实际输入形状）
input_spec = ct.TensorType(shape=(1, 16, 224, 224, 3))

# 执行转换
coreml_model = ct.convert(
    infer_signature,
    inputs=[input_spec],
    convert_to="mlprogram"
)

3. 检查并处理自定义层/状态组件

如果模型包含自定义状态化层，需要为其编写Core ML自定义转换器：

• 拆解自定义层的计算逻辑，替换为Core ML原生支持的操作组合；
• 使用coremltools.converters.tensorflow.convert.register_custom_op注册自定义转换逻辑。

额外提示 #

• 尝试升级Core ML Tools到最新版本（如7.x），新版本可能对StatefulPartitionedCall有更好的支持；
• 转换时明确指定inputs参数，可避免Core ML Tools解析模型签名时出现歧义。

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