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)

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, applied weights of analog layers 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 setting 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 with replaced digital layers with analog layers.

Return type

torch.nn.modules.module.Module