This page was generated from docs/developer/UML-diagrams.ipynb. Interactive online version:
UML diagrams for Rockpool
This notebook contains UML diagrams for the main packages and SW architecture of Rockpool.
[90]:
# - Load UML diagrams image
from PIL import Image
from pathlib import Path
im = Image.open(Path('uml', 'UML-diagrams.png'))
Module
base and core derived classes, and combinators
The Module
base class as well as all derived classes are under the nn.modules
package. All user-defined modules as well as back-end base classes derive from Module
or ModuleBase
. ModuleBase
is not designed as a user-facing class, and defines most of the features of Module
. Module
adds a facility for a post-initialisation callback on instantiation, used by some derived base classes.
Combinators used to build networks are in the package nn.combinators
. These are adapted within nn.combinators
for specific computational back-ends.
[93]:
im.crop([0, 0, 2450, 1800])
[93]:
Jax
and Torch
core modules
Derived Modules for specific back-ends are under nn.modules
, in sub-packages named by backend. Some backends require derived base classes to operate; e.g. JaxModule
and TorchModule
. All derived modules using jax
inherit from JaxModule
, and likewise for torch
.
TorchModule
also supports direct conversion of torch modules imported from torch.nn.Module
.
[75]:
im.crop([2700, 0, 4750, 1400])
[75]:
Rockpool Parameter
classes
Classes for managing parameters are in the package nn.parameters
. Parameter
, State
and SimulationParameter
derive from ParameterBase
.
Constant()
is a special class which wraps any other class, permitting Rockpool to identify that a constant argument should not be trainable.
[84]:
im.crop([7450, 0, 9000, 800])
[84]:
Hardware support packages
All hardware support packages are placed under devices
. At present only devices.xylo
is supported. See the documentation for Xylo for more information about these classes.
Different hardware versions are supported by providing sub-packages for each device. Within these subpackages are duplicate-named classes which support the corresponding specific device version.
e.g. devices.xylo.syns61201
provides XyloSim
, XyloSamna
, AFESamna
etc. for the SYNS61201 device. Likewise devices.xylo.syns61300
provides XyloSim
, XyloSamna
, AFESamna
etc. for the SYNS61300 device.
[87]:
im.crop([5000, 0, 7000, 1200])
[87]:
Computational graph representation
All classes used to serialise and represent Rockpool modules in a computational graph are under the package graph
.
graph
contains the base class GraphModule
, from which all derived GraphModule
classes should inherit. It also provides GraphHolder
and GraphNode
classes.
The sub-package graph.utils
provides a number of utility functions for manipulating graphs
[89]:
im.crop([0, 3500, 2300, 5750])
[89]:
Training and transformations
The package training
contains several sub-packages with useful functions for defining losses, regulariation, and training approaches.
The package transform
contains sub-packages for parameter and network transformations. Currently only the package transform.quantize_methods
is supported, providing methods for post-traning quantization for Xylo.
The beta implementation of QAT functionality for torch
modules is under transform.torch_transform
.
[79]:
im.crop([2400, 3500, 4600, 5200])
[79]: