adaptdl.torch.gradient_noise_scale module

class adaptdl.torch.gradient_noise_scale.GradientNoiseScale(adp, optimizer, mp_scaler=None, num_replicas=None, accum_scale=None)

Bases: object

This class tracks gradient related stats and takes care of gradient accumulation.

property accum_count

property accum_scale

gain(scale)

Current estimate of the GradientNoiseScale gain ratio.

Parameters

scale (float) – The total scale to estimate the gain ratio for.

Returns (float): Estimate of gain ratio.

get_progress()

property raw_sqr_avg

property raw_var_avg

reset_accumulation()

reset accumulation calculations and gradients.

set_accum_scale(accum_scale)

set_progress(progress)

property should_zero_grad

sqr_avg()

Current estimate of the squared l2-norm of the true gradient (sigma squared).

Returns (float): Estimate of squared l2-norm.

var_avg()

Current estimate of the trace of the covariance of the true gradient (mu squared).

Returns (float): Estimate of trace of the covariance.