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, tile_module_class=None, conversion_map=None, specific_rpu_config_fun=None, module_name='', ensure_analog_root=True, exclude_modules=None, inplace=False, verbose=False)

Convert a given digital model to its analog counterpart.

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 (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.

• tile_module_class (TileModule | None) – Custom tile module class

• conversion_map (Dict | None) –

  Dictionary of module classes to be replaced in case of custom replacement rules. By default all Conv and Linear 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.

• specific_rpu_config_fun (Callable | None) – Function that modifies the generic RPUConfig for specific modules. See specific_rpu_config_id() as an example how to specify it.

• module_name (str) – Explicitly given name of the base (root) module, given to specific_rpu_config_fun.

• ensure_analog_root (bool) –

  Whether to ensure that the root module is of layer type AnalogLayerBase so that custom analog are methods such as drift_analog_weigths are available. If set, it will wrap the model if AnalogWrapper if necessary.

  Note

  Since the module structure changes when wrapped, the checkpoint names will also change if this is enabled (for legacy load this might need to be disabled).

• exclude_modules (List[str] | None) – List of module names that are in the conversion map but should be excluded from the conversion

• inplace (bool) – Whether to perform in-place conversion (without deepcopy)

• verbose (bool) – Increase verbosity. Will print converted layers.

Returns:

Module where all the digital layers are replaced with analog mapped layers.

Raises:

ArgumentError – in case conversion map is not a dictionary

Return type:

Module

aihwkit.nn.conversion.convert_to_analog_mapped(module, rpu_config, tile_module_class=None, specific_rpu_config_fun=None, module_name='', ensure_analog_root=True, exclude_modules=None, inplace=False, verbose=False)

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 (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.

• tile_module_class (TileModule | None) – Custom tile module class

• specific_rpu_config_fun (Callable | None) – Function that modifies the generic RPUConfig for specific modules. See specific_rpu_config_id() as an example how to specify it.

• module_name (str) – Explicitly given name of the base (root) module, given to specific_rpu_config_fun.

• ensure_analog_root (bool) –

  Whether to ensure that the root module is of layer type AnalogLayerBase so that custom analog are methods such as drift_analog_weigths are available. If set, it will wrap the model if AnalogWrapper if necessary.

  Note

  Since the module structure changes when wrapped, the checkpoint names will also change if this is enabled (for legacy load this might need to be disabled).

• exclude_modules (List[str] | None) – List of modules names that are in the conversion map but should be excluded from the conversion

• inplace (bool) – Whether to for in place conversion (without deepcopy)

• verbose (bool) – Increase verbosity. Will print converted layers.

Returns:

Module where all the digital layers are replaced with analog mapped layers.

Return type:

Module

aihwkit.nn.conversion.convert_to_digital(module, conversion_set=None, module_name='', inplace=False)

Convert a given analog model to digital 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 (Module) – The torch module to convert having analog layers. All analog layer classes to be converted are defined in the conversion_set.

• conversion_set (Set | None) –

  Set of analog module classes to be replaced. By default all AnalogConv* and AnalogLinear* layers are replaced with their digital counterparts.

  Note

  The analog layer needs to have a class method to_digital which will be called during the conversion.

• module_name (str) – Explicitly given name of the base (root) module

• inplace (bool) – Whether to for in place conversion (without deepcopy)

Returns:

Module where all the analog layers are replaced with digital layers

Raises:

ArgumentError – in case conversion map is not a set or contains layers that cannot be converted

Return type:

Module

aihwkit.nn.conversion.specific_rpu_config_id(module_name, module, rpu_config)

ID default function for specifying the RPUConfig during conversion for specific layers.

A similar function can be given to the conversion.

Parameters:

• module_name (str) – The name of the module currently converted to analog

• module (Module) – the actual digital module to be converted

• rpu_config (RPUConfigGeneric) – a copy of the generic RPUConfig given to the overall conversion.

Returns:

modified RPUConfig

Return type:

RPUConfigGeneric