Segmentation tutorial

This example demonstrates image segmentation with an Akida-compatible model as illustrated through person segmentation using the Portrait128 dataset.

Using pre-trained models for quick runtime, this example shows the evolution of model performance for a trained TF-Keras floating-point model, a TF-Keras quantized and Quantization Aware Trained (QAT) model, and an Akida-converted model. Notice that the performance of the original TF-Keras floating-point model is maintained throughout the model conversion flow.

1. Load the dataset

import os
import numpy as np
from akida_models import fetch_file

# Download the validation set from the BrainChip data server, it contains 10% of the original
# dataset
data_path = fetch_file(fname="val.tar.gz",
                       origin="https://data.brainchip.com/dataset-mirror/portrait128/val.tar.gz",
                       cache_subdir=os.path.join("datasets", "portrait128"),
                       extract=True)

data_dir = os.path.join(os.path.dirname(data_path), "val")
x_val = np.load(os.path.join(data_dir, "val_img.npy"))
y_val = np.load(os.path.join(data_dir, "val_msk.npy")).astype('uint8')
batch_size = 32
steps = x_val.shape[0] // 32

# Visualize some data
import matplotlib.pyplot as plt

rng = np.random.default_rng()
id = rng.integers(0, x_val.shape[0] - 2)

fig, axs = plt.subplots(3, 3, constrained_layout=True)
for col in range(3):
    axs[0, col].imshow(x_val[id + col] / 255.)
    axs[0, col].axis('off')
    axs[1, col].imshow(1 - y_val[id + col], cmap='Greys')
    axs[1, col].axis('off')
    axs[2, col].imshow(x_val[id + col] / 255. * y_val[id + col])
    axs[2, col].axis('off')

fig.suptitle('Image, mask and masked image', fontsize=10)
plt.show()
Image, mask and masked image
Downloading data from https://data.brainchip.com/dataset-mirror/portrait128/val.tar.gz.

        0/267313385 [..............................] - ETA: 0s
   253952/267313385 [..............................] - ETA: 54s
  1417216/267313385 [..............................] - ETA: 19s
  2990080/267313385 [..............................] - ETA: 13s
  4628480/267313385 [..............................] - ETA: 11s
  5595136/267313385 [..............................] - ETA: 11s
  6873088/267313385 [..............................] - ETA: 11s
  8167424/267313385 [..............................] - ETA: 11s
  9248768/267313385 [>.............................] - ETA: 11s
 10797056/267313385 [>.............................] - ETA: 10s
 11886592/267313385 [>.............................] - ETA: 10s
 13230080/267313385 [>.............................] - ETA: 10s
 14524416/267313385 [>.............................] - ETA: 10s
 15294464/267313385 [>.............................] - ETA: 10s
 16785408/267313385 [>.............................] - ETA: 10s
 17899520/267313385 [=>............................] - ETA: 10s
 18898944/267313385 [=>............................] - ETA: 10s
 20127744/267313385 [=>............................] - ETA: 10s
 21454848/267313385 [=>............................] - ETA: 10s
 22781952/267313385 [=>............................] - ETA: 10s
 24010752/267313385 [=>............................] - ETA: 10s
 24952832/267313385 [=>............................] - ETA: 10s
 26222592/267313385 [=>............................] - ETA: 10s
 27598848/267313385 [==>...........................] - ETA: 10s
 28762112/267313385 [==>...........................] - ETA: 10s
 29843456/267313385 [==>...........................] - ETA: 10s
 31105024/267313385 [==>...........................] - ETA: 10s
 32350208/267313385 [==>...........................] - ETA: 10s
 33431552/267313385 [==>...........................] - ETA: 10s
 34709504/267313385 [==>...........................] - ETA: 9s 
 35921920/267313385 [===>..........................] - ETA: 9s
 37003264/267313385 [===>..........................] - ETA: 9s
 38313984/267313385 [===>..........................] - ETA: 9s
 39739392/267313385 [===>..........................] - ETA: 9s
 40886272/267313385 [===>..........................] - ETA: 9s
 41967616/267313385 [===>..........................] - ETA: 9s
 43294720/267313385 [===>..........................] - ETA: 9s
 44146688/267313385 [===>..........................] - ETA: 9s
 45490176/267313385 [====>.........................] - ETA: 9s
 46563328/267313385 [====>.........................] - ETA: 9s
 47718400/267313385 [====>.........................] - ETA: 9s
 49012736/267313385 [====>.........................] - ETA: 9s
 49963008/267313385 [====>.........................] - ETA: 9s
 51290112/267313385 [====>.........................] - ETA: 9s
 52502528/267313385 [====>.........................] - ETA: 9s
 52518912/267313385 [====>.........................] - ETA: 16s
 53895168/267313385 [=====>........................] - ETA: 16s
 54894592/267313385 [=====>........................] - ETA: 15s
 55713792/267313385 [=====>........................] - ETA: 15s
 57425920/267313385 [=====>........................] - ETA: 15s
 58662912/267313385 [=====>........................] - ETA: 15s
 59940864/267313385 [=====>........................] - ETA: 14s
 60989440/267313385 [=====>........................] - ETA: 14s
 62087168/267313385 [=====>........................] - ETA: 14s
 63414272/267313385 [======>.......................] - ETA: 14s
 64315392/267313385 [======>.......................] - ETA: 14s
 65609728/267313385 [======>.......................] - ETA: 14s
 66969600/267313385 [======>.......................] - ETA: 13s
 68313088/267313385 [======>.......................] - ETA: 13s
 69394432/267313385 [======>.......................] - ETA: 13s
 70721536/267313385 [======>.......................] - ETA: 13s
 71999488/267313385 [=======>......................] - ETA: 13s
 72867840/267313385 [=======>......................] - ETA: 12s
 74194944/267313385 [=======>......................] - ETA: 12s
 75423744/267313385 [=======>......................] - ETA: 12s
 76505088/267313385 [=======>......................] - ETA: 12s
 77864960/267313385 [=======>......................] - ETA: 12s
 78782464/267313385 [=======>......................] - ETA: 12s
 80076800/267313385 [=======>......................] - ETA: 12s
 81092608/267313385 [========>.....................] - ETA: 11s
 82452480/267313385 [========>.....................] - ETA: 11s
 83566592/267313385 [========>.....................] - ETA: 11s
 84893696/267313385 [========>.....................] - ETA: 11s
 86073344/267313385 [========>.....................] - ETA: 11s
 87318528/267313385 [========>.....................] - ETA: 11s
 88776704/267313385 [========>.....................] - ETA: 11s
 89858048/267313385 [=========>....................] - ETA: 11s
 91299840/267313385 [=========>....................] - ETA: 10s
 92741632/267313385 [=========>....................] - ETA: 10s
 93782016/267313385 [=========>....................] - ETA: 10s
 95100928/267313385 [=========>....................] - ETA: 10s
 96198656/267313385 [=========>....................] - ETA: 10s
 97525760/267313385 [=========>....................] - ETA: 10s
 98721792/267313385 [==========>...................] - ETA: 10s
