Skip to main content

AI Notebooks - Tutorial - Fine-Tune and export an AI model to ONNX

Stephanie R.
Updated

Learn how to fine-tune a PyTorch model and export it in ONNX (Open Neural Network Exchange) format.

For this purpose, we use an image classification model: DenseNet.

DenseNet models are pre-trained on the ImageNet dataset.

The goal is to fine-tune a DenseNet121 model to classify daily life images with the CIFAR-10 dataset. It is composed of 10 classes:

  • airplane
  • automobile
  • bird
  • cat
  • deer
  • dog
  • frog
  • horse
  • ship
  • truck

Find out more information about CIFAR-10 dataset on the following link.

At the end of training, the DenseNet model is saved in PyTorch format (.pth). It will then be transformed into ONNX.

Exporting your model in ONNX format allows you to optimize the inference of a Machine Learning model.

PaaS_AI_notebooks_exportONNX01.png

 

Requirements

 

Instructions

You can launch the notebook from the OVHcloud Control Panel or via the ovhai CLI.

 

Launching a Jupyter notebook with "Conda" via UI (Control Panel)

To launch your notebook from the OVHcloud Control Panel, refer to the following steps.

 

Code editor

Choose the Jupyterlab code editor.

 

Framework

In this tutorial, the conda framework is used.

 

Resources

Using GPUs is recommended to train the image classification model: densenet121.

Here, using 1 GPU is sufficient.

 

Launching a Jupyter notebook with "conda" via CLI

If you have not used our CLI yet, follow this guide to install it.

If you want to launch your notebook with the OVHcloud AI CLI, choose the jupyterlab editor and the conda framework.

To access the different versions of conda available, run the following command:

ovhai capabilities framework get conda -o yaml

You will also need to choose the number of GPUs to use in your notebook, using <nb-gpus>.

To launch your notebook, run the following command:

ovhai notebook run conda jupyterlab \
        --name <notebook-name> \
        --framework-version conda-py310-cuda11.8-v22-4 \
        --gpu <nb-gpus>

You can then reach your notebook’s URL once the notebook is running.

 

Accessing the notebooks

Once our AI examples repository has been cloned in your environment, find the fine-tuning notebook tutorial by following this path: ai-training-examples > notebooks > go-further > onnx > notebook_finetune_densenet_export_onnx.ipynb.

A preview of this notebook can be found on GitHub here.

 

Go further

For more information and tutorials, please see our other AI & Machine Learning support guides or explore the guides for other OVHcloud products and services.

If you need training or technical assistance to implement our solutions, contact your sales representative or click on this link to get a quote and ask our Professional Services experts for a custom analysis of your project.

Related articles