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, 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 (aihwkit.nn.conversion.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.

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

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)[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 (aihwkit.nn.conversion.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.

    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