aihwkit.nn.modules.linear module

Analog layers.

class aihwkit.nn.modules.linear.AnalogLinear(in_features, out_features, bias=True, rpu_config=None, tile_module_class=None)[source]

Bases: AnalogLayerBase, Linear

Linear layer that uses an analog tile.

Linear layer that uses an analog tile during its forward, backward and update passes.

Note

The tensor parameters of this layer (.weight and .bias) are not guaranteed to contain the same values as the internal weights and biases stored in the analog tile. Please use set_weights and get_weights when attempting to read or modify the weight/bias. This read/write process can simulate the (noisy and inexact) analog writing and reading of the resistive elements.

Parameters:

  • in_features (int) – input vector size (number of columns).

  • out_features (int) – output vector size (number of rows).

  • bias (bool) – whether to use a bias row on the analog tile or not. for setting initial weights and during reading of the weights.

  • rpu_config (RPUConfigBase | None) – resistive processing unit configuration.

  • tile_module_class (Type | None) – Class for the tile module (default will be specified from the RPUConfig).

forward(x_input)[source]

Compute the forward pass.

Parameters:

x_input (Tensor) –

Return type:

Tensor

classmethod from_digital(module, rpu_config, tile_module_class=None)[source]

Return an AnalogLinear layer from a torch Linear layer.

Parameters:

  • module (Linear) – The torch module to convert. All layers that are defined in the conversion_map.

  • rpu_config (RPUConfigBase) – RPU config to apply to all converted tiles. Applied to all converted tiles.

  • tile_module_class (Type | None) – Class of the underlying TileModule. If not given, will select based on the MappingParameter setting either TileModule or TileModuleArray

Returns:

an AnalogLinear layer based on the digital Linear module.

Return type:

AnalogLinear

reset_parameters()[source]

Reset the parameters (weight and bias).

Return type:

None

classmethod to_digital(module, realistic=False)[source]

Return an nn.Linear layer from an AnalogLinear layer.

Parameters:

  • module (AnalogLinear) – The analog module to convert.

  • realistic (bool) – whehter to estimate the weights with the non-ideal forward pass. If not set, analog weights are (unrealistically) copies exactly

Returns:

an torch Linear layer with the same dimension and weights as the analog linear layer.

Return type:

Linear