Function print_all

SwissArmyTransformer/sat/helpers.py:139–146

(msg, level=logging.INFO, flush=True)

        logger.log(level=level, msg=msg)

def print_all(msg, level=logging.INFO, flush=True):
    if isinstance(level, str):
        level = getattr(logging, level.upper())
    if torch.distributed.is_initialized():
        msg = f"[RANK {torch.distributed.get_rank()}] {msg}"
    logger.log(level=level, msg=msg)
    if flush:
        logger.handlers[0].flush()

def debug_param(name, param):
    print_all(f"param: {name}, min: {param.min()}, max: {param.max()}, mean: {param.mean()}, std: {param.std()}, scale: {param.abs().mean()}, first5: {param.flatten()[:5]}, last5: {param.flatten()[-5:]}")