100261888/267313385 [==========>...................] - ETA: 9s 
101277696/267313385 [==========>...................] - ETA: 9s
102768640/267313385 [==========>...................] - ETA: 9s
103915520/267313385 [==========>...................] - ETA: 9s
105340928/267313385 [==========>...................] - ETA: 9s
106749952/267313385 [==========>...................] - ETA: 9s
107945984/267313385 [===========>..................] - ETA: 9s
109273088/267313385 [===========>..................] - ETA: 9s
110501888/267313385 [===========>..................] - ETA: 9s
111616000/267313385 [===========>..................] - ETA: 9s
112975872/267313385 [===========>..................] - ETA: 8s
114188288/267313385 [===========>..................] - ETA: 8s
115662848/267313385 [===========>..................] - ETA: 8s
116760576/267313385 [============>.................] - ETA: 8s
118177792/267313385 [============>.................] - ETA: 8s
119414784/267313385 [============>.................] - ETA: 8s
120627200/267313385 [============>.................] - ETA: 8s
122019840/267313385 [============>.................] - ETA: 8s
123101184/267313385 [============>.................] - ETA: 8s
124542976/267313385 [============>.................] - ETA: 7s
125657088/267313385 [=============>................] - ETA: 7s
126984192/267313385 [=============>................] - ETA: 7s
128327680/267313385 [=============>................] - ETA: 7s
129441792/267313385 [=============>................] - ETA: 7s
130965504/267313385 [=============>................] - ETA: 7s
132145152/267313385 [=============>................] - ETA: 7s
133603328/267313385 [=============>................] - ETA: 7s
134586368/267313385 [==============>...............] - ETA: 7s
135290880/267313385 [==============>...............] - ETA: 7s
136896512/267313385 [==============>...............] - ETA: 7s
138076160/267313385 [==============>...............] - ETA: 7s
139304960/267313385 [==============>...............] - ETA: 6s
140582912/267313385 [==============>...............] - ETA: 6s
141877248/267313385 [==============>...............] - ETA: 6s
143138816/267313385 [===============>..............] - ETA: 6s
144318464/267313385 [===============>..............] - ETA: 6s
145301504/267313385 [===============>..............] - ETA: 6s
146497536/267313385 [===============>..............] - ETA: 6s
147759104/267313385 [===============>..............] - ETA: 6s
148889600/267313385 [===============>..............] - ETA: 6s
149921792/267313385 [===============>..............] - ETA: 6s
151150592/267313385 [===============>..............] - ETA: 6s
152477696/267313385 [================>.............] - ETA: 6s
153772032/267313385 [================>.............] - ETA: 6s
154836992/267313385 [================>.............] - ETA: 5s
155967488/267313385 [================>.............] - ETA: 5s
157245440/267313385 [================>.............] - ETA: 5s
158244864/267313385 [================>.............] - ETA: 5s
159424512/267313385 [================>.............] - ETA: 5s
160702464/267313385 [=================>............] - ETA: 5s
161734656/267313385 [=================>............] - ETA: 5s
162947072/267313385 [=================>............] - ETA: 5s
164257792/267313385 [=================>............] - ETA: 5s
165502976/267313385 [=================>............] - ETA: 5s
166567936/267313385 [=================>............] - ETA: 5s
167944192/267313385 [=================>............] - ETA: 5s
169271296/267313385 [=================>............] - ETA: 5s
170336256/267313385 [==================>...........] - ETA: 5s
171499520/267313385 [==================>...........] - ETA: 5s
172859392/267313385 [==================>...........] - ETA: 4s
174088192/267313385 [==================>...........] - ETA: 4s
175251456/267313385 [==================>...........] - ETA: 4s
176627712/267313385 [==================>...........] - ETA: 4s
177848320/267313385 [==================>...........] - ETA: 4s
178855936/267313385 [===================>..........] - ETA: 4s
180183040/267313385 [===================>..........] - ETA: 4s
181207040/267313385 [===================>..........] - ETA: 4s
182394880/267313385 [===================>..........] - ETA: 4s
183590912/267313385 [===================>..........] - ETA: 4s
184606720/267313385 [===================>..........] - ETA: 4s
185982976/267313385 [===================>..........] - ETA: 4s
186900480/267313385 [===================>..........] - ETA: 4s
188227584/267313385 [====================>.........] - ETA: 4s
189587456/267313385 [====================>.........] - ETA: 3s
190586880/267313385 [====================>.........] - ETA: 3s
191848448/267313385 [====================>.........] - ETA: 3s
192995328/267313385 [====================>.........] - ETA: 3s
194093056/267313385 [====================>.........] - ETA: 3s
195420160/267313385 [====================>.........] - ETA: 3s
196354048/267313385 [=====================>........] - ETA: 3s
197681152/267313385 [=====================>........] - ETA: 3s
198598656/267313385 [=====================>........] - ETA: 3s
199942144/267313385 [=====================>........] - ETA: 3s
201236480/267313385 [=====================>........] - ETA: 3s
202366976/267313385 [=====================>........] - ETA: 3s
203710464/267313385 [=====================>........] - ETA: 3s
204840960/267313385 [=====================>........] - ETA: 3s
206086144/267313385 [======================>.......] - ETA: 3s
207495168/267313385 [======================>.......] - ETA: 3s
208723968/267313385 [======================>.......] - ETA: 2s
209977344/267313385 [======================>.......] - ETA: 2s
211443712/267313385 [======================>.......] - ETA: 2s
212885504/267313385 [======================>.......] - ETA: 2s
214081536/267313385 [=======================>......] - ETA: 2s
215375872/267313385 [=======================>......] - ETA: 2s
216686592/267313385 [=======================>......] - ETA: 2s
217669632/267313385 [=======================>......] - ETA: 2s
219111424/267313385 [=======================>......] - ETA: 2s
220422144/267313385 [=======================>......] - ETA: 2s
221618176/267313385 [=======================>......] - ETA: 2s
223125504/267313385 [========================>.....] - ETA: 2s
224403456/267313385 [========================>.....] - ETA: 2s
225599488/267313385 [========================>.....] - ETA: 2s
226992128/267313385 [========================>.....] - ETA: 1s
228409344/267313385 [========================>.....] - ETA: 1s
229466112/267313385 [========================>.....] - ETA: 1s
230809600/267313385 [========================>.....] - ETA: 1s
232169472/267313385 [=========================>....] - ETA: 1s
233152512/267313385 [=========================>....] - ETA: 1s
234561536/267313385 [=========================>....] - ETA: 1s
235659264/267313385 [=========================>....] - ETA: 1s
236953600/267313385 [=========================>....] - ETA: 1s
238264320/267313385 [=========================>....] - ETA: 1s
239329280/267313385 [=========================>....] - ETA: 1s
240771072/267313385 [==========================>...] - ETA: 1s
241778688/267313385 [==========================>...] - ETA: 1s
243097600/267313385 [==========================>...] - ETA: 1s
244424704/267313385 [==========================>...] - ETA: 1s
245563392/267313385 [==========================>...] - ETA: 1s
247046144/267313385 [==========================>...] - ETA: 0s
248061952/267313385 [==========================>...] - ETA: 0s
249470976/267313385 [==========================>...] - ETA: 0s
250896384/267313385 [===========================>..] - ETA: 0s
251994112/267313385 [===========================>..] - ETA: 0s
253386752/267313385 [===========================>..] - ETA: 0s
254631936/267313385 [===========================>..] - ETA: 0s
255877120/267313385 [===========================>..] - ETA: 0s
257318912/267313385 [===========================>..] - ETA: 0s
258351104/267313385 [===========================>..] - ETA: 0s
259743744/267313385 [============================>.] - ETA: 0s
260808704/267313385 [============================>.] - ETA: 0s
262283264/267313385 [============================>.] - ETA: 0s
263446528/267313385 [============================>.] - ETA: 0s
264888320/267313385 [============================>.] - ETA: 0s
266264576/267313385 [============================>.] - ETA: 0s
267313385/267313385 [==============================] - 13s 0us/step
Download complete.

