aihwkit_lightning.inference.calibration.calibration module

Calibration for inference.

class aihwkit_lightning.inference.calibration.calibration.InputRangeCalibrationType(value)[source]

Bases: Enum

Input range post-training calibration type.

Different styles of calibrating the DAC ranges post-training.

CACHE_QUANTILE = 'CacheQuantile'

Caches inputs that are then used to compute the Xth quantile for the input range.

MAX = 'Max'

Takes the abs().max() over the inputs.

MOVING_QUANTILE = 'MovingQuantile'

Computes the moving average of the quantiles. Saves memory.

MOVING_STD = 'MovingStd'

Computes a moving average of x*standard deviation of the inputs.

NONE = 'None'

No Calibration.

aihwkit_lightning.inference.calibration.calibration.calibrate_input_ranges(model, calibration_type, dataloader, quantile=0.99995, max_samples=1000, std_alpha=None, verbose=True)[source]

Calibrate the input ranges according to the defined strategy.

Only tiles that support and have enabled input range learning will be calibrated. If noise management is turned on an error is raised.

Note

This implementation transiently registers a new forward_pre_hook on the analog tile level. It assumes that the user has not defined any other forward prehooks.

Parameters:

  • model (AnalogWrapper) – The analog model for which to calibrate the input ranges.

  • calibration_type (InputRangeCalibrationType) – Strategy of the calibration. See InputRangeCalibrationType

  • dataloader (Iterator) – Iterator that yields the next inputs. Is used like this x = next(dataloader); model(x)

  • quantile (float) – Quantile used for hard-coded quantile setting. Defaults to 0.99995.

  • max_samples (int) – Max batch samples to cache in each tile. Defaults to 1000.

  • std_alpha (float | None) – Number of standard deviations for moving standard deviation strategy. Defaults to init_std_alpha from RPUConfig

  • verbose (bool) – Whether to print verbose output.

Raises:

ConfigError – If RPUConfig does not support input range learning

Return type:

None