No public description

PiperOrigin-RevId: 770873511
Change-Id: I1219eeb60f5b4e1dee3c5f4ea82bd8a3c8d92a01

Author: lishanok

qkeras/quantizers.py

@@ -2260,6 +2260,11 @@ def get_config(self):
     return config
 
 
+@tf.function(jit_compile=True)
+def fast_relu_quantize(p, m_i, factor):
+  return m_i * tf.clip_by_value(tf.round(p) * factor, 0.0, 1.0 - factor)
+
+
 @quantizer_registry.register_quantizer
 class quantized_relu(base_quantizer.BaseQuantizer):  # pylint: disable=invalid-name
   """Computes a quantized relu to a number of bits.
@@ -2359,17 +2364,13 @@ def __str__(self):
       flags.append(str(int(self.use_stochastic_rounding)))
     return "quantized_relu(" + ",".join(flags) + ")"
 
-  @tf.function(jit_compile=True)
-  def fast_quantize(p, m_i, factor):
-    return m_i * tf.clip_by_value(tf.round(p) * factor, 0.0, 1.0 - factor)
-
   def __call__(self, x):
     if self.enable_fast_inference:
       # This is the fast inference version of the quantizer.
       m_i = 1 << self.integer
       p = x * (2 ** (self.bits - self.integer))
       factor = 2 ** -self.bits
-      return self.fast_quantize(p, m_i, factor)
+      return fast_relu_quantize(p, m_i, factor)
 
     if not self.built:
       self.build(var_name=self.var_name, use_variables=self.use_variables)

tests/qactivation_test.py

@@ -813,5 +813,12 @@ def test_quantized_hswish(bits, integer, symmetric, relu_shift,
   assert_allclose(result, expected_values, rtol=1e-05)
 
 
+def test_quantized_relu_fast_inference():
+  q1 = quantized_relu(10, 2, enable_fast_inference=False)
+  q2 = quantized_relu(10, 2, enable_fast_inference=True)
+  x = np.array([-2.1, 0.73, 2.36, 4.98])
+  np.testing.assert_array_equal(q1(x).numpy(), q2(x).numpy())
+
+
 if __name__ == '__main__':
   pytest.main([__file__])