Core ML uses too much memory when saving updated model

385 views Asked by At

For my thesis I'm using on-device training to update a model which uses images(CVPixelBuffers) as input to do its prediction.

For the training I'm using 200 images with the size of 224x224. There are two layers that are being updated. Both of them innerProducts (1000x12167 and 1x1000). Since these are a lot of parameters I tried reducing them to one layer(1x12167) but this only had little impact.

As soon as the MLUpdateTask starts, the memory usage jumps up to ⁓1,3 gigabytes (also with reduced parameters). The training is completed, but when the updated model is saved and loaded using

finalContext.model.write(to: fileURL)
model = model_updatable(contentsOf: fileURL)

it is terminated due to memory issue. What is causing this problem and is there a way to bypass it? Would it make a difference changing the input from images to multiArrays?

Edit: Using 20 images as input only results in a usage of ⁓1,1 gigabytes during training and allows the app to save/load the updated model with an overall max memory usage of 1,39 gigabytes. Sadly I do need the input of 200 images.

Kind regards Nils

1

There are 1 answers

1
Matthijs Hollemans On

Who knows what happens... Core ML is a black box and we can't look inside to see what is going on.

It's possible that your model is simply too large to update it on the device.

However, sometimes Core ML memory usage blows up because of a bug. I suggest filing a bug report with Apple.