aihwkit.nn.modules.base module

Base class for analog Modules.

class aihwkit.nn.modules.base.AnalogModuleBase

Bases: torch.nn.modules.module.Module

Base class for analog Modules.

Base Module for analog layers that use analog tiles. When subclassing, please note:

  • the _setup_tile() method is expected to be called by the subclass constructor, and it does not only create a tile, but also sets some instance attributes that are needed by the analog features (optimizer and others).

  • the weight and bias Parameters are not guaranteed to be in sync with the tile weights and biases during the lifetime of the instance, for performance reasons. The canonical way of reading and writing weights is via the set_weights() and get_weights() as opposed to using the attributes directly.

  • the BaseTile subclass that is created is retrieved from the rpu_config.tile_class attribute.

drift_analog_weights(t_inference=0.0)

(Program) and drift the analog weights.

Parameters

t_inference (float) – assumed time of inference (in sec)

Raises

ModuleError – if the layer is not in evaluation mode.

Return type

None

extra_repr()

Set the extra representation of the module.

Returns

A string with the extra representation.

Return type

str

get_weights(force_exact=False)

Get the weight (and bias) tensors.

This uses an realistic read if the property realistic_read_write of the layer is set, unless it is overwritten by force_exact. It scales the analog weights by the digital alpha scale if weight_scaling_omega is positive (see get_weights_scaled()).

Note

This is the recommended way for setting the weight/bias matrix from the analog tile, as it will correctly fetch the weights from the internal memory. Accessing self.weight and self.bias might yield wrong results as they are not always in sync with the analog tile library, for performance reasons.

Parameters

force_exact (bool) – forces an exact read to the analog tiles

Returns

weight matrix, bias vector

Return type

tuple

program_analog_weights()

Program the analog weights.

Raises

ModuleError – if the layer is not in evaluation mode.

Return type

None

register_analog_tile(tile)

Register the analog context of the tile.

Note

Needs to be called at the end init to register the tile for the analog optimizers.

Parameters

tile (BaseTile) – tile to register

Return type

None

set_weights(weight, bias=None, force_exact=False)

Set the weight (and bias) with given Tensors.

This uses an realistic write if the property realistic_read_write of the layer is set, unless it is overwritten by force_exact. It uses a scaled write if weight_scaling_omega is positive (see set_weights_scaled()).

Note

This is the recommended way for setting the weight/bias matrix of the analog tile, as it will correctly store the weights into the internal memory. Directly writing to self.weight and self.bias might yield wrong results as they are not always in sync with the analog tile Parameters, for performance reasons.

Parameters

  • weight (torch.Tensor) – weight matrix

  • bias (Optional[torch.Tensor]) – bias vector

  • force_exact (bool) – forces an exact write to the analog tiles

Return type

None

state_dict(destination=None, prefix='', keep_vars=False)

Return a dictionary containing a whole state of the module.

Parameters

  • destination (Any) –

  • prefix (str) –

  • keep_vars (bool) –

Return type

Dict

unregister_parameter(param_name)

Unregister module parameter from parameters.

Raises

ModuleError – In case parameter is not found

Parameters

param_name (str) –

Return type

None