Global Akida workflow

Using the MNIST dataset, this example shows the definition and training of a TF-Keras floating-point model, its quantization to 8-bit with the help of calibration, its quantization to 4-bit using QAT, its conversion to Akida and finally the measurement of the processor’s differentiator: power and energy per inference on an Akida device. Notice that the performance of the original TF-Keras floating-point model is maintained throughout the Akida flow. Please refer to the Akida user guide for further information.

If you train in PyTorch, the PyTorch to Akida workflow is the equivalent entry point, going through the ONNX format.

Note

Please refer to the TensorFlow tf_keras.models module for model creation/import details and the TensorFlow Guide for TensorFlow usage.

The MNIST example below is light enough so that a GPU is not needed for training.

Note

Power is measured on silicon and is currently supported on the AKD1000 SoC, which implements Akida 1.0, while this tutorial’s main flow targets Akida 2.0. The final section therefore switches to the Akida 1.0 version of an MNIST model from the model zoo to report the power figures. Everything else runs on the free software simulator.

Global Akida workflow

1. Create and train

1.1. Load and reshape MNIST dataset

import numpy as np
import tensorflow as tf

import matplotlib.cm as cm
import matplotlib.pyplot as plt

from tf_keras.datasets import mnist

# Load MNIST dataset
(x_train, y_train), (x_test, y_test) = mnist.load_data()

# Add a channels dimension to the image sets as Akida expects 4-D inputs corresponding to
# (num_samples, height, width, channels). Note: MNIST is a grayscale dataset and is unusual
# in this respect - most image data already includes a channel dimension, and this step will
# not be necessary.
x_train = np.expand_dims(x_train, -1)
x_test = np.expand_dims(x_test, -1)

# Display a few images from the test set
f, axarr = plt.subplots(1, 4)
for i in range(0, 4):
    axarr[i].imshow(x_test[i].reshape((28, 28)), cmap=cm.Greys_r)
    axarr[i].set_title('Class %d' % y_test[i])
plt.show()

Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz

    8192/11490434 [..............................] - ETA: 0s
   16384/11490434 [..............................] - ETA: 56s
   49152/11490434 [..............................] - ETA: 43s
   81920/11490434 [..............................] - ETA: 39s
  131072/11490434 [..............................] - ETA: 30s
  180224/11490434 [..............................] - ETA: 25s
  262144/11490434 [..............................] - ETA: 19s
  376832/11490434 [..............................] - ETA: 15s
  540672/11490434 [>.............................] - ETA: 11s
  802816/11490434 [=>............................] - ETA: 8s 
 1146880/11490434 [=>............................] - ETA: 6s
 1523712/11490434 [==>...........................] - ETA: 4s
 2334720/11490434 [=====>........................] - ETA: 3s
 3694592/11490434 [========>.....................] - ETA: 1s
 4767744/11490434 [===========>..................] - ETA: 1s
 5660672/11490434 [=============>................] - ETA: 0s
 5718016/11490434 [=============>................] - ETA: 0s
 5816320/11490434 [==============>...............] - ETA: 1s
 6070272/11490434 [==============>...............] - ETA: 0s
 7569408/11490434 [==================>...........] - ETA: 0s
 9224192/11490434 [=======================>......] - ETA: 0s
10698752/11490434 [==========================>...] - ETA: 0s
11490434/11490434 [==============================] - 1s 0us/step

1.2. Model definition

Note that at this stage, there is nothing specific to the Akida IP. The model constructed below, inspired by this example, is a completely standard TF-Keras CNN model.

import tf_keras as keras

model_keras = keras.models.Sequential([
    keras.layers.Input(shape=(28, 28, 1), name="input", dtype=tf.uint8),
    keras.layers.Rescaling(1. / 255),
    keras.layers.Conv2D(filters=32, kernel_size=3, strides=2),
    keras.layers.BatchNormalization(),
    keras.layers.ReLU(),
    # Separable layer
    keras.layers.DepthwiseConv2D(kernel_size=3, padding='same', strides=2),
    keras.layers.Conv2D(filters=64, kernel_size=1, padding='same'),
    keras.layers.BatchNormalization(),
    keras.layers.ReLU(),
    keras.layers.Flatten(),
    keras.layers.Dense(10)
], 'mnistnet')

