aihwkit.simulator.digital_low_precision.hijacker module

class aihwkit.simulator.digital_low_precision.hijacker.QuantizationHijacker(*args, activation=None, **kwargs)

Bases: QuantizedModule

Mixin class that ‘hijacks’ the forward pass in a module to perform quantization and dequantization on the weights and output distributions.

Usage: To make a quantized nn.Linear layer: ` >>> class QuantLinear(QuantizationHijacker, nn.Linear): ...     pass `

It is vital that QSchemeForwardHijacker is the first parent class, and that the second parent class derives from nn.Module, otherwise it will not be reached by a super(., .) call.

NB: this implementation (for now) assumes that there will always be some training involved, e.g. to estimate the activation ranges.

Parameters:

activation (Module)

forward(x, offsets=None)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

get_params()

get_weight_bias()

quantize_activations(activations)

Quantize a single activation tensor or all activations from a layer. I’m assuming that we should quantize all outputs for a layer with the same quantization scheme.

run_forward(x, weight, bias, offsets=None)