def build(self, input_shape):
super(QuantizeWrapper, self).build(input_shape)
self.optimizer_step = self.add_weight(
'optimizer_step',
initializer=tf.keras.initializers.Constant(-1),
dtype=tf.dtypes.int32,
trainable=False)
self._weight_vars = []
for weight, quantizer in \
self.quantize_config.get_weights_and_quantizers(self.layer):
quantizer_vars = quantizer.build(weight.shape,
self._weight_name(weight.name),
self)
self._weight_vars.append((weight, quantizer, quantizer_vars))
# Needed to ensure unquantized weights get trained as part of the wrapper.
self._trainable_weights.append(weight)
self._quantize_activations = []
for activation, quantizer in \
self.quantize_config.get_activations_and_quantizers(self.layer):
quantize_activation = quantize_aware_activation.QuantizeAwareActivation(
activation, quantizer, self.optimizer_step, self)
self._quantize_activations.append(quantize_activation)
self._output_quantizers = self.quantize_config.get_output_quantizers(
self.layer)
if self._output_quantizers:
self._output_quantizer_vars = self._output_quantizers[0].build(
self.layer.compute_output_shape(input_shape), 'output', self)
def build(self, input_shape):
# responsible for trainable self.kernel weights
super(_ConvBatchNorm2D, self).build(input_shape)
# resposible for trainable gamma and beta weights
self.batchnorm.build(self.compute_output_shape(input_shape))
if self.is_quantized:
self._weight_min_var = self.add_variable(
'weight_min',
initializer=initializers.Constant(-6.0),
trainable=False)
self._weight_max_var = self.add_variable(
'weight_max',
initializer=initializers.Constant(6.0),
trainable=False)
self.optimizer_step = self.add_weight(
'optimizer_step',
initializer=initializers.Constant(-1),
dtype=dtypes.int32,
trainable=False)
self.post_activation = quantize_aware_activation.QuantizeAwareActivation(
self.post_activation, self.activation_quantizer,
self.optimizer_step, self)
def build(self, input_shape):
super(QuantizeWrapper, self).build(input_shape)
self.optimizer_step = self.add_weight(
'optimizer_step',
initializer=initializers.Constant(-1),
dtype=dtypes.int32,
trainable=False)
self._weight_vars = []
for weight, quantizer in \
self.quantize_provider.get_weights_and_quantizers(self.layer):
min_var, max_var = self._add_range_weights(
self._weight_name(weight.name))
self._weight_vars.append((weight, quantizer, min_var, max_var))
# Needed to ensure unquantized weights get trained as part of the wrapper.
self._trainable_weights.append(weight)
self._quantize_activations = []
for activation, quantizer in \
self.quantize_provider.get_activations_and_quantizers(self.layer):
quantize_activation = quantize_aware_activation.QuantizeAwareActivation(
activation, quantizer, self.optimizer_step, self)
self._quantize_activations.append(quantize_activation)
self._output_quantizers = self.quantize_provider.get_output_quantizers(
self.layer)
if self._output_quantizers:
self._output_min_max = self._add_range_weights('output')
def _quantize_activation(activation, parent_class, quantize_params):
try:
return quantize_aware_activation.QuantizeAwareActivation(
activation.__name__, parent_class, **quantize_params)
except TypeError:
# Non-standard activation. Could be a custom callable, or an advanced
# activation. Simply return the original activation for now.
# TODO(pulkitb): Determine how to handle custom activations and advanced
# activations.
return activation
