# ******************************************************************************
# Copyright 2020 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.
# ******************************************************************************
"""Quantized model calibration tools"""
import keras
import tensorflow as tf
from ..quantization_layers import (QuantizedReLU, QuantizedDense,
QuantizedConv2D, QuantizedSeparableConv2D)
from ..quantization_ops import WeightQuantizer
from ..transforms.clone import clone_model_with_weights
def create_submodel(model, first_layer, last_layer):
"""Create a submodel of a Keras Model
Args:
model (:obj:`keras.Model`): a Keras Model
first_layer (:obj:`keras.Layer`): the submodel first layer
last_layer (:obj:`keras.Layer`): the submodel last layer
Returns:
:obj:`keras.Model`: a Keras sub-model
"""
submodel = keras.Sequential()
input_shape = first_layer.input_shape[1:]
submodel.add(keras.layers.Input(shape=input_shape))
add_layer = False
for layer in model.layers:
if layer == first_layer:
add_layer = True
if add_layer:
submodel.add(layer)
if layer == last_layer:
break
return submodel
class NoneQuantizer(WeightQuantizer):
"""A WeightQuantizer that returns float weights
"""
def scale_factor(self, _):
return 1
def quantize(self, w):
return w
def disable_quantizers(layer):
"""Disable quantizers for the specified layer
Args:
layer (:obj:`keras.Layer`): the layer
Returns:
[:obj:`cnn2snn.quantization_ops.WeightQuantizer`]: the list of disabled
quantizers.
"""
quantizers = [layer.quantizer]
layer.quantizer = NoneQuantizer(layer.quantizer.bitwidth)
if hasattr(layer, 'quantizer_dw'):
quantizers.append(layer.quantizer_dw)
layer.quantizer_dw = NoneQuantizer(layer.quantizer_dw.bitwidth)
return quantizers
def restore_quantizers(layer, quantizers):
"""Restores quantizers for the specified layer
Args:
layer (:obj:`keras.Layer`): the layer
quantizers ([:obj:`cnn2snn.quantization_ops.WeightQuantizer`]): the list of disabled
quantizers.
"""
layer.quantizer = quantizers[0]
if hasattr(layer, 'quantizer_dw'):
layer.quantizer_dw = quantizers[1]
def next_activation(layer):
"""Find the next activation after the specified layer
Args:
layer (:obj:`keras.Layer`): the layer
Returns:
:obj:`keras.Layer`: the next activation layer or None
"""
outbound_nodes = layer.outbound_nodes
for node in outbound_nodes:
next_layer = node.layer
if isinstance(next_layer, QuantizedReLU):
return next_layer
activation = next_activation(next_layer)
if activation is not None:
return activation
return None
def is_on_top(layer, bottom_layer):
"""Checks if a layer follows another layer
Args:
layer (:obj:`keras.Layer`): the top layer
bottom_layer (:obj:`keras.Layer`): the previous layer
Returns:
bool: True if the layer is on top
"""
outbound_nodes = bottom_layer.outbound_nodes
if len(outbound_nodes) == 0:
# Bottom layer has no successors
return False
for node in outbound_nodes:
# Check if layer is on top of one of the outbound layers
next_layer = node.layer
if next_layer == layer or is_on_top(layer, next_layer):
return True
return False
def is_quantized_neural(layer):
"""Check is a layer is a quantized neural layer
These layers are the only ones eligible for calibration.
Args:
layer (:obj:`keras.Layer`): the top layer
Returns:
bool: True if the layer can be calibrated
"""
return isinstance(
layer, (QuantizedDense, QuantizedConv2D, QuantizedSeparableConv2D))
[docs]
class QuantizationSampler():
"""A tool to inspect the layer outputs of a quantized model
The sampler is initialized with a quantized model and a set of samples used
for the evaluation of the layer outputs. An optional batch size can be
specified to avoid out-of-memory errors when using a GPU.
To evaluate the outputs of a specific layer, it must first be selected
using the `select_layer` member.
Once done, three methods are available to inspect the layer outputs:
- `quantized_outputs` returns the actual outputs of the layer,
- `float_outputs` returns the outputs of the layer if its weights were not
quantized,
- `quantization_error` applies the `keras.metrics.Metric` passed as
arguments to the difference between the float and quantized outputs.
Example:
>>> # Evaluate the quantization MSE of a quantized layer
>>> model = tf.keras.Sequential([
... tf.keras.layers.Dense(5, input_shape=(3,)),
... tf.keras.layers.ReLU()])
>>> model_quantized = cnn2snn.quantize(model,
... weight_quantization=4,
... activ_quantization=4)
>>> # Instantiate a QuantizationSampler with a few dataset samples
>>> sampler = QuantizationSampler(model_quantized, samples)
>>> # Select the quantized layer
>>> sampler.select_layer(model_quantized.layers[0])
>>> # Evaluate the Mean Squared Error
>>> m = keras.metrics.MeanSquaredError()
>>> mse = sampler.quantization_error(m)
Args:
model (:obj:`keras.Model`): a quantized Keras model
samples (:obj:`np.ndarray`): a set of calibration samples
batch_size (int): the batch size used for evaluation.
"""
def __init__(self, model, samples, batch_size=None):
self.model = model
n_inputs = samples.shape[0]
if batch_size is None:
batch_size = n_inputs
# We only process full batches
self.n_batches = n_inputs // batch_size
# Split samples by batches
self.samples = []
for i in range(self.n_batches):
b_from = i * batch_size
b_to = (i + 1) * batch_size
self.samples.append(samples[b_from:b_to, :, :, :])
self.inputs = None
self.submodel = None
self.layer = None
def select_layer(self, layer, include_activation=False):
"""Select the layer to inspect
Args:
layer (:obj:`keras.Layer`): a quantized keras layer
include_activation (bool): evaluate the outputs after activation
"""
if not is_quantized_neural(layer):
raise ValueError("Output layer must be a Quantized layer")
# Evaluate the outputs immediately before the selected layer to save
# computation time when repeatedly evaluating the layer outputs
if self.submodel is None or not is_on_top(layer, self.layer):
# Evaluate layer inputs from the beginning of the model
input_layer = self.model.layers[0]
self.inputs = []
for i in range(self.n_batches):
self.inputs.append(self.samples[i])
else:
# Evaluate layer inputs using the outputs of the current submodel
input_layer = self.submodel.layers[-1].outbound_nodes[0].layer
for i in range(self.n_batches):
self.inputs[i] = self.quantized_outputs(i)
if layer != input_layer:
# Advance inputs
prev_layer = layer.inbound_nodes[0].inbound_layers
submodel = create_submodel(self.model, input_layer, prev_layer)
for i in range(self.n_batches):
self.inputs[i] = submodel(self.inputs[i])
# Create the submodel including the layer (and optionally its activation)
activation = next_activation(layer) if include_activation else None
if activation is not None:
self.submodel = create_submodel(self.model, layer, activation)
else:
self.submodel = create_submodel(self.model, layer, layer)
# Remember current layer
self.layer = layer
def quantized_outputs(self, batch_index=0):
"""Evaluates the quantized outputs of the selected layer
Args:
batch_index (int): the samples batch index
Returns:
:obj:`tf.Tensor`: the quantized outputs
"""
if self.submodel is None:
raise SystemError("Set output layer before requesting data")
if batch_index < 0 or batch_index >= self.n_batches:
raise ValueError("Invalid batch index")
return self.submodel(self.inputs[batch_index])
def float_outputs(self, batch_index=0):
"""Evaluates the 'float' outputs of the selected layer
This disables the weights quantization and evaluates the outputs of the
layer as if it was using float weights.
Args:
batch_index (int): the samples batch index
Returns:
:obj:`tf.Tensor`: the quantized outputs
"""
if self.submodel is None:
raise SystemError("Set output layer before requesting data")
if batch_index < 0 or batch_index >= self.n_batches:
raise ValueError("Invalid batch index")
quantizers = disable_quantizers(self.layer)
outputs = self.submodel(self.inputs[batch_index])
restore_quantizers(self.layer, quantizers)
return outputs
def quantization_error(self, metric):
"""Evaluates the quantization error of the selected layer
This evaluates the difference between the layer quantized outputs and
the outputs if the layer was using float weights using the Metric
passed as parameter.
Args:
metric (:obj:`keras.metrics.Metric`): the Metric for evaluation
Returns:
the Metric results (usually a single float)
"""
metric.reset_state()
for i in range(self.n_batches):
metric.update_state(self.float_outputs(i),
self.quantized_outputs(i))
return metric.result().numpy()
def previous_neural(layer):
"""Find the previous neural layer before the specified layer
Args:
layer (:obj:`keras.Layer`): the layer
Returns:
:obj:`keras.Layer`: the next activation layer or None
"""
inbound_nodes = layer.inbound_nodes
for node in inbound_nodes:
previous_layer = node.inbound_layers
if is_quantized_neural(previous_layer):
return previous_layer
neural_layer = previous_neural(previous_layer)
if neural_layer is not None:
return neural_layer
return None
def activations_rescaling(model, samples, batch_size=None):
"""iterate over the layers of a quantized model, evaluating
the pre-activations before each QuantizedReLu layer,
and adjusting the max_value parameter of each QuantizedReLu
accordingly.
Args:
model (:obj:`keras.Model`): a Keras Model
samples (:obj:`np.ndarray`): a set of calibration samples
batch_size (int): the batch size used for evaluation.
"""
clone_model = clone_model_with_weights(model)
sampler = QuantizationSampler(clone_model, samples, batch_size)
for layer in clone_model.layers:
if not isinstance(layer, QuantizedReLU):
continue
previous_neural_layer = previous_neural(layer)
sampler.select_layer(previous_neural_layer, True)
float_output = sampler.float_outputs()
max_val = tf.reduce_max(float_output)
if max_val < layer.max_value_:
layer.max_value_ = tf.constant(max_val, dtype=tf.float32)
return clone_model