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:
- classmethod to_digital(module, *args, **kwargs)[source]
Construct Sequential in-place from AnalogSequential.
- Parameters:
module (AnalogSequential) –
args (Any) –
kwargs (Any) –
- 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:
- 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:
- classmethod to_digital(module, *args, **kwargs)[source]
Construct Sequential in-place from AnalogSequential.
- Parameters:
module (AnalogWrapper) –
args (Any) –
kwargs (Any) –
- Return type:
Module