2. Load a pre-trained native TF-Keras model

The model used in this example is AkidaUNet. It has an AkidaNet (0.5) backbone to extract features combined with a succession of separable transposed convolutional blocks to build an image segmentation map. A pre-trained floating-point TF-Keras model is downloaded to save training time.

Note

  • The “transposed” convolutional feature is new in Akida 2.0.

  • The “separable transposed” operation is realized through the combination of a QuantizeML custom DepthwiseConv2DTranspose layer with a standard pointwise convolution.

The performance of the model is evaluated using both pixel accuracy and Binary IoU. The pixel accuracy describes how well the model can predict the segmentation mask pixel by pixel, and the Binary IoU takes into account how close the predicted mask is to the ground truth.

from akida_models.model_io import load_model

# Retrieve the model file from the BrainChip data server
model_file = fetch_file(fname="akida_unet_portrait128.h5",
                        origin="https://data.brainchip.com/models/AkidaV2/akida_unet/akida_unet_portrait128.h5",
                        cache_subdir='models')

# Load the native TF-Keras pre-trained model
model_keras = load_model(model_file)
model_keras.summary()
Downloading data from https://data.brainchip.com/models/AkidaV2/akida_unet/akida_unet_portrait128.h5.

      0/4501976 [..............................] - ETA: 0s
 204800/4501976 [>.............................] - ETA: 1s
