aihwkit.nn.low_precision_modules.quantized_base_modules module

Basic quantized modules

class aihwkit.nn.low_precision_modules.quantized_base_modules.QuantBatchNorm2d(org_model, **quant_params)[source]

Bases: QuantizedModel

Quantization of the BatchNorm2d module. output activations are quantized.

Parameters:

  • org_model (Module)

  • quant_params (Any)

forward(x)[source]

Execute BatchNorm2d and then quantize its output

Parameters:

x (Tensor)

Return type:

Tensor

class aihwkit.nn.low_precision_modules.quantized_base_modules.QuantConv2d(*args, activation=None, **kwargs)[source]

Bases: QuantizationHijacker, Conv2d

Quantized layer of torch.nn.Conv2d with weight/act quantization

Parameters:

activation (Module)

run_forward(x, weight, bias, offsets=None)[source]
Parameters:

  • x (Tensor)

  • weight (Tensor)

  • bias (Tensor)

  • offsets (Any | None)

Return type:

Tensor

class aihwkit.nn.low_precision_modules.quantized_base_modules.QuantEmbedding(*args, activation=None, **kwargs)[source]

Bases: QuantizationHijacker, Embedding

Quantization of the Embedding, weight quantization. Note: Embedding should not quantize activations, as it is simply a lookup table, which is already quantized.

Parameters:

  • args (Any)

  • activation (Any | None)

  • kwargs (Any)

run_forward(x, weight, bias, offsets=None)[source]
Parameters:

  • x (Tensor)

  • weight (Tensor)

  • bias (Tensor)

  • offsets (Any | None)

Return type:

Tensor

class aihwkit.nn.low_precision_modules.quantized_base_modules.QuantLayerNorm(*args, activation=None, **kwargs)[source]

Bases: QuantizationHijacker, LayerNorm

Quantized layer of torch.nn.LayerNorm with input and weight quantization

Parameters:

activation (Module)

run_forward(x, weight, bias, offsets=None)[source]
Parameters:

  • x (Tensor)

  • weight (Tensor)

  • bias (Tensor)

  • offsets (Any | None)

Return type:

Tensor

class aihwkit.nn.low_precision_modules.quantized_base_modules.QuantLinear(*args, activation=None, **kwargs)[source]

Bases: QuantizationHijacker, Linear

Quantized layer of torch.nn.Linear with weight/act quantization

Parameters:

activation (Module)

run_forward(x, weight, bias, offsets=None)[source]
Parameters:

  • x (Tensor)

  • weight (Tensor)

  • bias (Tensor)

  • offsets (Any | None)

Return type:

Tensor