fastreid.solver.optim.lamb Namespace Reference

Classes
class	Lamb

Functions
	log_lamb_rs (Optimizer optimizer, SummaryWriter event_writer, int token_count)

Function Documentation

log_lamb_rs()

fastreid.solver.optim.lamb.log_lamb_rs	(	Optimizer	optimizer,
		SummaryWriter	event_writer,
		int	token_count )

Log a histogram of trust ratio scalars in across layers.

Definition at line 12 of file lamb.py.

def log_lamb_rs(optimizer: Optimizer, event_writer: SummaryWriter, token_count: int):
    """Log a histogram of trust ratio scalars in across layers."""
    results = collections.defaultdict(list)
    for group in optimizer.param_groups:
        for p in group['params']:
            state = optimizer.state[p]
            for i in ('weight_norm', 'adam_norm', 'trust_ratio'):
                if i in state:
                    results[i].append(state[i])
 
    for k, v in results.items():
        event_writer.add_histogram(f'lamb/{k}', torch.tensor(v), token_count)