aihwkit.simulator.tiles.base module

High level analog tiles (base).

class aihwkit.simulator.tiles.base.BaseTile(*args, **kwds)

Bases: typing.Generic

Base class for tiles.

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)

backward(d_input)

Perform the backward pass.

Parameters

d_input (torch.Tensor) – [N, out_size] tensor. If out_trans is set, transposed.

Returns

[N, in_size] tensor. If in_trans is set, transposed.

Return type

torch.Tensor

backward_indexed(d_input)

Perform the backward pass for convolutions.

Parameters

d_input (torch.Tensor) – [N, out_size] tensor. If out_trans is set, transposed.

Returns

[N, in_size] tensor. If in_trans is set, transposed.

Return type

torch.Tensor

cuda(device=None)

Return a copy of this tile in CUDA memory.

Parameters

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

Return type

aihwkit.simulator.tiles.base.BaseTile

decay_weights(alpha=1.0)

Decays the weights once.

Parameters

alpha (float) – additional decay scale (such as LR). The base decay rate is set during tile init.

Return type

None

diffuse_weights()

Diffuses the weights once.

The base diffusion rate is set during tile init.

Return type

None

forward(x_input, is_test=False)

Perform the forward pass.

Parameters

  • x_input – [N, in_size] tensor. If in_trans is set, transposed.

  • is_test – whether to assume testing mode.

Returns

[N, out_size] tensor. If out_trans is set, transposed.

Return type

torch.Tensor

forward_indexed(x_input, is_test=False)

Perform the forward pass for convolutions.

Parameters

  • x_input – [N, in_size] tensor. If in_trans is set, transposed.

  • is_test – whether to assume testing mode.

Returns

[N, out_size] tensor. If out_trans is set, transposed.

Return type

torch.Tensor

get_hidden_parameters()

Get the hidden parameters of the tile.

Returns

Ordered dictionary of hidden parameter tensors.

Return type

collections.OrderedDict

get_learning_rate()

Return the tile learning rate.

Returns

the tile learning rate.

Return type

float

get_weights(realistic=False)

Get the tile weights (and biases).

Gets the tile weights and extracts the mathematical weight matrix and biases (if present, by determined by the self.bias parameter).

Note

By default this is not hardware realistic. Use set realistic to True for a realistic transfer.

Parameters

realistic (bool) – Whether to use the forward pass to read out the tile weights iteratively, using get_weights_realistic()

Returns

a tuple where the first item is the [out_size, in_size] weight matrix; and the second item is either the [out_size] bias vector or None if the tile is set not to use bias.

Return type

Tuple[torch.Tensor, Optional[torch.Tensor]]

is_cuda = False
post_update_step()

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

Return type

None

reset_columns(start_column_idx=0, num_columns=1, reset_prob=1.0)

Reset (a number of) columns.

Resets the weights with device-to-device and cycle-to-cycle variability (depending on device type), typically:

\[W_{ij} = \xi*\sigma_\text{reset} + b^\text{reset}_{ij}\]
Parameters

  • start_column_idx (int) – a start index of columns (0..x_size-1)

  • num_columns (int) – how many consecutive columns to reset (with circular warping)

  • reset_prob (float) – individual probability of reset.

Return type

None

The reset parameter are set during tile init.

set_hidden_parameters(ordered_parameters)

Set the hidden parameters of the tile.

Parameters

ordered_parameters (collections.OrderedDict) – Ordered dictionary of hidden parameter tensors.

Return type

None

set_indexed(indices, image_sizes)

Sets the index matrix for convolutions ans switches to indexed forward/backward/update versions.

Parameters

  • indices (torch.Tensor) – torch.tensor with int indices

  • image_sizes (List) – [C_in, H_in, W_in, H_out, W_out] sizes

Return type

None

set_learning_rate(learning_rate)

Set the tile learning rate.

Set the tile learning rate to -learning_rate. Note that the learning rate is always taken to be negative (because of the meaning in gradient descent) and positive learning rates are not supported.

Parameters

learning_rate (float) – the desired learning rate.

Return type

None

set_weights(weights, biases=None, realistic=False, n_loops=10)

Set the tile weights (and biases).

Sets the internal tile weights to the specified values, and also the internal tile biases if the tile was set to use bias (via self.bias).

Note

By default this is not hardware realistic. You can set the realistic parameter to True for a realistic transfer.

Parameters

  • weights (torch.Tensor) – [out_size, in_size] weight matrix.

  • biases (Optional[torch.Tensor]) – [out_size] bias vector. This parameter is required if self.bias is True, and ignored otherwise.

  • realistic (bool) – whether to use the forward and update pass to program the weights iteratively, using set_weights_realistic().

  • n_loops (int) – number of times the columns of the weights are set in a closed-loop manner. A value of 1 means that all columns in principle receive enough pulses to change from w_min to w_max.

Return type

None

update(x_input, d_input)

Perform the update pass.

Parameters

  • x_input (torch.Tensor) – [N, in_size] tensor. If in_trans is set, transposed.

  • d_input (torch.Tensor) – [N, out_size] tensor. If out_trans is set, transposed.

Return type

None

update_indexed(x_input, d_input)

Perform the update pass for convolutions.

Parameters

  • x_input (torch.Tensor) – [N, in_size] tensor. If in_trans is set, transposed.

  • d_input (torch.Tensor) – [N, out_size] tensor. If out_trans is set, transposed.

Return type

None