aihwkit.nn.conversion module¶
Digital/analog model conversion utilities.
This module includes tools for converting a given torch model to a model containing analog layers.
-
aihwkit.nn.conversion.
convert_to_analog
(module, rpu_config, realistic_read_write=False, weight_scaling_omega=0.0, conversion_map=None)[source]¶ Convert a given digital model to analog counter parts.
Note
The torch device (cuda/cpu) is inferred from the original models parameters, however, if multiple torch devices are used in a given module, the corresponding analog module is not moved to any device.
- Parameters
module (torch.nn.modules.module.Module) – The torch module to convert. All layers that are defined in the
conversion_map
.rpu_config (RPUConfigGeneric) – RPU config to apply to all converted tiles. Applied to all converted tiles.
realistic_read_write (bool) – Whether to use closed-loop programming when setting the weights. Applied to all converted tiles.
weight_scaling_omega (float) –
If non-zero, the analog weights will be scaled by
weight_scaling_omega
divided by the absolute maximum value of the original weight matrix.Note
Make sure that the weight max and min settings of the device support the desired analog weight range.
conversion_map (Optional[Dict]) –
Dictionary of module classes to be replaced in case of custom replacement rules. By default all
Conv
andLinear
layers are replaced with their analog counterparts.Note
The analog layer needs to have a class method
from_digital
which will be called during the conversion.
- Returns
Module where all the digital layers are replaced with analog mapped layers.
- Return type
torch.nn.modules.module.Module
-
aihwkit.nn.conversion.
convert_to_analog_mapped
(module, rpu_config, realistic_read_write=False, weight_scaling_omega=0.0)[source]¶ Convert a given digital model to its analog counterpart with tile mapping support.
Note
The torch device (cuda/cpu) is inferred from the original models parameters, however, if multiple torch devices are used in a given module, the corresponding analog module is not moved to any device.
- Parameters
module (torch.nn.modules.module.Module) – The torch module to convert. All layers that are defined in the
conversion_map
.rpu_config (RPUConfigGeneric) – RPU config to apply to all converted tiles.
realistic_read_write (bool) – Whether to use closed-loop programming when setting the weights. Applied to all converted tiles.
weight_scaling_omega (float) –
If non-zero, the analog weights will be scaled by
weight_scaling_omega
divided by the absolute maximum value of the original weight matrix.Note
Make sure that the weight max and min settings of the device support the desired analog weight range.
- Returns
Module where all the digital layers are replaced with analog mapped layers.
- Return type
torch.nn.modules.module.Module