CoreMLTools中FP16模型转换时的激活函数兼容性问题分析

问题背景

在使用CoreMLTools将PyTorch模型转换为Core ML格式时，当模型被转换为FP16精度后，某些激活函数如LeakyReLU、RReLU、PReLU和ELU会出现导出失败的情况。这是由于这些激活函数的alpha参数未能正确跟随模型整体转换为FP16精度所导致的。

问题现象

具体表现为转换过程中会抛出类型不匹配的错误信息："alpha has dtype fp32 whereas x has dtype fp16"。这是因为虽然模型主体已经转换为FP16精度，但激活函数的alpha参数仍然保持FP32精度，导致数据类型不一致。

技术分析

根本原因

1. 历史遗留问题：在早期使用torch.jit.trace进行模型转换时，Torch CPU运算仅支持FP32精度，因此CoreMLTools设计时默认只接受FP32精度的Torch模型，通过compute_precision参数来控制最终Core ML模型的精度。

2. 数据类型传播不一致：当使用model.to(torch.float16)将整个模型转换为FP16时，某些参数如激活函数的alpha值未能正确跟随转换，导致数据类型不匹配。

解决方案探讨

目前有两种可行的解决方案：

1. 保持Torch模型为FP32：不强制使用model.to(torch.float16)，而是通过CoreMLTools的convert函数的compute_precision参数控制输出模型的精度。对于需要FP16输入输出的情况，可以在ExecuTorch CoreML委托中增加标志来显式指定输入输出为FP16。

2. 支持FP16 Torch模型：允许模型整体转换为FP16，在CoreMLTools转换过程中遇到数据类型不匹配时，自动调用promote_input_dtypes函数进行类型提升和统一。

最佳实践建议

对于大多数使用场景，建议采用第一种方案：

1. 保持PyTorch模型为FP32精度
2. 在CoreMLTools转换时明确指定：

   compute_precision=ct.precision.FLOAT16
   

3. 通过inputs参数显式指定输入输出的数据类型：

   inputs=[ct.TensorType(dtype=np.float16), ...]
   

这种方法更加稳定，且能避免因数据类型不匹配导致的各种问题。同时，Core ML运行时仍会以FP16精度执行计算，不会影响最终的性能表现。

技术细节补充

对于确实需要将PyTorch模型转换为FP16的特殊场景，可以在CoreMLTools的转换代码中添加类型提升逻辑。例如对于LeakyReLU激活函数，可以在处理节点前添加：

alpha, x = promote_input_dtypes([alpha, x])
res = mb.leaky_relu(x=x, alpha=negative_slope, name=node.name)

这样可以确保所有输入参数的数据类型一致，避免转换失败。

总结

CoreMLTools在FP16模型转换时的激活函数兼容性问题源于历史设计决策和数据类型传播机制。通过理解问题本质并采用适当的解决方案，开发者可以顺利完成模型转换工作。对于大多数应用场景，推荐保持PyTorch模型为FP32精度，而通过CoreMLTools的参数控制最终模型的精度，这是最稳定可靠的实践方案。