CoreML: Failed in 2nd reshape after missing custom layer info

Viewed 96

I our very big Tensorflow (2.x) model there is a part where we use InceptionResNetV2 and create multi-pooled model from it:

model_base = InceptionResNetV2(weights     = 'imagenet',
                                   include_top = False,
                                   input_shape = input_shape)

ImgResizer = Lambda(lambda x: tf.image.resize(x, pool_size, method='area'),
                          name='feature_resizer')

feature_layers = [l for l in model_base.layers if 'mixed' in l.name]
feature_layers = [feature_layers[i] for i in indexes]
pools = [ImgResizer(l.output) for l in feature_layers]
conc_pools = Concatenate(name='conc_pools', axis=3)(pools)

model = Model(inputs  = model_base.input, outputs = conc_pools)

Here we use Lambda function to resize 4D tensor ([batch, height, width, channels]) to a new 4D tensor with pool_size (5,5), input_shape=(None, None, 3) and indexes=list(range(43)). Unfortunately, this operation (ResizeArea) is not supported in the current version of coremltools (5.0) and I have to do a custom operation wrapper and then implement it on a swift.

@register_op(doc_str='Custom ResizeArea Layer', is_custom_op=True)
class custom_resize_area(Operation):
    input_spec = InputSpec(
             x = TensorInputType(),
             s = ScalarOrTensorInputType()
    )

    bindings = { 'class_name'  : 'CustomResizeArea',
                 'input_order' : ['x', 's'],
                 'parameters'  : [],
                 'description' : "Resize area custom layer"
                }

    def __init__(self, **kwargs):
        super(custom_resize_area, self).__init__(**kwargs)

    def type_inference(self):
        x_type = self.x.dtype
        x_shape = self.x.shape
        s = list(self.s.val)

        ret_shape = list(x_shape)
        ret_shape[1] = s[0]
        ret_shape[2] = s[1]
        #print(x_shape, ret_shape)
        return types.tensor(x_type, ret_shape)

# Override ResizeArea op with override=True flag
@register_tf_op(tf_alias=['ResizeArea'], override=True)
def CustomResizeArea(context, node):
    #input: "model_2/mixed_5b/concat"
    #input: "model_2/lambda/resize/size"
    x = context[node.inputs[0]]
    s = context[node.inputs[1]]
    x = mb.custom_resize_area(x=x, s=s, name=node.name)
    context.add(node.name, x)

The conversion was successful and now I started to implement this custom layer on swift. Here is a part from my code:

@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]] {
        print(#function, inputShapes)
        let outputShape = inputShapes[0]
        
        // [ sequence, batch, channel, height, width ] - why?
        if outputShape.count == 5 {
            print(outputShape)
            return [outputShape]
        }
        print([outputShape[0], 5, 5, outputShape[3]])
        return [[outputShape[0], 5, 5, outputShape[3]]]
    }
    ...
}

I am publishing only outputShapes function since the CoreML gives me an error before evaluate function will actually be executed. Here the last logs (with some output shapes parameters inside):

outputShapes(forInputShapes:) [[0, 0, 0, 0, 0], [0, 0, 0, 0, 0]]
[0, 0, 0, 0, 0]
outputShapes(forInputShapes:) [[0, 0, 0, 0, 0], [0, 0, 0, 0, 0]]
[0, 0, 0, 0, 0]
.....
outputShapes(forInputShapes:) [[1, 0, 0, 448], [2]]
[1, 5, 5, 448]
outputShapes(forInputShapes:) [[1, 0, 0, 448], [2]]
[1, 5, 5, 448]
outputShapes(forInputShapes:) [[1, 0, 0, 448], [2]]
[1, 5, 5, 448]
2021-11-06 10:15:55.085694+0100 snafu[11263:4717917] [espresso] [Espresso::handle_ex_plan] exception=Failed in 2nd reshape after missing custom layer info.
2021-11-06 10:15:55.086108+0100 snafu[11263:4717917] [coreml] Error in adding network -1.
2021-11-06 10:15:55.086477+0100 snafu[11263:4717917] [coreml] MLModelAsset: load failed with error Error Domain=com.apple.CoreML Code=0 "Error in declaring network." UserInfo={NSLocalizedDescription=Error in declaring network.}

I have no idea why I get [[0, 0, 0, 0, 0], [0, 0, 0, 0, 0]] as an inout, and why are there zeros in, for example, [[1, 0, 0, 448], 2]. Actually my model should support any image size on input, that's why? I think I have tried everything...

P.S.: I also tried to replace area reize method with bilinear in the python code (Lambda) because there is a BilinearResize converter in coremltools. But I can not convert it, because I get the following error:

File "/home/alex/anaconda3/envs/tfgpu/lib/python3.8/site-packages/coremltools/converters/mil/mil/builder.py", line 75, in _add_const
    raise ValueError("Cannot add const {}".format(val))
ValueError: Cannot add const 5.0001/is57

I suppose that 5.0001 here is my pool_size, but I don't understand why is it incorrect.

0 Answers
Related