Trenowanie modelu RetinaNet do wykrywania obrazów

Otwórz wersję notebooka tej strony

Wytrenuj od podstaw model wykrywania obiektów RetinaNet w AI Runtime za pomocą PyTorch i torchvision. RetinaNet jest jednoetapowym modelem wykrywania obiektów, który używa sieci piramidy cech (FPN) i straty ogniskowej w celu zarządzania nierównowagą klas.

Zeszyt obejmuje:

Ładowanie i przekształcanie zestawu danych COCO na potrzeby wykrywania obiektów
Trenowanie modelu RetinaNet z backendem ResNet-50 na jednym GPU
Skalowanie trenowania na wielu procesorach GPU przy użyciu rozproszonego równoległego przetwarzania danych (DDP)
Rejestrowanie metryk trenowania za pomocą biblioteki MLflow

Note

Ten przykład wymaga środowiska AI Databricks w wersji 5 lub wyższej.

Nawiązywanie połączenia z bezserwerowym przetwarzaniem GPU

Aby uruchomić ten notatnik, połącz się z bezserwerowym środowiskiem obliczeniowym GPU z konfiguracją 1xA10 do trenowania na pojedynczym GPU lub 8xH100 do sekcji treningu rozproszonego.

Kliknij selektor obliczeniowy notesu w prawym górnym rogu i wybierz pozycję Bezserwerowy procesor GPU.
Po prawej stronie kliknij przycisk środowiska.
Wybierz 1xA10 lub 8xH100 jako Akcelerator.
Wybierz AI v6 jako środowisko, a następnie kliknij Apply.

Instalowanie wymaganych pakietów

Zainstaluj narzędzia pycocotools dla narzędzi zestawu danych COCO i uruchom ponownie środowisko języka Python, aby załadować nowy pakiet.

%pip install pycocotools
dbutils.library.restartPython()

Konfigurowanie ścieżek Unity Catalog za pomocą widżetów

Zdefiniuj widżety, aby określić katalog Unity Catalog, schemat i wolumin, gdzie przechowywany jest zestaw danych COCO.

dbutils.widgets.text("uc_catalog", "main")
dbutils.widgets.text("uc_schema", "default")
dbutils.widgets.text("uc_volume", "coco_data")

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

print(f"UC_CATALOG: {UC_CATALOG}")
print(f"UC_SCHEMA: {UC_SCHEMA}")
print(f"UC_VOLUME: {UC_VOLUME}")

Importowanie bibliotek PyTorch

Zaimportuj torch i torchvision na potrzeby tworzenia i trenowania modelu wykrywania obrazów.

import torch
import torchvision

Importowanie klas modelu i zestawu danych

Zaimportuj architekturę modelu RetinaNet i narzędzia zestawu danych COCO z torchvision.

import os
from torchvision.models.detection import retinanet_resnet50_fpn_v2

# For this example we will be using a default Dataset from torch
from torchvision.datasets import CocoDetection

Importowanie rozproszonych narzędzi szkoleniowych

Zaimportuj rozproszone moduły szkoleniowe PyTorch i bezserwerowy dekorator rozproszonych procesorów GPU na potrzeby trenowania z wieloma procesorami GPU.

Zdefiniuj hiperparametry trenowania i ścieżki danych

Skonfiguruj ścieżki danych, rozmiar partii, liczbę klas, szybkość nauki i inne parametry trenowania. Dostosuj BATCH_SIZE i NUM_EPOCHS na podstawie typu GPU i wymagań dotyczących trenowania.

import torch.distributed as dist
from serverless_gpu import distributed

DATA_PATH = f"/Volumes/{UC_CATALOG}/{UC_SCHEMA}/{UC_VOLUME}/"
TRAIN_IMG_PATH = os.path.join(DATA_PATH, "val2017")
TRAIN_ANN_PATH = os.path.join(DATA_PATH, "annotations", "instances_val2017.json")

