AI Training - Tutoriel - Entraîner un modèle PyTorch et l'exporter au format ONNX

Objectif

L'objectif de ce tutoriel est de vous montrer comment entraîner un modèle PyTorch personnalisé et l'exporter au format ONNX (Open Neural Network Exchange).

L'objectif est d'entraîner votre propre modèle de classification d'images sur le célèbre dataset MNIST. À la fin de l'entraînement du modèle, celui-ci est sauvegardé au format PyTorch (.pt) puis transformé en ONNX.

Exporter votre modèle au format ONNX vous permet d'optimiser l'inférence d'un modèle de Machine Learning.

Prérequis

• Un projet Public Cloud créé.
• L'interface CLI ovhai installée sur votre système (plus d'informations ici).
• Docker installé et configuré pour construire des images.
• Un registre d'images OCI / Docker. Vous pouvez utiliser un registre public (comme Docker Hub par exemple) ou un registre privé. Consultez la documentation Créer un registre privé pour créer un registre privé basé sur Harbor. Pour rendre votre registre compatible avec l'utilisation d'AI Solutions, suivez le guide Utiliser et gérer vos registres.
• Des connaissances sur la construction d'images avec Dockerfile.

Accès à l'espace client OVHcloud

• Lien direct : Tous mes projets Public Cloud
• Pour accéder à vos services : Public Cloud > Sélectionnez votre projet

En pratique

Créer un bucket Object Storage pour votre modèle ONNX

Pour pouvoir récupérer et utiliser le modèle ONNX à la fin de l'entraînement, vous devez créer un bucket vide pour le stocker.

Créer votre bucket via l'interface (Control Panel)

Si vous ne vous sentez pas à l'aise avec les commandes, cette méthode peut être plus intuitive.

Cliquez sur ce lien pour accéder à votre projet Public Cloud, puis sélectionnez la section Object Storage (dans la catégorie Storage) et créez un nouveau conteneur d'objets en cliquant sur Storage > Object Storage > Créer un conteneur d'objets.

Vous pouvez créer le bucket qui stockera votre modèle ONNX à la fin de l'entraînement. Sélectionnez le type de conteneur et la région qui correspondent à vos besoins.

Créer votre bucket via la CLI ovhai

Pour suivre cette partie, assurez-vous d'avoir installé la CLI ovhai sur votre ordinateur ou sur une instance.

Comme dans le Control Panel, vous devrez préciser la region et le name (cnn-model-onnx) de votre bucket. Créez votre bucket Object Storage comme suit :

ovhai bucket create <region> cnn-model-onnx

Écrire le code Python d'entraînement du modèle

Pour entraîner le modèle, nous allons utiliser AI Training. Cet outil puissant vous permet d'entraîner vos modèles d'IA à partir de vos propres images Docker.

Vous devez créer un script Python chargé de réaliser l'entraînement : train_image_classification.py.

Tout d'abord, importez les dépendances Python.

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable
from torchvision import datasets, transforms

Ensuite, définissez l'architecture du réseau de neurones.

class Network(nn.Module):

    def __init__(self):
        super(Network, self).__init__()
        self.conv1 = nn.Conv2d(1, 20, 5, 1)
        self.conv2 = nn.Conv2d(20, 50, 5, 1)
        self.fc1 = nn.Linear(4*4*50, 500)
        self.fc2 = nn.Linear(500, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, 2, 2)
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2, 2)
        x = x.view(-1, 4*4*50)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)

        return F.log_softmax(x, dim=1)

Maintenant, définissez la fonction load_data.

def load_data():

    train_loader = torch.utils.data.DataLoader(
    datasets.MNIST('/workspace/data', train=True, download=True,
                   transform=transforms.Compose([
                       transforms.ToTensor(),
                       transforms.Normalize((0.1307,), (0.3081,))
                   ])),
        batch_size=64, shuffle=True)

    test_loader = torch.utils.data.DataLoader(
    datasets.MNIST('/workspace/data', train=False, transform=transforms.Compose([
                       transforms.ToTensor(),
                       transforms.Normalize((0.1307,), (0.3081,))
                   ])),
    batch_size=64, shuffle=True)

    return train_loader, test_loader

Vérifiez la disponibilité du GPU.

def check_gpu():

    if torch.cuda.is_available():
        device = torch.device('cuda:0')
        print('Running on GPU...')
    else:
        device = torch.device('cpu')
        print('Running on CPU...')

    return device

Ensuite, définissez la fonction qui va entraîner votre modèle.

def train_model(model, device, train_loader, optimizer, epoch):

    model.train()

    running_loss=0
    correct=0
    total=0

    for data in train_loader:

        inputs, labels = data[0].to(device), data[1].to(device)

        optimizer.zero_grad()
        outputs = model(inputs)
        loss = F.nll_loss(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

        _, predicted = outputs.max(1)
        total += labels.size(0)
        correct += predicted.eq(labels).sum().item()

    train_loss = running_loss/len(train_loader)
    accu = 100.*correct/total

    train_losses.append(train_loss)
    train_accu.append(accu)

    print('Training - Loss: %.3f | Accuracy: %.3f'%(train_loss, accu))

    return

Vous pouvez également définir la fonction qui va tester votre modèle.

def test_model(model, device, test_loader):

    model.eval()

    running_loss=0
    correct=0
    total=0

    with torch.no_grad():

        for data in test_loader:

            images,labels=data[0].to(device), data[1].to(device)

            outputs = model(images)
            loss = F.nll_loss(outputs, labels)
            running_loss+=loss.item()

            _, predicted = outputs.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()

    test_loss = running_loss/len(test_loader)
    accu = 100.*correct/total

    test_losses.append(test_loss)
    test_accu.append(accu)

    print('Validation - Loss: %.3f | Accuracy: %.3f'%(test_loss,accu))

    return

Enfin, exportez le modèle au format ONNX grâce à la fonction suivante.

def export_onnx(model):

    model.load_state_dict(torch.load('/workspace/models/model_mnist_classification.pt'))

    dummy_input = torch.randn(1, 1, 28, 28, device="cuda")
    torch.onnx.export(model, dummy_input, "/workspace/models/model_mnist_classification.onnx")

    return

Il est maintenant temps d'appeler ces fonctions dans le main !

if __name__ == '__main__':

    train_loader, test_loader = load_data()   

    device =  check_gpu()
    model = Network().to(device)
    optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5)

    train_losses, train_accu, test_losses, test_accu = ([] for i in range(4))

    for epoch in range(0, 10):
        print('\nEpoch %d/%d:'%(epoch, 10))
        train_model(model, device, train_loader, optimizer, epoch)
        test_model(model, device, test_loader)

    torch.save(model.state_dict(),"/workspace/models/model_mnist_classification.pt")
    export_onnx(model)

