Spaces:
Sleeping
Sleeping
File size: 1,370 Bytes
801501a |
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 |
import gc
import rich
from rich.tree import Tree
from rich.syntax import Syntax
import torch
from omegaconf import DictConfig, OmegaConf
def print_config(
config,
fields=(
"trainer",
"model",
"callbacks",
"logger",
"seed",
"name",
),
resolve: bool = True,
save_config: bool = False,
) -> None:
"""Prints content of DictConfig using Rich library and its tree structure.
Args:
config (DictConfig): Configuration composed by Hydra.
fields (Sequence[str], optional): Determines which main fields from config will
be printed and in what order.
resolve (bool, optional): Whether to resolve reference fields of DictConfig.
"""
style = "dim"
tree = Tree("CONFIG", style=style, guide_style=style)
for field in fields:
branch = tree.add(field, style=style, guide_style=style)
config_section = config.get(field)
branch_content = str(config_section)
if isinstance(config_section, DictConfig):
branch_content = OmegaConf.to_yaml(config_section, resolve=resolve)
branch.add(Syntax(branch_content, "yaml"))
rich.print(tree)
if save_config:
with open("config_tree.log", "w") as fp:
rich.print(tree, file=fp)
def clean_gpu():
gc.collect()
torch.cuda.empty_cache()
|