Träna en RetinaNet-bildidentifieringsmodell

Träna en RetinaNet-objektidentifieringsmodell från grunden på AI Runtime med PyTorch och torchvision. RetinaNet är en objektidentifieringsmodell med ett steg som använder ett funktionspyramidnätverk (FPN) och fokuserad förlust för att hantera obalans mellan klasser.

Anteckningsboken omfattar:

• Läsa in och transformera COCO-datamängden för objektidentifiering
• Träna en RetinaNet-modell med ResNet-50-stamnät på en enda GPU
• Skala träning över flera GPU:er med DDP (Distributed Data Parallel)
• Logga träningsmått med MLflow

Ansluta till serverlös GPU-beräkning

För att köra den här anteckningsboken ansluter du till serverlösa GPU-beräkningsresurser med 1xA10 för träning med en GPU eller 8xH100 för den distribuerade delen.

1. Klicka på notebook-filens beräkningsväljare längst upp till höger och välj Serverlös GPU.
2. Klicka på miljöknappen till höger.
3. Välj 1xA10 eller 8xH100 som accelerator.
4. Välj AI v5 som miljö och klicka sedan på Använd.

Installera de paket som krävs

Installera pycocotools för COCO-datauppsättningsverktyg och starta om Python-miljön för att läsa in det nya paketet.

%pip install pycocotools
dbutils.library.restartPython()

Konfigurera Unity-katalogsökvägar med widgetar

Definiera widgetar för att ange katalogen, schemat och volymen för Unity-katalogen där COCO-datauppsättningen lagras.

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

Importera PyTorch-bibliotek

Importera torch och torchvision för att bygga och träna modellen för bilddetektering.

import torch
import torchvision

Importera modell- och datamängdsklasser

Importera RetinaNet-modellarkitekturen och COCO-datauppsättningsverktygen från 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

Importera distribuerade träningsverktyg

Importera de distribuerade träningsmodulerna i PyTorch och den serverlösa distribuerade GPU-dekoreraren för träning över flera GPU:er.

Definiera träningshyperparametrar och datasökvägar

Konfigurera datasökvägar, batchstorlek, antal klasser, utbildningshastighet och andra träningsparametrar. Justera BATCH_SIZE och NUM_EPOCHS baserat på GPU-typen och träningskraven.

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

Initiera RetinaNet-modellen

Skapa en RetinaNet-modell med ResNet-50-ryggrad utan förträningsvikter, konfigurerad för antalet klasser i COCO-datasetet.

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

Transformera bilder och anteckningar till modellindata

Modellen kräver indata som tensorer med form (C, H, W), float32-datatyp och normaliserat intervall (0,0 till 1,0). Funktionen get_transform konverterar PIL-bilder och tillämpar dataförstoring. Klassen CocoWrapper omsluter COCO-datamängden för att formatera avgränsningsrutor och etiketter korrekt.

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

Skapa datainläsaren

Definiera en DataLoader med anpassad sortering för batchbilder och mål för träning.

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
)

Konfigurera optimeringen

Konfigurera SGD-optimeraren med parametrar för inlärningshastighet, momentum och viktförfall.

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
)

Träna modellen på en enda GPU

Kör träningsloopen för det angivna antalet epoker och logga förlustmått till 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.")

Träna med distribuerad dataparallellism (DDP)

Skala träning över flera GPU:er med hjälp av dekoratören @distributed . Den här metoden läser datauppsättningen direkt från Unity Catalog-volymen och använder DistributedSampler för att partitionera datamängden mellan GPU:er. Alla importer, transformeringar, datamängden och DataLoader omdefinieras i träningsfunktionen, vilket krävs av den distribuerade körningsmodellen.

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

Nästa steg

• Metodtips för serverlös GPU-beräkning
• Felsöka problem med serverlös GPU-beräkning
• Distribuerad träning med flera GPU:n och flera noder
• Träna modeller med PyTorch
• Loggmodeller med MLflow

Exempelanteckningsbok

Träna en RetinaNet-bildidentifieringsmodell

Hämta anteckningsbok