Classification d'images à l'aide de réseaux de neurones convolutifs

Ouvrir dans Databricks

Entraînez un réseau de neurones convolutif (CNN) pour la classification d'images sur AI Runtime à l'aide de PyTorch et du dataset MNIST. MNIST contient 70 000 images en niveaux de gris de chiffres manuscrits (0-9), ce qui en fait un ensemble idéal pour apprendre les techniques de classification d'images.

Vous apprendrez comment :

• Connectez votre notebook à un compute GPU serverless avec un GPU A10G
• Définir une architecture simple de réseau de neurones convolutif
• Entraînez le modèle sur un seul GPU et enregistrez les métriques dans MLflow
• Enregistrer les points de contrôle du modèle dans un volume Unity Catalog
• Chargez et évaluez le modèle entraîné

Se connecter au compute GPU serverless

Ce Notebook nécessite un GPU pour entraîner le réseau de neurones efficacement. Suivez ces étapes pour vous connecter au compute GPU Serverless :

1. Cliquez sur le menu déroulant **Connecter** en haut du Notebook.
2. Sélectionnez Serverless GPU .
3. Ouvrez le panneau latéral Environnement sur le côté droit du notebook.
4. Définissez l' Accélérateur sur 1xA10 pour cette démo.
5. Sélectionnez **AI v5** dans le menu déroulant **Environnement**.
6. Sélectionnez Appliquer et cliquez sur Confirmer pour appliquer cet environnement à votre Notebook.

Pour plus d'informations, consultez Serverless GPU compute.

Configurer l'emplacement de stockage des points de contrôle.

La cellule suivante crée des paramètres de widget pour spécifier où les points de contrôle du modèle seront enregistrés dans Unity Catalog. Ces paramètres définissent :

• uc_catalog: Le nom du catalogue Unity Catalog
• uc_schema: le schéma (base de données) dans le catalogue
• uc_volume: Le volume pour stocker les fichiers de point de contrôle
• uc_model_name: Le sous-répertoire dans le volume pour ce modèle spécifique

Ces valeurs sont utilisées dans tout le notebook pour construire le chemin de point de contrôle : /Volumes/{uc_catalog}/{uc_schema}/{uc_volume}/{uc_model_name}

La cellule suivante utilise des valeurs d'espace réservé par default. Mettez à jour les valeurs à l'aide des widgets en haut du Notebook. Ou, mettez à jour les valeurs default directement dans la cellule suivante.

Python

dbutils.widgets.text("uc_catalog", "main")
dbutils.widgets.text("uc_schema", "default")
dbutils.widgets.text("uc_volume", "checkpoints")
dbutils.widgets.text("uc_model_name", "cnn_mnist")

Définir le réseau de neurones convolutif

La cellule suivante définit une architecture CNN simple pour la classification d'images. Le réseau se compose de :

• Deux couches de convolution avec max pooling pour extraire les caractéristiques des images
• Deux couches entièrement connectées pour classer les caractéristiques extraites
• Couches de dropout pour éviter le surapprentissage

Le code définit également des classes utilitaires pour le jalonnement de l'état du modèle et de l'optimiseur dans un volume Unity Catalog, et des fonctions pour configurer l'entraînement distribué (utilisé pour les scénarios multi-GPU).

Cette implémentation est adaptée de l'exemple MNIST Horovod PyTorch.

Python

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.distributed as dist
import torch.distributed.checkpoint as dcp
from datetime import timedelta
import os

from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict
from torch.distributed.checkpoint.stateful import Stateful

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

UC_CATALOG = dbutils.widgets.get("uc_catalog")
UC_SCHEMA = dbutils.widgets.get("uc_schema")
UC_VOLUME = dbutils.widgets.get("uc_volume")
UC_MODEL_NAME = dbutils.widgets.get("uc_model_name")

# Ensure that the UC Volume directory exists first
CHECKPOINT_DIR = f"/Volumes/{UC_CATALOG}/{UC_SCHEMA}/{UC_VOLUME}/{UC_MODEL_NAME}"

class AppState(Stateful):
    """This is a useful wrapper for checkpointing the Application State. Since this object is compliant
    with the Stateful protocol, DCP will automatically call state_dict/load_stat_dict as needed in the
    dcp.save/load APIs.

    Note: We take advantage of this wrapper to hande calling distributed state dict methods on the model
    and optimizer.
    """

    def __init__(self, model, optimizer=None):
        self.model = model
        self.optimizer = optimizer

    def state_dict(self):
        # this line automatically manages FSDP FQN's, as well as sets the default state dict type to FSDP.SHARDED_STATE_DICT
        model_state_dict, optimizer_state_dict = get_state_dict(self.model, self.optimizer)
        return {
            "model": model_state_dict,
            "optim": optimizer_state_dict
        }

    def load_state_dict(self, state_dict):
        # sets our state dicts on the model and optimizer, now that we've loaded
        set_state_dict(
            self.model,
            self.optimizer,
            model_state_dict=state_dict["model"],
            optim_state_dict=state_dict["optim"]
        )

def setup():
    rank = int(os.environ["RANK"])
    world_size = int(os.environ["WORLD_SIZE"])
    # Shorter timeouts help surface failures quickly instead of hanging
    dist.init_process_group(
        backend="nccl",
        timeout=timedelta(seconds=120),
        init_method="env://",
        rank=rank,
        world_size=world_size,
    )
    torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", 0)))
    dist.barrier()
    if rank == 0:
        print("PG up; all ranks reached barrier")


def cleanup():
    try:
        dist.barrier()
    finally:
        dist.destroy_process_group()

