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, specific_rpu_config_fun=None, module_name='')[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.

  • specific_rpu_config_fun (Optional[Callable]) – Function that modifies the generic RPUConfig for specific modules. See :fun:`~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.

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, specific_rpu_config_fun=None, module_name='')[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.

  • specific_rpu_config_fun (Optional[Callable]) – Function that modifies the generic RPUConfig for specific modules. See :fun:`~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.

Returns

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

Return type

torch.nn.modules.module.Module

aihwkit.nn.conversion.specific_rpu_config_id(module_name, module, rpu_config)[source]

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 (torch.nn.modules.module.Module) – the actual digital module to be converted

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

Returns

modified RPUConfig

Return type

aihwkit.nn.conversion.RPUConfigGeneric