pytorch_svgrender/libs/engine/model_state.py

# -*- coding: utf-8 -*-
# Copyright (c) XiMing Xing. All rights reserved.
# Author: XiMing Xing
# Description:

from typing import Union, List
from pathlib import Path
from datetime import datetime
import logging

from omegaconf import OmegaConf, DictConfig
from pprint import pprint
import torch
from accelerate.utils import LoggerType
from accelerate import Accelerator

from ..utils.logging import get_logger


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

    features:
        - Precision
        - Device
        - Optimizer
        - Logger (default: python system print and logging)
        - Monitor (default: wandb, tensorboard)
    """

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

        """check valid"""
        mixed_precision = self.state_cfg.get("mprec")
        # Bug: omegaconf convert 'no' to false
        mixed_precision = "no" if type(mixed_precision) == bool else mixed_precision

        """create working space"""
        # rule: ['./config'. 'method_name', 'exp_name.yaml']
        # -> result_path: ./runs/{method_name}-{exp_name}, as a base folder
        now_time = datetime.now().strftime('%Y-%m-%d-%H-%M')
        results_folder = self.args.get("result_path", None)
        if results_folder is None:
            self.result_path = Path("./workdir") / f"{self.x_cfg.method}-{now_time}"
        else:
            self.result_path = Path(results_folder) / f"{self.x_cfg.method}-{now_time}"

        # update result_path: ./runs/{method_name}-{exp_name}/{log_path_suffix}
        # noting: can be understood as "results dir / methods / ablation study / your result"
        config_name_only = str(self.x_cfg.method).split(".")[0]
        if log_path_suffix is not None:
            self.result_path = self.result_path / f"{config_name_only}-{log_path_suffix}"
        else:
            self.result_path = self.result_path / f"{config_name_only}"

        """init visualized tracker"""
        # TODO: monitor with WANDB or TENSORBOARD
        self.log_with = []
        # if self.state_cfg.wandb:
        #     self.log_with.append(LoggerType.WANDB)
        # if self.state_cfg.tensorboard:
        #     self.log_with.append(LoggerType.TENSORBOARD)

        """HuggingFace Accelerator"""
        self.accelerator = Accelerator(
            device_placement=True,
            mixed_precision=mixed_precision,
            cpu=True if self.state_cfg.cpu else False,
            log_with=None if len(self.log_with) == 0 else self.log_with,
            project_dir=self.result_path / "vis",
        )

        """logs"""
        if self.accelerator.is_local_main_process:
            # logging
            self.log = logging.getLogger(__name__)

            # log results in a folder periodically
            self.result_path.mkdir(parents=True, exist_ok=True)
            if not ignore_log:
                self.logger = get_logger(
                    logs_dir=self.result_path.as_posix(),
                    file_name=f"{now_time}-{args.seed}-log.txt"
                )

            print("==> system args: ")
            sys_cfg = OmegaConf.masked_copy(args, ["x"])
            print(sys_cfg)
            print("==> yaml config args: ")
            print(self.x_cfg)

            print("\n***** Model State *****")
            print(f"-> Mixed Precision: {mixed_precision}, AMP: {self.accelerator.native_amp}")
            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()}")

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

        """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")

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

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

    @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 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:
            >>> params_optimized = self.optimized_params(uvit, verbose=True)
            >>> optimizer = torch.optim.AdamW(params_optimized, lr=1e-3)

        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.result_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.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_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 len(self.log_with) > 0:
            self.close_tracker()
        self.print(msg)