Configurer les paramètres d'entraînement

La cellule suivante définit les hyperparamètres d'entraînement :

• batch_size: Nombre d'images traitées à chaque itération de formation
• num_epochs: Nombre de passages complets à travers le dataset d'entraînement
• momentum: Facteur de momentum pour l'optimiseur SGD
• log_interval: Fréquence de la journalisation de la progression de l'entraînement

Python

# Specify training parameters
batch_size = 100
num_epochs = 5
momentum = 0.5
log_interval = 100

Définir la boucle d'entraînement

La cellule suivante définit la fonction train_one_epoch, qui :

• Itère sur des batchs de données d'entraînement
• Effectue une propagation avant et arrière
• Met à jour les pondérations du modèle à l'aide de l'optimiseur
• Logs la perte d'entraînement vers MLflow à intervalles réguliers.

Python

def train_one_epoch(model, device, data_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(data_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % log_interval == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(data_loader) * len(data),
                100. * batch_idx / len(data_loader), loss.item()))
            # Log metrics
            mlflow.log_metric('loss', loss.item(), step=epoch * len(data_loader) + batch_idx)

Entraînez le modèle sur un seul GPU

La cellule suivante définit la fonction principale d'entraînement qui :

• Charge le dataset d'entraînement MNIST.
• Initialise le modèle et l'optimiseur.
• Entraîne le modèle pour le nombre d'époques spécifié.
• Enregistre les points de contrôle dans le volume Unity Catalog après chaque époque
• Logs metrics to MLflow pour le suivi d'expérimentation.

Python

import mlflow
import torch.optim as optim
from torchvision import datasets, transforms

def train(learning_rate):

  with mlflow.start_run() as run:
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

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

    model = Net().to(device)
    optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=momentum)
    with torch.no_grad():
      input_example, _ = next(iter(data_loader))
      output_example = model(input_example.to(device))

    for epoch in range(1, num_epochs + 1):
      train_one_epoch(model, device, data_loader, optimizer, epoch)

      state_dict = { "app": AppState(model, optimizer) }
      dcp.save(state_dict, checkpoint_id=CHECKPOINT_DIR)
      print(f"saved checkpoint to {CHECKPOINT_DIR}")

Exécutez la fonction d'entraînement

La cellule suivante exécute la fonction train avec un taux d'apprentissage de 0,001. Le processus de formation :

• Download le dataset MNIST (s'il n'est pas déjà mis en cache).
• Entraîner le modèle pour 5 époques
• Afficher la progression de l’entraînement et les valeurs de perte
• Enregistrez les points de contrôle de modèle dans le volume Unity Catalog
• Journalisez les métriques dans MLflow.

L'entraînement prend généralement quelques minutes sur un GPU A10G.

Python

train(learning_rate = 0.001)

Charger et évaluer le modèle entraîné

Après l'entraînement, vous pouvez charger le modèle à partir du point de contrôle et évaluer ses performances sur le dataset de test.

La cellule suivante définit une fonction test qui :

• Charge l'état du modèle depuis le point de contrôle du volume Unity Catalog.
• download le dataset de test MNIST
• Évalue le modèle sur les données de test.
• Calcule et affiche la perte moyenne de test

Python

def test():
  # Load model state from checkpoint using dcp
  model = Net()
  optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=momentum)
  app_state = AppState(model, optimizer)
  state_dict = { "app": app_state }
  dcp.load(state_dict, checkpoint_id=CHECKPOINT_DIR)
  model.eval()

  device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
  model.to(device)
  test_dataset = datasets.MNIST(
    'data',
    train=False,
    download=True,
    transform=transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]))
  data_loader = torch.utils.data.DataLoader(test_dataset)

  test_loss = 0
  for data, target in data_loader:
      data, target = data.to(device), target.to(device)
      output = model(data)
      test_loss += F.nll_loss(output, target)

  test_loss /= len(data_loader.dataset)
  print("Average test loss: {}".format(test_loss.item()))

Exécuter l’évaluation

La cellule suivante exécute la fonction test pour évaluer le modèle entraîné sur le test dataset MNIST. Une perte de test inférieure indique de meilleures performances du modèle.

Python

test()

Conclusion

Félicitations ! Vous avez entraîné avec succès un modèle de classification d'images en utilisant le compute GPU serverless. Vous avez appris à :

• Configurez et connectez-vous au compute GPU Serverless
• Définir une architecture de réseau de neurones convolutif
• Entraînez un modèle avec PyTorch et enregistrez les métriques dans MLflow
• Enregistrer les points de contrôle de modèle dans les volumes Unity Catalog
• Charger et évaluer un modèle entraîné

Se déconnecter du compute GPU

Pour éviter une utilisation inutile du GPU, déconnectez-vous manuellement de votre GPU :

1. Sélectionnez Connecté en haut du Notebook
2. Survolez Serverless
3. Sélectionnez Terminate dans le menu déroulant
4. Sélectionnez Confirmer pour terminer.

Remarque : Si vous ne vous déconnectez pas manuellement, votre connexion se termine automatiquement après 60 minutes d'inactivité.

Étapes suivantes

Explorez ces ressources pour en savoir plus sur le Machine Learning sur Databricks :

• Bonnes pratiques pour le compute GPU Serverless
• Dépanner les problèmes sur le compute GPU serverless
• Formation distribuée multi-GPU et multi-nœuds
• Suivi MLflow
• Entraîner des modèles avec PyTorch

Exemple de Notebook

Classification d'images à l'aide de réseaux de neurones convolutionnels

Ouvrir le notebook dans un nouvel ongletOuvrir dans Databricks