model_keras.summary()

Model: "mnistnet"
_________________________________________________________________
 Layer (type)                Output Shape              Param #
=================================================================
 rescaling (Rescaling)       (None, 28, 28, 1)         0

 conv2d (Conv2D)             (None, 13, 13, 32)        320

 batch_normalization (Batch  (None, 13, 13, 32)        128
 Normalization)

 re_lu (ReLU)                (None, 13, 13, 32)        0

 depthwise_conv2d (Depthwis  (None, 7, 7, 32)          320
 eConv2D)

 conv2d_1 (Conv2D)           (None, 7, 7, 64)          2112

 batch_normalization_1 (Bat  (None, 7, 7, 64)          256
 chNormalization)

 re_lu_1 (ReLU)              (None, 7, 7, 64)          0

 flatten (Flatten)           (None, 3136)              0

 dense (Dense)               (None, 10)                31370

=================================================================
Total params: 34506 (134.79 KB)
Trainable params: 34314 (134.04 KB)
Non-trainable params: 192 (768.00 Byte)
_________________________________________________________________

1.3. Model training

Given the model created above, train the model and check its accuracy. The model should achieve a test accuracy over 98% after 10 epochs.

from tf_keras.optimizers import Adam

model_keras.compile(
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    optimizer=Adam(learning_rate=1e-3),
    metrics=['accuracy'])

_ = model_keras.fit(x_train, y_train, epochs=10, validation_split=0.1, verbose=2)

Epoch 1/10
1688/1688 - 6s - loss: 0.1738 - accuracy: 0.9473 - val_loss: 0.0793 - val_accuracy: 0.9783 - 6s/epoch - 4ms/step
Epoch 2/10
1688/1688 - 4s - loss: 0.0678 - accuracy: 0.9789 - val_loss: 0.0531 - val_accuracy: 0.9860 - 4s/epoch - 2ms/step
Epoch 3/10
1688/1688 - 4s - loss: 0.0509 - accuracy: 0.9843 - val_loss: 0.0547 - val_accuracy: 0.9830 - 4s/epoch - 2ms/step
Epoch 4/10
1688/1688 - 4s - loss: 0.0395 - accuracy: 0.9872 - val_loss: 0.0512 - val_accuracy: 0.9868 - 4s/epoch - 2ms/step
Epoch 5/10
1688/1688 - 4s - loss: 0.0321 - accuracy: 0.9895 - val_loss: 0.0485 - val_accuracy: 0.9878 - 4s/epoch - 2ms/step
Epoch 6/10
1688/1688 - 4s - loss: 0.0279 - accuracy: 0.9908 - val_loss: 0.0553 - val_accuracy: 0.9872 - 4s/epoch - 2ms/step
Epoch 7/10
1688/1688 - 4s - loss: 0.0218 - accuracy: 0.9930 - val_loss: 0.0584 - val_accuracy: 0.9853 - 4s/epoch - 2ms/step
Epoch 8/10
1688/1688 - 4s - loss: 0.0192 - accuracy: 0.9935 - val_loss: 0.0653 - val_accuracy: 0.9862 - 4s/epoch - 2ms/step
Epoch 9/10
1688/1688 - 4s - loss: 0.0170 - accuracy: 0.9942 - val_loss: 0.0590 - val_accuracy: 0.9872 - 4s/epoch - 2ms/step
Epoch 10/10
1688/1688 - 4s - loss: 0.0150 - accuracy: 0.9948 - val_loss: 0.0601 - val_accuracy: 0.9873 - 4s/epoch - 2ms/step

score = model_keras.evaluate(x_test, y_test, verbose=0)
print('Test accuracy:', score[1])

Test accuracy: 0.9850999712944031

2. Quantize

2.1. 8-bit quantization

