Source code for akida_models.tenn_recurrent.model_tenn_recurrent

#!/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)