File size: 3,232 Bytes
4947683 | 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 | import os
from pathlib import Path
from typing import Optional
import dotenv
import pytorch_lightning as pl
from omegaconf import DictConfig, OmegaConf
def get_env(env_name: str, default: Optional[str] = None) -> str:
"""
Safely read an environment variable.
Raises errors if it is not defined or it is empty.
:param env_name: the name of the environment variable
:param default: the default (optional) value for the environment variable
:return: the value of the environment variable
"""
if env_name not in os.environ:
if default is None:
raise KeyError(
f"{env_name} not defined and no default value is present!")
return default
env_value: str = os.environ[env_name]
if not env_value:
if default is None:
raise ValueError(
f"{env_name} has yet to be configured and no default value is present!"
)
return default
return env_value
def load_envs(env_file: Optional[str] = None) -> None:
"""
Load all the environment variables defined in the `env_file`.
This is equivalent to `. env_file` in bash.
It is possible to define all the system specific variables in the `env_file`.
:param env_file: the file that defines the environment variables to use. If None
it searches for a `.env` file in the project.
"""
dotenv.load_dotenv(dotenv_path=env_file, override=True)
STATS_KEY: str = "stats"
# Adapted from https://github.com/hobogalaxy/lightning-hydra-template/blob/6bf03035107e12568e3e576e82f83da0f91d6a11/src/utils/template_utils.py#L125
def log_hyperparameters(
cfg: DictConfig,
model: pl.LightningModule,
trainer: pl.Trainer,
) -> None:
"""This method controls which parameters from Hydra config are saved by Lightning loggers.
Additionally saves:
- sizes of train, val, test dataset
- number of trainable model parameters
Args:
cfg (DictConfig): [description]
model (pl.LightningModule): [description]
trainer (pl.Trainer): [description]
"""
hparams = OmegaConf.to_container(cfg, resolve=True)
# save number of model parameters
hparams[f"{STATS_KEY}/params_total"] = sum(p.numel()
for p in model.parameters())
hparams[f"{STATS_KEY}/params_trainable"] = sum(
p.numel() for p in model.parameters() if p.requires_grad
)
hparams[f"{STATS_KEY}/params_not_trainable"] = sum(
p.numel() for p in model.parameters() if not p.requires_grad
)
if trainer.logger is not None:
# 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 = lambda params: None
# Load environment variables
load_envs()
# Set the cwd to the project root
PROJECT_ROOT: Path = Path(get_env("PROJECT_ROOT"))
assert (
PROJECT_ROOT.exists()
), "You must configure the PROJECT_ROOT environment variable in a .env file!"
os.chdir(PROJECT_ROOT)
|