1236992/4501976 [=======>......................] - ETA: 0s
2629632/4501976 [================>.............] - ETA: 0s
3317760/4501976 [=====================>........] - ETA: 0s
4501976/4501976 [==============================] - 0s 0us/step
Download complete.
Model: "akida_unet"
_________________________________________________________________
 Layer (type)                Output Shape              Param #
=================================================================
 input (InputLayer)          [(None, 128, 128, 3)]     0

 rescaling (Rescaling)       (None, 128, 128, 3)       0

 conv_0 (Conv2D)             (None, 64, 64, 16)        432

 conv_0/BN (BatchNormalizat  (None, 64, 64, 16)        64
 ion)

 conv_0/relu (ReLU)          (None, 64, 64, 16)        0

 conv_1 (Conv2D)             (None, 64, 64, 32)        4608

 conv_1/BN (BatchNormalizat  (None, 64, 64, 32)        128
 ion)

 conv_1/relu (ReLU)          (None, 64, 64, 32)        0

 conv_2 (Conv2D)             (None, 32, 32, 64)        18432

 conv_2/BN (BatchNormalizat  (None, 32, 32, 64)        256
 ion)

 conv_2/relu (ReLU)          (None, 32, 32, 64)        0

 conv_3 (Conv2D)             (None, 32, 32, 64)        36864

 conv_3/BN (BatchNormalizat  (None, 32, 32, 64)        256
 ion)

 conv_3/relu (ReLU)          (None, 32, 32, 64)        0

 dw_separable_4 (DepthwiseC  (None, 16, 16, 64)        576
 onv2D)

 pw_separable_4 (Conv2D)     (None, 16, 16, 128)       8192

 pw_separable_4/BN (BatchNo  (None, 16, 16, 128)       512
 rmalization)

 pw_separable_4/relu (ReLU)  (None, 16, 16, 128)       0

 dw_separable_5 (DepthwiseC  (None, 16, 16, 128)       1152
 onv2D)

 pw_separable_5 (Conv2D)     (None, 16, 16, 128)       16384

 pw_separable_5/BN (BatchNo  (None, 16, 16, 128)       512
 rmalization)

 pw_separable_5/relu (ReLU)  (None, 16, 16, 128)       0

 dw_separable_6 (DepthwiseC  (None, 8, 8, 128)         1152
 onv2D)

 pw_separable_6 (Conv2D)     (None, 8, 8, 256)         32768

 pw_separable_6/BN (BatchNo  (None, 8, 8, 256)         1024
 rmalization)

 pw_separable_6/relu (ReLU)  (None, 8, 8, 256)         0

 dw_separable_7 (DepthwiseC  (None, 8, 8, 256)         2304
 onv2D)

 pw_separable_7 (Conv2D)     (None, 8, 8, 256)         65536

 pw_separable_7/BN (BatchNo  (None, 8, 8, 256)         1024
 rmalization)

 pw_separable_7/relu (ReLU)  (None, 8, 8, 256)         0

 dw_separable_8 (DepthwiseC  (None, 8, 8, 256)         2304
 onv2D)

 pw_separable_8 (Conv2D)     (None, 8, 8, 256)         65536

 pw_separable_8/BN (BatchNo  (None, 8, 8, 256)         1024
 rmalization)

 pw_separable_8/relu (ReLU)  (None, 8, 8, 256)         0

 dw_separable_9 (DepthwiseC  (None, 8, 8, 256)         2304
 onv2D)

 pw_separable_9 (Conv2D)     (None, 8, 8, 256)         65536

 pw_separable_9/BN (BatchNo  (None, 8, 8, 256)         1024
 rmalization)

 pw_separable_9/relu (ReLU)  (None, 8, 8, 256)         0

 dw_separable_10 (Depthwise  (None, 8, 8, 256)         2304
 Conv2D)

 pw_separable_10 (Conv2D)    (None, 8, 8, 256)         65536

 pw_separable_10/BN (BatchN  (None, 8, 8, 256)         1024
 ormalization)

 pw_separable_10/relu (ReLU  (None, 8, 8, 256)         0
 )

 dw_separable_11 (Depthwise  (None, 8, 8, 256)         2304
 Conv2D)

 pw_separable_11 (Conv2D)    (None, 8, 8, 256)         65536

 pw_separable_11/BN (BatchN  (None, 8, 8, 256)         1024
 ormalization)

 pw_separable_11/relu (ReLU  (None, 8, 8, 256)         0
 )

 dw_separable_12 (Depthwise  (None, 4, 4, 256)         2304
 Conv2D)

 pw_separable_12 (Conv2D)    (None, 4, 4, 512)         131072

 pw_separable_12/BN (BatchN  (None, 4, 4, 512)         2048
 ormalization)

 pw_separable_12/relu (ReLU  (None, 4, 4, 512)         0
 )

 dw_separable_13 (Depthwise  (None, 4, 4, 512)         4608
 Conv2D)

 pw_separable_13 (Conv2D)    (None, 4, 4, 512)         262144

 pw_separable_13/BN (BatchN  (None, 4, 4, 512)         2048
 ormalization)

 pw_separable_13/relu (ReLU  (None, 4, 4, 512)         0
 )

 dw_sepconv_t_0 (DepthwiseC  (None, 8, 8, 512)         5120
 onv2DTranspose)

 pw_sepconv_t_0 (Conv2D)     (None, 8, 8, 256)         131328

 pw_sepconv_t_0/BN (BatchNo  (None, 8, 8, 256)         1024
 rmalization)

 pw_sepconv_t_0/relu (ReLU)  (None, 8, 8, 256)         0

 dropout (Dropout)           (None, 8, 8, 256)         0

 dw_sepconv_t_1 (DepthwiseC  (None, 16, 16, 256)       2560
 onv2DTranspose)

 pw_sepconv_t_1 (Conv2D)     (None, 16, 16, 128)       32896

 pw_sepconv_t_1/BN (BatchNo  (None, 16, 16, 128)       512
 rmalization)

 pw_sepconv_t_1/relu (ReLU)  (None, 16, 16, 128)       0

 dropout_1 (Dropout)         (None, 16, 16, 128)       0

 dw_sepconv_t_2 (DepthwiseC  (None, 32, 32, 128)       1280
 onv2DTranspose)

 pw_sepconv_t_2 (Conv2D)     (None, 32, 32, 64)        8256

 pw_sepconv_t_2/BN (BatchNo  (None, 32, 32, 64)        256
 rmalization)

 pw_sepconv_t_2/relu (ReLU)  (None, 32, 32, 64)        0

 dropout_2 (Dropout)         (None, 32, 32, 64)        0

 dw_sepconv_t_3 (DepthwiseC  (None, 64, 64, 64)        640
 onv2DTranspose)

 pw_sepconv_t_3 (Conv2D)     (None, 64, 64, 32)        2080

 pw_sepconv_t_3/BN (BatchNo  (None, 64, 64, 32)        128
 rmalization)

 pw_sepconv_t_3/relu (ReLU)  (None, 64, 64, 32)        0

 dropout_3 (Dropout)         (None, 64, 64, 32)        0

 dw_sepconv_t_4 (DepthwiseC  (None, 128, 128, 32)      320
 onv2DTranspose)

 pw_sepconv_t_4 (Conv2D)     (None, 128, 128, 16)      528

 pw_sepconv_t_4/BN (BatchNo  (None, 128, 128, 16)      64
 rmalization)

 pw_sepconv_t_4/relu (ReLU)  (None, 128, 128, 16)      0

 dropout_4 (Dropout)         (None, 128, 128, 16)      0

 head (Conv2D)               (None, 128, 128, 1)       17

