UnSAMFlow/trainer/base_trainer.py

"""
Copyright (c) Meta Platforms, Inc. and affiliates.
"""

import os

from abc import abstractmethod

import numpy as np
import torch

from utils.manifold_utils import pathmgr
from utils.torch_utils import (
    AdamW,
    bias_parameters,
    load_checkpoint,
    other_parameters,
    save_checkpoint,
    weight_parameters,
)

from .object_cache import ObjectCache


class BaseTrainer:
    """
    Base class for all trainers
    """

    def __init__(
        self,
        train_loaders,
        valid_loaders,
        model,
        loss_func,
        save_root,
        config,
        resume=False,
        train_sets_epoches=None,
        summary_writer=None,
        rank=0,
        world_size=1,
    ):
        self.cfg = config
        self.save_root = save_root
        self.summary_writer = summary_writer
        self.train_loaders, self.valid_loaders = train_loaders, valid_loaders
        self.train_sets_epoches = train_sets_epoches

        self.rank, self.world_size = rank, world_size
        self.device = model.device
        self.loss_func = loss_func

        if resume:  # load all states
            self._load_resume_ckpt(model)
        else:
            self.model = self._init_model(model)
            self.i_epoch, self.i_iter = 0, 0
            self.i_train_set = 0
            while (
                self.train_sets_epoches[self.i_train_set] == 0
            ):  # skip the datasets of 0 epoches
                self.i_train_set += 1

            self.optimizer = self._create_optimizer()
            self.scheduler = self._create_scheduler(
                self.optimizer, self.train_sets_epoches[self.i_train_set]
            )

            self.best_error = np.inf

    @abstractmethod
    def _run_one_epoch(self):
        ...

    @abstractmethod
    def _validate_with_gt(self):
        ...

    def log(self, s):
        if self.rank == 0:
            print(s)

    def set_up_obj_cache(self, cache_size=500):
        self.obj_cache = ObjectCache(cache_size=cache_size)

    def train(self):

        if (
            self.cfg.pretrained_model is not None
        ):  # if using a pretrained model, evaluate that first to compare
            if self.rank == 0:
                self._validate_with_gt()
            torch.distributed.barrier()

        for _epoch in range(self.i_epoch, self.cfg.epoch_num):
            self._run_one_epoch()

            if self.i_epoch >= sum(self.train_sets_epoches[: (self.i_train_set + 1)]):
                self.i_train_set += 1
                self.optimizer = (
                    self._create_optimizer()
                )  # reset the states of optimizer as well
                self.scheduler = self._create_scheduler(
                    self.optimizer, self.train_sets_epoches[self.i_train_set]
                )

            if self.rank == 0:
                if self.i_epoch % self.cfg.val_epoch_size == 0:
                    self._validate_with_gt()
                    self.log(" * Epoch {} validation complete.".format(self.i_epoch))

            torch.distributed.barrier()

    # def zero_grad(self):
    #     # One Pytorch tutorial suggests clearing the gradients this way for faster speed
    #     # https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html
    #     for param in self.model.parameters():
    #         param.grad = None

    def _init_model(self, model):
        model = model.to(self.device)

        if self.cfg.pretrained_model:
            self.log(
                "=> using pre-trained weights {}.".format(self.cfg.pretrained_model)
            )
            epoch, weights = load_checkpoint(self.cfg.pretrained_model)
            model.module.load_state_dict(weights)
        else:
            self.log("=> Train from scratch.")
            model.module.init_weights()

        self.log("number of parameters: {}".format(self.count_parameters(model)))
        self.log(
            "gpu memory allocated (model parameters only): {} Bytes".format(
                torch.cuda.memory_allocated()
            )
        )
        return model

    def _create_optimizer(self):
        self.log("=> setting {} optimizer".format(self.cfg.optim))
        param_groups = [
            {
                "params": bias_parameters(self.model.module),
                "weight_decay": self.cfg.bias_decay,
            },
            {
                "params": weight_parameters(self.model.module),
                "weight_decay": self.cfg.weight_decay,
            },
            {"params": other_parameters(self.model.module), "weight_decay": 0},
        ]

        if self.cfg.optim == "adamw":
            optimizer = AdamW(
                param_groups, self.cfg.lr, betas=(self.cfg.momentum, self.cfg.beta)
            )
        elif self.cfg.optim == "adam":
            optimizer = torch.optim.Adam(
                param_groups,
                self.cfg.lr,
                betas=(self.cfg.momentum, self.cfg.beta),
                eps=1e-7,
            )
        else:
            raise NotImplementedError(self.cfg.optim)

        return optimizer

    def _create_scheduler(self, optimizer, epoches=np.inf):

        if (
            self.i_train_set < len(self.train_sets_epoches) - 1
        ):  # try only the last loader uses onecyclelr
            scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=1)
            return scheduler

        if "lr_scheduler" in self.cfg.keys():
            self.log("=> setting {} scheduler".format(self.cfg.lr_scheduler.module))

            params = self.cfg.lr_scheduler.params

            if self.cfg.lr_scheduler.module == "OneCycleLR":
                params["epochs"] = min(epoches, self.cfg.epoch_num - self.i_epoch)
                params["steps_per_epoch"] = self.cfg.epoch_size

            scheduler = getattr(torch.optim.lr_scheduler, self.cfg.lr_scheduler.module)(
                optimizer, **params
            )
        else:  # a dummy scheduler by default
            scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=1)

        return scheduler

    def _load_resume_ckpt(self, model):
        self.log("==> resuming")

        with pathmgr.open(
            os.path.join(self.save_root, "model_ckpt.pth.tar"), "rb"
        ) as f:
            ckpt_dict = torch.load(f)

        if "iter" not in ckpt_dict.keys():
            ckpt_dict["iter"] = ckpt_dict["epoch"] * self.cfg.epoch_size
        if "best_error" not in ckpt_dict.keys():
            ckpt_dict["best_error"] = np.inf
        self.i_epoch, self.i_iter, self.best_error = (
            ckpt_dict["epoch"],
            ckpt_dict["iter"],
            ckpt_dict["best_error"],
        )
        self.i_train_set = np.where(self.i_epoch < np.cumsum(self.train_sets_epoches))[
            0
        ][0]

        model = model.to(self.device)
        model.module.load_state_dict(ckpt_dict["state_dict"])
        # self.model = torch.nn.DataParallel(model, device_ids=self.device_ids)

        self.optimizer = self._create_optimizer()
        self.scheduler = self._create_scheduler(
            self.optimizer, self.train_sets_epoches[self.i_train_set]
        )

        if "optimizer_dict" in ckpt_dict.keys():
            self.optimizer.load_state_dict(ckpt_dict["optimizer_dict"])
        if "scheduler_dict" in ckpt_dict.keys():
            self.scheduler.load_state_dict(ckpt_dict["scheduler_dict"])

        return

    # def _prepare_device(self, n_gpu_use):
    #     """
    #     setup GPU device if available, move model into configured device
    #     """
    #     n_gpu = torch.cuda.device_count()
    #     if n_gpu_use > 0 and n_gpu == 0:
    #         self.log(
    #             "Warning: There's no GPU available on this machine,"
    #             "training will be performed on CPU."
    #         )
    #         n_gpu_use = 0
    #     if n_gpu_use > n_gpu:
    #         self.log(
    #             "Warning: The number of GPU's configured to use is {}, "
    #             "but only {} are available.".format(n_gpu_use, n_gpu)
    #         )
    #         n_gpu_use = n_gpu
    #     device = torch.device("cuda:0" if n_gpu_use > 0 else "cpu")
    #     list_ids = list(range(n_gpu_use))
    #     self.log("=> gpu in use: {} gpu(s)".format(n_gpu_use))
    #     self.log(
    #         "device names: {}".format([torch.cuda.get_device_name(i) for i in list_ids])
    #     )
    #     return device, list_ids

    def count_parameters(self, model):
        return sum(p.numel() for p in model.parameters() if p.requires_grad)

    def save_model(self, name, save_with_runtime=True):
        if save_with_runtime:
            models = {
                "epoch": self.i_epoch,
                "iter": self.i_iter,
                "best_error": self.best_error,
                "state_dict": self.model.module.state_dict(),
                "optimizer_dict": self.optimizer.state_dict(),
                "scheduler_dict": self.scheduler.state_dict(),
            }
        else:
            models = {"state_dict": self.model.module.state_dict()}

        save_checkpoint(self.save_root, models, name, is_best=False)