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
andbias
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 theset_weights()
andget_weights()
as opposed to using the attributes directly.
-
TILE_CLASS_ANALOG
¶
-
TILE_CLASS_FLOATING_POINT
¶ alias of
aihwkit.simulator.tiles.floating_point.FloatingPointTile
-
TILE_CLASS_INFERENCE
¶
-
cpu
()¶ Moves all model parameters, buffers and tiles to the CPU.
Note
Please be aware that moving analog layers from GPU to CPU is currently not supported.
- Returns
This layer with its parameters, buffers and tiles in CPU.
- Return type
-
cuda
(device=None)¶ Moves all model parameters, buffers and tiles to the GPU.
This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on GPU while being optimized.
- Parameters
device (int, optional) – if specified, all parameters will be copied to that GPU device
- Returns
This layer with its parameters, buffers and tiles in GPU.
- Return type
-
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 byforce_exact
. It scales the analog weights by the digital alpha scale ifweight_scaling_omega
is positive (seeget_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
andself.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
-
set_weights
(weight, bias=None, force_exact=False)¶ Sets 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 byforce_exact
. It uses a scaled write ifweight_scaling_omega
is positive (seeset_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
andself.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)¶ Returns a dictionary containing a whole state of the module.
- Parameters
destination (Any) –
prefix (str) –
keep_vars (bool) –
- Return type
Dict