MobileNetV3Large BatchNorm Incompatible With Expected Resource

[binocla]

❓Question

Intro

Hi! I'm having the following MobileNetV3Large model for Semantic Human Segmentation:

def create_mobilenet_segmentation_model(input_shape):
    base_model = tf.keras.applications.MobileNetV3Large(input_shape=input_shape[1:], include_top=False, weights='imagenet')
    base_model.trainable = True 

    inputs = tf.keras.Input(batch_shape=input_shape)

    # Encoder - Using MobileNetV3 as the encoder
    x = base_model(inputs, training=False)
    x = layers.Conv2D(256, (3, 3), padding='same', activation='selu')(x)
    x = layers.UpSampling2D(size=(2, 2))(x)

    # ...

    x = layers.Conv2D(16, (3, 3), padding='same', activation='selu')(x)
    x = layers.UpSampling2D(size=(2, 2))(x)

    # Output layer for binary segmentation mask
    output = layers.Conv2D(1, kernel_size=1, strides=1, padding='same', activation='sigmoid')(x)

    output = tf.image.resize(output, (input_shape[1], input_shape[2]))

    model = tf.keras.models.Model(inputs=[inputs], outputs=[output])
    return model

input_shape = (1, IMAGE_SIZE, IMAGE_SIZE, 3)
model = create_mobilenet_segmentation_model(input_shape)
model.summary()

Model summary results in the following:

Layer (type)                Output Shape              Param #   
=================================================================
 input_6 (InputLayer)        [(1, 512, 512, 3)]        0         
                                                                 
 MobilenetV3large (Function  (None, 16, 16, 960)       2996352   
 al)                                                             
                                                                 
 conv2d_6 (Conv2D)           (1, 16, 16, 256)          2212096   
                                                                 
 up_sampling2d_5 (UpSamplin  (1, 32, 32, 256)          0         
 g2D)                                                            
                                                                 
# ... 
                                                                 
 up_sampling2d_9 (UpSamplin  (1, 512, 512, 16)         0         
 g2D)                                                            
                                                                 
 conv2d_11 (Conv2D)          (1, 512, 512, 1)          17        
                                                                 
 tf.image.resize_1 (TFOpLam  (1, 512, 512, 1)          0         
 bda)                                          

Conversion

import tensorflow as tf
from keras.models import load_model
import coremltools as ct

# Load the Keras model
model_path = '/mnt/host_home/checkpoints/mobileNet_large_8.keras'
model = load_model(model_path)

# Fix the batch size to 1 for the model
new_input = tf.keras.Input(shape=(512, 512, 3), batch_size=1, name="input")
new_output = model(new_input)
fixed_batch_model = tf.keras.Model(inputs=new_input, outputs=new_output)

# Convert the model to Core ML
mlmodel = ct.convert(
    fixed_batch_model,
    source="tensorflow",
    inputs=[ct.ImageType(name="input", shape=(1, 512, 512, 3), scale=1/255.0, bias=[0, 0, 0])]
)

# Save the Core ML model
coreml_model_path = '/mnt/host_home/mobileNet_large_8.mlmodel'
mlmodel.save(coreml_model_path)

Exception:

Error
		ValueError
		---------------------------------------------------------------------------
		InvalidArgumentError                      Traceback (most recent call last)
		/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/importer.py in _import_graph_def_internal(graph_def, input_map, return_elements, validate_colocation_constraints, name, producer_op_list, propagate_device_spec)
		    510         with graph._c_graph.get() as c_graph:  # pylint: disable=protected-access
		--> 511           results = c_api.TF_GraphImportGraphDefWithResults(
		    512               c_graph, serialized, options)
		
		InvalidArgumentError: Input 0 of node model_7/model_5/MobilenetV3large/Conv/BatchNorm/AssignNewValue was passed float from model_7/model_5/MobilenetV3large/Conv/BatchNorm/FusedBatchNormV3/ReadVariableOp/resource:0 incompatible with expected resource.
		
		During handling of the above exception, another exception occurred:
		
		ValueError                                Traceback (most recent call last)
		<ipython-input-13-f12f82c8401f> in <cell line: 15>()
		     13 
		     14 # Convert the model to Core ML
		---> 15 mlmodel = ct.convert(
		     16     fixed_batch_model,
		     17     source="tensorflow",

I tried to convert to TF Saved_model and after that to CoreML, to ONNX and then to CoreML, to freeze model, but still the same exception

Working conversion to .tflite:

import tensorflow as tf
from tensorflow.keras import layers, models
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt

IMAGE_SIZE = 512

model = tf.keras.models.load_model('/mnt/host_home/checkpoints/mobileNet_large_8.keras')

saved_model_path = 'saved_model/human_segmentation_model_1'
model.save(saved_model_path)

converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_path)
converter.target_spec.supported_ops = [
    tf.lite.OpsSet.TFLITE_BUILTINS,
    tf.lite.OpsSet.SELECT_TF_OPS 
]
tflite_model = converter.convert()

tflite_model_path = '/mnt/host_home/human_segmentation_model_marko.tflite'
with open(tflite_model_path, 'wb') as f:
    f.write(tflite_model)

I don't know how to fix this BatchNormalization layer so that it's possible to convert trained model to CoreML

I look forward to any advices or hints in due course. Thank you!

[TobyRoseman]

The error seems to be coming from TensorFlow (tensorflow/python/framework/importer.py). I think it can't load your model. This doesn't seem like a coremltools issue.

If you still believe this to be a coremltools issue, please give us a standalone, minimal example, which doesn't include loading an external model. Also let us know which version of TensorFlow and Coremltools you are using.