Search code examples
swiftcoremlcoremltools

CoreML automatic cast int to float

I am new user of coremltool, currently I am able to convert my TF model, and run it in python. However the output layer of my model is automatically converted to float instead of int32. I have to perform myself in python a rollback conversion.

model.predict({"Input": image}).astype(int)

Does somebody knows how to cancel this automatic conversion ? And keep my output as a multidimensional array - int ?

Thank you,

++t

Solution

  • You can change the datatype of the model's outputs to INT32 using the coremltools Python library. That will automatically cast those numbers from floats to integers.

    Something like this:

    import coremltools
model = coremltools.models.MLModel("YourModel.mlmodel")
spec = model._spec
spec.description.output[0].type.multiArrayType.dataType = coremltools.proto.FeatureTypes_pb2.ArrayFeatureType.INT32
coremltools.models.utils.save_spec(spec, "YourNewModel.mlmodel")

    The reason this is not done automatically, is because usually floats are the correct datatype.