Quantization aware training for executorch deployment on ethos-U

[Steph71]

Hi,

I am banging my head against this since a wile with only partly success.

Is there any example, that shows how to

• use torch.ao for quantization aware training of a simple, but not trivial NN (a cnn that classifies an MNSIT dataset should be good enough)
• deploy this NN on arm ethosU using executorch, while preserving the quantization parameters determined during QAT

A notebook, that goes through the individual steps, including some lines of explanation, what is going on would be highly appreciated.

Even better: Some information how to inspect the intermediate results for correctness and some debugging hints ...

Best Regards
Steph

[cccclai]

tagging the PoC I know from ARM @zingo

[AdrianLundell]

Hi Steph, we have a pytest which performs a qat training loop using the TOSAQuantizer here: https://github.com/pytorch/executorch/blob/main/backends/arm/test/misc/test_qat_training_loop.py. You should be able to run it using pytest backends/arm/test/misc/test_qat_training_loop.py or you could try to copy parts of it and use it together with our ethos_u_minimal example notebook https://github.com/pytorch/executorch/blob/main/examples/arm/ethos_u_minimal_example.ipynb.

You do raise a good point that we should have better documentation for this though and we should definitely upstream a notebook on this in the long term! Feel free to post any issues you have here in the meanwhile.

[Steph71]

This was exactly what I am looking for, thanks for that. A notebook, implementing the end to end flow including qat and execution on fvp is available here:
https://gist.github.com/Steph71/3c45dcd753accc0ea4ad28298be1308c

[AdrianLundell]

That's great, thanks for putting it together! If you're interested in upstreaming it feel free to open a PR, we're happy to review and accept contributions.