Mempercepat dan mengurangi penggunaan memori tanpa mengorbankan akurasi model

Mempercepat dan mengurangi penggunaan memori tanpa mengorbankan akurasi model menggunakan GradScaler. torch.cuda.amp.GradScaler adalah kelas yang disediakan oleh PyTorch untuk memfasilitasi pelatihan model deep learning dengan precision campuran (mixed precision). Mixed precision adalah teknik yang menggunakan kombinasi angka floating point 16-bit (FP16) dan 32-bit (FP32) untuk mempercepat pelatihan dan mengurangi penggunaan memori tanpa mengorbankan akurasi model.

Berikut adalah penjelasan tentang beberapa komponen dan cara kerja torch.cuda.amp.GradScaler:

Fungsi dan Kegunaan

• Mempercepat Pelatihan: Dengan menggunakan FP16, komputasi dapat dilakukan lebih cepat karena angka-angka lebih kecil dan lebih sedikit memori yang digunakan.
• Mengurangi Penggunaan Memori: FP16 menggunakan setengah dari memori yang dibutuhkan oleh FP32, sehingga memungkinkan batch size yang lebih besar atau model yang lebih kompleks untuk muat dalam GPU.
• Mencegah Underflow: FP16 memiliki range nilai yang lebih kecil dibandingkan FP32, sehingga lebih rentan terhadap underflow (nilai yang sangat kecil menjadi nol). GradScaler membantu mencegah hal ini dengan secara dinamis menskalakan gradien selama pelatihan.

Cara Kerja

1. Scaling Gradien: GradScaler secara otomatis menskalakan loss sebelum backward pass dengan faktor scaling tertentu. Ini dilakukan untuk memastikan bahwa gradien cukup besar untuk direpresentasikan dalam FP16.
2. Backward Pass: Selama backward pass, gradien yang dihitung akan diskalakan balik dengan faktor scaling yang digunakan sebelumnya.
3. Update Parameter: GradScaler memeriksa apakah gradien mengandung NaN (Not a Number) atau inf (tak terhingga) setelah backward pass. Jika ada, maka step optimasi (update parameter) akan dilewati untuk iterasi tersebut dan faktor scaling akan diturunkan. Jika tidak ada NaN atau inf, step optimasi akan dijalankan dan faktor scaling mungkin dinaikkan untuk iterasi berikutnya.

GradScaler() ditemukan pada tutorial https://pytorch.org/tutorials/intermediate/torchvision_tutorial.html yaitu berada di https://raw.githubusercontent.com/pytorch/vision/main/references/detection/engine.py Berikut kode yang tertera pada engine.py di function train_one_epoch

import math
import sys
import time

import torch
import torchvision.models.detection.mask_rcnn
import utils
from coco_eval import CocoEvaluator
from coco_utils import get_coco_api_from_dataset


def train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq, scaler=None):
    model.train()
    metric_logger = utils.MetricLogger(delimiter="  ")
    metric_logger.add_meter("lr", utils.SmoothedValue(window_size=1, fmt="{value:.6f}"))
    header = f"Epoch: [{epoch}]"

    lr_scheduler = None
    if epoch == 0:
        warmup_factor = 1.0 / 1000
        warmup_iters = min(1000, len(data_loader) - 1)

        lr_scheduler = torch.optim.lr_scheduler.LinearLR(
            optimizer, 
            start_factor = warmup_factor, 
            total_iters = warmup_iters
        )

    for images, targets in metric_logger.log_every(data_loader, print_freq, header):
        images = list(image.to(device) for image in images)
        targets = [{k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in t.items()} for t in targets]
        with torch.cuda.amp.autocast(enabled=scaler is not None):
            loss_dict = model(images, targets)
            losses = sum(loss for loss in loss_dict.values())

        # reduce losses over all GPUs for logging purposes
        loss_dict_reduced = utils.reduce_dict(loss_dict)
        losses_reduced = sum(loss for loss in loss_dict_reduced.values())

        loss_value = losses_reduced.item()

        if not math.isfinite(loss_value):
            print(f"Loss is {loss_value}, stopping training")
            print(loss_dict_reduced)
            sys.exit(1)

        optimizer.zero_grad()
        if scaler is not None:
            scaler.scale(losses).backward()
            scaler.step(optimizer)
            scaler.update()
        else:
            losses.backward()
            optimizer.step()

        if lr_scheduler is not None:
            lr_scheduler.step()

        metric_logger.update(loss=losses_reduced, **loss_dict_reduced)
        metric_logger.update(lr=optimizer.param_groups[0]["lr"])

    return metric_logger