=================================================================
Total params: 1058865 (4.04 MB)
Trainable params: 1051889 (4.01 MB)
Non-trainable params: 6976 (27.25 KB)
_________________________________________________________________
from tf_keras.metrics import BinaryIoU

# Compile the native TF-Keras model (required to evaluate the metrics)
model_keras.compile(loss='binary_crossentropy', metrics=[BinaryIoU(), 'accuracy'])

# Check Keras model performance
_, biou, acc = model_keras.evaluate(x_val, y_val, steps=steps, verbose=0)

print(f"TF-Keras binary IoU / pixel accuracy: {biou:.4f} / {100*acc:.2f}%")
TF-Keras binary IoU / pixel accuracy: 0.9356 / 96.78%

3. Load a pre-trained quantized TF-Keras model

The next step is to quantize and potentially perform Quantization Aware Training (QAT) on the TF-Keras model from the previous step. After the TF-Keras model is quantized to 8-bit for all weights and activations, QAT is used to maintain the performance of the quantized model. Again, a pre-trained model is downloaded to save runtime.

from akida_models import akida_unet_portrait128_pretrained

# Load the pre-trained quantized model
model_quantized_keras = akida_unet_portrait128_pretrained()
model_quantized_keras.summary()
Downloading data from https://data.brainchip.com/models/AkidaV2/akida_unet/akida_unet_portrait128_i8_w8_a8.h5.

      0/4527816 [..............................] - ETA: 0s
 253952/4527816 [>.............................] - ETA: 0s
