将模型转换为 Core ML
Core ML 是 Apple 用于将机器学习模型集成到 iOS、macOS、watchOS 和 tvOS 应用程序的框架。Core ML 的主要功能之一是它能够将各种流行的机器学习框架中的模型转换为 Core ML 格式。在这篇博客文章中，我们将探讨将模型转换为 Core ML 的过程，重点关注 PyTorch 模型。
理解 Core ML Tools
Core ML Tools 是一个 Python 软件包，用于将机器学习模型转换为 Core ML 格式。它支持多种模型类型和框架，包括 TensorFlow、Keras、scikit-learn、XGBoost 和 PyTorch。
Core ML Tools 的主要特性：
- 统一的模型转换 API
- 支持各种机器学习框架
- 针对 Apple 设备进行自动模型优化
- 内置对常见预处理和后处理步骤的支持
将 PyTorch 模型转换为 Core ML
将 PyTorch 模型转换为 Core ML 的过程涉及两个主要步骤：
- 将 PyTorch 模型转换为 TorchScript
- 使用 Core ML Tools 将 TorchScript 模型转换为 Core ML 格式
步骤 1：将 PyTorch 转换为 TorchScript
TorchScript 是 PyTorch 模型的一种中间表示，可以在 C++ 等高性能环境中运行。有两种方法可以将 PyTorch 模型转换为 TorchScript：
1. 追踪
追踪通过运行示例输入并记录执行的操作来工作。此方法适用于具有固定控制流的模型。
import torch
def trace_model(model, example_input):
return torch.jit.trace(model, example_input)
# Example usage
traced_model = trace_model(my_pytorch_model, torch.rand(1, 3, 224, 224))
2. 脚本化
脚本化分析模型的 Python 代码并将其转换为 TorchScript。此方法更灵活，可以处理具有动态控制流的模型。
import torch
def script_model(model):
return torch.jit.script(model)
# Example usage
scripted_model = script_model(my_pytorch_model)
步骤 2：将 TorchScript 转换为 Core ML
一旦您拥有 TorchScript 模型，您就可以使用 Core ML Tools 将其转换为 Core ML 格式。
import coremltools as ct
def convert_to_coreml(torchscript_model, input_shape):
mlmodel = ct.convert(
torchscript_model,
inputs=[ct.TensorType(shape=input_shape)]
)
return mlmodel
# Example usage
coreml_model = convert_to_coreml(traced_model, (1, 3, 224, 224))
coreml_model.save("my_model.mlmodel")
模型追踪与模型脚本化
了解何时使用追踪与脚本化对于成功进行模型转换至关重要。让我们深入探讨这两种方法：
模型追踪
追踪通常更简单，并且通常会生成更优化的 TorchScript 代码。它非常适合具有静态计算图的模型。
追踪的优势：
- 通常生成更快的模型
- 对于简单的模型来说，使用更简单
- 适用于具有固定结构的模型
追踪的局限性：
- 无法捕获动态控制流
- 如果模型行为因输入而异，可能无法很好地泛化
模型脚本化
脚本化更灵活，可以处理具有动态行为的模型，但在某些情况下可能会生成优化程度较低的代码。
脚本化的优势：
- 可以捕获动态控制流（if 语句、循环）
- 适用于更广泛的 PyTorch 模型
- 保留了更多的原始 Python 代码结构
脚本化的局限性：
- 可能生成优化程度较低的 TorchScript 代码
- 如果发生错误，调试可能更复杂
追踪与脚本化的选择
以下是一些帮助您选择合适方法的指南：
- 如果您的模型具有固定结构且不依赖于动态控制流，请使用追踪
- 如果您的模型包含依赖于输入的条件语句或循环，请使用脚本化
- 对于复杂模型，可以考虑使用混合方法（追踪某些组件，脚本化其他组件）
案例研究：转换分割模型
让我们通过一个真实的例子来演示将 PyTorch 分割模型转换为 Core ML 的过程。我们将使用带有 ResNet-101 主干的 DeepLabV3 模型。
import torch
import torchvision
import coremltools as ct
from PIL import Image
# Load the pre-trained model
model = torchvision.models.segmentation.deeplabv3_resnet101(pretrained=True).eval()
# Prepare a sample input
input_image = Image.open("sample_image.jpg")
preprocess = torchvision.transforms.Compose([
torchvision.transforms.Resize((256, 256)),
torchvision.transforms.ToTensor(),
])
input_tensor = preprocess(input_image).unsqueeze(0)
# Attempt to trace the model (this will fail)
try:
traced_model = torch.jit.trace(model, input_tensor)
except RuntimeError as e:
print(f"Tracing failed: {e}")
# Create a wrapper class to handle dictionary output
class WrappedDeepLabV3(torch.nn.Module):
def __init__(self, model):
super().__init__()
self.model = model
def forward(self, x):
return self.model(x)['out']
# Wrap the model and trace it
wrapped_model = WrappedDeepLabV3(model)
traced_model = torch.jit.trace(wrapped_model, input_tensor)
# Convert to Core ML
mlmodel = ct.convert(
traced_model,
inputs=[ct.ImageType(name="input", shape=input_tensor.shape, scale=1/255.0, bias=[0, 0, 0])],
outputs=[ct.TensorType(name="output")]
)
# Set metadata
mlmodel.author = "Your Name"
mlmodel.license = "Your License"
mlmodel.short_description = "DeepLabV3 Segmentation Model"
mlmodel.version = "1.0"
# Save the model
mlmodel.save("deeplabv3_segmentation.mlmodel")
案例研究：CLIP-Finder 图像编码器转换
现在，让我们看一个更复杂的例子：转换 CLIP-Finder 项目中使用的 CLIP 模型中的图像编码器。本案例研究展示了将最先进的多模态模型转换为 Core ML 的过程。欲了解更多详情，请查看：
🤗 MobileCLIP 在 Hugging Face 上已转换
import coremltools
import torch
import mobileclip
from mobileclip.modules.common.mobileone import reparameterize_model
from mobileclip.clip import CLIP
# Define the model configuration
model_cfg = {
"embed_dim": 512,
"image_cfg": {
"image_size": 256,
"model_name": "mci0"
},
"text_cfg": {
"context_length": 77,
"vocab_size": 49408,
"dim": 512,
"ffn_multiplier_per_layer": 4.0,
"n_heads_per_layer": 8,
"n_transformer_layers": 4,
"norm_layer": "layer_norm_fp32",
"causal_masking": False,
"model_name": "mct"
}
}
# Create a custom CLIP class for image encoding
class CLIP_encode_image(CLIP):
def __init__(self, cfg, output_dict=False, *args, **kwargs):
super().__init__(cfg, output_dict, *args, **kwargs)
def forward(self, image):
return self.encode_image(image, normalize=True)
# Initialize and load the model
model_ie = CLIP_encode_image(cfg=model_cfg)
model_ie.eval()
chkpt = torch.load("checkpoints/mobileclip_s0.pt")
model_ie.load_state_dict(chkpt)
# Reparameterize the model for inference
reparameterized_model = reparameterize_model(model_ie)
reparameterized_model.eval()
# Trace the model
image = torch.rand(1, 3, 256, 256)
traced_model = torch.jit.trace(reparameterized_model, image)
# Convert to Core ML
input_image = coremltools.ImageType(name="input_image", shape=(1, 3, 256, 256), color_layout=coremltools.colorlayout.RGB, scale=1/255.0, bias=[0, 0, 0])
output_tensor = [coremltools.TensorType(name="output_embeddings")]
ml_model = coremltools.convert(
model=traced_model,
outputs=output_tensor,
inputs=[input_image],
convert_to="mlprogram",
minimum_deployment_target=coremltools.target.iOS17,
compute_units=coremltools.ComputeUnit.ALL,
debug=True,
)
# Save the model
ml_model.save("clip_mci_image_s0.mlpackage")
最佳实践和故障排除技巧
将模型转换为 Core ML 时，请记住以下最佳实践：
- 在追踪或脚本化之前，务必将模型置于评估模式（`model.eval()`）
- 使用代表性的输入数据进行追踪，以确保捕获所有代码路径
- 准备好创建包装类或修改模型以处理复杂的输出或输入
- 验证转换后的模型输出与原始 PyTorch 模型是否一致
- 使用 Core ML Tools 的调试功能来识别和解决转换问题
常见故障排除步骤包括：
- 如果追踪失败，请尝试脚本化或混合方法
- 对于具有动态形状的模型，请使用 `coremltools.EnumeratedShapes` 来指定可能的输入维度
- 如果您遇到不支持的操作，请考虑将其实现为自定义层或复合操作
- 使用最新版本的 PyTorch 和 Core ML Tools，以确保与较新模型架构的兼容性
结论
将 PyTorch 模型转换为 Core ML 为在 Apple 设备上部署复杂的机器学习模型提供了无限可能。通过理解追踪和脚本化的细微差别，并遵循最佳实践，您可以成功转换各种模型，从简单的分类器到复杂的混合模态架构（如 CLIP）。
请记住，机器学习和模型转换领域正在不断发展。请随时关注 PyTorch 和 Core ML Tools 的最新发展，以确保您使用最有效、最高效的转换技术来完成您的项目。
完整脚本
CLIPImageModel 到 CoreML
本 Notebook 演示了将 CLIP 图像模型转换为 CoreML 格式的过程。CLIPTextModel 到 CoreML
本 Notebook 演示了将 CLIP 文本模型转换为 CoreML 格式的过程。
相关项目
CLIP-Finder GitHub 仓库
https://github.com/fguzman82/CLIP-Finder2转换后的 CLIP Core ML 模型
🤗 MobileCLIP 在 Hugging Face 上
