Classificação de imagens usando redes neurais convolucionais (CNN)
Este Notebook demonstra como ensinar uma rede neural convolucional (CNN) para classificação de imagens usando o datasetMNIST e PyTorch. O dataset MNIST contém 70.000 imagens em tons de cinza de dígitos manuscritos (0-9), o que o torna ideal para o aprendizado de técnicas de classificação de imagens.
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
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 :
- Clique no dropdown Conectar na parte superior do Bloco de Anotações.
- Selecione GPU sem servidor .
- Abra o painel lateral Ambiente , localizado no lado direito do Notebook.
- Configure o acelerador para A10 para esta demonstração.
- Selecione Aplicar e clique em Confirmar para aplicar este ambiente ao seu Notebook.
Para obter mais informações, consulte computeGPU sem servidor.
Atualize o MLflow para a versão 3.0 ou superior.
MLflow 3.0 ou superior é recomendado para aprendizagem de fluxo de trabalho profundo. A célula a seguir atualiza o MLflow e reinicia o ambiente Python para aplicar as alterações.
Para obter mais informações, consulte Melhores práticas para MLflow 3.0+ aprendizagem profundo fluxo de trabalho.
%pip install -U mlflow>=3
%restart_python
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 Cataloguc_schemaO esquema (banco de dados) dentro do catálogouc_volume: O volume para armazenar arquivos de ponto de verificaçãouc_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.
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.
import torch
import torch.nn as nn
import torch.nn.functional as F
import os
import torch.distributed as dist
import torch.distributed.checkpoint as dcp
import torch.multiprocessing as mp
from datetime import timedelta
import os
from torch.distributed.fsdp import fully_shard
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 treinamentonum_epochsNúmero de passagens completas pelo datasetde treinamentomomentumFator de momentum para o otimizador SGDlog_intervalFrequência de registro do progresso do treinamento
# Specify training parameters
batch_size = 100
num_epochs = 3
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
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
import mlflow
import mlflow.pyfunc
import torch.optim as optim
from torchvision import datasets, transforms
from time import time
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).
- Ensinar o modelo por 3 é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.
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 seguinte 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.
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.
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:
- Selecione Conectado na parte superior do Notebook.
- Passe o cursor sobre "sem servidor"
- Selecione "Encerrar" no menu dropdown .
- 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