Spaces:
Runtime error
Runtime error
| import logging | |
| from importlib.util import find_spec | |
| from typing import List, Optional, Union | |
| from omegaconf import DictConfig, OmegaConf | |
| from pie_modules.models.interface import RequiresTaskmoduleConfig | |
| from pytorch_ie import PyTorchIEModel, TaskModule | |
| from pytorch_lightning.loggers import Logger | |
| from pytorch_lightning.utilities import rank_zero_only | |
| def get_pylogger(name=__name__) -> logging.Logger: | |
| """Initializes multi-GPU-friendly python command line logger.""" | |
| logger = logging.getLogger(name) | |
| # this ensures all logging levels get marked with the rank zero decorator | |
| # otherwise logs would get multiplied for each GPU process in multi-GPU setup | |
| logging_levels = ("debug", "info", "warning", "error", "exception", "fatal", "critical") | |
| for level in logging_levels: | |
| setattr(logger, level, rank_zero_only(getattr(logger, level))) | |
| return logger | |
| log = get_pylogger(__name__) | |
| def log_hyperparameters( | |
| logger: Optional[List[Logger]] = None, | |
| config: Optional[Union[dict, DictConfig]] = None, | |
| model: Optional[PyTorchIEModel] = None, | |
| taskmodule: Optional[TaskModule] = None, | |
| key_prefix: str = "_", | |
| **kwargs, | |
| ) -> None: | |
| """Controls which config parts are saved by lightning loggers. | |
| Additional saves: | |
| - Number of model parameters | |
| """ | |
| hparams = {} | |
| if not logger: | |
| log.warning("Logger not found! Skipping hyperparameter logging...") | |
| return | |
| # this is just for backwards compatibility: usually, the taskmodule_config should be passed to | |
| # the model and, thus, be logged there automatically | |
| if model is not None and not isinstance(model, RequiresTaskmoduleConfig): | |
| if taskmodule is None: | |
| raise ValueError( | |
| "If model is not an instance of RequiresTaskmoduleConfig, taskmodule must be passed!" | |
| ) | |
| # here we use the taskmodule/model config how it is after preparation/initialization | |
| hparams["taskmodule_config"] = taskmodule.config | |
| if model is not None: | |
| # save number of model parameters | |
| hparams[f"{key_prefix}num_params/total"] = sum(p.numel() for p in model.parameters()) | |
| hparams[f"{key_prefix}num_params/trainable"] = sum( | |
| p.numel() for p in model.parameters() if p.requires_grad | |
| ) | |
| hparams[f"{key_prefix}num_params/non_trainable"] = sum( | |
| p.numel() for p in model.parameters() if not p.requires_grad | |
| ) | |
| if config is not None: | |
| hparams[f"{key_prefix}config"] = ( | |
| OmegaConf.to_container(config, resolve=True) if OmegaConf.is_config(config) else config | |
| ) | |
| # add additional hparams | |
| for k, v in kwargs.items(): | |
| hparams[f"{key_prefix}{k}"] = v | |
| # send hparams to all loggers | |
| for current_logger in logger: | |
| current_logger.log_hyperparams(hparams) | |
| def close_loggers() -> None: | |
| """Makes sure all loggers closed properly (prevents logging failure during multirun).""" | |
| log.info("Closing loggers...") | |
| if find_spec("wandb"): # if wandb is installed | |
| import wandb | |
| if wandb.run: | |
| log.info("Closing wandb!") | |
| wandb.finish() | |