BATCH_SIZE = 2 # Please use batch size of 8 with H100 for best performance
NUM_CLASSES = 91
LEARNING_RATE = 0.005
MOMENTUM = 0.9
WEIGHT_DECAY = 0.0005
NUM_EPOCHS = 1 # Update num_epochs accordingly

Zainicjuj model RetinaNet

Utwórz model RetinaNet z siecią szkieletową ResNet-50 bez wstępnie wytrenowanych wag, skonfigurowany na liczbę klas w zbiorze danych COCO.

# Since we are training the model from scratch, we need to initialize weights to None
model = retinanet_resnet50_fpn_v2(weights=None, num_classes=NUM_CLASSES)

Przekształcanie obrazów i adnotacji w dane wejściowe modelu

Model wymaga danych wejściowych jako tensorów z kształtem (C, H, W), typem danych float32 i znormalizowanymi zakresami (od 0,0 do 1,0). Funkcja get_transform konwertuje obrazy PIL i stosuje rozszerzenie danych. Klasa CocoWrapper opakowuje zestaw danych COCO w celu poprawnego formatowania pól ograniczenia i etykiet.

from torchvision.transforms import v2
from torchvision import tv_tensors

def get_transform(train):
  transforms = []
  transforms.append(v2.ToImage())
  transforms.append(v2.ToDtype(torch.float32, scale=True))
  if train:
    transforms.append(v2.RandomHorizontalFlip())
  return v2.Compose(transforms)

class CocoWrapper(CocoDetection):
  def __init__(self, root, annFile, transforms=None):
    super().__init__(root, annFile)
    self._transforms = transforms

  def __getitem__(self, idx):
    img, target = super().__getitem__(idx)
    image_id = self.ids[idx]

    boxes = []
    labels = []

    for obj in target:
      x, y, w, h = obj["bbox"]

      boxes.append([x, y, x + w, y + h])
      labels.append(obj["category_id"])

    if len(boxes) == 0:
      boxes = torch.zeros((0, 4), dtype=torch.float32)
      labels = torch.zeros((0,), dtype=torch.int64)
    else:
      boxes = torch.as_tensor(boxes, dtype=torch.float32)
      labels = torch.as_tensor(labels, dtype=torch.int64)

    w, h = img.size
    boxes = torchvision.tv_tensors.BoundingBoxes(
      data=boxes,
      format=torchvision.tv_tensors.BoundingBoxFormat.XYXY,
      canvas_size=(h, w)
    )

    final_target = {
      "boxes": boxes,
      "labels": labels,
      "image_id": torch.tensor([image_id])
    }

    if self._transforms is not None:
      img, final_target = self._transforms(img, final_target)

    return img, final_target

dataset = CocoWrapper(
  root = TRAIN_IMG_PATH,
  annFile=TRAIN_ANN_PATH,
  transforms=get_transform(train=True)
)

# Sanity Check
img, target = dataset[0]

print("Image type:", type(img))
print("Image shape:", img.shape)        # should be [3, H, W]
print("Image dtype:", img.dtype)

print("\nTarget keys:", target.keys())
print("Boxes shape:", target["boxes"].shape)
print("Labels shape:", target["labels"].shape)
print("Image ID:", target["image_id"])

Tworzenie modułu ładującego dane

Zdefiniuj obiekt DataLoader z niestandardowym łączeniem do grupowania obrazów i celów do trenowania.

from torch.utils.data import DataLoader

def collate_fn(batch):
    images, targets = list(zip(*batch))
    return list(images), list(targets)

train_loader = DataLoader(
    dataset,
    batch_size=BATCH_SIZE,
    shuffle=True,
    num_workers=16,
    collate_fn=collate_fn,
    pin_memory=True,
    prefetch_factor=2 # Please use a prefetch_factor of 4 with H100 for best performance
)

Konfigurowanie optymalizatora

Skonfiguruj optymalizator SGD z parametrami tempa uczenia, tempa i rozkładu masy.

device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
model.to(device)

params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(
    params,
    lr=LEARNING_RATE,
    momentum=MOMENTUM,
    weight_decay=WEIGHT_DECAY
)

Trenowanie modelu na jednym procesorze GPU

