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: Generic[aihwkit.simulator.tiles.base.RPUConfigGeneric]

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

Parameters

• out_size – output size

• in_size – input size

• rpu_config – resistive processing unit configuration.

• bias – whether to add a bias column to the tile.

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

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

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

cuda(device=None)

Return a copy of this tile in CUDA memory.

Parameters

device (Optional[Union[torch.device, str, int]]) – CUDA device

Returns

Self with the underlying C++ tile moved to CUDA memory.

Raises

CudaError – if the library has not been compiled with CUDA.

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.

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 (torch.Tensor) –

• dim (int) –

• is_test (bool) –

• ctx (Any) –

Return type

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

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 compuate a new drift compensation baseline

Parameters

from_reference (bool) – Whether to use weights from reference

Return type

None