Help! I am unable to get the Apple Neural Engine(ANE) to run the DenseNet121 model after fine-tuning it using PyTorch. When downloading the model directly from torchvision with retrained weights, it managed to run with the ANE, but after fine-tuning the model, I am unable to run the model using ANE.

This is the model:

#Downloading densenet model

import torchvision
model = torchvision.models.densenet121()
model.classifier=nn.Linear(in_features=1024,out_features=100,bias=True)


#Converting to CoreML model package

import coremltools as ct
model.to("cpu").eval()
example_input = torch.rand(1, 3, 32, 32) 
traced_model = torch.jit.trace(model, example_input)
class_labels=[item for item in range(100)]
classifier_config=ct.ClassifierConfig(class_labels)
converted_model = ct.convert(
    traced_model,
    source='pytorch',
    convert_to="mlprogram",
    inputs=[ct.ImageType(name="input_1",shape=example_input.shape,color_layout=ct.colorlayout.RGB,)],
    classifier_config=classifier_config,
 )
converted_model.save("./converted_models/densenet121.mlpackage")

The photo below shows the model using the neural network.

enter image description here

After loading in the weights for the network and trying the same thing again, it will only use the GPU.

#Loading model checkpoint
checkpoint=torch.load("./trained_model/model_best.pt")
torch_model.load_state_dict(checkpoint['model_state_dict'])


import coremltools as ct
model.to("cpu").eval()
example_input = torch.rand(1, 3, 32, 32) 
traced_model = torch.jit.trace(model, example_input)
class_labels=[item for item in range(100)]
classifier_config=ct.ClassifierConfig(class_labels)
converted_model = ct.convert(
    traced_model,
    source='pytorch',
    convert_to="mlprogram",
    inputs=[ct.ImageType(name="input_1",shape=example_input.shape,color_layout=ct.colorlayout.RGB,)],
    classifier_config=classifier_config,
    compute_precision=ct.precision.FLOAT16,
 )
converted_model.save("./converted_models/densenet121.mlpackage")

I have also tried to convert the model directly after training on a linux system. It will also only use the GPU for inferencing.

I have also added an additional argument for compute precision to be FLOAT16 as Neural engine only takes in Float16 for computation, but it still fails to use ANE.

Is there anyway to load a PyTorch model, convert it to mlpackage and be able to use ANE for the model for inferencing?

New contributor

Chin Kin Mun is a new contributor to this site. Take care in asking for clarification, commenting, and answering.
Check out our Code of Conduct.