#!/usr/bin/env python
# ******************************************************************************
# Copyright 2021 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.
# ******************************************************************************
"""
Common utility methods used by the training scripts.
"""
import argparse
import numpy as np
from keras.preprocessing.image import DirectoryIterator
from keras.optimizer_v2.adam import Adam
from keras.layers import Activation
[docs]def compile_model(model,
learning_rate=1e-3,
loss='categorical_crossentropy',
metrics=None):
""" Compiles the model using Adam optimizer.
Args:
model (keras.Model): the model to compile
learning_rate (float, optional): the learning rate. Defaults to 1e-3.
loss (str or function, optional): the loss function. Defaults to
'categorical_crossentropy'.
metrics (list, optional): list of metrics to be evaluated during
training and testing. Defaults to None.
"""
if metrics is None:
metrics = ['accuracy']
model.compile(optimizer=Adam(learning_rate=learning_rate),
loss=loss,
metrics=metrics)
[docs]def freeze_model_before(model, freeze_before):
""" Freezes the model before the given layer name.
Args:
model (keras.Model): the model to freeze
freeze_before (str): name of the layer from which the model will not be
frozen
Raises:
ValueError: if the provided layer name was not found in the model
"""
if freeze_before is not None:
trainable = False
for layer in model.layers:
if layer.name == freeze_before:
trainable = True
layer.trainable = trainable
if not trainable:
raise ValueError(f"No such layer {freeze_before} in model.")
[docs]def evaluate_model(model,
x,
y=None,
batch_size=None,
steps=None,
print_history=False):
""" Evaluates model performances.
Args:
model (keras.Model): the model to evaluate
x (tf.Dataset, np.array or generator): evaluation input data
y (tf.Dataset, np.array or generator, optional): evaluation target
data. Defaults to None.
batch_size (int, optional): the batch size. Defaults to None.
steps (int, optional): total number of steps before declaring the
evaluation round finished. Defaults to None.
print_history (bool, optional): either to print all history or just
accuracy. Defaults to False.
"""
history = model.evaluate(x=x,
y=y,
batch_size=batch_size,
steps=steps,
verbose=0)
if print_history:
print('Validation history: ', history)
else:
print('Validation accuracy: ', history[1])
[docs]def evaluate_akida_model(model, x, activation):
""" Evaluates Akida model and return predictions and labels to compute
accuracy.
Args:
model (akida.Model): the model to evaluate
x (tf.Dataset, np.array or generator): evaluation input data
activation (str): activation function to apply to potentials
Returns:
np.array, np.array: predictions and labels
"""
# Initialize to None to allow different shapes depending on the caller
labels = None
pots = None
# Datasets created as ImageDataGenerator will be looped over indefinitely,
# they thus need a special treatment
steps = None
if isinstance(x, DirectoryIterator):
steps = np.math.ceil(x.samples / x.batch_size)
for batch, label_batch in x:
if not isinstance(batch, np.ndarray):
batch = batch.numpy()
pots_batch = model.predict(batch.astype('uint8'))
if labels is None:
labels = label_batch
pots = pots_batch.squeeze(axis=(1, 2))
else:
labels = np.concatenate((labels, label_batch))
pots = np.concatenate((pots, pots_batch.squeeze(axis=(1, 2))))
# End the for loop if the number of batches has reached the number of
# steps
if steps is not None and x.total_batches_seen == steps:
break
return Activation(activation)(pots), labels
def get_training_parser(batch_size,
tune=False,
freeze_before=False,
global_batch_size=True,
global_parser=None):
""" Instantiates a base arguments parser for training scripts.
The returned parser comes with train, tune and eval actions as subparsers
(with default parameters) for which one can add arguments if required.
Args:
batch_size (int): batch size default value
tune (bool, optional): either to add a tune action or not. Default to
False.
freeze_before (bool, optional): either to add a "--freeze_before"
parameter in train and tune actions or not. Default to False.
global_batch_size (bool, optional): either to add batch size as a
global parser parameter or for train/tune actions only. Default to
True.
global_parser (argparse.ArgumentParser, optional): global parser with
custom parameters. Defaults to None.
Returns:
argparse.ArgumentParser, argparse.ArgumentParser,
argparse.ArgumentParser, argparse.ArgumentParser,
argparse._SubParsersAction: main parser and train, tune and eval actions
subparsers and the subparser object.
"""
# Define a strictly positive int type for parameters
def positive_int(value):
ivalue = int(value)
if ivalue <= 0:
raise argparse.ArgumentTypeError(
"%s is an invalid value, must be >0." % value)
return ivalue
parser = argparse.ArgumentParser()
# Base arguments
if global_parser is None:
global_parser = argparse.ArgumentParser(add_help=False)
global_parser.add_argument("-m",
"--model",
type=str,
required=True,
help="Model to load")
if global_batch_size:
global_parser.add_argument("-b",
"--batch_size",
type=positive_int,
default=batch_size,
help="The batch size")
# Create an action subparser
subparsers = parser.add_subparsers(help="Desired action to perform",
dest="action")
# Create parent subparser for arguments shared between train and tune
# actions
parent_parser = argparse.ArgumentParser(add_help=False,
parents=[global_parser])
parent_parser.add_argument("-s",
"--savemodel",
type=str,
default=None,
help="Save model with the specified name")
parent_parser.add_argument("-e",
"--epochs",
type=positive_int,
required=True,
help="The number of training epochs")
if not global_batch_size:
parent_parser.add_argument("-b",
"--batch_size",
type=positive_int,
default=batch_size,
help="The batch size")
if freeze_before:
parent_parser.add_argument(
"-fb",
"--freeze_before",
type=str,
default=None,
help="Freeze the layers of the model before this \
layer")
# Subparsers based on parent
train_parser = subparsers.add_parser("train",
parents=[parent_parser],
help="Train a Keras model")
tune_parser = None
if tune:
tune_parser = subparsers.add_parser("tune",
parents=[parent_parser],
help="Tune a Keras model")
# Other subparsers
eval_parser = subparsers.add_parser("eval",
help="Evaluate a model",
parents=[global_parser])
eval_parser.add_argument("-ak",
"--akida",
action='store_true',
help="Converts to an Akida model and evaluate")
return parser, train_parser, tune_parser, eval_parser, subparsers