Coremltools: Split a pre-trained CoreML model into two CoreML Models ?

[anujdutt9]

I have a Sound Classification model from turicreate example here: https://apple.github.io/turicreate/docs/userguide/sound_classifier/

I am trying to split this model into two and save the two parts as separate CoreML Models using coremltools library. Can anyone please guide me on how to do this?

I am able to load the model and even print out the spec of the model. But don't know where to go from here.

import coremltools

# Load the Saved CoreML Model
mlmodel = coremltools.models.MLModel('./EnvSceneClassification.mlmodel')

# Get spec from the model
spec = mlmodel.get_spec()

Thanks.

[aseemw]

Once the spec is loaded it corresponds to the protobuf message message Model (https://github.com/apple/coremltools/blob/6af076f9f169a2903741b551c4d467ae32826bd6/mlmodel/format/Model.proto#L213)

Since its a protobuf object it can be edited in Python using standard python protobuf APIs, as long as the structure of the proto is followed, which is provided in here

For example, if you want to split the model generated by the example in https://apple.github.io/turicreate/docs/userguide/sound_classifier/ , the code would look like:

import coremltools

mlmodel = coremltools.models.MLModel('mymodel.mlmodel')

spec = mlmodel.get_spec()

# get the type of the model : PipelineClassifier

print('Type of mlmodel: ', spec.WhichOneof('Type'))

print('***** input description *****: ')
print(spec.description.input)
print('***** output description *****: ')
print(spec.description.output)

pipeline_spec = spec.pipelineClassifier.pipeline
print('number of models in pipeline :', len(pipeline_spec.models))
if len(pipeline_spec.names) > 0:
    print('names of the models in the pipeline: ')
    print(pipeline_spec.names)


print('type of each model present in the pipeline: ')
for i in range(len(pipeline_spec.models)):
    print('-' * 40)
    spec_i = pipeline_spec.models[i]
    print('{}-th model type : {}'.format(i, spec_i.WhichOneof('Type')))
    print('***** input description *****: ')
    print(spec_i.description.input)
    print('***** output description *****: ')
    print(spec_i.description.output)


last_model_spec = pipeline_spec.models[2]
del pipeline_spec.models[-1]

# delete outputs from the spec, as these are produced from the second model now
del spec.description.output[:]
# copy the output which is the input to the last model
out = spec.description.output.add()
out.CopyFrom(last_model_spec.description.input[0])

model_beginning = coremltools.models.MLModel(spec)
model_last = coremltools.models.MLModel(last_model_spec)

model_beginning.save('mymodel_broken_1st_part.mlmodel')
model_last.save('mymodel_broken_2nd_part.mlmodel')

[anujdutt9]

Awesome. Thanks a lot for this @aseemw : )