Uruchom pętlę trenowania dla określonej liczby epok, rejestruj metryki straty w systemie MLflow.

import time
import mlflow

model.train()

lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)

with mlflow.start_run():
    for epoch in range(NUM_EPOCHS):
        start_time = time.time()
        epoch_loss = 0

        for i, (images, targets) in enumerate(train_loader):

            images = list(image.to(device) for image in images)
            targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
            loss_dict = model(images, targets)
            losses = sum(loss for loss in loss_dict.values())
            optimizer.zero_grad()
            losses.backward()
            optimizer.step()

            epoch_loss += losses.item()
            mlflow.log_metric("loss", losses.item(), step=epoch * len(train_loader) + i)

            if i % 50 == 0:
                print(f"Epoch {epoch+1} | Step {i}/{len(train_loader)} | Loss: {losses.item():.4f}")

        lr_scheduler.step()

        end_time = time.time()
        avg_loss = epoch_loss / len(train_loader)
        mlflow.log_metric("epoch_avg_loss", avg_loss, step=epoch)
        print(f"Epoch {epoch+1} Finished! Avg Loss: {avg_loss:.4f} | Time: {(end_time - start_time)/60:.2f} min")

print("Training Complete.")

Trenuj z równoległym przetwarzaniem danych rozproszonych (DDP)

Skaluj trenowanie na wielu procesorach GPU przy użyciu dekoratora @distributed. Takie podejście odczytuje zestaw danych bezpośrednio z woluminu Unity Catalog i używa narzędzia DistributedSampler do podziału zestawu danych pomiędzy procesory GPU. Wszystkie importy, przekształcenia, zestaw danych i moduł DataLoader są ponownie definiowane wewnątrz funkcji trenowania, zgodnie z wymaganiami modelu rozproszonego wykonywania.

from datetime import timedelta

BATCH_SIZE_PER_GPU = 8  # for better performance with H100

