PyTorch Image Models终极导出指南：ONNX、TorchScript与CoreML完整教程
PyTorch Image Models终极导出指南：ONNX、TorchScript与CoreML完整教程
项目地址: https://gitcode.com/GitHub_Trending/py/pytorch-image-models PyTorch Image Models（timm）是一个由Hugging Face开发维护的PyTorch视觉模型库，包含多个高性能的预训练模型，适用于图像识别、分类等视觉任务。本教程将详细介绍如何将timm模型导出为ONNX、TorchScript和CoreML格式，帮助开发者在不同平台上部署和应用这些强大的视觉模型。
准备工作：环境搭建与模型选择
在开始导出之前，需要确保你的开发环境中已经安装了必要的依赖库。除了PyTorch和timm库外，导出ONNX格式需要安装onnx和onnxruntime，导出TorchScript无需额外依赖，而导出CoreML则需要coremltools。你可以通过以下命令安装这些依赖：
pip install onnx onnxruntime coremltools
接下来，选择一个你想要导出的模型。timm库提供了丰富的预训练模型，例如MobileNet、ResNet、EfficientNet等。你可以通过timm.list_models()查看所有可用模型，并使用timm.create_model()加载模型。
ONNX导出：跨平台部署的最佳选择
ONNX（Open Neural Network Exchange）是一种开放的神经网络模型格式，支持多种深度学习框架和部署平台。timm库提供了专门的ONNX导出脚本onnx_export.py，简化了导出过程。
基本导出命令
使用以下命令可以将模型导出为ONNX格式：
python onnx_export.py output.onnx --model mobilenetv3_large_100 --batch-size 1 --img-size 224
这个命令将MobileNetV3 Large模型导出为output.onnx文件，输入批次大小为1，图像尺寸为224x224。
高级导出选项
onnx_export.py提供了多种高级选项，以满足不同的需求：
--opset：指定ONNX opset版本，默认为10。不同的opset版本支持不同的操作，较高的版本可能提供更好的性能。--dynamic-size：导出支持动态宽度和高度的模型，适用于需要处理不同尺寸输入的场景。--check-forward：在导出后进行前向传播检查，确保导出的模型与原始PyTorch模型输出一致。--reparam：对模型进行重参数化，融合层以提高推理性能。
例如，导出支持动态尺寸的模型：
python onnx_export.py output_dynamic.onnx --model resnet50 --dynamic-size --check-forward
导出后的验证
导出ONNX模型后，可以使用timm提供的onnx_validate.py脚本验证模型的正确性：
python onnx_validate.py output.onnx --model mobilenetv3_large_100
该脚本会加载ONNX模型并与原始PyTorch模型的输出进行比较，确保两者的结果一致。
TorchScript导出：PyTorch生态系统的优化部署
TorchScript是PyTorch的模型序列化格式，允许你将PyTorch模型转换为可序列化和可优化的中间表示。这对于在生产环境中部署PyTorch模型非常有用，特别是在需要高性能推理的场景。
使用训练脚本导出
timm的train.py脚本支持在训练过程中导出TorchScript模型。通过添加--torchscript参数，可以在训练结束时自动导出模型：
python train.py --model resnet50 --data-path ./data --torchscript --epochs 10
这将在训练结束后生成一个TorchScript模型文件，通常以.pt或.pth为扩展名。
手动导出TorchScript模型
你也可以手动编写代码导出TorchScript模型。以下是一个简单的示例：
import torch
import timm
model = timm.create_model('resnet50', pretrained=True)
model.eval()
# 创建一个示例输入
example_input = torch.randn(1, 3, 224, 224)
# 跟踪模型
traced_model = torch.jit.trace(model, example_input)
# 保存模型
traced_model.save('resnet50_torchscript.pt')
在导出过程中，需要注意模型是否支持TorchScript。timm中的大多数模型都经过优化，可以很好地支持TorchScript，但某些复杂的模型可能需要额外的适配。
TorchScript模型的优化
TorchScript提供了多种优化选项，可以进一步提高模型的推理性能。例如，可以使用torch.jit.optimize_for_inference()函数对模型进行优化：
optimized_model = torch.jit.optimize_for_inference(traced_model)
optimized_model.save('resnet50_torchscript_optimized.pt')
CoreML导出：iOS和macOS平台的部署
CoreML是Apple开发的机器学习框架，允许在iOS、macOS等Apple设备上部署机器学习模型。虽然timm库没有专门的CoreML导出脚本，但可以通过PyTorch的CoreML导出功能实现。
使用coremltools导出
首先，确保你已经安装了coremltools库。然后，可以使用以下代码将PyTorch模型转换为CoreML格式：
import torch
import timm
import coremltools as ct
model = timm.create_model('mobilenetv3_large_100', pretrained=True)
model.eval()
# 创建示例输入
example_input = torch.randn(1, 3, 224, 224)
# 跟踪模型
traced_model = torch.jit.trace(model, example_input)
# 转换为CoreML模型
coreml_model = ct.convert(
traced_model,
inputs=[ct.ImageType(name="input", shape=example_input.shape, scale=1/255.0, bias=[-0.485, -0.456, -0.406])]
)
# 保存CoreML模型
coreml_model.save('mobilenetv3_coreml.mlmodel')
在转换过程中，需要注意输入图像的预处理步骤，如缩放和归一化。上述代码中，scale=1/255.0和bias=[-0.485, -0.456, -0.406]对应于ImageNet数据集的预处理参数。
MobileViT模型的CoreML优化
timm中的MobileViT模型特别针对移动设备进行了优化，并提供了CoreML导出支持。在创建MobileViT模型时，可以设置coreml_exportable=True以确保模型兼容CoreML导出：
model = timm.create_model('mobilevit_xxs', pretrained=True, coreml_exportable=True)
这将禁用一些CoreML不支持的操作，确保模型能够顺利导出。
常见问题与解决方案
导出时出现不支持的操作
不同的导出格式支持的操作有所不同。如果在导出过程中遇到不支持的操作，可以尝试以下解决方案：
- 更新PyTorch和相关库到最新版本，新版本通常支持更多操作。
- 使用
--aten-fallback选项（ONNX导出），回退到ATEN操作。 - 重写模型中不支持的操作，使用更通用的实现。
导出后的模型精度下降
如果导出后的模型精度下降，可以检查以下几点：
- 确保在导出前将模型设置为评估模式（
model.eval()）。 - 检查输入数据的预处理步骤是否与训练时一致。
- 使用
--check-forward选项（ONNX导出）验证模型输出。
模型文件过大
模型文件过大可能会影响部署和推理速度。可以尝试以下方法减小模型大小：
- 使用模型剪枝（
--reparam选项）融合层和参数。 - 量化模型，将浮点数参数转换为低精度格式（如INT8）。
- 选择更小的模型架构，如MobileNet、EfficientNet等。
总结
本教程详细介绍了如何将timm模型导出为ONNX、TorchScript和CoreML格式，涵盖了基本导出命令、高级选项、验证方法以及常见问题的解决方案。通过选择合适的导出格式，你可以在不同的平台上高效部署timm模型，充分发挥其在图像识别和分类任务中的强大性能。
无论是跨平台部署、PyTorch生态系统优化还是Apple设备上的应用，timm库都提供了灵活且高效的解决方案。希望本教程能够帮助你顺利完成模型导出，并在实际应用中取得良好的效果！
智能硬件社区聚焦AI智能硬件技术生态，汇聚嵌入式AI、物联网硬件开发者，打造交流分享平台，同步全国赛事资讯、开展 OPC 核心人才招募，助力技术落地与开发者成长。
更多推荐
13
0- 0
已为社区贡献7条内容
所有评论(0)