1515520/4527816 [=========>....................] - ETA: 0s
2973696/4527816 [==================>...........] - ETA: 0s
4527816/4527816 [==============================] - 0s 0us/step
Download complete.
Model: "akida_unet"
_________________________________________________________________
 Layer (type)                Output Shape              Param #
=================================================================
 input (InputLayer)          [(None, 128, 128, 3)]     0

 rescaling (QuantizedRescal  (None, 128, 128, 3)       0
 ing)

 conv_0 (QuantizedConv2D)    (None, 64, 64, 16)        448

 conv_0/relu (QuantizedReLU  (None, 64, 64, 16)        32
 )

 conv_1 (QuantizedConv2D)    (None, 64, 64, 32)        4640

 conv_1/relu (QuantizedReLU  (None, 64, 64, 32)        64
 )

 conv_2 (QuantizedConv2D)    (None, 32, 32, 64)        18496

 conv_2/relu (QuantizedReLU  (None, 32, 32, 64)        128
 )

 conv_3 (QuantizedConv2D)    (None, 32, 32, 64)        36928

 conv_3/relu (QuantizedReLU  (None, 32, 32, 64)        128
 )

 dw_separable_4 (QuantizedD  (None, 16, 16, 64)        704
 epthwiseConv2D)

 pw_separable_4 (QuantizedC  (None, 16, 16, 128)       8320
 onv2D)

 pw_separable_4/relu (Quant  (None, 16, 16, 128)       256
 izedReLU)

 dw_separable_5 (QuantizedD  (None, 16, 16, 128)       1408
 epthwiseConv2D)

 pw_separable_5 (QuantizedC  (None, 16, 16, 128)       16512
 onv2D)

 pw_separable_5/relu (Quant  (None, 16, 16, 128)       256
 izedReLU)

 dw_separable_6 (QuantizedD  (None, 8, 8, 128)         1408
 epthwiseConv2D)

 pw_separable_6 (QuantizedC  (None, 8, 8, 256)         33024
 onv2D)

 pw_separable_6/relu (Quant  (None, 8, 8, 256)         512
 izedReLU)

 dw_separable_7 (QuantizedD  (None, 8, 8, 256)         2816
 epthwiseConv2D)

 pw_separable_7 (QuantizedC  (None, 8, 8, 256)         65792
 onv2D)

 pw_separable_7/relu (Quant  (None, 8, 8, 256)         512
 izedReLU)

 dw_separable_8 (QuantizedD  (None, 8, 8, 256)         2816
 epthwiseConv2D)

 pw_separable_8 (QuantizedC  (None, 8, 8, 256)         65792
 onv2D)

 pw_separable_8/relu (Quant  (None, 8, 8, 256)         512
 izedReLU)

 dw_separable_9 (QuantizedD  (None, 8, 8, 256)         2816
 epthwiseConv2D)

 pw_separable_9 (QuantizedC  (None, 8, 8, 256)         65792
 onv2D)

 pw_separable_9/relu (Quant  (None, 8, 8, 256)         512
 izedReLU)

 dw_separable_10 (Quantized  (None, 8, 8, 256)         2816
 DepthwiseConv2D)

 pw_separable_10 (Quantized  (None, 8, 8, 256)         65792
 Conv2D)

 pw_separable_10/relu (Quan  (None, 8, 8, 256)         512
 tizedReLU)

 dw_separable_11 (Quantized  (None, 8, 8, 256)         2816
 DepthwiseConv2D)

 pw_separable_11 (Quantized  (None, 8, 8, 256)         65792
 Conv2D)

 pw_separable_11/relu (Quan  (None, 8, 8, 256)         512
 tizedReLU)

 dw_separable_12 (Quantized  (None, 4, 4, 256)         2816
 DepthwiseConv2D)

 pw_separable_12 (Quantized  (None, 4, 4, 512)         131584
 Conv2D)

 pw_separable_12/relu (Quan  (None, 4, 4, 512)         1024
 tizedReLU)

 dw_separable_13 (Quantized  (None, 4, 4, 512)         5632
 DepthwiseConv2D)

 pw_separable_13 (Quantized  (None, 4, 4, 512)         262656
 Conv2D)

 pw_separable_13/relu (Quan  (None, 4, 4, 512)         1024
 tizedReLU)

 dw_sepconv_t_0 (QuantizedD  (None, 8, 8, 512)         6144
 epthwiseConv2DTranspose)

 pw_sepconv_t_0 (QuantizedC  (None, 8, 8, 256)         131328
 onv2D)

 pw_sepconv_t_0/relu (Quant  (None, 8, 8, 256)         512
 izedReLU)

 dropout (QuantizedDropout)  (None, 8, 8, 256)         0

 dw_sepconv_t_1 (QuantizedD  (None, 16, 16, 256)       3072
 epthwiseConv2DTranspose)

 pw_sepconv_t_1 (QuantizedC  (None, 16, 16, 128)       32896
 onv2D)

 pw_sepconv_t_1/relu (Quant  (None, 16, 16, 128)       256
 izedReLU)

 dropout_1 (QuantizedDropou  (None, 16, 16, 128)       0
 t)

 dw_sepconv_t_2 (QuantizedD  (None, 32, 32, 128)       1536
 epthwiseConv2DTranspose)

 pw_sepconv_t_2 (QuantizedC  (None, 32, 32, 64)        8256
 onv2D)

 pw_sepconv_t_2/relu (Quant  (None, 32, 32, 64)        128
 izedReLU)

 dropout_2 (QuantizedDropou  (None, 32, 32, 64)        0
 t)

 dw_sepconv_t_3 (QuantizedD  (None, 64, 64, 64)        768
 epthwiseConv2DTranspose)

 pw_sepconv_t_3 (QuantizedC  (None, 64, 64, 32)        2080
 onv2D)

 pw_sepconv_t_3/relu (Quant  (None, 64, 64, 32)        64
 izedReLU)

 dropout_3 (QuantizedDropou  (None, 64, 64, 32)        0
 t)

 dw_sepconv_t_4 (QuantizedD  (None, 128, 128, 32)      384
 epthwiseConv2DTranspose)

 pw_sepconv_t_4 (QuantizedC  (None, 128, 128, 16)      528
 onv2D)

 pw_sepconv_t_4/relu (Quant  (None, 128, 128, 16)      32
 izedReLU)

 dropout_4 (QuantizedDropou  (None, 128, 128, 16)      0
 t)

 head (QuantizedConv2D)      (None, 128, 128, 1)       19

 dequantizer (Dequantizer)   (None, 128, 128, 1)       0

