summaryrefslogtreecommitdiff
path: root/training/utils.py
blob: 6f57e4d9e31a51d942458cec11e2e0e0014a7586 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
"""Util functions for training with hydra and pytorch lightning."""
from typing import Any, List, Type
import warnings

import hydra
from loguru import logger as log
from omegaconf import DictConfig, OmegaConf
from pytorch_lightning import (
    Callback,
    LightningModule,
    Trainer,
)
from pytorch_lightning.loggers import LightningLoggerBase
from pytorch_lightning.loggers.wandb import WandbLogger
from pytorch_lightning.utilities import rank_zero_only
from tqdm import tqdm
import wandb


def save_config(config: DictConfig) -> None:
    """Save config to experiment directory."""
    with open("config.yaml", "w") as f:
        OmegaConf.save(config, f=f)


def print_config(config: DictConfig) -> None:
    """Prints config."""
    print(OmegaConf.to_yaml(config))


@rank_zero_only
def configure_logging(config: DictConfig) -> None:
    """Configure the loguru logger for output to terminal and disk."""
    # Remove default logger to get tqdm to work properly.
    log.remove()
    log.add(lambda msg: tqdm.write(msg, end=""), colorize=True, level=config.logging)


def configure_callbacks(config: DictConfig,) -> List[Type[Callback]]:
    """Configures Lightning callbacks."""
    callbacks = []
    if config.get("callbacks"):
        for callback_config in config.callbacks.values():
            if callback_config.get("_target_"):
                log.info(f"Instantiating callback <{callback_config._target_}>")
                callbacks.append(hydra.utils.instantiate(callback_config))
    return callbacks


def configure_logger(config: DictConfig) -> List[Type[LightningLoggerBase]]:
    """Configures Lightning loggers."""
    logger = []
    if config.get("logger"):
        for logger_config in config.logger.values():
            if logger_config.get("_target_"):
                log.info(f"Instantiating logger <{logger_config._target_}>")
                logger.append(hydra.utils.instantiate(logger_config))
    return logger


def extras(config: DictConfig) -> None:
    """Sets optional utilities."""
    # Enable adding new keys.
    OmegaConf.set_struct(config, False)

    if config.get("ignore_warnings"):
        log.info("Disabling python warnings! <config.ignore_warnings=True>")
        warnings.filterwarnings("ignore")

    if config.get("debug"):
        log.info("Running in debug mode! <config.debug=True>")
        config.trainer.fast_dev_run = True

    if config.trainer.get("fast_dev_run"):
        log.info(
            "Forcing debugger friendly configuration! <config.trainer.fast_dev_run=True>"
        )
        # Debuggers do not like GPUs and multiprocessing.
        if config.trainer.get("gpus"):
            config.trainer.gpus = 0
        if config.trainer.get("precision"):
            config.trainer.precision = 32
        if config.datamodule.get("pin_memory"):
            config.datamodule.pin_memory = False
        if config.datamodule.get("num_workers"):
            config.datamodule.num_workers = 0

    # Disable adding new keys to config
    OmegaConf.set_struct(config, True)


def empty(*args: Any, **kwargs: Any) -> None:
    pass


@rank_zero_only
def log_hyperparameters(
    config: DictConfig, model: LightningModule, trainer: Trainer,
) -> None:
    """This method saves hyperparameters with the logger."""
    hparams = {}

    # choose which parts of hydra config will be saved to loggers
    hparams["trainer"] = config["trainer"]
    hparams["model"] = config["model"]
    hparams["datamodule"] = config["datamodule"]
    if "callbacks" in config:
        hparams["callbacks"] = config["callbacks"]

    # save number of model parameters
    hparams["model/params_total"] = sum(p.numel() for p in model.parameters())
    hparams["model/params_trainable"] = sum(
        p.numel() for p in model.parameters() if p.requires_grad
    )
    hparams["model/params_not_trainable"] = sum(
        p.numel() for p in model.parameters() if not p.requires_grad
    )

    # send hparams to all loggers
    trainer.logger.log_hyperparams(hparams)

    # disable logging any more hyperparameters for all loggers
    # this is just a trick to prevent trainer from logging hparams of model,
    # since we already did that above
    trainer.logger.log_hyperparams = empty


def finish(logger: List[Type[LightningLoggerBase]],) -> None:
    """Makes sure everything closed properly."""
    for lg in logger:
        if isinstance(lg, WandbLogger):
            wandb.finish()