train_znet.py

from datetime import datetime
from pathlib import Path
import cv2

import numpy as np
import timm.utils
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import tqdm
import typer
from torch.utils.tensorboard import SummaryWriter

from effdet.data.dataset import ImagePoseDatset
from effdet.data.parsers import PoseMeParserCfg
from effdet.data.transforms import transforms_z_train, transforms_z_val
from effdet.znet import ZRegressionNet

app = typer.Typer(no_args_is_help=True)


def weighted_z_loss(input, target):
    expanded_input, expanded_target = torch.broadcast_tensors(input, target)
    loss = torch.pow(expanded_input - expanded_target, 2)
    # weight the loss by distance to mean z
    # target is already normalized
    weights = torch.pow(torch.abs(target) + 1, 2)
    weighted_loss = loss * weights
    return weighted_loss.mean()


@app.command()
def train(
    data_dir: Path = typer.Argument(
        Path("../../data/fleckenzwerg_dataset_imagepose_cleaned"),
        help="Path to data directory",
    ),
    batch_size: int = typer.Option(48, help="Batch size"),
    num_epochs: int = typer.Option(200, help="Number of epochs"),
    lr: float = typer.Option(1e-3, help="Learning rate"),
    min_lr: float = typer.Option(1e-5, help="Minimum learning rate"),
    num_workers: int = typer.Option(4, help="Number of workers"),
    output_dir: Path = typer.Option("output", help="Output directory"),
):
    model = ZRegressionNet().to("cuda")
    optimizer = torch.optim.SGD(
        model.parameters(), lr=lr, momentum=0.9, weight_decay=1e-4
    )
    lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
        optimizer=optimizer,
        T_max=num_epochs,
        eta_min=min_lr,
        last_epoch=-1,
    )
    parser_train_cfg = PoseMeParserCfg(
        data_dir=data_dir / "train", has_labels=True, extension=".png"
    )
    parser_val_cfg = PoseMeParserCfg(
        data_dir=data_dir / "val", has_labels=True, extension=".png"
    )
    dataset_train = ImagePoseDatset(
        data_dir=data_dir / "train",
        parser="poseme",
        transform=transforms_z_train(),
        parser_kwargs={"cfg": parser_train_cfg},
    )
    dataset_val = ImagePoseDatset(
        data_dir=data_dir / "val",
        parser="poseme",
        transform=transforms_z_val(),
        parser_kwargs={"cfg": parser_val_cfg},
    )
    loader_train = torch.utils.data.DataLoader(
        dataset_train,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        persistent_workers=True if num_workers > 0 else False,
    )
    loader_val = torch.utils.data.DataLoader(
        dataset_val,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        persistent_workers=True if num_workers > 0 else False,
    )
    # loss_fn = F.mse_loss
    loss_fn = weighted_z_loss
    metric_fn = F.l1_loss

    exp_name = f"znet-{datetime.now().strftime('%Y%m%d-%H%M%S')}"
    run_dir = timm.utils.get_outdir(output_dir, "train_znet", exp_name)

    best_metric = None
    writer = SummaryWriter(log_dir=run_dir)
    saver = timm.utils.CheckpointSaver(
        model=model,
        optimizer=optimizer,
        max_history=5,
        decreasing=True,
        checkpoint_dir=run_dir,
    )

    try:
        for epoch in tqdm.trange(num_epochs, desc="Train Epoch", position=0):
            train_loss = train_epoch(
                epoch,
                model,
                loader_train,
                loss_fn,
                metric_fn,
                optimizer,
                saver,
                writer,
            )

            eval_loss = validate_epoch(model, loader_val, metric_fn)

            lr_scheduler.step()

            lr = optimizer.param_groups[0]["lr"]
            writer.add_scalar("lr/train/epochs", lr, epoch)
            writer.add_scalar("loss/train/epochs", train_loss, epoch)
            writer.add_scalar("loss/eval/epochs", eval_loss, epoch)

            # save proper checkpoint with eval metric
            best_metric, best_epoch = saver.save_checkpoint(
                epoch=epoch, metric=eval_loss
            )
    finally:
        writer.close()
        print(f"Best metric: {best_metric:.4f} at epoch {best_epoch}")


