Coremltools模型转换中的浮点精度问题解析

问题背景

在使用Coremltools进行模型转换时，开发者可能会遇到一个典型问题：成功将PyTorch模型转换为.mlpackage格式后，使用原生PyTorch模型推理能得到正确结果，但通过CoreML加载模型进行预测时却输出全NaN值。这种情况通常与模型计算精度设置有关。

问题现象分析

当开发者完成模型转换后，可能出现以下现象：

1. 模型能够成功转换并保存为.mlpackage格式
2. 原始PyTorch模型(traced_model)推理结果正常
3. 加载CoreML模型(coreml_model.predict())后输出全为NaN值
4. 转换过程中可能出现"overflow encountered in cast"的运行时警告

根本原因

这个问题的主要原因是CoreML默认使用16位浮点(FP16)精度进行模型转换和计算。当模型包含较大数值范围的计算时，FP16精度可能导致数值溢出，从而产生NaN结果。

解决方案

通过明确指定计算精度为32位浮点(FP32)可以解决此问题。具体方法是在转换函数ct.convert()中添加compute_precision参数：

coreml_model = ct.convert(
    traced_model,
    compute_precision=ct.precision.FLOAT32,
    inputs=[...],
    outputs=[...]
)

技术原理

FP16(半精度浮点)和FP32(单精度浮点)的主要区别在于：

1. FP16使用16位存储，指数位5位，尾数位10位
2. FP32使用32位存储，指数位8位，尾数位23位
3. FP16的数值范围约为±65,504，而FP32约为±3.4×10³⁸

当模型中的数值超出FP16的表示范围时，就会发生溢出，导致NaN结果。这在深度学习模型中较为常见，特别是当模型包含：

• 大数值的输入或权重
• 指数运算
• 累积求和操作

最佳实践建议

1. 对于大多数现代神经网络模型，建议优先尝试FP32精度转换
2. 如果模型较小且对部署效率要求高，可以尝试FP16精度，但需验证输出结果
3. 转换后务必进行结果验证，比较原始框架和CoreML的输出差异
4. 注意监控转换过程中的警告信息，特别是数值溢出相关的警告

总结

Coremltools在模型转换过程中默认使用FP16精度可能导致数值溢出问题。通过显式设置compute_precision=ct.precision.FLOAT32参数，可以确保模型使用FP32精度进行计算，避免NaN结果。这一解决方案简单有效，是处理CoreML模型转换中数值精度问题的首选方法。