#!/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.
# ******************************************************************************
"""Equalization tools for Keras/CNN2SNN Sequential models.
These transformations take models with heterogeneities in their weights or
biases that are identified as harmful for the model quantization of for their
conversion to akida. They produce models with reduced heterogeneities, yet
returning nearly identical results.
"""
import numpy as np
import keras
from .clone import clone_model_with_weights
def _get_filter_max_values(layer):
if isinstance(layer, keras.layers.Conv2D):
# Layer "weights" are filters, biases
filters = layer.get_weights()[0]
# Filters are HWCN, i.e. the filter index is the last dimension
return np.amax(np.abs(filters), axis=(0, 1, 2))
if isinstance(layer, keras.layers.SeparableConv2D):
# Layer "weights" are depthwise filters, pointwise filters, biases
# We are only interested in rescaling pointwise filters
filters = layer.get_weights()[1]
# Pointwise filters are HWCN, i.e. the filter index is the last
# dimension
return np.amax(np.abs(filters), axis=(0, 1, 2))
if isinstance(layer, keras.layers.Dense):
# Layer "weights" are filters, biases
filters = layer.get_weights()[0]
# Filters are CN, i.e. the filter index is the last dimension
return np.amax(np.abs(filters), axis=0)
return None
def _get_channel_max_values(layer):
if isinstance(layer, keras.layers.Conv2D):
# Layer "weights" are filters, biases
filters = layer.get_weights()[0]
# Filters are HWCN, i.e. the channel index is the third dimension
return np.amax(np.abs(filters), axis=(0, 1, 3))
if isinstance(layer, keras.layers.SeparableConv2D):
# Layer "weights" are depthwise filters, pointwise filters, biases
dw_filters = layer.get_weights()[0]
# Depthwise filters are HWCN, i.e. channel index is the third dimension
return np.amax(np.abs(dw_filters), axis=(0, 1, 3))
if isinstance(layer, keras.layers.Dense):
# Layer "weights" are filters, biases
filters = layer.get_weights()[0]
# Filters are CN, i.e. the channel index is the first dimension
return np.amax(np.abs(filters), axis=1)
return None
def _rescale_fused_layer(layer, scales):
weights = layer.get_weights()
if isinstance(layer, keras.layers.SeparableConv2D):
# Layer "weights" are depthwise filters, pointwise filters, biases
dw_filters = weights[0]
pw_filters = weights[1]
# Divide dw_filters and multiply pw_filters by the same scale
for i in range(scales.shape[0]):
dw_filters[:, :, i, :] = dw_filters[:, :, i, :] / scales[i]
pw_filters[:, :, i, :] = pw_filters[:, :, i, :] * scales[i]
weights[0] = dw_filters
weights[1] = pw_filters
# Update layer weights
layer.set_weights(weights)
def _get_homogeneity_rate(layer):
filter_max = _get_filter_max_values(layer)
if filter_max is None:
return None
global_max = np.max(filter_max)
# Evaluate the relative score of each filter
filter_rate = filter_max / global_max
return np.mean(filter_rate)
[docs]def weights_homogeneity(model):
"""Give an estimation of the homogeneity of layer weights
For each Conv or Dense layer in the model, this compares the ranges of
the weights for each filter with the range of the tensor.
The score for each filter is expressed as an homogeneity rate (1 is the
maximum), and the layer homogeneity rate is the mean of all filter rates.
Args:
model (:obj:`tf.keras.Model`): a Keras model.
Returns:
dict: rates indexed by layer names.
"""
scores = {}
for layer in model.layers:
score = _get_homogeneity_rate(layer)
if score is not None:
scores[layer.name] = score
return scores
[docs]def normalize_separable_layer(layer):
"""This normalizes the depthwise weights of a SeparableConv2D.
In order to limit the quantization error when using a per-tensor
quantization of depthwise weights, this rescales all depthwise weights
to fit within the [-1, 1] range.
To preserve the output of the layer, each depthwise kernel is rescaled
independently to the [-1, 1] interval by dividing all weights by the
absolute maximum value, and inversely, all pointwise filters 'looking'
at these kernels are multiplied by the same value.
Args:
layer (:obj:`tf.keras.layers.SeparableConv2D`): a Keras SeparableConv2D
layer.
"""
if not isinstance(layer, keras.layers.SeparableConv2D):
raise ValueError("The layer is not a SeparableConv2D")
# Get maximum ranges per channel
dw_max = _get_channel_max_values(layer)
# Rescale depthwise to [-1, 1] and adjust pointwise accordingly
_rescale_fused_layer(layer, dw_max)
[docs]def normalize_separable_model(model):
"""This normalizes the depthwise weights of all SeparableConv2D in a Model.
Args:
model (:obj:`tf.keras.Model`): a Keras model.
Returns:
:obj:`tf.keras.Model`: a new Keras model with normalized depthwise
weights in SeparableConv2D layers.
"""
# Clone model
new_model = clone_model_with_weights(model)
# Normalize SeparableConv2D depthwise weights
for layer in new_model.layers:
if isinstance(layer, keras.layers.SeparableConv2D):
normalize_separable_layer(layer)
return new_model