#!/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_sc12", "tenn_recurrent_sc12_pretrained",
"tenn_recurrent_uored", "tenn_recurrent_uored_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_sc12(input_shape=(16384, 1), num_classes=12, 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 12.
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 SC12
"""
# 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_sc12")
[docs]
def tenn_recurrent_uored(input_shape=(4096, 1), num_classes=4, input_scaling=(2**15, 0)):
""" Instantiates a TENN recurrent architecture for UORED-VAFCLS bearing fault classification.
Outputs raw logits (no activation) for use with BinaryCrossentropy(from_logits=True).
Args:
input_shape (tuple, optional): the input shape. Defaults to (4096, 1).
num_classes (int, optional): number of fault classes. Defaults to 4
(inner race, outer race, ball, cage).
input_scaling (tuple, optional): scale factor and offset for int16 input rescaling.
Defaults to (2**15, 0).
Returns:
keras.Model: a TENN recurrent model for UORED-VAFCLS
"""
# architecture parameters
num_coeffs = 32
channels = [8, 16, 32, 32, 64]
subsampling_pattern = [8, 4, 2, 2, 2] # total 256x -> 4096/256 = 16 final temporal steps
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 (raw logits, no batchnorm or activation)
x = kernelized_block(x, num_coeffs, num_classes, subsampling=False, add_batchnorm=False,
relu_activation=False, name='ssm_layer_head')
x = GlobalAveragePooling1D(name='gap')(x)
return Model(inputs, x, name="tenn_recurrent_uored")
[docs]
def tenn_recurrent_sc12_pretrained(quantized=True):
""" Helper method to retrieve an `tenn_recurrent_sc12` model that was trained on SC12 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_sc12_stateful_i8_w8_a8.h5'
file_hash_v2 = 'bb844379bc357bc5aa1084ad5bf59a9bb956290cb6371f40c06bcd6cde53159a'
else:
model_name_v2 = 'tenn_recurrent_sc12_stateful.h5'
file_hash_v2 = 'aefa050ecc923b92486ae9888eb9e7a23e73ba2eea007ce18ef2ab37e3e21cf5'
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)
[docs]
def tenn_recurrent_uored_pretrained(quantized=True):
""" Helper method to retrieve an `tenn_recurrent_uored` model that was trained
on UORED-VAFCLS 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_uored_stateful_i8_w8_a8.h5'
file_hash_v2 = 'fc16370d2c5de11e3021d7df8069a6399f72b859cc19117c12578b570d40222b'
else:
model_name_v2 = 'tenn_recurrent_uored_stateful.h5'
file_hash_v2 = 'a71d725fa222191164695ed45f667654bc3cae11b39310ee4d7cc0f04150aef5'
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)