aihwkit.nn.modules.linear module

Analog layers.

class aihwkit.nn.modules.linear.AnalogLinear(in_features, out_features, bias=True, rpu_config=None, realistic_read_write=False, weight_scaling_omega=0.0)

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.

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

  • in_features – input vector size (number of columns).

  • out_features – output vector size (number of rows).

  • rpu_config – resistive processing unit configuration.

  • bias – whether to use a bias row on the analog tile or not

  • realistic_read_write – whether to enable realistic read/write for setting initial weights and read out of weights

  • weight_scaling_omega – the weight value where the max weight will be scaled to. If zero, no weight scaling will be performed

forward(x_input)

Computes the forward pass.

Parameters

x_input (torch.Tensor) –

Return type

torch.Tensor

in_features: int
out_features: int
realistic_read_write: bool
reset_parameters()

Reset the parameters (weight and bias).

Return type

None

weight_scaling_omega: float