Créer le fichier requirements.txt

Ensuite, créez un fichier requirements.txt pour déclarer les dépendances Python et leurs versions.

torch==2.0.1
torchvision==0.15.2
onnx==1.14.1

Construire votre propre image Docker

Créer un Dockerfile compatible avec AI Training

Le Dockerfile doit commencer par l'instruction FROM indiquant l'image parente à utiliser. Dans notre cas, nous choisissons de partir d'une image python:3.10.

Ensuite, précisez le chemin du workspace, installez les dépendances Python et lancez l'entraînement du modèle à l'aide de CMD.

FROM python:3.10

WORKDIR /workspace
ADD . /workspace

RUN pip install -r requirements.txt

CMD [ "python3"  , "/workspace/train_image_classification.py"]

RUN chown -R 42420:42420 /workspace
ENV HOME=/workspace

Construire l'image Docker depuis le Dockerfile

Depuis le répertoire contenant votre Dockerfile, exécutez l'une des commandes suivantes pour construire l'image de votre application :

# Build the image using your machine's default architecture
docker build . -t train-cnn-model-export-onnx:latest

# Build image targeting the linux/amd64 architecture
docker buildx build --platform linux/amd64 -t train-cnn-model-export-onnx:latest .

• La première commande construit l'image en utilisant l'architecture par défaut de votre système. Cela peut fonctionner si votre machine utilise déjà l'architecture linux/amd64, qui est requise pour exécuter des containers avec nos produits AI. Cependant, sur des systèmes avec une architecture différente (par exemple ARM64 sur Apple Silicon), l'image résultante ne sera pas compatible et ne pourra pas être déployée.

• La seconde commande cible explicitement l'architecture linux/amd64 afin de garantir la compatibilité avec nos services AI. Elle nécessite buildx, qui n'est pas installé par défaut. Si vous n'avez jamais utilisé buildx, vous pouvez l'installer en exécutant : docker buildx install

Envoyer l'image vers le registre partagé

Trouvez l'adresse de votre registre partagé en lançant cette commande :

ovhai registry list

Connectez-vous au registre partagé avec vos identifiants habituels d'utilisateur de la plateforme AI :

docker login -u <user> -p <password> <shared-registry-address>

Envoyez l'image compilée vers le registre partagé :

docker tag train-cnn-model-export-onnx:latest <shared-registry-address>/train-cnn-model-export-onnx:latest
docker push <shared-registry-address>/train-cnn-model-export-onnx:latest

Une fois votre image Docker créée et envoyée vers le registre, vous pouvez directement utiliser la commande ovhai pour créer votre entraînement de modèle. Vous pouvez lancer l'entraînement en spécifiant plus ou moins de GPU selon la vitesse souhaitée pour votre entraînement.

Lancer le job AI Training

Vous pouvez lancer le job d'entraînement via l'interface ou la CLI.

Créer votre job d'entraînement via l'interface (Control Panel)

Si vous ne vous sentez pas à l'aise avec les commandes, cette méthode peut être plus intuitive.

Cliquez sur ce lien pour accéder à votre projet Public Cloud, puis sélectionnez la section AI Training (dans la catégorie AI & Machine Learning) et créez un nouveau job en cliquant sur AI Training > Lancer un job.

Vous pouvez créer le job qui entraînera votre modèle et l'exportera au format ONNX. Sélectionnez la région et ajoutez votre image Docker personnalisée (<shared-registry-address>/train-cnn-model-export-onnx:latest).

Ensuite, attachez votre conteneur Object Storage cnn-model-onnx et définissez le répertoire de montage : /workspace/models.

Enfin, configurez votre job et choisissez au moins 1 GPU.

Créer votre job d'entraînement via la CLI ovhai

La commande suivante démarre l'entraînement :

ovhai job run \
	--name <job_name> \
	--gpu 1 \
	--volume cnn-model-onnx@<region>/:/workspace/models:RW \
	<shared-registry-address>/train-cnn-model-export-onnx:latest

Vous pouvez ensuite suivre l'évolution de l'entraînement via les logs du job :

ovhai job logs <job_id>

Aller plus loin

Consultez nos autres tutoriels pour apprendre à :

• Utiliser le Transfer Learning avec ResNet50 pour la classification d'images

• Fine-tuner et exporter un modèle d'IA au format ONNX via un AI Notebook

Votre avis nous intéresse !

N’hésitez pas à nous faire part de vos questions, retours et suggestions pour améliorer le service :

• Sur le serveur Discord OVHcloud