@distributed(gpus=8, gpu_type='H100')
def train_distributed():
    import os
    import torch
    import torch.distributed as dist
    import time
    import torchvision
    import mlflow
    from torch.nn.parallel import DistributedDataParallel as DDP
    from torch.utils.data import DataLoader
    from torch.utils.data.distributed import DistributedSampler
    from torchvision.models.detection import retinanet_resnet50_fpn_v2
    from torchvision.transforms import v2
    from torchvision import tv_tensors
    from torchvision.datasets import CocoDetection

    def get_transform(train):
        transforms = []
        transforms.append(v2.ToImage())
        transforms.append(v2.ToDtype(torch.float32, scale=True))
        if train:
            transforms.append(v2.RandomHorizontalFlip())
        return v2.Compose(transforms)

    def collate_fn(batch):
        images, targets = list(zip(*batch))
        return list(images), list(targets)

    class CocoWrapper(CocoDetection):
        def __init__(self, root, annFile, transforms=None):
            super().__init__(root, annFile)
            self._transforms = transforms

        def __getitem__(self, idx):
            img, target = super().__getitem__(idx)
            image_id = self.ids[idx]

            boxes = []
            labels = []

            for obj in target:
                x, y, w, h = obj["bbox"]

                boxes.append([x, y, x + w, y + h])
                labels.append(obj["category_id"])

            if len(boxes) == 0:
                boxes = torch.zeros((0, 4), dtype=torch.float32)
                labels = torch.zeros((0,), dtype=torch.int64)
            else:
                boxes = torch.as_tensor(boxes, dtype=torch.float32)
                labels = torch.as_tensor(labels, dtype=torch.int64)

            w, h = img.size
            boxes = torchvision.tv_tensors.BoundingBoxes(
                data=boxes,
                format=torchvision.tv_tensors.BoundingBoxFormat.XYXY,
                canvas_size=(h, w)
            )

            final_target = {
                "boxes": boxes,
                "labels": labels,
                "image_id": torch.tensor([image_id])
            }

            if self._transforms is not None:
                img, final_target = self._transforms(img, final_target)

            return img, final_target

    dist.init_process_group(backend="nccl", timeout=timedelta(minutes=30))

    rank = int(os.environ["RANK"])
    local_rank = int(os.environ["LOCAL_RANK"])
    world_size = int(os.environ["WORLD_SIZE"])

    torch.cuda.set_device(local_rank)
    device = torch.device(f"cuda:{local_rank}")

    # Read the dataset directly from the Unity Catalog volume on every rank.
    train_img_path = os.path.join(DATA_PATH, "val2017")
    train_ann_path = os.path.join(DATA_PATH, "annotations", "instances_val2017.json")

    dataset = CocoWrapper(
        root=train_img_path,
        annFile=train_ann_path,
        transforms=get_transform(train=True)
    )

    train_sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True)

    train_loader = DataLoader(
        dataset,
        batch_size=BATCH_SIZE_PER_GPU,
        shuffle=False,
        num_workers=8, # 8 workers * 8 GPUs = 64 total CPU threads
        collate_fn=collate_fn,
        pin_memory=True,
        prefetch_factor=4,
        sampler=train_sampler
    )

    model = retinanet_resnet50_fpn_v2(weights=None, num_classes=NUM_CLASSES)
    model.to(device)

    model = DDP(model, device_ids=[local_rank])

    params = [p for p in model.parameters() if p.requires_grad]
    optimizer = torch.optim.SGD(params, lr=0.04, momentum=0.9, weight_decay=0.0005)
    lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)

    model.train()
    if rank == 0:
        print(f"Training on {world_size} GPUs. Global Batch Size: {BATCH_SIZE_PER_GPU * world_size}")

    with mlflow.start_run():
        for epoch in range(NUM_EPOCHS):
            train_sampler.set_epoch(epoch)

            start_time = time.time()
            epoch_loss = 0

            for i, (images, targets) in enumerate(train_loader):
                images = [image.to(device) for image in images]
                targets = [{k: v.to(device) for k, v in t.items()} for t in targets]

                loss_dict = model(images, targets)
                losses = sum(loss for loss in loss_dict.values())

                optimizer.zero_grad()
                losses.backward()
                optimizer.step()

                epoch_loss += losses.item()

                if rank == 0:
                    mlflow.log_metric("loss", losses.item(), step=epoch * len(train_loader) + i)

                if rank == 0 and i % 50 == 0:
                    print(f"Rank 0 | Step {i}/{len(train_loader)} | Loss: {losses.item():.4f}")

            lr_scheduler.step()

            if rank == 0:
                avg_loss = epoch_loss / len(train_loader)
                mlflow.log_metric("epoch_avg_loss", avg_loss, step=epoch)
                print(f"Epoch {epoch+1} Finished! Avg Loss: {avg_loss:.4f} | Time: {(time.time() - start_time)/60:.2f} min")

    dist.destroy_process_group()

train_distributed.distributed()

Następne kroki

Przykładowy notatnik

Trenowanie modelu RetinaNet do wykrywania obrazów

Pobierz laptopa

Opinia

Czy ta strona była pomocna?

Ostatnia aktualizacja o 2026-09-11

Trenowanie modelu RetinaNet do wykrywania obrazów

Nawiązywanie połączenia z bezserwerowym przetwarzaniem GPU

Instalowanie wymaganych pakietów

Konfigurowanie ścieżek Unity Catalog za pomocą widżetów

Importowanie bibliotek PyTorch

Importowanie klas modelu i zestawu danych

Importowanie rozproszonych narzędzi szkoleniowych

Zdefiniuj hiperparametry trenowania i ścieżki danych

Zainicjuj model RetinaNet

Przekształcanie obrazów i adnotacji w dane wejściowe modelu

Tworzenie modułu ładującego dane

Konfigurowanie optymalizatora

Trenowanie modelu na jednym procesorze GPU

Trenuj z równoległym przetwarzaniem danych rozproszonych (DDP)

Następne kroki

Przykładowy notatnik

Trenowanie modelu RetinaNet do wykrywania obrazów

Opinia

Dodatkowe zasoby