finetune.1.py

# %%
import os

# limit resource usage
os.environ["OMP_NUM_THREADS"] = "4"  # export OMP_NUM_THREADS=4
os.environ["OPENBLAS_NUM_THREADS"] = "4"  # export OPENBLAS_NUM_THREADS=4
os.environ["MKL_NUM_THREADS"] = "4"  # export MKL_NUM_THREADS=6
os.environ["VECLIB_MAXIMUM_THREADS"] = "4"  # export VECLIB_MAXIMUM_THREADS=4
os.environ["NUMEXPR_NUM_THREADS"] = "4"  # export NUMEXPR_NUM_THREADS=6
os.environ["GDAL_NUM_THREADS"] = "4"

from pathlib import Path
import random
import sys

from accelerate import Accelerator, load_checkpoint_in_model
from accelerate.utils import ProjectConfiguration, set_seed
from addict import Dict
import einops as ein
import evaluate
import numpy as np
from rich.progress import track
import tensorwatch as tw
import torch
from torchmetrics import Accuracy
import torchview as tv
from tqdm import tqdm

from DeepHyperX.models import get_model
from src.utils import get_finetune_config, get_supervised_data
from src.vit_original import ViTRGB
from src.vit_spatial_spectral import ViTSpatialSpectral

SEED = 5
set_seed(SEED)


# %%
def load_checkpoint(config, model, classifier_name):
    print("Initializing pre-trained weights...")
    checkpoint = torch.load(config.checkpoint_path)

    encoder_weights = checkpoint["model_state_dict"]
    for k in list(encoder_weights.keys()):
        encoder_weights[k.replace("encoder.", "")] = encoder_weights[k]

        # delete old keys and those that are not part of the encoder
        del encoder_weights[k]

    # wrong output shape in pre-trained mpl_head
    if model.pixelwise:
        w, b = model.mlp_head[2].weight, model.mlp_head[2].bias
        linear_idx = 2
    else:
        w, b = model.mlp_head[1].weight, model.mlp_head[1].bias
        linear_idx = 1

    if (
        config.patch_sub != 0
        and isinstance(model, ViTSpatialSpectral)
        and "pos_embed" in encoder_weights
    ):
        # pre_trained with different image_size
        print(f"{encoder_weights['pos_embed'].shape=}")
        assert model.pos_embed.shape[1] == (config.image_size - config.patch_sub) ** 2
        encoder_weights["pos_embed"] = encoder_weights["pos_embed"][
            :, : model.pos_embed.shape[1], :
        ]
        print(f"{encoder_weights['pos_embed'].shape=}")

    del encoder_weights[f"{classifier_name}.1.bias"]
    del encoder_weights[f"{classifier_name}.1.weight"]
    encoder_weights[f"{classifier_name}.{linear_idx}.bias"] = b
    encoder_weights[f"{classifier_name}.{linear_idx}.weight"] = w
    print(model.load_state_dict(encoder_weights))

    return model


# %%
def get_val_epochs(config, dataloader):
    """fix the number of validation runs
    training will last for `epochs` or `max_steps`, whatever takes longer
    for small data_fraction and fixed batch_size, epochs will be very short
    and and validation time will dominate"""

    steps_per_epoch = len(dataloader)
    total_steps = steps_per_epoch * config.epoch
    if total_steps > config.max_steps:
        # max epochs is reached first, eval after each epoch
        validation_epochs = torch.arange(config.epoch)
    else:
        # run will stop when max_steps is reached, still only eval `epoch` many times
        total_epochs = config.max_steps // steps_per_epoch
        validation_epochs = list(map(int, np.linspace(0, total_epochs, config.epoch)))

    return validation_epochs


## %%
def stack_image_batch(config, img, label):
    """tile image into multiple image_size,image_size patches
    and stack along batch dimension"""
    cutoff_h = img.shape[2] % (config.image_size - config.patch_sub)
    cutoff_w = img.shape[3] % (config.image_size - config.patch_sub)
    assert cutoff_h == cutoff_w
    if cutoff_h != 0:
        # remove border pixels s.t. image is divisible by patch size
        img = img[:, :, :-cutoff_h, :-cutoff_w]
        label = label[:, :-cutoff_h, :-cutoff_w]
    img = ein.rearrange(
        img,
        "b c (h p1) (w p2) -> (b h w) c p1 p2",
        p1=config.image_size - config.patch_sub,
        p2=config.image_size - config.patch_sub,
    )
    label = ein.rearrange(
        label,
        "b (h p1) (w p2) -> (b h w) p1 p2",
        p1=config.image_size - config.patch_sub,
        p2=config.image_size - config.patch_sub,
    )

    return img, label


