aihwkit.nn.modules.container module

Analog Modules that contain children Modules.

class aihwkit.nn.modules.container.AnalogContainerBase[source]

Bases: AnalogLayerBase

Base class for analog containers.

IS_CONTAINER: bool = True

Class constant indicating whether sub-layers exist or whether this layer is a leave node (that is only having tile modules)

get_weights(**kwargs)[source]

Returns all weights, bias tuples in an ordered dictionary.

Parameters:

kwargs (Any) – passed to the TileModule get_weights call

Returns:

All analog weight of all layers

Return type:

OrderedDict[str, Tuple[Tensor, Tensor | None]]

set_weights(weights_dic, **kwargs)[source]

Set all analog weights part of this parent module.

Parameters:
  • weights_dic (OrderedDict[str, Tuple[Tensor, Tensor | None]]) – Ordered dictionary of weight data

  • kwargs (Any) – passed to the TileModule set_weights call

Raises:

ModuleError – in case tile name cannot be found

Return type:

None

class aihwkit.nn.modules.container.AnalogSequential(*args: Module)[source]
class aihwkit.nn.modules.container.AnalogSequential(arg: OrderedDict[str, Module])

Bases: AnalogContainerBase, Sequential

An analog-aware sequential container.

Specialization of torch nn.Sequential with extra functionality for handling analog layers:

  • apply analog-specific functions to all its children (drift and program weights).

Note

This class is recommended to be used in place of nn.Sequential in order to correctly propagate the actions to all the children analog layers. If using regular containers, please be aware that operations need to be applied manually to the children analog layers when needed.

classmethod from_digital(module, *args, **kwargs)[source]

Construct AnalogSequential in-place from Sequential.

Parameters:
  • module (Sequential) –

  • args (Any) –

  • kwargs (Any) –

Return type:

AnalogSequential

classmethod to_digital(module, *args, **kwargs)[source]

Construct Sequential in-place from AnalogSequential.

Parameters:
Return type:

Sequential

class aihwkit.nn.modules.container.AnalogWrapper(module=None, **__)[source]

Bases: AnalogContainerBase

Generic wrapper over an given Module.

Will add the AnalogLayerBase functionality to the given Module (as an added subclass).

Note

Here the state dictionary of the give module will be simply copied by reference. The original model therefore should not be used any more as the underlying tensor data is shared.

Parameters:
  • model – model to wrap with the analog wrapper.

  • module (Module | None) –

  • __ (Any) –

Return type:

AnalogWrapper

SUBCLASSES: Dict[str, Type] = {}

Registry of the created subclasses.

classmethod from_digital(module, *args, **kwargs)[source]

Construct AnalogSequential in-place from any module.

Parameters:
  • module (Module) –

  • args (Any) –

  • kwargs (Any) –

Return type:

AnalogWrapper

classmethod to_digital(module, *args, **kwargs)[source]

Construct Sequential in-place from AnalogSequential.

Parameters:
Return type:

Module