Source code for quantizeml.layers.separable_convolution

#!/usr/bin/env python
# ******************************************************************************
# Copyright 2022 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
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# See the License for the specific language governing permissions and
# limitations under the License.
# ******************************************************************************
QuantizedSeparableConv2D layer definition.

__all__ = ["QuantizedSeparableConv2D"]

import tensorflow as tf
from keras.layers import SeparableConv2D
from keras.utils import conv_utils
from keras import backend

from .layers_base import (register_quantize_target, rescale_outputs, tensor_inputs,
                          neural_layer_init, register_aligned_inputs)
from ..tensors import FixedPoint

[docs]@register_quantize_target(SeparableConv2D) @register_aligned_inputs @tf.keras.utils.register_keras_serializable() class QuantizedSeparableConv2D(SeparableConv2D): """ A separable convolutional layer that operates on quantized inputs and weights. Args: quant_config (dict, optional): the serialized quantization configuration. Defaults to None. """ @neural_layer_init(True) def __init__(self, *args, quant_config=None, **kwargs): if 'depth_multiplier' in kwargs: if kwargs['depth_multiplier'] != 1: raise ValueError("Keyword argument 'depth_multiplier' is not supported in \ QuantizedSeparableConv2D.") @tensor_inputs([FixedPoint, tf.Tensor]) @rescale_outputs def call(self, inputs): # Although the dephwise operation does not require it, we only accept inputs quantized # per-tensor to avoid increasing too much the fractional bits of the depthwise outputs. inputs.assert_per_tensor() # Quantize the weights depthwise_kernel = self.dw_weight_quantizer(self.depthwise_kernel) pointwise_kernel = self.pw_weight_quantizer(self.pointwise_kernel) dw_outputs_q = backend.depthwise_conv2d( inputs, depthwise_kernel, strides=self.strides, padding=self.padding, dilation_rate=self.dilation_rate, data_format=self.data_format) outputs = tf.nn.convolution( dw_outputs_q, pointwise_kernel, strides=[1, 1, 1, 1], padding='VALID', data_format=conv_utils.convert_data_format(self.data_format, ndim=4)) if self.use_bias: # Quantize bias and align it on the outputs bias = self.bias_quantizer(self.bias, outputs) outputs = tf.add(outputs, bias) return outputs def get_config(self): config = super().get_config() config["quant_config"] = self.quant_config return config