Classificação de imagens usando redes neurais convolucionais (CNN)

Open in Databricks

Ensine uma Rede Neural Convolucional (CNN) para classificação de imagem no Runtime de AI usando PyTorch e o dataset MNIST. O MNIST contém 70.000 imagens em escala de cinza de dígitos manuscritos (0-9), tornando-o ideal para aprender técnicas de classificação de imagem.

Você aprenderá como:

• Conecte seu notebook à compute GPU serverless com uma GPU A10G.
• Definir uma arquitetura simples de redes neurais convolucionais (CNN)
• Ensine o modelo em uma única GPU e log as métricas no MLflow
• Salvar pontos de verificação do modelo em um volume Unity Catalog
• Carregar e avaliar o modelo treinado

nota

Este exemplo requer a versão 5 ou superior do ambiente Databricks AI.

Conecte-se à computeGPU serverless

Este notebook requer uma GPU para treinar a rede neural de forma eficiente. Siga estes passos para se conectar à compute GPU serverless :

1. Clique no dropdown Conectar na parte superior do Bloco de Anotações.
2. Selecione GPU sem servidor .
3. Abra o painel lateral Ambiente , localizado no lado direito do Notebook.
4. Defina o **acelerador** como **1xA10** para esta demonstração.
5. Selecione AI v6 no dropdown Environment.
6. Selecione Aplicar e clique em Confirmar para aplicar este ambiente ao seu Notebook.

Para obter mais informações, consulte computeGPU sem servidor.

Configurar local de armazenamento do ponto de verificação

A célula a seguir cria parâmetros de widget para especificar onde os pontos de verificação do modelo serão salvos no Unity Catalog. Esses parâmetros definem:

• uc_catalog: O nome do Unity Catalog
• uc_schemaO esquema (banco de dados) dentro do catálogo
• uc_volume: O volume para armazenar arquivos de ponto de verificação
• uc_model_name: O subdiretório dentro do volume para este modelo específico

Esses valores são usados em todo o Notebook para construir o caminho do ponto de verificação: /Volumes/{uc_catalog}/{uc_schema}/{uc_volume}/{uc_model_name}

A célula a seguir usa valores de espaço reservado como padrão. Atualize os valores usando os widgets na parte superior do Notebook. Ou, atualize os valores default diretamente na próxima célula.

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")

Definir as redes neurais convolucionais (CNN)

A célula a seguir define uma arquitetura CNN simples para classificação de imagens. A rede é composta por:

• Convolucional de duas camadas com max pooling para extrair recursos de imagens.
• Duas camadas totalmente conectadas para classificar o recurso extraído.
• Camadas de dropout para evitar sobreajuste

O código também define classes auxiliares para salvar o estado do modelo e do otimizador em um Volume Unity Catalog , além de funções para configurar o treinamento distribuído (usado em cenários com múltiplas GPUs).

Esta implementação foi adaptada do exemplo MNIST do 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()

Configurar parâmetros de treinamento

A célula seguinte define os hiperparâmetros para o treinamento:

• batch_sizeNúmero de imagens processadas em cada iteração de treinamento
• num_epochsNúmero de passagens completas pelo datasetde treinamento
• momentumFator de momentum para o otimizador SGD
• log_intervalFrequência de registro do progresso do treinamento

Python

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

Defina o ciclo de treinamento

A célula seguinte define a função train_one_epoch , que:

• Itera por lotes de dados de treinamento.
• Realiza propagação direta e inversa.
• Atualiza os pesos do modelo usando o otimizador.
• Registros treinados de perda para MLflow em intervalos regulares

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)

ensinar o modelo em uma única GPU

A célula seguinte define a principal função de treinamento que:

• Carrega o datasetMNIST de treinamento.
• Inicializa o modelo e o otimizador.
• Ensinar o modelo para o número especificado de épocas.
• Salva pontos de verificação no Volume Unity Catalog após cada época.
• registra métricas no MLflow para acompanhamento do experimento

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}")

execução da função de treinamento

A célula seguinte executa a função train com uma taxa de aprendizagem de 0,001. O processo de treinamento irá:

• Baixe o dataset MNIST (caso ainda não esteja em cache).
• Ensine o modelo por 5 épocas
• Exibir valores de progresso e perda de treinamento.
• Salve os pontos de verificação do modelo no volume Unity Catalog
• logar métricas para MLflow

O treinamento normalmente leva alguns minutos em uma GPU A10G.

Python

train(learning_rate = 0.001)

Carregar e avaliar o modelo treinado

Após o treinamento, você pode carregar o modelo a partir do ponto de verificação e avaliar seu desempenho no dataset de teste.

A célula a seguir define uma função test que:

• Carrega o estado do modelo a partir do ponto de verificação do volume Unity Catalog
• baixa o datasetde teste MNIST
• Avalia o modelo em dados de teste.
• Calcula e exibe a perda média de teste.

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()))

execução da avaliação

A célula seguinte executa a função test para avaliar o modelo treinado no dataset de teste MNIST. Uma menor perda de teste indica um melhor desempenho do modelo.

Python

test()

Conclusão

Parabéns! Você treinou com sucesso um modelo de classificação de imagens usando compute em GPU serverless . Você aprendeu como:

• Configure e conecte-se à computeGPU serverless
• Definir uma arquitetura de redes neurais convolucionais (CNN)
• ensinando um modelo com PyTorch e log métricas no MLflow
• Salvar pontos de verificação do modelo em volumes Unity Catalog
• Carregar e avaliar um modelo treinado

Desconectar-se do computeda GPU

Para evitar o uso desnecessário da GPU, desconecte-se manualmente da sua GPU:

1. Selecione Conectado na parte superior do Notebook.
2. Passe o cursor sobre "sem servidor"
3. Selecione "Encerrar" no menu dropdown .
4. Selecione Confirmar para encerrar.

Observação : Se você não desconectar manualmente, sua conexão será encerrada automaticamente após 60 minutos de inatividade.

Próximos passos

Explore estes recursos para aprender mais sobre machine learning no Databricks:

• Melhores práticas para computede GPU serverless
• Solucionar problemas em computeGPU serverless
• Treinamento distribuído com múltiplas GPUs e múltiplos nós
• MLflow
• ensinar modelos com PyTorch

Exemplo de caderno

Classificação de imagens usando redes neurais convolucionais (CNN)

Abrir notebook em uma nova abaOpen in Databricks