=================================================================
Total params: 1061603 (4.05 MB)
Trainable params: 1047905 (4.00 MB)
Non-trainable params: 13698 (53.51 KB)
_________________________________________________________________
# Compile the quantized TF-Keras model (required to evaluate the metrics)
model_quantized_keras.compile(loss='binary_crossentropy', metrics=[BinaryIoU(), 'accuracy'])

# Check Keras model performance
_, biou, acc = model_quantized_keras.evaluate(x_val, y_val, steps=steps, verbose=0)

print(f"TF-Keras quantized binary IoU / pixel accuracy: {biou:.4f} / {100*acc:.2f}%")
TF-Keras quantized binary IoU / pixel accuracy: 0.9344 / 96.77%

4. Conversion to Akida

Finally, the quantized TF-Keras model from the previous step is converted into an Akida model and its performance is evaluated. Note that the original performance of the TF-Keras floating-point model is maintained throughout the conversion process in this example.

from cnn2snn import convert

# Convert the model
model_akida = convert(model_quantized_keras)
model_akida.summary()
                  Model Summary
_________________________________________________
Input shape    Output shape   Sequences  Layers
=================================================
[128, 128, 3]  [128, 128, 1]  1          36
_________________________________________________

_____________________________________________________________________________
Layer (type)                               Output shape    Kernel shape

====================== SW/conv_0-dequantizer (Software) =====================

