Function train_one_epoch

tools/test_linear_probing_hf.py:257–299

Train for one epoch.

(
    feature_model: nn.Module,
    linear_classifiers: AllClassifiers,
    optimizer: torch.optim.Optimizer,
    scheduler: torch.optim.lr_scheduler._LRScheduler,
    criterion: nn.Module,
    train_loader,
    epoch: int,
    epoch_length: int,
    device: torch.device,
)

def train_one_epoch(
    feature_model: nn.Module,
    linear_classifiers: AllClassifiers,
    optimizer: torch.optim.Optimizer,
    scheduler: torch.optim.lr_scheduler._LRScheduler,
    criterion: nn.Module,
    train_loader,
    epoch: int,
    epoch_length: int,
    device: torch.device,
) -> float:
    """Train for one epoch."""
    linear_classifiers.train()
    total_loss = 0.0
    num_batches = 0

    progress_bar = tqdm(train_loader, total=epoch_length, desc=f"Epoch {epoch}") if is_main_process() else train_loader

    for batch_idx, (images, labels) in enumerate(progress_bar):
        if batch_idx >= epoch_length:
            break

        images = images.to(device, non_blocking=True)
        labels = labels.to(device, non_blocking=True)

        features = feature_model(images)
        outputs = linear_classifiers(features)

        losses = {f"loss_{k}": criterion(v, labels) for k, v in outputs.items()}
        loss = sum(losses.values())

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        scheduler.step()

        total_loss += loss.item()
        num_batches += 1

        if is_main_process() and batch_idx % 50 == 0:
            progress_bar.set_postfix(loss=loss.item(), lr=optimizer.param_groups[0]["lr"])

    return total_loss / max(num_batches, 1)


@torch.no_grad()