CoreMLTools模型转换中的数值溢出问题分析与解决方案

问题背景

在使用CoreMLTools将PyTorch模型转换为Core ML格式时，开发者可能会遇到数值溢出问题。具体表现为转换过程中出现RuntimeWarning警告，如"overflow encountered in true_divide"和"overflow encountered in subtract"，最终导致转换后的模型输出结果与原始PyTorch模型不一致，出现inf或nan值。

问题分析

该问题主要发生在处理复数除法运算时。在原始PyTorch模型中，复数除法通过分解实部和虚部进行计算，公式如下：

real_result = (real_a * real_b + imag_a * imag_b) / denominator
imag_result = (imag_a * real_b - real_a * imag_b) / denominator

其中denominator = real_b² + imag_b²

当CoreMLTools进行模型转换时，默认会使用float16精度进行计算，而PyTorch模型通常使用float32精度。这种精度差异会导致：

1. float16的数值范围(-65504 ~ +65504)远小于float32(约-3.4e38 ~ +3.4e38)
2. 在计算大数值的平方或除法时，float16更容易出现溢出
3. 当denominator过小时，float16的除法运算会产生inf值

解决方案

要解决这个问题，可以在CoreMLTools转换时显式指定计算精度：

mlmodel = ct.convert(
    trace,
    inputs=[input_a],
    minimum_deployment_target=ct.target.iOS16,
    compute_precision=ct.precision.FLOAT32  # 强制使用float32精度
)

技术建议

1. 精度选择：对于涉及大数值范围或敏感数值计算(如复数运算)的模型，建议始终使用FLOAT32精度

2. 数值稳定性检查：

   • 在模型转换前后都应检查输出结果的有效性
   • 使用torch.isnan()和torch.isinf()检测异常值
   • 对Core ML模型输出使用numpy.isnan()和numpy.isinf()检查

3. 复数运算优化：

   • 考虑对输入数据进行归一化处理
   • 在除法运算前添加小的epsilon值防止除以零
   • 对极端值进行裁剪(clipping)

4. 模型验证：

   • 建立自动化测试验证转换前后模型的一致性
   • 对关键运算层进行单独测试
   • 使用代表性输入数据验证模型行为

总结

CoreMLTools在模型转换过程中默认使用float16精度可能导致数值溢出问题，特别是在处理复数运算等敏感计算时。通过显式指定FLOAT32计算精度可以有效解决这一问题。开发者应当充分了解不同精度格式的数值特性，并在模型转换过程中进行充分的验证测试，确保转换后的模型保持与原始模型一致的数值行为。