conv_0 (InputConv2D)                       [64, 64, 16]    (3, 3, 3, 16)
_____________________________________________________________________________
conv_1 (Conv2D)                            [64, 64, 32]    (3, 3, 16, 32)
_____________________________________________________________________________
conv_2 (Conv2D)                            [32, 32, 64]    (3, 3, 32, 64)
_____________________________________________________________________________
conv_3 (Conv2D)                            [32, 32, 64]    (3, 3, 64, 64)
_____________________________________________________________________________
dw_separable_4 (DepthwiseConv2D)           [16, 16, 64]    (3, 3, 64, 1)
_____________________________________________________________________________
pw_separable_4 (Conv2D)                    [16, 16, 128]   (1, 1, 64, 128)
_____________________________________________________________________________
dw_separable_5 (DepthwiseConv2D)           [16, 16, 128]   (3, 3, 128, 1)
_____________________________________________________________________________
pw_separable_5 (Conv2D)                    [16, 16, 128]   (1, 1, 128, 128)
_____________________________________________________________________________
dw_separable_6 (DepthwiseConv2D)           [8, 8, 128]     (3, 3, 128, 1)
_____________________________________________________________________________
pw_separable_6 (Conv2D)                    [8, 8, 256]     (1, 1, 128, 256)
_____________________________________________________________________________
dw_separable_7 (DepthwiseConv2D)           [8, 8, 256]     (3, 3, 256, 1)
_____________________________________________________________________________
pw_separable_7 (Conv2D)                    [8, 8, 256]     (1, 1, 256, 256)
_____________________________________________________________________________
dw_separable_8 (DepthwiseConv2D)           [8, 8, 256]     (3, 3, 256, 1)
_____________________________________________________________________________
pw_separable_8 (Conv2D)                    [8, 8, 256]     (1, 1, 256, 256)
_____________________________________________________________________________
dw_separable_9 (DepthwiseConv2D)           [8, 8, 256]     (3, 3, 256, 1)
_____________________________________________________________________________
pw_separable_9 (Conv2D)                    [8, 8, 256]     (1, 1, 256, 256)
_____________________________________________________________________________
dw_separable_10 (DepthwiseConv2D)          [8, 8, 256]     (3, 3, 256, 1)
_____________________________________________________________________________
pw_separable_10 (Conv2D)                   [8, 8, 256]     (1, 1, 256, 256)
_____________________________________________________________________________
dw_separable_11 (DepthwiseConv2D)          [8, 8, 256]     (3, 3, 256, 1)
_____________________________________________________________________________
pw_separable_11 (Conv2D)                   [8, 8, 256]     (1, 1, 256, 256)
_____________________________________________________________________________
dw_separable_12 (DepthwiseConv2D)          [4, 4, 256]     (3, 3, 256, 1)
_____________________________________________________________________________
pw_separable_12 (Conv2D)                   [4, 4, 512]     (1, 1, 256, 512)
_____________________________________________________________________________
dw_separable_13 (DepthwiseConv2D)          [4, 4, 512]     (3, 3, 512, 1)
_____________________________________________________________________________
pw_separable_13 (Conv2D)                   [4, 4, 512]     (1, 1, 512, 512)
_____________________________________________________________________________
dw_sepconv_t_0 (DepthwiseConv2DTranspose)  [8, 8, 512]     (3, 3, 512, 1)
_____________________________________________________________________________
pw_sepconv_t_0 (Conv2D)                    [8, 8, 256]     (1, 1, 512, 256)
_____________________________________________________________________________
dw_sepconv_t_1 (DepthwiseConv2DTranspose)  [16, 16, 256]   (3, 3, 256, 1)
_____________________________________________________________________________
pw_sepconv_t_1 (Conv2D)                    [16, 16, 128]   (1, 1, 256, 128)
_____________________________________________________________________________
dw_sepconv_t_2 (DepthwiseConv2DTranspose)  [32, 32, 128]   (3, 3, 128, 1)
_____________________________________________________________________________
pw_sepconv_t_2 (Conv2D)                    [32, 32, 64]    (1, 1, 128, 64)
_____________________________________________________________________________
dw_sepconv_t_3 (DepthwiseConv2DTranspose)  [64, 64, 64]    (3, 3, 64, 1)
_____________________________________________________________________________
pw_sepconv_t_3 (Conv2D)                    [64, 64, 32]    (1, 1, 64, 32)
_____________________________________________________________________________
dw_sepconv_t_4 (DepthwiseConv2DTranspose)  [128, 128, 32]  (3, 3, 32, 1)
_____________________________________________________________________________
pw_sepconv_t_4 (Conv2D)                    [128, 128, 16]  (1, 1, 32, 16)
_____________________________________________________________________________
head (Conv2D)                              [128, 128, 1]   (1, 1, 16, 1)
_____________________________________________________________________________
dequantizer (Dequantizer)                  [128, 128, 1]   N/A
_____________________________________________________________________________
import tf_keras as keras

# Check Akida model performance
labels, pots = None, None

for s in range(steps):
    batch = x_val[s * batch_size: (s + 1) * batch_size, :]
    label_batch = y_val[s * batch_size: (s + 1) * batch_size, :]
    pots_batch = model_akida.predict(batch.astype('uint8'))

    if labels is None:
        labels = label_batch
        pots = pots_batch
    else:
        labels = np.concatenate((labels, label_batch))
        pots = np.concatenate((pots, pots_batch))
preds = keras.activations.sigmoid(pots)

m_binary_iou = keras.metrics.BinaryIoU(target_class_ids=[0, 1], threshold=0.5)
m_binary_iou.update_state(labels, preds)
binary_iou = m_binary_iou.result().numpy()

m_accuracy = keras.metrics.Accuracy()
m_accuracy.update_state(labels, preds > 0.5)
accuracy = m_accuracy.result().numpy()
print(f"Akida binary IoU / pixel accuracy: {binary_iou:.4f} / {100*accuracy:.2f}%")

# For non-regression purposes
assert binary_iou > 0.9
Akida binary IoU / pixel accuracy: 0.9249 / 96.71%

5. Segment a single image

For visualization of the person segmentation performed by the Akida model, display a single image along with the segmentation produced by the original floating-point model and the ground truth segmentation.

import matplotlib.pyplot as plt

# Run inference on a random single image and display TF-Keras and Akida outputs
sample = np.expand_dims(x_val[id, :], 0)
keras_out = model_keras(sample)
akida_out = keras.activations.sigmoid(model_akida.forward(sample.astype('uint8')))

fig, axs = plt.subplots(1, 3, constrained_layout=True)
axs[0].imshow(keras_out[0] * sample[0] / 255.)
axs[0].set_title('Keras segmentation', fontsize=10)
axs[0].axis('off')

axs[1].imshow(akida_out[0] * sample[0] / 255.)
axs[1].set_title('Akida segmentation', fontsize=10)
axs[1].axis('off')

axs[2].imshow(y_val[id] * sample[0] / 255.)
axs[2].set_title('Expected segmentation', fontsize=10)
axs[2].axis('off')

plt.show()
Keras segmentation, Akida segmentation, Expected segmentation

Total running time of the script: (2 minutes 4.373 seconds)

Gallery generated by Sphinx-Gallery