## %%
def train_step(
    img,
    label,
    model,
    config,
    criterion,
    optimizer,
    acc_criterion,
    accelerator: Accelerator,
):
    if config.image_size != 64 and config.dataset in ["dfc", "worldcover"]:
        if config.shifting_window:
            # divide image into non-overlapping patches and stack them
            img, label = stack_image_batch(config, img, label)
        else:
            # train with one smaller random crop
            x, y = torch.randint(
                0, 64 - config.image_size - config.patch_sub, size=(2,)
            )
            img = img[
                :,
                :,
                x : x + config.image_size - config.patch_sub,
                y : y + config.image_size - config.patch_sub,
            ]
            label = label[
                :,
                x : x + config.image_size - config.patch_sub,
                y : y + config.image_size - config.patch_sub,
            ]

    if config.method_name == "li" or config.pixelwise:
        # baseline model only predicts class for the center pixel of the patch
        center_idx = (config.image_size - config.patch_sub) // 2
        if config.dataset in ["dfc", "worldcover"]:
            label = label[:, center_idx, center_idx]  # .unsqueeze(1).unsqueeze(1)

        # extra dim for 3D conv model
        if config.method_name == "li":
            img = img.unsqueeze(1)

    optimizer.zero_grad()

    output = model(img)
    loss = criterion(output, label)

    if torch.isnan(loss):
        ValueError("Loss is NaN")

    pred = output.argmax(dim=1)

    accelerator.backward(loss)

    optimizer.step()

    valid_idx = label != config.ignored_label
    pred, label, loss, valid_idx = accelerator.gather_for_metrics(
        (pred, label, loss, valid_idx)
    )

    acc = (pred[valid_idx] == label[valid_idx]).sum() / pred[valid_idx].numel()
    macro_acc = (
        acc_criterion(pred[valid_idx].to(int), label[valid_idx])
        if valid_idx.sum() != 0
        else acc
    )

    return loss.mean(), acc.mean(), macro_acc.mean()


# %%
def validate(
    config,
    epoch,
    model,
    val_dataloader,
    criterion,
    acc_criterion,
    step,
    best_val_acc,
    accelerator: Accelerator,
    pixelwise=False,
):
    with torch.no_grad():
        val_losses = []
        val_accs = []
        val_macro_accs = []
        model.eval()
        val_pbar = track(
            val_dataloader,
            description=f"Validation {epoch}/{config.epoch}",
            transient=True,
            disable=not accelerator.is_local_main_process,
        )
        for batch in val_pbar:
            # val_pbar.set_description(f"Validation {step:,}")
            img_whole = batch["img"]
            label_whole = batch["label"]

            if config.image_size != 64 and config.dataset in ["dfc", "worldcover"]:
                # validate each tile sub-patch
                for x in range(0, 64, config.image_size - config.patch_sub):
                    for y in range(0, 64, config.image_size - config.patch_sub):
                        img = img_whole[
                            :,
                            :,
                            x : x + config.image_size - config.patch_sub,
                            y : y + config.image_size - config.patch_sub,
                        ]
                        label = label_whole[
                            :,
                            x : x + config.image_size - config.patch_sub,
                            y : y + config.image_size - config.patch_sub,
                        ]
                        if x + config.image_size >= 64 or y + config.image_size > 64:
                            continue

                        if config.method_name == "li" or pixelwise:
                            # baseline model only predicts class for the center pixel of the patch
                            center_idx = (config.image_size - config.patch_sub) // 2

                            label = label[
                                :, center_idx, center_idx
                            ]  # .unsqueeze(1).unsqueeze(1)
                            if config.method_name == "li":
                                img = img.unsqueeze(1)

                        output = model(img)
                        loss = criterion(output, label)

                        pred = output.argmax(dim=1)
                        valid_idx = label != config.ignored_label
                        acc = (pred[valid_idx] == label[valid_idx]).sum() / pred[
                            valid_idx
                        ].numel()
                        macro_acc = acc_criterion(
                            pred[valid_idx].to(int), label[valid_idx]
                        )

            else:
                img = img_whole
                label = label_whole

                if config.method_name == "li" or pixelwise:
                    # baseline model only predicts class for the center pixel of the patch
                    if config.dataset != "houston2018":
                        center_idx = (config.image_size - config.patch_sub) // 2
                        label = label[
                            :, center_idx, center_idx
                        ]  # .unsqueeze(1).unsqueeze(1)
                    if config.method_name == "li":
                        img = img.unsqueeze(1)

                output = model(img)
                loss = criterion(output, label)
                pred = output.argmax(dim=1)
                valid_idx = label != config.ignored_label
                acc = (pred[valid_idx] == label[valid_idx]).sum() / pred[
                    valid_idx
                ].numel()
                if valid_idx.sum() != 0:
                    macro_acc = acc_criterion(pred[valid_idx].to(int), label[valid_idx])
                else:
                    macro_acc = acc

            val_losses.append(loss.detach().item())
            val_accs.append(acc.detach().item())
            val_macro_accs.append(macro_acc.detach().item())

        current_val_acc = torch.tensor(val_accs).mean().item()
        aim_kwargs = {"epoch": epoch, "context": {"subset": "val"}}
        accelerator.log(
            {
                "val_acc": current_val_acc,
                "val_macro_acc": torch.tensor(val_macro_accs).mean().item(),
                "val_loss": torch.tensor(val_losses).mean().item(),
            },
            step,
            log_kwargs={"aim": aim_kwargs},
        )

    if (
        epoch == config.epoch
        or current_val_acc > best_val_acc
        or epoch in config.checkpoint_save_epochs
    ):
        if epoch == config.epoch or epoch % 10 == 0:
            accelerator.save_state(
                f"{config.checkpoints_dir}/{config.method_name}_ep{epoch}",
                safe_serialization=False,
            )
        if current_val_acc > best_val_acc:
            accelerator.save_state(
                f"{config.checkpoints_dir}/best_{config.method_name}_ep{epoch}",
                safe_serialization=False,
            )

            accelerator.wait_for_everyone()
            accelerator.save_model(
                model,
                f"{config.models_dir}/finetune_{config.dataset}",
                safe_serialization=False,
            )

            best_val_acc = current_val_acc

    return val_losses, best_val_acc


