aihwkit_lightning.optim.analog_optimizer module

Analog-aware inference optimizer.

class aihwkit_lightning.optim.analog_optimizer.AnalogOptimizer(optimizer_cls, *_, **__)[source]

Bases: Optimizer

Generic optimizer that wraps an existing Optimizer for analog inference.

This class wraps an existing Optimizer, customizing the optimization step for triggering the analog update needed for analog tiles. All other (digital) parameters are governed by the given torch optimizer. In case of hardware-aware training (InferenceTile) the tile weight update is also governed by the given optimizer, otherwise it is using the internal analog update as defined in the rpu_config.

The AnalogOptimizer constructor expects the wrapped optimizer class as the first parameter, followed by any arguments required by the wrapped optimizer.

Note

The instances returned are of a new type that is a subclass of:

  • the wrapped Optimizer (allowing access to all their methods and attributes).

  • this AnalogOptimizer.

Example

The following block illustrate how to create an optimizer that wraps standard SGD:

>>> from torch.optim import SGD
>>> from torch.nn import Linear
>>> from aihwkit.simulator.configs.configs import InferenceRPUConfig
>>> from aihwkit.optim import AnalogOptimizer
>>> model = AnalogLinear(3, 4, rpu_config=InferenceRPUConfig)
>>> optimizer = AnalogOptimizer(SGD, model.parameters(), lr=0.02)
Parameters:

  • optimizer_cls (Type) –

  • _ (Any) –

  • __ (Any) –

Return type:

AnalogOptimizer

SUBCLASSES: Dict[str, Type] = {}

Registry of the created subclasses.