#!/usr/bin/env python
# ******************************************************************************
# Copyright 2023 Brainchip Holdings Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ******************************************************************************
"""
TENN recurrent architecture definition.
"""
__all__ = ["tenn_recurrent_sc10", "tenn_recurrent_sc10_pretrained"]
import tensorflow as tf
from tf_keras.models import Model
from tf_keras.layers import Input, GlobalAveragePooling1D, SpatialDropout1D, Rescaling
from ..layer_blocks import kernelized_block
from ..utils import fetch_file
from ..model_io import load_model, get_model_path
[docs]
def tenn_recurrent_sc10(input_shape=(16384, 1), num_classes=10, input_scaling=(2**15, 0)):
""" Instantiates a TENN recurrent architecture.
Args:
input_shape (tuple, optional): the input shape. Defaults to (16384, 1).
num_classes (int, optional): number of classes. Defaults to 10.
input_scaling (tuple, optional): scale factor set to the max value of a 16bits unsigned
inputs and offset set to 0. Note that following Akida convention, the scale factor
is a number used as a divisor.
Returns:
keras.Model: a TENN recurrent model for SC10
"""
# architecture parameters
num_coeffs = 32
channels = [8, 16, 32, 64, 96, 128]
subsampling_pattern = [4, 4, 2, 2, 2, 2]
inputs = Input(shape=input_shape, dtype=tf.int16, name="input")
scale, offset = input_scaling
x = Rescaling(1. / scale, offset, name="rescaling")(inputs)
# main network stack
for i, (channel, subsample) in enumerate(zip(channels, subsampling_pattern)):
x = kernelized_block(x, num_coeffs, channel, subsampling=subsample,
add_batchnorm=True, relu_activation='ReLU', name=f'ssm_layer_{i}')
x = SpatialDropout1D(0.1)(x)
# the classification head
x = kernelized_block(x, num_coeffs, num_classes, subsampling=False, add_batchnorm=False,
relu_activation=False, name='ssm_layer_head')
# for benchmarking a GAP is applied to the entire sequence to make a single prediction on the
# 1sec sample (output filtering).
x = GlobalAveragePooling1D(name='gap')(x)
return Model(inputs, x, name="tenn_recurrent_sc10")
[docs]
def tenn_recurrent_sc10_pretrained(quantized=True):
""" Helper method to retrieve an `tenn_recurrent_sc10` model that was trained on SC10 dataset.
Args:
quantized (bool, optional): a boolean indicating whether the model should be loaded
quantized or not. Defaults to True.
Returns:
keras.Model: a Keras Model instance.
"""
if quantized:
model_name_v2 = 'tenn_recurrent_sc10_stateful_i8_w8_a8.h5'
file_hash_v2 = 'fcc7f3f9f626c1be5c0a05e7c34531db8c36ed4816be6e9521e36fd7245c40fa'
else:
model_name_v2 = 'tenn_recurrent_sc10_stateful.h5'
file_hash_v2 = 'fe8459926fe2270369dd25f32cae982dff9d4488176d9444db8f27995094187e'
model_path, model_name, file_hash = get_model_path(
"tenn_recurrent", model_name_v2=model_name_v2, file_hash_v2=file_hash_v2)
model_path = fetch_file(model_path, fname=model_name, file_hash=file_hash,
cache_subdir='models')
return load_model(model_path)