aihwkit.nn.modules.linear_mapped module

Analog mapped layers.

class aihwkit.nn.modules.linear_mapped.AnalogLinearMapped(in_features, out_features, bias=True, rpu_config=None, realistic_read_write=False, weight_scaling_omega=None)[source]

Bases: aihwkit.nn.modules.base.AnalogModuleBase, torch.nn.modules.linear.Linear

Linear layer that uses an analog tile.

Linear layer that uses an analog tile during its forward, backward and update passes. In contrast to Linear the maximal in and/or out dimension can be restricted, in which case the linear layer is split into multiple parts and computed on multiple tiles of given max sizes.

In contrast to Linear, the bias vector (if requested) is always handled in digital (floating point).

Note

Mapping is controlled by the aihwkit.simulator.configs.utils.MappingParameter.

Note

The tensor parameters of this layer (.weight and .bias) are not guaranteed to contain the same values as the internal weights and biases stored in the analog tile. Please use set_weights and get_weights when attempting to read or modify the weight/bias. This read/write process can simulate the (noisy and inexact) analog writing and reading of the resistive elements.

Parameters
analog_bias: bool
digital_bias: bool
extra_repr()[source]

Set the extra representation of the module.

Returns

A string with the extra representation.

Return type

str

forward(x_input)[source]

Compute the forward pass.

Parameters

x_input (torch.Tensor) –

Return type

torch.Tensor

classmethod from_digital(module, rpu_config=None, realistic_read_write=False)[source]

Return an AnalogLinearMapped layer from a torch Linear layer.

Parameters
Returns

an AnalogLinearMapped layer based on the digital Linear module.

Return type

aihwkit.nn.modules.linear_mapped.AnalogLinearMapped

get_split_sizes(size, split_max_size)[source]

Computed the split sizes.

Parameters

  • size (int) – number of elements of the layer in one dimension

  • split_max_size (int) – max size of the split

Returns

List of split sizes

Return type

List[int]

get_weights(force_exact=False, apply_out_scales=True)[source]

Get the weight (and bias) tensors.

This uses an realistic read if the property realistic_read_write of the layer is set, unless it is overwritten by force_exact. It scales the analog weights by the digital alpha scale if weight_scaling_omega is positive (see get_weights_scaled()).

Note

This is the recommended way for setting the weight/bias matrix from the analog tile, as it will correctly fetch the weights from the internal memory. Accessing self.weight and self.bias might yield wrong results as they are not always in sync with the analog tile library, for performance reasons.

Parameters

  • force_exact (bool) – forces an exact read to the analog tiles

  • apply_out_scales (bool) – Whether to return the weights with the (digital) output scaling factors applied. Note the “logical” weights of the layer which the DNN is effectively using are those with the output scales applied. If apply_out_scales is set to False, then only the weight values that is programmed onto the crossbar array are returned, without applying the digital scales.

Returns

weight matrix, bias vector

Return type

tuple

in_features: int
in_sizes: List[int]
out_features: int
out_sizes: List[int]
realistic_read_write: bool
reset_parameters()[source]

Reset the parameters (weight and bias).

Return type

None

set_weights(weight, bias=None, force_exact=False, remap_weights=True, weight_scaling_omega=None)[source]

Set the weight (and bias) with given Tensors.

This uses an realistic write if the property realistic_read_write of the layer is set, unless it is overwritten by force_exact. It uses a scaled write if weight_scaling_omega is positive (see set_weights_scaled()).

Note

This is the recommended way for setting the weight/bias matrix of the analog tile, as it will correctly store the weights into the internal memory. Directly writing to self.weight and self.bias might yield wrong results as they are not always in sync with the analog tile Parameters, for performance reasons.

Parameters

  • weight (torch.Tensor) – weight matrix

  • bias (Optional[torch.Tensor]) – bias vector

  • force_exact (bool) – forces an exact write to the analog tiles

  • remap_weights (bool) –

    Whether to rescale the given weight matrix and populate the digital output scaling factors as specified in the configuration MappingParameter. A new weight_scaling_omega can be given. Note that this will overwrite the existing digital out scaling factors.

    Note that each tile (in case of multiple mapped tiles) has it separate out scaling factors.

  • weight_scaling_omega (Optional[float]) – The weight scaling omega factor (see MappingParameter). If given explicitly here, it will overwrite the value in the mapping field.

Return type

None

use_bias: bool
weight_scaling_omega: float