The library is comprised by several layers:
The PyTorch layer is the high-level layer that provides primitives to users for using the features on the library from PyTorch, in particular layers and optimizers.
Overall, the elements on this layer take advantage of PyTorch facilities
(inheriting from existing PyTorch classes and integrating with the rest of
PyTorch features), replacing the default functionality with calls to a
object from the simulator abstraction layer.
Python simulator abstraction layer¶
This layer provides a series of Python objects that can be transparently manipulated and used as any other existing Python functionality, without requiring explicit references to the lower level constructs. By providing this separate Python interface, this allows us for greater flexibility when defining it, keeping all the extra operations and calls to the real bindings internal and performing any translations on behalf of the user.
The main purpose of this layer is to abstract away the implementation-specific complexities of the simulator layers, and map the structures and classes into an interface that caters to the needs of the PyTorch layer. This also provides benefits in regards to serialization and separating concerns overall.
Pybind Python layer¶
This layer is the bridge between C++ and Python. The Python classes and functions in this layer are built using Pybind, and in general consist of exposing selected classes and methods from the C++ simulator, handling the conversion between specific types.
As a results, using the classes from this layer is very similar to how using the C++ classes would be. This is purposeful: by keeping the mapping close to 1:1 on this layer, we (and users that are experimenting directly with the simulator) benefit from being able to translate code almost directly. However, in general users are encouraged to not use the objects from this layer direcly, as it involves an extra overhead and precautions when using them that is managed by the upper classes.
C++ layer (
Ultimately, this is the layer where the real operations over
place, and the one that implements the actual simulation and most of the
features. It is not directly accesible from Python - however, it can be actually
used directly from other C++ programs by using the provided headers.
Layer interaction example¶
For example, using this excerpt of code:
1model = AnalogLinear(2, 1) 2opt = AnalogSGD(model.parameters(), lr=0.5) 3... 4 5for epoch in range(100): 6 pred = model(x_b) 7 loss = mse_loss(pred, y_b) 8 loss.backward() 9 opt.step()
AnalogLinearconstructor (line 1) will:
aihwkit.simulator.tiles.FloatingPointTile. As no extra arguments are passed to the constructor, it will also create as a default a
FloatingPointResistiveDevicethat uses the default
FloatingPointResistiveDeviceParametersparameters. These three objects are the ones from the pure-python layer.
aihwkit.simulator.tiles.FloatingPointTileconstructor will create a
aihwkit.simulator.rpu_base.tiles.FloatingPointTileinstance, along with other objects. These objects are not exposed to the PyTorch layer, and are the ones from the Pybind bindings layer at
instantiating the bindings classes will create the C++ objects internally.
AnalogSGDconstructor (line 2) will:
setup the optimizer, using the attributes of the
AnalogLinearlayer in order to identify which Parameters are to be handled differently during the optimization.
During the training loop (lines 6-8), the forward and backward steps will be performed in the analog tile:
AnalogLinearlayer, PyTorch will call the function defined at
these functions will call the
backward()functions defined in the
aihwkit.simulator.tiles.FloatingPointTileof the layer.
in turn, they will delegate on the
backward()functions defined in the bindings, which in turn delegate on the C++ methods.
The optimizer (line 9) will perform the update step in the analog tile:
using the information constructed during its initialization, the
AnalogSGDwill retrieve the reference to the
aihwkit.simulator.tiles.FloatingPointTile, calling its
in turn, it will delegate on the
update()function defined in the bindings object, which in turn delegate on the C++ method.