An Akida accelerator processes 8- or 4-bit integer activations and weights. Therefore, the floating-point TF-Keras model must be quantized in preparation for running on an Akida accelerator.

The QuantizeML quantize function can be used to quantize a TF-Keras model for Akida. For this step in this example, an “8/8/8” quantization scheme will be applied to the floating-point TF-Keras model to produce 8-bit weights in the first layer, 8-bit weights in all other layers, and 8-bit activations.

The quantization process results in a TF-Keras model with custom QuantizeML quantized layers substituted for the original TF-Keras layers. All TF-Keras API functions can be applied to this new model: summary(), compile(), fit(), etc.

Note

The quantize function applies several transformations to the original model. For example, it folds the batch normalization layers into the corresponding neural layers. The new weights are computed according to this folding operation.

from quantizeml.models import quantize, QuantizationParams

qparams = QuantizationParams(input_weight_bits=8, weight_bits=8, activation_bits=8)
model_quantized = quantize(model_keras, qparams=qparams)

/usr/local/lib/python3.11/dist-packages/quantizeml/models/quantize.py:577: UserWarning: Quantizing per-axis with random calibration samples is not accurate. Set QuantizationParams.per_tensor_activations=True when calibrating with random samples. Continuing execution.
  warnings.warn("Quantizing per-axis with random calibration samples is not accurate. "

   1/1024 [..............................] - ETA: 2:28
  52/1024 [>.............................] - ETA: 0s  
 103/1024 [==>...........................] - ETA: 0s
 156/1024 [===>..........................] - ETA: 0s
 208/1024 [=====>........................] - ETA: 0s
 259/1024 [======>.......................] - ETA: 0s
 310/1024 [========>.....................] - ETA: 0s
 361/1024 [=========>....................] - ETA: 0s
 412/1024 [===========>..................] - ETA: 0s
 463/1024 [============>.................] - ETA: 0s
 515/1024 [==============>...............] - ETA: 0s
 567/1024 [===============>..............] - ETA: 0s
 618/1024 [=================>............] - ETA: 0s
 670/1024 [==================>...........] - ETA: 0s
 722/1024 [====================>.........] - ETA: 0s
 773/1024 [=====================>........] - ETA: 0s
 824/1024 [=======================>......] - ETA: 0s
 875/1024 [========================>.....] - ETA: 0s
 926/1024 [==========================>...] - ETA: 0s
 978/1024 [===========================>..] - ETA: 0s
1024/1024 [==============================] - 1s 982us/step

model_quantized.summary()

Model: "mnistnet"
_________________________________________________________________
 Layer (type)                Output Shape              Param #
=================================================================
 input (InputLayer)          [(None, 28, 28, 1)]       0

 rescaling (QuantizedRescal  (None, 28, 28, 1)         0
 ing)

 conv2d (QuantizedConv2D)    (None, 13, 13, 32)        320

 re_lu (QuantizedReLU)       (None, 13, 13, 32)        64

 depthwise_conv2d (Quantize  (None, 7, 7, 32)          384
 dDepthwiseConv2D)

 conv2d_1 (QuantizedConv2D)  (None, 7, 7, 64)          2112

 re_lu_1 (QuantizedReLU)     (None, 7, 7, 64)          128

 flatten (QuantizedFlatten)  (None, 3136)              0

 dense (QuantizedDense)      (None, 10)                31370

 dequantizer (Dequantizer)   (None, 10)                0

=================================================================
Total params: 34378 (134.29 KB)
Trainable params: 34122 (133.29 KB)
Non-trainable params: 256 (1.00 KB)
_________________________________________________________________

Note

The number of parameters for the floating and quantized models differs, a consequence of the BatchNormalization folding and the additional parameters added for quantization. For further details, please refer to their respective summaries.

Check the quantized model accuracy.

def compile_evaluate(model):
    """ Compiles and evaluates the model, then return accuracy score. """
    model.compile(metrics=['accuracy'])
    return model.evaluate(x_test, y_test, verbose=0)[1]


print('Test accuracy after 8-bit quantization:', compile_evaluate(model_quantized))

Test accuracy after 8-bit quantization: 0.9807000160217285

2.2. Effect of calibration

The previous call to quantize was made with random samples for calibration (default parameters). While the observed drop in accuracy is minimal, that is, around 1%, it can be worse on more complex models. Therefore, it is advised to use a set of real samples from the training set for calibration during a call to quantize. Note that this remains a calibration step rather than a training step in that no output labels are required. Furthermore, any relevant data could be used for calibration. The recommended settings for calibration that are widely used to obtain the zoo performance are:

• 1024 samples

• a batch size of 100

• 2 epochs

model_quantized = quantize(model_keras, qparams=qparams,
                           samples=x_train, num_samples=1024, batch_size=100, epochs=2)

 1/11 [=>............................] - ETA: 1s
11/11 [==============================] - 0s 1ms/step

 1/11 [=>............................] - ETA: 0s
11/11 [==============================] - 0s 1ms/step

Check the accuracy for the quantized and calibrated model.

print('Test accuracy after calibration:', compile_evaluate(model_quantized))

Test accuracy after calibration: 0.9850000143051147

Calibrating with real samples on this model recovers the initial float accuracy.

2.3. 4-bit quantization

The accuracy of the 8/8/8 quantized model is equal to that of the TF-Keras floating-point model. In some cases, a smaller memory size for the model is required. This can be accomplished through quantization of the model to smaller bitwidths.

The model will now be quantized to 8/4/4, that is, 8-bit weights in the first layer with 4-bit weights and activations in all other layers. Such a quantization scheme will usually introduce a performance drop.

qparams = QuantizationParams(input_weight_bits=8, weight_bits=4, activation_bits=4)
model_quantized = quantize(model_keras, qparams=qparams,
                           samples=x_train, num_samples=1024, batch_size=100, epochs=2)

 1/11 [=>............................] - ETA: 1s
11/11 [==============================] - 0s 1ms/step

 1/11 [=>............................] - ETA: 0s
11/11 [==============================] - 0s 1ms/step

Check the 4-bit quantized accuracy.

print('Test accuracy after 4-bit quantization:', compile_evaluate(model_quantized))

Test accuracy after 4-bit quantization: 0.9836999773979187

2.4. Model fine-tuning (Quantization Aware Training)

When a model suffers from an accuracy drop after quantization, fine-tuning or Quantization Aware Training (QAT) may recover some or all of the original performance.

Note that since this is a fine-tuning step, both the number of epochs and the learning rate are expected to be lower than during the initial float training.

model_quantized.compile(
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    optimizer=Adam(learning_rate=1e-4),
    metrics=['accuracy'])

model_quantized.fit(x_train, y_train, epochs=5, validation_split=0.1, verbose=2)

Epoch 1/5
1688/1688 - 15s - loss: 0.0104 - accuracy: 0.9968 - val_loss: 0.0568 - val_accuracy: 0.9888 - 15s/epoch - 9ms/step
Epoch 2/5
1688/1688 - 10s - loss: 0.0068 - accuracy: 0.9983 - val_loss: 0.0547 - val_accuracy: 0.9888 - 10s/epoch - 6ms/step
Epoch 3/5
1688/1688 - 11s - loss: 0.0052 - accuracy: 0.9989 - val_loss: 0.0557 - val_accuracy: 0.9892 - 11s/epoch - 6ms/step
Epoch 4/5
1688/1688 - 11s - loss: 0.0043 - accuracy: 0.9992 - val_loss: 0.0561 - val_accuracy: 0.9900 - 11s/epoch - 6ms/step
Epoch 5/5
1688/1688 - 11s - loss: 0.0038 - accuracy: 0.9994 - val_loss: 0.0582 - val_accuracy: 0.9900 - 11s/epoch - 6ms/step

<tf_keras.src.callbacks.History object at 0x7aebf706c290>

score = model_quantized.evaluate(x_test, y_test, verbose=0)[1]
print('Test accuracy after fine-tuning:', score)

Test accuracy after fine-tuning: 0.9879999756813049

3. Convert

3.1 Convert to Akida model

When the quantized model produces satisfactory performance, it can be converted to the native Akida format. The convert function returns a model in Akida format ready for inference.

As with TF-Keras, the summary() method provides a textual representation of the Akida model.

from cnn2snn import convert

model_akida = convert(model_quantized)
model_akida.summary()

                Model Summary
______________________________________________
Input shape  Output shape  Sequences  Layers
==============================================
[28, 28, 1]  [1, 1, 10]    1          5
______________________________________________

__________________________________________________________________
Layer (type)                        Output shape  Kernel shape

=============== SW/conv2d-dequantizer_2 (Software) ===============

conv2d (InputConv2D)                [13, 13, 32]  (3, 3, 1, 32)
__________________________________________________________________
depthwise_conv2d (DepthwiseConv2D)  [7, 7, 32]    (3, 3, 32, 1)
__________________________________________________________________
conv2d_1 (Conv2D)                   [7, 7, 64]    (1, 1, 32, 64)
__________________________________________________________________
dense (Dense1D)                     [1, 1, 10]    (3136, 10)
__________________________________________________________________
dequantizer_2 (Dequantizer)         [1, 1, 10]    N/A
__________________________________________________________________

3.2. Check performance

accuracy = model_akida.evaluate(x_test, y_test.astype(np.int32))
print('Test accuracy after conversion:', accuracy)

# For non-regression purposes
assert accuracy > 0.96

Test accuracy after conversion: 0.9846000075340271

3.3 Show predictions for a single image

Display one of the test images, such as the first image in the dataset from above, to visualize the output of the model.

# Test a single example
sample_image = 0
image = x_test[sample_image]
outputs = model_akida.predict(image.reshape(1, 28, 28, 1))
print('Input Label: %i' % y_test[sample_image])

f, axarr = plt.subplots(1, 2)
axarr[0].imshow(x_test[sample_image].reshape((28, 28)), cmap=cm.Greys_r)
axarr[0].set_title('Class %d' % y_test[sample_image])
axarr[1].bar(range(10), outputs.squeeze())
axarr[1].set_xticks(range(10))
plt.show()

print(outputs.squeeze())

Input Label: 7
[-15.541361  -10.673538   -3.2178562  -1.458841  -17.463495  -10.089144
 -32.559746   10.078885   -9.606353   -1.546738 ]

Consider the output from the model above. As is typical in backprop-trained models, the final layer is a Dense layer with one neuron for each of the 10 classes in the dataset. The goal of training is to maximize the response of the neuron corresponding to the label of each training sample while minimizing the responses of the other neurons.

In the bar chart above, you can see the outputs from all 10 neurons. It is easy to see that neuron 7 responds much more strongly than the others. The first sample is indeed a number 7.

4. Measure power and energy on hardware

Accuracy parity is only half of the story: the reason to run a model on Akida is efficiency. This last section measures the power the processor draws while classifying the MNIST test set, read from a sensor on the chip.

Power is measured on silicon, currently on the AKD1000 SoC that implements Akida 1.0 (see the See the power number page). Since the model trained above uses the default Akida 2.0 quantization, this section switches to the Akida 1.0 MNIST model from the model zoo: the GXNOR/MNIST CNN, pretrained and quantized to 2-bit weights and 1-bit activations.

4.1. Load and convert a pretrained Akida 1.0 model

The set_akida_version context selects the Akida 1.0 version of the pretrained model, which is then converted exactly like the trained model above.

from cnn2snn import set_akida_version, AkidaVersion
from akida_models import gxnor_mnist_pretrained

# Use a quantized model with pretrained quantized weights
with set_akida_version(AkidaVersion.v1):
    model_quantized_1_0 = gxnor_mnist_pretrained()

model_akida_1_0 = convert(model_quantized_1_0)

/usr/local/lib/python3.11/dist-packages/akida_models/model_io.py:147: UserWarning: Model gxnor_mnist_iq2_wq2_aq1.h5 has been trained with akida_models 1.1.10 which is the last version supporting 1.0 models training. Continuing execution.
  warnings.warn(f'Model {model_name_v1} has been trained with akida_models 1.1.10 which is '
Downloading data from https://data.brainchip.com/models/AkidaV1/gxnor/gxnor_mnist_iq2_wq2_aq1.h5.

      0/6556040 [..............................] - ETA: 0s
 245760/6556040 [>.............................] - ETA: 1s
1228800/6556040 [====>.........................] - ETA: 0s
2621440/6556040 [==========>...................] - ETA: 0s
3702784/6556040 [===============>..............] - ETA: 0s
4816896/6556040 [=====================>........] - ETA: 0s
6488064/6556040 [============================>.] - ETA: 0s
6556040/6556040 [==============================] - 0s 0us/step
Download complete.
WARNING:tensorflow:No training configuration found in the save file, so the model was *not* compiled. Compile it manually.

4.2. Map on hardware

List available Akida devices and check that an NSoC V2, Akida 1.0 production chip is available.

import akida

devices = akida.devices()
print(f'Available devices: {[dev.desc for dev in devices]}')
assert len(devices), "No device found, this example needs an Akida NSoC_v2 device."
device = devices[0]
assert device.version == akida.NSoC_v2, "Wrong device found, this example needs an Akida NSoC_v2."

Available devices: ['PCIe/NSoC_v2/0']

Map the model on the device

model_akida_1_0.map(device)

# Check model mapping: NP allocation and binary size
model_akida_1_0.summary()

                                     Model Summary
_______________________________________________________________________________________
Input shape  Output shape  Sequences  Layers  NPs  Skip DMAs  External Memory (Bytes)
=======================================================================================
[28, 28, 1]  [1, 1, 10]    1          4       66   0          409600
_______________________________________________________________________________________

_________________________
Component (type)  Count
=========================
HRC               1
_________________________
CNP1              64
_________________________
FNP2              1
_________________________
FNP3              1
_________________________

        External Memory Summary
________________________________________
Layer (type)   External Memory (Bytes)
========================================
fc_1 (Fully.)  409600
________________________________________

__________________________________________________________________________
Layer (type)                 Output shape  Kernel shape       Components

====== HW/block_1/conv_1-predictions (Hardware) - size: 559748 bytes =====

block_1/conv_1 (InputConv.)  [14, 14, 32]  (5, 5, 1, 32)      1 HRC
__________________________________________________________________________
block_2/conv_1 (Conv.)       [7, 7, 64]    (3, 3, 32, 64)     64 CNP1
__________________________________________________________________________
fc_1 (Fully.)                [1, 1, 512]   (1, 1, 3136, 512)  1 FNP2
__________________________________________________________________________
predictions (Fully.)         [1, 1, 10]    (1, 1, 512, 10)    1 FNP3
__________________________________________________________________________

4.3. Performance measurement

Power measurement must be enabled on the device’s SoC (disabled by default). After sending data for inference, performance measurements are available in the model statistics.

# Enable power measurement
device.soc.power_measurement_enabled = True

# Send the test set for inference
_ = model_akida_1_0.forward(x_test)

# Display floor power
floor_power = device.soc.power_meter.floor
print(f'Floor power: {floor_power:.2f} mW')

# Retrieve statistics
print(model_akida_1_0.statistics)

Floor power: 909.58 mW
Average framerate = 1797.59 fps
Last inference power range (mW):  Avg 1198.42 / Min 909.00 / Max 1208.00 / Std 36.15
Last inference energy consumed (mJ/frame): 0.67
Last inference clock: 1657293495
Last program clock: 130144

The floor power is the idle draw of the board. In the statistics above, next to the average framerate, you should see the two numbers this workflow was building toward:

• Last inference power range (mW): the power drawn while classifying, measured on-chip,

• Last inference energy consumed (mJ/frame): the energy cost of classifying one digit.

Both figures include the floor power.

Note

Power is read from a sensor on the silicon: the software simulator cannot measure it. Without a device, the model zoo performance page publishes the measured reference figure for this model — 0.34 mJ per frame on an AKD1500 device.