CoreML加载含自定义ResizeArea层模型报2nd reshape失败求助

问题原因说明 #

你遇到的两个报错分别对应自定义算子形状推演不兼容、coremltools 5.x版本常量解析bug两类问题，先明确两个现象的正常性：

1. 加载阶段出现的全零维度、5维张量都是CoreML的正常行为：全零是动态维度的占位值，5维是CoreML内部处理动态尺寸图像模型的临时序列布局，不是你的模型逻辑出错。
2. BilinearResize转换报错是coremltools 5.x对float类型的尺寸参数解析有bug，不是你的pool_size值非法。

解决方案 #

方案1：修复自定义ResizeArea层 #

第一步：修正Python侧自定义算子的形状推演逻辑

CoreML的MIL中间表示默认使用NCHW（[batch, channel, height, width]）布局，和TensorFlow的NHWC布局不同，你原来的type_inference修改了错误的维度索引，修正后代码如下：

def type_inference(self):
    x_type = self.x.dtype
    x_shape = self.x.shape
    s = list(self.s.val)
    ret_shape = list(x_shape)
    # 适配CoreML NCHW布局，height对应索引2，width对应索引3
    ret_shape[2] = s[0]
    ret_shape[3] = s[1]
    return types.tensor(x_type, ret_shape)

第二步：修正Swift侧outputShapes逻辑

兼容4维、5维输入，正确替换高宽维度：

@objc(CustomResizeArea) class CustomResizeArea: NSObject, MLCustomLayer {
    required init(parameters: [String : Any]) throws {
        super.init()
    }
    
    func setWeightData(_ weights: [Data]) throws {}
    
    func outputShapes(forInputShapes inputShapes: [[NSNumber]]) throws -> [[NSNumber]] {
        let inputXShape = inputShapes[0]
        // 固定目标尺寸为5x5，也可从inputShapes[1]解析动态尺寸
        let targetH: NSNumber = 5
        let targetW: NSNumber = 5
        
        var outputShape = inputXShape
        if inputXShape.count == 5 {
            // 5维结构：[sequence, batch, channel, height, width]
            outputShape[3] = targetH
            outputShape[4] = targetW
        } else if inputXShape.count == 4 {
            // 4维结构：[batch, channel, height, width]
            outputShape[2] = targetH
            outputShape[3] = targetW
        } else {
            throw NSError(domain: "CustomResizeArea", code: -1, userInfo: [
                NSLocalizedDescriptionKey: "Unsupported input dimension"
            ])
        }
        return [outputShape]
    }
    
    // 补充实现evaluate方法即可完成自定义层
    func evaluate(_ inputs: [MLCustomLayerInput], outputs: [MLCustomLayerOutput]) throws {
        // 此处实现ResizeArea的运算逻辑即可
    }
}

方案2：绕开自定义层，用内置算子实现相同逻辑 #

如果你不想维护自定义层，可以选择以下两种替代实现，coremltools 5.x完全支持：

替代1：修复BilinearResize转换问题

将resize的目标尺寸改为int32类型的常量，绕开coremltools的常量解析bug：

# 原Lambda层修改为：
ImgResizer = Lambda(
    lambda x: tf.image.resize(x, tf.constant([5,5], dtype=tf.int32), method='bilinear'),
    name='feature_resizer'
)
# 转换时指定compute_units避免精度问题
mlmodel = ct.convert(model, compute_units=ct.ComputeUnit.CPU_AND_GPU)

替代2：用自适应平均池化替换ResizeArea

对于下采样场景，AdaptiveAvgPool2D((5,5))的效果和ResizeArea几乎一致，完全不需要处理算子兼容问题：

ImgResizer = tf.keras.layers.AdaptiveAvgPool2D((5,5), name='feature_resizer')

内容的提问来源于stack exchange，提问作者mazy