Source code for labml.utils.pytorch

from typing import Optional, Dict, TYPE_CHECKING

import torch

from labml import tracker
from labml.configs import BaseConfigs

    from torch.optim.optimizer import Optimizer

def store_l1_l2(name: str, tensor: torch.Tensor):
    if tensor.is_floating_point():
        tracker.add(f"{name}.mean", tensor.mean())
        tracker.add(f"{name}.l1", tensor.abs().mean())
        tracker.add(f"{name}.l2", (tensor ** 2).mean().sqrt())

def store_var(name: str, tensor: torch.Tensor):
    if tensor.is_floating_point():
        dims = tuple(i for i in range(len(tensor.shape)))
        tracker.add(f"{name}.mean", tensor.mean())
        var = (tensor ** 2).mean(dim=dims) - tensor.mean(dim=dims) ** 2
        tracker.add(f"{name}.var", var.mean())

[docs]def store_model_indicators(model: torch.nn.Module, *, model_name: str = "model"): """ Track model parameters and gradients. Arguments: model (Optimizer): PyTorch model Keyword Arguments: model_name (str, optional): name of the model """ for name, param in model.named_parameters(): if param.requires_grad: with torch.no_grad(): store_l1_l2(f"param.{model_name}.{name}", param) if param.grad is not None: store_l1_l2(f"grad.{model_name}.{name}", param.grad)
[docs]def store_optimizer_indicators(optimizer: 'Optimizer', *, models: Optional[Dict[str, torch.nn.Module]] = None, optimizer_name: str = "optimizer"): """ Track optimizer stats such as moments. Arguments: optimizer (Optimizer): PyTorch optimizer Keyword Arguments: models (Dict[str, torch.nn.Module], optional): a dictionary of modules being optimized. This is used to get the proper parameter names. optimizer_name (str, optional): name of the optimizer """ if models is None: models = {} names = {} for model_name, model in models.items(): for name, p in model.named_parameters(): names[p] = f'{model_name}.{name}' unknown = 0 for group in optimizer.param_groups: for p in group['params']: if p.grad is None: continue state = optimizer.state[p] if len(state) == 0: continue name = names.get(p, None) if name is None: name = f'unknown.{unknown}' unknown += 1 for k, v in state.items(): if isinstance(v, float) or isinstance(v, int): tracker.add(f'optim.{optimizer_name}.{name}.{k}', v) if isinstance(v, torch.Tensor): store_l1_l2(f'optim.{optimizer_name}.{name}.{k}', v)
[docs]def get_modules(configs: BaseConfigs): """ Get all the PyTorch modules in ``configs`` object. Arguments: configs (labml.configs.BaseConfigs): configurations object """ keys = dir(configs) modules = {} for k in keys: type_ = configs._get_type(k) try: if issubclass(type_, torch.nn.Module): modules[k] = getattr(configs, k) except TypeError: pass return modules
[docs]def get_device(module: torch.nn.Module): """ Get the device the ``module`` is in. Arguments: module (torch.nn.Module): PyTorch module """ params = module.parameters() try: sample_param = next(params) return sample_param.device except StopIteration: raise RuntimeError(f"Unable to determine" f" device of {module.__class__.__name__}") from None