def train_epoch(
    epoch,
    model,
    loader,
    loss_fn,
    metric_fn,
    optimizer,
    saver=None,
    writer=None,
):
    losses_m = timm.utils.AverageMeter()
    metric_m = timm.utils.AverageMeter()
    z_error_m = timm.utils.AverageMeter()
    max_z_error = 0

    last_idx = len(loader) - 1
    num_updates = epoch * len(loader)

    z_mean = 0.460 * 1000
    z_std = 0.050 * 1000

    tqdmloader = tqdm.tqdm(loader, desc="Train Batch", position=1)

    model.train()
    for batch_idx, (input, target) in enumerate(tqdmloader):
        last_batch = batch_idx == last_idx

        # image_grid = torchvision.utils.make_grid(input, nrow=8, normalize=True, scale_each=True)
        # cv2.imshow("train_grid", image_grid.numpy().transpose(1, 2, 0))
        # cv2.waitKey(1)

        input = input.type(torch.float32).cuda()
        target_z = target["translation"][:, 0, [2]].type(torch.float32).cuda()

        output_z = model(input)
        loss = loss_fn(output_z, target_z)
        losses_m.update(loss.item(), input.size(0))

        metric = metric_fn(output_z, target_z)
        metric_m.update(metric.item(), input.size(0))

        denormalized_output = output_z.detach().cpu().numpy() * z_std + z_mean
        denormalized_target = target_z.detach().cpu().numpy() * z_std + z_mean
        z_error = np.abs(denormalized_output - denormalized_target)

        z_error_m.update(z_error.mean(), input.size(0))
        max_z_error = np.max([max_z_error, z_error.max()])

        tqdmloader.set_postfix(
            loss=losses_m.avg, metric=metric_m.avg, z_error=z_error_m.avg, max_z_error=max_z_error
        )

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

        num_updates += 1

        if writer:
            lr = optimizer.param_groups[0]["lr"]
            writer.add_scalar("loss/train/minibatches", losses_m.val, num_updates)
            writer.add_scalar("lr/train/minibatches", lr, num_updates)

        if saver is not None and last_batch:
            saver.save_recovery(epoch, batch_idx=batch_idx)

    return losses_m.avg


def validate_epoch(
    model,
    loader,
    metric_fn,
):
    metric_m = timm.utils.AverageMeter()
    z_error_m = timm.utils.AverageMeter()
    max_z_error = 0

    z_mean = 0.460 * 1000
    z_std = 0.050 * 1000

    tqdmloader = tqdm.tqdm(loader, desc="Val Batch", position=2, postfix={"z_error": 0})

    model.eval()
    with torch.no_grad():
        for input, target in tqdmloader:
            # image_grid = torchvision.utils.make_grid(input, nrow=8, normalize=True, scale_each=True)
            # cv2.imshow("val_grid", image_grid.numpy().transpose(1, 2, 0))
            # cv2.waitKey(1)

            input = input.type(torch.float32).cuda()
            target_z = target["translation"][:, 0, [2]].type(torch.float32).cuda()

            output_z = model(input)
            metric = metric_fn(output_z, target_z)
            metric_m.update(metric.item(), input.size(0))

            denormalized_output = output_z.detach().cpu().numpy() * z_std + z_mean
            denormalized_target = target_z.detach().cpu().numpy() * z_std + z_mean
            z_error = np.abs(denormalized_output - denormalized_target)
            z_error_m.update(z_error.mean(), input.size(0))

            max_z_error = np.max([max_z_error, z_error.max()])

            tqdmloader.set_postfix(
                metric=metric_m.avg, z_error=z_error_m.avg, max_z_error=max_z_error
            )

    return metric_m.avg


if __name__ == "__main__":
    app()