libs/engine/model_state.py

import os
from functools import partial
from typing import Union, List
from pathlib import Path
from datetime import datetime, timedelta

from omegaconf import DictConfig
from pprint import pprint
import torch
from accelerate.utils import LoggerType
from accelerate import (
    Accelerator,
    GradScalerKwargs,
    DistributedDataParallelKwargs,
    InitProcessGroupKwargs
)

from ..modules.ema import EMA
from ..utils.logging import get_logger


class ModelState:
    """
    Handling logger and `hugging face` accelerate training

    features:
        - Mixed Precision
        - Gradient Scaler
        - Gradient Accumulation
        - Optimizer
        - EMA
        - Logger (default: python print)
        - Monitor (default: wandb, tensorboard)
    """

    def __init__(
            self,
            args,
            log_path_suffix: str = None,
            ignore_log=False,  # whether to create log file or not
    ) -> None:
        self.args: DictConfig = args

        """check valid"""
        mixed_precision = self.args.get("mixed_precision")
        # Bug: omegaconf convert 'no' to false
        mixed_precision = "no" if type(mixed_precision) == bool else mixed_precision
        split_batches = self.args.get("split_batches", False)
        gradient_accumulate_step = self.args.get("gradient_accumulate_step", 1)
        assert gradient_accumulate_step >= 1, f"except gradient_accumulate_step >= 1, get {gradient_accumulate_step}"

        """create working space"""
        # rule: ['./config'. 'method_name', 'exp_name.yaml']
        # -> results_path: ./runs/{method_name}-{exp_name}, as a base folder
        # config_prefix, config_name = str(self.args.get("config")).split('/')
        # config_name_only = str(config_name).split(".")[0]

        config_name_only = str(self.args.get("config")).split(".")[0]
        results_folder = self.args.get("results_path", None)
        if results_folder is None:
            # self.results_path = Path("./workdir") / f"{config_prefix}-{config_name_only}"
            self.results_path = Path("./workdir")
        else:
            # self.results_path = Path(results_folder) / f"{config_prefix}-{config_name_only}"
            self.results_path = Path(os.path.join(results_folder, self.args.get("edit_type"), ))

        # update results_path: ./runs/{method_name}-{exp_name}/{log_path_suffix}
        # noting: can be understood as "results dir / methods / ablation study / your result"
        if log_path_suffix is not None:
            self.results_path = self.results_path / log_path_suffix

        kwargs_handlers = []
        """mixed precision training"""
        if args.mixed_precision == "no":
            scaler_handler = GradScalerKwargs(
                init_scale=args.init_scale,
                growth_factor=args.growth_factor,
                backoff_factor=args.backoff_factor,
                growth_interval=args.growth_interval,
                enabled=True
            )
            kwargs_handlers.append(scaler_handler)

        """distributed training"""
        ddp_handler = DistributedDataParallelKwargs(
            dim=0,
            broadcast_buffers=True,
            static_graph=False,
            bucket_cap_mb=25,
            find_unused_parameters=False,
            check_reduction=False,
            gradient_as_bucket_view=False
        )
        kwargs_handlers.append(ddp_handler)

        init_handler = InitProcessGroupKwargs(timeout=timedelta(seconds=1200))
        kwargs_handlers.append(init_handler)

        """init visualized tracker"""
        log_with = []
        self.args.visual = False
        if args.use_wandb:
            log_with.append(LoggerType.WANDB)
        if args.tensorboard:
            log_with.append(LoggerType.TENSORBOARD)

        """hugging face Accelerator"""
        self.accelerator = Accelerator(
            device_placement=True,
            split_batches=split_batches,
            mixed_precision=mixed_precision,
            gradient_accumulation_steps=args.gradient_accumulate_step,
            cpu=True if args.use_cpu else False,
            log_with=None if len(log_with) == 0 else log_with,
            project_dir=self.results_path / "vis",
            kwargs_handlers=kwargs_handlers,
        )

        """logs"""
        if self.accelerator.is_local_main_process:
            # for logging results in a folder periodically
            self.results_path.mkdir(parents=True, exist_ok=True)
            if not ignore_log:
                now_time = datetime.now().strftime('%Y-%m-%d-%H-%M')
                # self.logger = get_logger(
                #     logs_dir=self.results_path.as_posix(),
                #     file_name=f"log.txt"
                # )

            print("==> command line args: ")
            print(args.cmd_args)
            print("==> yaml config args: ")
            print(args.yaml_config)

            print("\n***** Model State *****")
            if self.accelerator.distributed_type != "NO":
                print(f"-> Distributed Type: {self.accelerator.distributed_type}")
            print(f"-> Split Batch Size: {split_batches}, Total Batch Size: {self.actual_batch_size}")
            print(f"-> Mixed Precision: {mixed_precision}, AMP: {self.accelerator.native_amp},"
                  f" Gradient Accumulate Step: {gradient_accumulate_step}")
            print(f"-> Weight dtype:  {self.weight_dtype}")

            if self.accelerator.scaler_handler is not None and self.accelerator.scaler_handler.enabled:
                print(f"-> Enabled GradScaler: {self.accelerator.scaler_handler.to_kwargs()}")

            if args.use_wandb:
                print(f"-> Init trackers: 'wandb' ")
                self.args.visual = True
                self.__init_tracker(project_name="my_project", tags=None, entity="")

            print(f"-> Working Space: '{self.results_path}'")

        """EMA"""
        self.use_ema = args.get('ema', False)
        self.ema_wrapper = self.__build_ema_wrapper()

        """glob step"""
        self.step = 0

        """log process"""
        self.accelerator.wait_for_everyone()
        print(f'Process {self.accelerator.process_index} using device: {self.accelerator.device}')

        self.print("-> state initialization complete \n")

    def __init_tracker(self, project_name, tags, entity):
        self.accelerator.init_trackers(
            project_name=project_name,
            config=dict(self.args),
            init_kwargs={
                "wandb": {
                    "notes": "accelerate trainer pipeline",
                    "tags": [
                        f"total batch_size: {self.actual_batch_size}"
                    ],
                    "entity": entity,
                }}
        )

    def __build_ema_wrapper(self):
        if self.use_ema:
            self.print(f"-> EMA: {self.use_ema}, decay: {self.args.ema_decay}, "
                       f"update_after_step: {self.args.ema_update_after_step}, "
                       f"update_every: {self.args.ema_update_every}")
            ema_wrapper = partial(
                EMA, beta=self.args.ema_decay,
                update_after_step=self.args.ema_update_after_step,
                update_every=self.args.ema_update_every
            )
        else:
            ema_wrapper = None

        return ema_wrapper

    @property
    def device(self):
        return self.accelerator.device

    @property
    def weight_dtype(self):
        weight_dtype = torch.float32
        if self.accelerator.mixed_precision == "fp16":
            weight_dtype = torch.float16
        elif self.accelerator.mixed_precision == "bf16":
            weight_dtype = torch.bfloat16
        return weight_dtype

    @property
    def actual_batch_size(self):
        if self.accelerator.split_batches is False:
            actual_batch_size = self.args.batch_size * self.accelerator.num_processes * self.accelerator.gradient_accumulation_steps
        else:
            assert self.actual_batch_size % self.accelerator.num_processes == 0
            actual_batch_size = self.args.batch_size
        return actual_batch_size

    @property
    def n_gpus(self):
        return self.accelerator.num_processes

    @property
    def no_decay_params_names(self):
        no_decay = [
            "bn", "LayerNorm", "GroupNorm",
        ]
        return no_decay

    def no_decay_params(self, model, weight_decay):
        """optimization tricks"""
        optimizer_grouped_parameters = [
            {
                "params": [
                    p for n, p in model.named_parameters()
                    if not any(nd in n for nd in self.no_decay_params_names)
                ],
                "weight_decay": weight_decay,
            },
            {
                "params": [
                    p for n, p in model.named_parameters()
                    if any(nd in n for nd in self.no_decay_params_names)
                ],
                "weight_decay": 0.0,
            },
        ]
        return optimizer_grouped_parameters

    def optimized_params(self, model: torch.nn.Module, verbose=True) -> List:
        """return parameters if `requires_grad` is True

        Args:
            model: pytorch models
            verbose: log optimized parameters

        Examples:
            >>> self.params_optimized = self.optimized_params(uvit, verbose=True)
            >>> optimizer = torch.optim.AdamW(self.params_optimized, lr=args.lr)

        Returns:
                a list of parameters
        """
        params_optimized = []
        for key, value in model.named_parameters():
            if value.requires_grad:
                params_optimized.append(value)
                if verbose:
                    self.print("\t {}, {}, {}".format(key, value.numel(), value.shape))
        return params_optimized

    def save_everything(self, fpath: str):
        """Saving and loading the model, optimizer, RNG generators, and the GradScaler."""
        if not self.accelerator.is_main_process:
            return
        self.accelerator.save_state(fpath)

    def load_save_everything(self, fpath: str):
        """Loading the model, optimizer, RNG generators, and the GradScaler."""
        self.accelerator.load_state(fpath)

    def save(self, milestone: Union[str, float, int], checkpoint: object) -> None:
        if not self.accelerator.is_main_process:
            return

        torch.save(checkpoint, self.results_path / f'model-{milestone}.pt')

    def save_in(self, root: Union[str, Path], checkpoint: object) -> None:
        if not self.accelerator.is_main_process:
            return

        torch.save(checkpoint, root)

    def load_ckpt_model_only(self, model: torch.nn.Module, path: Union[str, Path], rm_module_prefix: bool = False):
        ckpt = torch.load(path, map_location=self.accelerator.device)

        unwrapped_model = self.accelerator.unwrap_model(model)
        if rm_module_prefix:
            unwrapped_model.load_state_dict({k.replace('module.', ''): v for k, v in ckpt.items()})
        else:
            unwrapped_model.load_state_dict(ckpt)
        return unwrapped_model

    def load_shared_weights(self, model: torch.nn.Module, path: Union[str, Path]):
        ckpt = torch.load(path, map_location=self.accelerator.device)
        self.print(f"pretrained_dict len: {len(ckpt)}")
        unwrapped_model = self.accelerator.unwrap_model(model)
        model_dict = unwrapped_model.state_dict()
        pretrained_dict = {k: v for k, v in ckpt.items() if k in model_dict}
        model_dict.update(pretrained_dict)
        unwrapped_model.load_state_dict(model_dict, strict=False)
        self.print(f"selected pretrained_dict: {len(model_dict)}")
        return unwrapped_model

    def print(self, *args, **kwargs):
        """Use in replacement of `print()` to only print once per server."""
        self.accelerator.print(*args, **kwargs)

    def pretty_print(self, msg):
        if self.accelerator.is_local_main_process:
            pprint(dict(msg))

    def close_tracker(self):
        self.accelerator.end_training()

    def free_memory(self):
        self.accelerator.clear()

    def close(self, msg: str = "Training complete."):
        """Use in end of training."""
        self.free_memory()

        if torch.cuda.is_available():
            self.print(f'\nGPU memory usage: {torch.cuda.max_memory_reserved() / 1024 ** 3:.2f} GB')
        if self.args.visual:
            self.close_tracker()
        self.print(msg)