bytecover/models/train_module.py

import logging
import os
from copy import deepcopy
from typing import Dict, List

import numpy as np
import torch
import torch.nn as nn
import wandb
from tqdm import tqdm, trange

from bytecover.models.data_model import BatchDict, Postfix, TestResults, ValDict
from bytecover.models.early_stopper import EarlyStopper
from bytecover.models.modules import Bottleneck, Resnet50
from bytecover.models.utils import (
    calculate_ranking_metrics,
    dataloader_factory,
    dir_checker,
    save_best_log,
    save_logs,
    save_predictions,
)

logger: logging.Logger = logging.getLogger()  # The logger used to log output


class TrainModule:
    def __init__(self, config: Dict) -> None:
        self.config = config
        self.state = "initializing"
        self.best_model_path: str = None
        self.num_classes = self.config["train"]["num_classes"]
        self.max_len = self.config["train"]["max_seq_len"][0]

        self.model = Resnet50(
            Bottleneck,
            num_channels=self.config["num_channels"],
            num_classes=self.num_classes,
            compress_ratio=self.config["train"]["compress_ratio"],
            tempo_factors=self.config["train"]["tempo_factors"],
        )
        self.model.to(self.config["device"])
        if self.config["wandb"]:
            wandb.watch(self.model)
        self.postfix: Postfix = {}

        self.triplet_loss = nn.TripletMarginLoss(margin=config["train"]["triplet_margin"])
        self.cls_loss = nn.CrossEntropyLoss(label_smoothing=config["train"]["smooth_factor"])

        self.early_stop = EarlyStopper(patience=self.config["train"]["patience"])
        self.optimizer = self.configure_optimizers()
        if self.config["device"] != "cpu":
            self.scaler = torch.cuda.amp.GradScaler(enabled=self.config["train"]["mixed_precision"])

    def pipeline(self) -> None:
        self.config["val"]["output_dir"] = dir_checker(self.config["val"]["output_dir"])

        if self.config["train"]["model_ckpt"] is not None:
            self.model.load_state_dict(torch.load(self.config["train"]["model_ckpt"]), strict=False)
            logger.info(f'Model loaded from checkpoint: {self.config["train"]["model_ckpt"]}')

        self.t_loaders = dataloader_factory(config=self.config, data_split="TRAIN")
        self.v_loader = dataloader_factory(config=self.config, data_split="VAL")[0]

        self.state = "running"

        self.pbar = trange(
            self.config["train"]["epochs"], disable=(not self.config["progress_bar"]), position=0, leave=True
        )
        for epoch in self.pbar:
            if self.state in ["early_stopped", "interrupted", "finished"]:
                return

            self.postfix["Epoch"] = epoch
            self.pbar.set_postfix(self.postfix)

            try:
                self.train_procedure()
            except KeyboardInterrupt:
                logger.warning("\nKeyboard Interrupt detected. Attempting gracefull shutdown...")
                self.state = "interrupted"
            except Exception as err:
                raise (err)

            if self.state == "interrupted":
                self.validation_procedure()
                self.pbar.set_postfix(
                    {k: self.postfix[k] for k in self.postfix.keys() & {"train_loss_step", "mr1", "mAP"}}
                )

        self.state = "finished"

    def validate(self) -> None:
        self.v_loader = dataloader_factory(config=self.config, data_split="VAL")[0]
        self.state = "running"
        self.validation_procedure()
        self.state = "finished"

    def test(self) -> None:
        self.test_loader = dataloader_factory(config=self.config, data_split="TEST")[0]
        self.test_results: TestResults = {}

        if self.best_model_path is not None:
            self.model.load_state_dict(torch.load(self.best_model_path), strict=False)
            print(f"Best model loaded from checkpoint: {self.best_model_path}")
        elif self.config["test"]["model_ckpt"] is not None:
            self.model.load_state_dict(torch.load(self.config["test"]["model_ckpt"], map_location='cpu'), strict=False)
            print(f'Model loaded from checkpoint: {self.config["test"]["model_ckpt"]}')
        elif self.state == "initializing":
            print("Warning: Testing with random weights")

        self.state = "running"
        self.test_procedure()
        self.state = "finished"

    def train_procedure(self) -> None:
        self.model.train()
        pbar_loaders = tqdm(self.t_loaders, disable=(not self.config["progress_bar"]), position=1, leave=False)
        for _, t_loader in enumerate(pbar_loaders):
            train_loss_list = []
            train_cls_loss_list = []
            train_triplet_loss_list = []
            self.max_len = t_loader.dataset.max_len
            pbar_loaders.set_postfix_str(f"max_seq_len={self.max_len}")
            for step, batch in tqdm(
                enumerate(t_loader),
                total=len(t_loader),
                disable=(not self.config["progress_bar"]),
                position=2,
                leave=False,
            ):
                train_step = self.training_step(batch)
                self.postfix["train_loss_step"] = float(f"{train_step['train_loss_step']:.3f}")
                train_loss_list.append(train_step["train_loss_step"])
                self.postfix["train_cls_loss_step"] = float(f"{train_step['train_cls_loss']:.3f}")
                train_cls_loss_list.append(train_step["train_cls_loss"])
                self.postfix["train_triplet_loss_step"] = float(f"{train_step['train_triplet_loss']:.3f}")
                train_triplet_loss_list.append(train_step["train_triplet_loss"])
                self.pbar.set_postfix(
                    {k: self.postfix[k] for k in self.postfix.keys() & {"train_loss_step", "mr1", "mAP"}}
                )
                if self.config["wandb"]:
                    wandb.log(self.postfix)
                if step % self.config["train"]["log_steps"] == 0:
                    save_logs(
                        dict(
                            epoch=self.postfix["Epoch"],
                            seq_len=self.max_len,
                            step=step,
                            train_loss_step=f"{train_step['train_loss_step']:.3f}",
                            train_cls_loss_step=f"{train_step['train_cls_loss']:.3f}",
                            train_triplet_loss_step=f"{train_step['train_triplet_loss']:.3f}",
                        ),
                        output_dir=self.config["val"]["output_dir"],
                        name="log_steps",
                    )
            train_loss = torch.tensor(train_loss_list)
            train_cls_loss = torch.tensor(train_cls_loss_list)
            train_triplet_loss = torch.tensor(train_triplet_loss_list)
            self.postfix["train_loss"] = train_loss.mean().item()
            self.postfix["train_cls_loss"] = train_cls_loss.mean().item()
            self.postfix["train_triplet_loss"] = train_triplet_loss.mean().item()
            if self.config["wandb"]:
                wandb.log(self.postfix)
            self.validation_procedure()
            if self.config["wandb"]:
                wandb.log(self.postfix)
            self.overfit_check()
            self.pbar.set_postfix({k: self.postfix[k] for k in self.postfix.keys() & {"train_loss_step", "mr1", "mAP"}})

    def training_step(self, batch: BatchDict) -> Dict[str, float]:
        with torch.autocast(
            device_type=self.config["device"].split(":")[0], enabled=self.config["train"]["mixed_precision"]
        ):
            anchor = self.model.forward(batch["anchor"].to(self.config["device"]))
            positive = self.model.forward(batch["positive"].to(self.config["device"]))
            negative = self.model.forward(batch["negative"].to(self.config["device"]))
            l1 = self.triplet_loss(anchor["f_t"], positive["f_t"], negative["f_t"])
            labels = nn.functional.one_hot(batch["anchor_label"].long(), num_classes=self.num_classes)
            l2 = self.cls_loss(anchor["cls"], labels.float().to(self.config["device"]))
            loss = l1 + l2

        self.optimizer.zero_grad()
        if self.config["device"] != "cpu":
            self.scaler.scale(loss).backward()
            self.scaler.step(self.optimizer)
            self.scaler.update()
        else:
            loss.backward()
            self.optimizer.step()

        return {"train_loss_step": loss.item(), "train_triplet_loss": l1.item(), "train_cls_loss": l2.item()}

    def validation_procedure(self) -> None:
        self.model.eval()
        embeddings: Dict[str, torch.Tensor] = {}
        for batch in tqdm(self.v_loader, disable=(not self.config["progress_bar"]), position=1, leave=False):
            val_dict = self.validation_step(batch)
            if val_dict["f_t"].ndim == 1:
                val_dict["f_c"] = val_dict["f_c"].unsqueeze(0)
                val_dict["f_t"] = val_dict["f_t"].unsqueeze(0)
            for anchor_id, triplet_embedding, embedding in zip(val_dict["anchor_id"], val_dict["f_t"], val_dict["f_c"]):
                embeddings[anchor_id] = torch.stack([triplet_embedding, embedding])

        val_outputs = self.validation_epoch_end(embeddings)
        logger.info(
            f"\n{' Validation Results ':=^50}\n"
            + "\n".join([f'"{key}": {value}' for key, value in self.postfix.items()])
            + f"\n{' End of Validation ':=^50}\n"
        )

        if self.config["val"]["save_val_outputs"]:
            val_outputs["val_embeddings"] = torch.stack(list(embeddings.values()))[:, 1].numpy()
            save_predictions(val_outputs, output_dir=self.config["val"]["output_dir"])
            save_logs(self.postfix, output_dir=self.config["val"]["output_dir"])
        self.model.train()

    def validation_epoch_end(self, outputs: Dict[str, torch.Tensor]) -> Dict[str, np.ndarray]:
        val_loss = torch.zeros(len(outputs))
        pos_ids = []
        neg_ids = []
        clique_ids = []
        for k, (anchor_id, embeddings) in enumerate(outputs.items()):
            clique_id, pos_id, neg_id = self.v_loader.dataset._triplet_sampling(anchor_id)
            val_loss[k] = self.triplet_loss(embeddings[0], outputs[pos_id][0], outputs[neg_id][0]).item()
            pos_ids.append(pos_id)
            neg_ids.append(neg_id)
            clique_ids.append(clique_id)
        anchor_ids = np.stack(list(outputs.keys()))
        preds = torch.stack(list(outputs.values()))[:, 1]
        self.postfix["val_loss"] = val_loss.mean().item()
        ranks, average_precisions = calculate_ranking_metrics(embeddings=preds.numpy(), cliques=clique_ids)
        self.postfix["mr1"] = ranks.mean()
        self.postfix["mAP"] = average_precisions.mean()
        return {
            "triplet_ids": np.stack(list(zip(clique_ids, anchor_ids, pos_ids, neg_ids))),
            "ranks": ranks,
            "average_precisions": average_precisions,
        }

    def validation_step(self, batch: BatchDict) -> ValDict:
        anchor_id = batch["anchor_id"]
        positive_id = batch["positive_id"]
        negative_id = batch["negative_id"]

        features = self.model.forward(batch["anchor"].to(self.config["device"]))

        return {
            "anchor_id": anchor_id,
            "positive_id": positive_id,
            "negative_id": negative_id,
            "f_t": features["f_t"].squeeze(0).detach().cpu(),
            "f_c": features["f_c"].squeeze(0).detach().cpu(),
        }

    def test_procedure(self) -> None:
        self.model.eval()
        clique_ids = []
        embeddings: Dict[str, torch.Tensor] = {}
        for batch in tqdm(self.test_loader, disable=(not self.config["progress_bar"])):
            clique_ids_batch = self.test_loader.dataset.labels.loc[batch["anchor_id"], "clique"]
            test_dict = self.validation_step(batch)
            if test_dict["f_c"].ndim == 1:
                test_dict["f_c"] = test_dict["f_c"].unsqueeze(0)
            for anchor_id, clique_id, embedding in zip(test_dict["anchor_id"], clique_ids_batch, test_dict["f_c"]):
                embeddings[anchor_id] = embedding
                clique_ids.append(clique_id)

        test_outputs = self.test_epoch_end(embeddings, clique_ids)
        logger.info(
            f"\n{' Test Results ':=^50}\n"
            + "\n".join([f'"{key}": {value}' for key, value in self.test_results.items()])
            + f"\n{' End of Testing ':=^50}\n"
        )

        if self.config["test"]["save_test_outputs"]:
            test_outputs["test_embeddings"] = torch.stack(list(embeddings.values())).numpy()
            save_predictions(test_outputs, output_dir=self.config["test"]["output_dir"])
            save_logs(self.test_results, output_dir=self.config["test"]["output_dir"])

    def test_epoch_end(self, outputs: Dict[str, torch.Tensor], clique_ids: List[int]) -> Dict[str, np.ndarray]:
        anchor_ids = np.stack(list(outputs.keys()))
        preds = torch.stack(list(outputs.values()))
        ranks, average_precisions = calculate_ranking_metrics(embeddings=preds.numpy(), cliques=clique_ids)
        self.test_results["test_mr1"] = ranks.mean()
        self.test_results["test_mAP"] = average_precisions.mean()
        return {
            "anchor_ids": np.stack(list(zip(clique_ids, anchor_ids))),
            "ranks": ranks,
            "average_precisions": average_precisions,
        }

    def overfit_check(self) -> None:
        if self.early_stop(self.postfix["val_loss"]):
            logger.info(f"\nValidation not improved for {self.early_stop.patience} consecutive epochs. Stopping...")
            self.state = "early_stopped"

        if self.early_stop.counter > 0:
            logger.info("\nValidation loss was not improved")
        else:
            logger.info(f"\nMetric improved. New best score: {self.early_stop.min_validation_loss:.3f}")
            save_best_log(self.postfix, output_dir=self.config["val"]["output_dir"])

            logger.info("Saving model...")
            epoch = self.postfix["Epoch"]
            max_secs = self.max_len
            prev_model = deepcopy(self.best_model_path)
            self.best_model_path = os.path.join(
                self.config["val"]["output_dir"], "model", f"best-model-{epoch=}-{max_secs=}.pt"
            )
            os.makedirs(os.path.dirname(self.best_model_path), exist_ok=True)
            torch.save(deepcopy(self.model.state_dict()), self.best_model_path)
            if prev_model is not None:
                os.remove(prev_model)

    def configure_optimizers(self) -> torch.optim.Optimizer:
        optimizer = torch.optim.Adam(self.model.parameters(), lr=self.config["train"]["learning_rate"])

        return optimizer