# %%
if __name__ == "__main__":
    dataset_name = "houston2018"

    config = get_finetune_config(
        f"configs/finetune_config_{dataset_name}.yaml", "configs/config.yaml", SEED
    )

    # %%
    project_dir_path = Path(f"./finetune-{config.dataset}/").expanduser().resolve()
    logging_dir_path = project_dir_path / "logs"
    checkpoints_dir_path = project_dir_path / "checkpoints" / config.version
    models_dir_path = project_dir_path / "models" / config.version
    graphs_dir_path = project_dir_path / "graphs" / config.version
    project_name = f"finetune-{config.dataset}"

    config.project_dir = str(project_dir_path)
    config.logging_dir = str(logging_dir_path)
    config.checkpoints_dir = str(checkpoints_dir_path)
    config.models_dir = str(models_dir_path)
    config.graphs_dir = str(graphs_dir_path)
    config.project_name = project_name

    # %%
    project_config = ProjectConfiguration(
        project_dir=str(project_dir_path), logging_dir=str(logging_dir_path)
    )
    accelerator = Accelerator(log_with="all", project_config=project_config)

    # %%
    model = ViTSpatialSpectral(
        image_size=config.image_size - config.patch_sub,
        spatial_patch_size=config.patch_size,
        spectral_patch_size=config.band_patch_size,
        num_classes=config.n_classes,
        dim=config.transformer_dim,
        depth=config.transformer_depth,
        heads=config.transformer_n_heads,
        mlp_dim=config.transformer_mlp_dim,
        dropout=config.transformer_dropout,
        emb_dropout=config.transformer_emb_dropout,
        channels=config.n_bands,
        spectral_pos=config.spectral_pos,
        spectral_pos_embed=config.spectral_pos_embed,
        blockwise_patch_embed=config.blockwise_patch_embed,
        spectral_only=config.spectral_only,
        pixelwise=config.pixelwise,
        pos_embed_len=config.pos_embed_len,
    )

    classifier_name = "mlp_head"
    # %%
    if config.checkpoint_path:
        model = load_checkpoint(config, model, classifier_name)

    if accelerator.is_main_process:
        # draw model architecture
        tv.draw_graph(
            model,
            input_size=(
                config.batch_size,
                config.n_bands,
                config.image_size,
                config.image_size,
            ),
            depth=2,
            graph_name=f"finetune-{config.dataset}-{config.version}",
            save_graph=True,
            directory=config.graphs_dir,
        )

    if config.linear_eval:
        print("Linear evaluation... only training mlp_head")
        for n, p in model.named_parameters():
            if classifier_name not in n:
                p.requires_grad = False
        params = list(getattr(model, classifier_name).parameters())
    else:
        # fine-tuning
        params = list(model.parameters())
        # set different LR for transformer and MLP head
        if config.lr != config.mlp_head_lr:
            mlp_param_list = [
                p for n, p in model.named_parameters() if classifier_name in n
            ]
            rest_param_list = [
                p for n, p in model.named_parameters() if classifier_name not in n
            ]
            params = [
                {"params": mlp_param_list, "lr": config.mlp_head_lr},
                {"params": rest_param_list},
            ]

    if config.method_name != "li" or config.overwrite_li_optim:
        optimizer = torch.optim.Adam(
            params, lr=config.lr, weight_decay=config.weight_decay
        )
        criterion = torch.nn.CrossEntropyLoss(ignore_index=config.ignored_label)

    # %%
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, factor=0.9, patience=5, verbose=True
    )

    acc_criterion = Accuracy(
        "multiclass", num_classes=config.n_classes, average="macro"
    )
    acc_criterion.to(accelerator.device)
    model_params = sum([p.numel() for p in model.parameters()])
    config.num_params = model_params

    print(f"Model name: {config.method_name}")
    print(f"Model parameters: {model_params:,}")

    dataloader, val_dataloader = get_supervised_data(config)

    model, optimizer, scheduler, dataloader, val_dataloader = accelerator.prepare(
        model, optimizer, scheduler, dataloader, val_dataloader
    )

    # %% register checkpointing
    accelerator.register_for_checkpointing(model, optimizer, scheduler)

    # %%
    aim_kwargs = {"log_system_params": True}
    accelerator.init_trackers(
        project_name, config.to_dict(), init_kwargs={"aim": aim_kwargs}
    )

    # os.mkdir(f"models/{config.run_id}/")
    Path(f"models/{config.run_id}/").mkdir(parents=True, exist_ok=True)

    losses = []
    accs = []
    macro_accs = []
    acc_per_epoch = []
    current_val_acc = 0
    best_val_acc = 0
    step = 0
    epoch = 0
    validation_epochs = get_val_epochs(config, dataloader)

    while epoch < config.epoch + 1 or step < config.max_steps + 1:
        model.train()

        train_pbar = track(
            dataloader,
            description=f"Epoch {epoch}/{config.epoch}",
            transient=True,
            disable=not accelerator.is_local_main_process or sys.gettrace(),
        )
        for batch in train_pbar:
            img = batch["img"]
            label = batch["label"]

            loss, acc, macro_acc = train_step(
                img,
                label,
                model,
                config,
                criterion,
                optimizer,
                acc_criterion,
                accelerator,
            )
            step += 1

            losses.append(loss.detach().item())
            accs.append(acc.detach().item())
            macro_accs.append(macro_acc.detach().item())

            if step % config.logging_freq == 0:
                aim_kwargs = {"epoch": epoch, "context": {"subset": "train"}}
                accelerator.log(
                    {
                        "acc": np.array(accs[-1 * config.logging_freq :]).mean(),
                        "macro_acc": np.array(
                            macro_accs[-1 * config.logging_freq :]
                        ).mean(),
                        "loss": np.array(losses[-1 * config.logging_freq :]).mean(),
                        "lr": optimizer.param_groups[0]["lr"],
                    },
                    step=step,
                    log_kwargs={"aim": aim_kwargs},
                )
                accs = []
                macro_accs = []
                losses = []

        # log at end of training epoch (to same step as validation stats below)
        aim_kwargs = {"epoch": epoch, "context": {"subset": "train"}}
        accelerator.log(
            {"acc": acc.item(), "loss": loss.item()},
            step=step,
            log_kwargs={"aim": aim_kwargs},
        )

        if epoch in validation_epochs:
            val_losses, best_val_acc = validate(
                config,
                epoch,
                model,
                val_dataloader,
                criterion,
                acc_criterion,
                step,
                best_val_acc,
                accelerator,
                pixelwise=config.pixelwise,
            )

        scheduler.step(torch.tensor(val_losses).mean().item())
        epoch += 1

    accelerator.end_training()