aihwkit.nn.modules.container module

Analog Modules that contain children Modules.

class aihwkit.nn.modules.container.AnalogSequential(*args)

Bases: torch.nn.modules.container.Sequential

An analog-aware sequential container.

Specialization of torch nn.Sequential with extra functionality for handling analog layers:

• correct handling of .cuda() for children modules.

• apply analog-specific functions to all its children (drift and program weights).

Note

This class is recommended to be used in place of nn.Sequential in order to correctly propagate the actions to all the children analog layers. If using regular containers, please be aware that operations need to be applied manually to the children analog layers when needed.

cpu()

Moves all model parameters and buffers to the CPU.

Note

This method modifies the module in-place.

Returns

self

Return type

Module

cuda(device=None)

Moves all model parameters and buffers to the GPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on GPU while being optimized.

Note

This method modifies the module in-place.

Parameters

device (int, optional) – if specified, all parameters will be copied to that device

Returns

self

Return type

Module

drift_analog_weights(t_inference=0.0)

(Program) and drift all analog inference layers of a given model.

Parameters

t_inference (float) – assumed time of inference (in sec)

Raises

ModuleError – if the layer is not in evaluation mode.

Return type

None

classmethod from_digital(module, *args, **kwargs)

Construct AnalogSequential in-place from Sequential.

Parameters

• module (torch.nn.modules.container.Sequential) –

• args (Any) –

• kwargs (Any) –

Return type

aihwkit.nn.modules.container.AnalogSequential

load_state_dict(state_dict, strict=True, load_rpu_config=True)

Specializes torch’s load_state_dict to add a flag whether to load the RPU config from the saved state.

Parameters

• state_dict (OrderedDict[str, Tensor]) – see torch’s load_state_dict

• strict (bool) – see torch’s load_state_dict

• load_rpu_config (bool) –

  Whether to load the saved RPU config or use the current RPU config of the model.

  Caution

  If load_rpu_config=False the RPU config can be changed from the stored model. However, the user has to make sure that the changed RPU config makes sense.

  For instance, changing the device type might change the expected fields in the hidden parameters and result in an error.

Returns

see torch’s load_state_dict

Return type

NamedTuple

Raises: ModuleError: in case the rpu_config class mismatches

for load_rpu_config=False.

prepare_for_ddp()

Adds ignores to avoid broadcasting the analog tile states in case of distributed training.

Note

Call this function before the mode is converted with DDP.

Important

Only InferenceTile supports DDP.

Raises

ModuleError – In case analog tiles are used that do not support data-parallel model, ie. all analog training tiles.

Return type

None

program_analog_weights()

Program all analog inference layers of a given model.

Raises

ModuleError – if the layer is not in evaluation mode.

Return type

None

to(device=None)

Move and/or cast the parameters, buffers and analog tiles.

Note

Please be aware that moving analog layers from GPU to CPU is currently not supported.

Parameters

device (Optional[Union[torch.device, str, int]]) – the desired device of the parameters, buffers and analog tiles in this module.

Returns

This module in the specified device.

Return type

aihwkit.nn.modules.container.AnalogSequential