def _get_iou_types(model):
    model_without_ddp = model
    if isinstance(model, torch.nn.parallel.DistributedDataParallel):
        model_without_ddp = model.module
    iou_types = ["bbox"]
    if isinstance(model_without_ddp, torchvision.models.detection.MaskRCNN):
        iou_types.append("segm")
    if isinstance(model_without_ddp, torchvision.models.detection.KeypointRCNN):
        iou_types.append("keypoints")
    return iou_types


@torch.inference_mode()
def evaluate(model, data_loader, device):
    n_threads = torch.get_num_threads()
    # FIXME remove this and make paste_masks_in_image run on the GPU
    torch.set_num_threads(1)
    cpu_device = torch.device("cpu")
    model.eval()
    metric_logger = utils.MetricLogger(delimiter="  ")
    header = "Test:"

    coco = get_coco_api_from_dataset(data_loader.dataset)
    iou_types = _get_iou_types(model)
    coco_evaluator = CocoEvaluator(coco, iou_types)

    for images, targets in metric_logger.log_every(data_loader, 100, header):
        images = list(img.to(device) for img in images)

        if torch.cuda.is_available():
            torch.cuda.synchronize()
        model_time = time.time()
        outputs = model(images)

        outputs = [{k: v.to(cpu_device) for k, v in t.items()} for t in outputs]
        model_time = time.time() - model_time

        res = {target["image_id"]: output for target, output in zip(targets, outputs)}
        evaluator_time = time.time()
        coco_evaluator.update(res)
        evaluator_time = time.time() - evaluator_time
        metric_logger.update(model_time=model_time, evaluator_time=evaluator_time)

    # gather the stats from all processes
    metric_logger.synchronize_between_processes()
    print("Averaged stats:", metric_logger)
    coco_evaluator.synchronize_between_processes()

    # accumulate predictions from all images
    coco_evaluator.accumulate()
    coco_evaluator.summarize()
    torch.set_num_threads(n_threads)
    return coco_evaluator

Mari kita buat lebih simple saja

import torch
from torch.cuda.amp import GradScaler, autocast

# Inisialisasi model, loss function, optimizer, dan scaler
model = MyModel().cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
scaler = GradScaler()

for epoch in range(num_epochs):
    for inputs, targets in dataloader:
        inputs, targets = inputs.cuda(), targets.cuda()

        # Forward pass dengan autocast untuk mixed precision
        with autocast():
            outputs = model(inputs)
            loss = loss_function(outputs, targets)

        # Backward pass dengan GradScaler
        scaler.scale(loss).backward()

        # Update parameter
        scaler.step(optimizer)
        scaler.update()

        # Clear gradients
        optimizer.zero_grad()

dengan penjelasan sebagai berikut

Penjelasan Kode

1. Inisialisasi: model, optimizer, dan GradScaler diinisialisasi.
2. autocast: Digunakan untuk melakukan forward pass dengan mixed precision. Semua operasi di dalam blok ini akan dijalankan dengan FP16 jika memungkinkan.
3. scale: Loss diskalakan sebelum backward pass.
4. backward: Backward pass dijalankan pada loss yang telah diskalakan.
5. step: Optimizer melakukan update parameter.
6. update: GradScaler mengupdate faktor scaling untuk iterasi berikutnya.
7. zero_grad: Gradien dibersihkan untuk iterasi berikutnya.

Dengan menggunakan torch.cuda.amp.GradScaler, Anda dapat memanfaatkan keuntungan dari mixed precision training dengan cara yang aman dan efisien.