aihwkit.simulator.tiles.inference module

High level analog tiles (inference).

class aihwkit.simulator.tiles.inference.InferenceTile(out_size, in_size, rpu_config=None, bias=False, in_trans=False, out_trans=False, shared_weights=True)

Bases: TileModule, InferenceTileWithPeriphery, RPUCudaSimulatorTileWrapper

Tile used for analog inference and hardware-aware training for inference.

Note

This tile uses RPUCuda library with backward and update set to perfect.

Parameters:

• out_size (int) – output size

• in_size (int) – input size

• rpu_config (InferenceRPUConfig | None) – resistive processing unit configuration.

• bias (bool) – whether to add a bias column to the tile.

• in_trans (bool) – Whether to assume an transposed input (batch first)

• out_trans (bool) – Whether to assume an transposed output (batch first)

• shared_weights (bool) – Whether to keep the weight in torch’s memory space

forward(x_input, tensor_view=None)

Torch forward function that calls the analog forward

Parameters:

• x_input (Tensor) –

• tensor_view (Tuple | None) –

Return type:

Tensor

supports_ddp: bool = True

class aihwkit.simulator.tiles.inference.InferenceTileWithPeriphery

Bases: TileWithPeriphery

Additional (peripheral) functionality for hardware-aware training and inference.

Note

Here it is assumed that the training is done in software and only the inference pass is done on analog hardware.

cpu()

Return a copy of this tile in CPU memory.

Returns:

Self with the underlying buffers moved to CPU memory.

Return type:

BaseTile

cuda(device=None)

Return a copy of this tile in CUDA memory.

Parameters:

device (str | device | int | None) – CUDA device

Returns:

Self with the underlying buffers to CUDA memory.

Return type:

BaseTile

drift_weights(t_inference=0.0)

Programs and drifts the current reference weights.

The current weight reference is either the current weights or the ones at the time when initialize_drift_reference() was called, which then would overwrite the current weights with the drifted ones.

Parameters:

t_inference (float) – Time (in sec) of assumed inference time. Programming ends at t=0s. The rest is waiting time, where the devices might drift and accumulate noise. See noise model used for details.

Raises:

ConfigError – in case of noise_model is not defined in the RPUConfig

Return type:

None

init_mapping_scales()

Helper function to initialize the mapping scales used to scale the weights in digital and determine the conductance conversion.

Note

This method is called from the constructor.

Return type:

None

post_forward(x_output, dim, is_test=False, ctx=None)

Operations after the actual forward step for post processing

Parameters:

• x_output (Tensor) –

• dim (int) –

• is_test (bool) –

• ctx (Any | None) –

Return type:

Tensor

post_update_step()

Operators that need to be called once per mini-batch.

In the InferenceTile, the following calls are made (if enabled in the rpu_config settings). First, the post update step of the parent is called, then the weight clipping is done, subsequently then remapping is done (if enforced), and finally the forward-backward weight modifier is called. The latter will modify the weights that are used during forward and backward (but not update) until the next time this function is called.

Return type:

None

program_weights(from_reference=True, noise_model=None)

Apply weights noise to the current tile weights and saves these for repeated drift experiments.

This method also establishes the drift coefficients for each conductance slice.

Will also reset the drift readout tensor and compute a new drift compensation baseline

Parameters:

• from_reference (bool) – Whether to use weights from reference

• noise_model (BaseNoiseModel | None) –

  Optional defining the noise model to be used. If not given, it will use the noise model defined in the RPUConfig of this tile.

  Caution

  If given a noise model here it will overwrite the stored rpu_config.noise_model definition.

Raises:

ConfigError – in case of noise_model is not defined in the RPUConfig or the given noise model is of the wrong type

Return type:

None