Porting TF fake_quant_with_min_max functions #20641

doncarlos999 · 2024-12-13T09:57:27Z

Based on the discussion here: #20319 I started porting the fake_quant_with_min_max functions from tensorflow to keras3.
This PR contains those ported functions and the relevant tests from https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/tests/fake_quant_ops_test.py.

I didn't implement tf.quantization.fake_quant_with_min_max_vars as it looks the same as tf.quantization.fake_quant_with_min_max_args. But, I can add this one too if required.

~~For the CLA I am waiting on our CTO to add me to the Edge Impulse <-> Google CLA. But I figured that I can work on revisions to the PR in the meantime.~~

CC: @matpalm, @dansitu, @james77777778

* adds fake_quant_with_min_max functions from TF to keras3

google-cla · 2024-12-13T09:57:31Z

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

codecov-commenter · 2024-12-13T10:06:12Z

Codecov Report

Attention: Patch coverage is 89.80892% with 16 lines in your changes missing coverage. Please review.

Project coverage is 72.50%. Comparing base (84b531c) to head (5c48be2).

Files with missing lines	Patch %	Lines
keras/src/quantizers/quantizers.py	91.72%	7 Missing and 5 partials ⚠️
keras/api/_tf_keras/keras/quantizers/__init__.py	0.00%	4 Missing ⚠️

❗ There is a different number of reports uploaded between BASE (84b531c) and HEAD (5c48be2). Click for more details.

HEAD has 4 uploads less than BASE

Flag BASE (84b531c) HEAD (5c48be2)

keras 5 3

keras-numpy 1 0

keras-jax 1 0

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #20641      +/-   ##
==========================================
- Coverage   81.95%   72.50%   -9.46%     
==========================================
  Files         543      543              
  Lines       50663    50820     +157     
  Branches     7828     7842      +14     
==========================================
- Hits        41523    36849    -4674     
- Misses       7246    12204    +4958     
+ Partials     1894     1767     -127

Flag	Coverage Δ
keras	`72.37% <89.17%> (-9.42%)`	⬇️
keras-jax	`?`
keras-numpy	`?`
keras-openvino	`29.89% <12.10%> (-0.06%)`	⬇️
keras-tensorflow	`64.73% <89.17%> (+0.07%)`	⬆️
keras-torch	`63.80% <88.53%> (+0.07%)`	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

fchollet

Thanks for the PR!

keras/src/quantizers/quantizers.py

james77777778

Hi @doncarlos999
I have left some comments.

Additionally, I think we still need fake_quant_with_min_max_vars, as it is used in TFMOT:
https://github.com/tensorflow/model-optimization/blob/master/tensorflow_model_optimization/python/core/quantization/keras/quant_ops.py#L340

james77777778 · 2024-12-19T05:20:44Z

keras/api/_tf_keras/keras/quantizers/__init__.py

@@ -12,4 +12,14 @@
 from keras.src.quantizers.quantizers import abs_max_quantize
 from keras.src.quantizers.quantizers import compute_float8_amax_history
 from keras.src.quantizers.quantizers import compute_float8_scale
+from keras.src.quantizers.quantizers import fake_quant_with_min_max_args
+from keras.src.quantizers.quantizers import (
+    fake_quant_with_min_max_args_gradient,


For QAT purposes, I don't think we need *_gradient ops.

james77777778 · 2024-12-19T05:21:04Z

keras/api/_tf_keras/keras/quantizers/__init__.py

+    fake_quant_with_min_max_vars_per_channel,
+)
+from keras.src.quantizers.quantizers import (
+    fake_quant_with_min_max_vars_per_channel_gradient,


For QAT purposes, I don't think we need *_gradient ops.

james77777778 · 2024-12-19T05:21:22Z

keras/api/quantizers/__init__.py

@@ -12,4 +12,14 @@
 from keras.src.quantizers.quantizers import abs_max_quantize
 from keras.src.quantizers.quantizers import compute_float8_amax_history
 from keras.src.quantizers.quantizers import compute_float8_scale
+from keras.src.quantizers.quantizers import fake_quant_with_min_max_args
+from keras.src.quantizers.quantizers import (
+    fake_quant_with_min_max_args_gradient,


james77777778 · 2024-12-19T05:21:29Z

keras/api/quantizers/__init__.py

+    fake_quant_with_min_max_vars_per_channel,
+)
+from keras.src.quantizers.quantizers import (
+    fake_quant_with_min_max_vars_per_channel_gradient,


james77777778 · 2024-12-19T05:21:56Z

keras/src/quantizers/__init__.py

@@ -6,6 +6,16 @@
 from keras.src.quantizers.quantizers import abs_max_quantize
 from keras.src.quantizers.quantizers import compute_float8_amax_history
 from keras.src.quantizers.quantizers import compute_float8_scale
+from keras.src.quantizers.quantizers import fake_quant_with_min_max_args
+from keras.src.quantizers.quantizers import (
+    fake_quant_with_min_max_args_gradient,


james77777778 · 2024-12-19T05:46:45Z

keras/src/quantizers/quantizers.py

+@keras_export(
+    "keras.quantizers.fake_quant_with_min_max_vars_per_channel_gradient"
+)
+def fake_quant_with_min_max_vars_per_channel_gradient(


For QAT purposes, I don't think we need *_gradient ops.

james77777778 · 2024-12-19T05:50:03Z

keras/src/quantizers/quantizers_test.py

@@ -100,3 +100,759 @@ def test_quantize_and_dequantize(self):
        )
        # A loose assertion due to an expected quantization error
        self.assertAllClose(qdq_values, values, atol=5e-1)
+
+    def _TestOp(


We can use @parameterized.named_parameters and named_product to organize similar tests like this one:
https://github.com/keras-team/keras/blob/master/keras/src/ops/nn_test.py#L2355-L2365

james77777778 · 2024-12-19T05:52:43Z

keras/src/quantizers/quantizers_test.py

+            num_bits=num_bits,
+            narrow_range=narrow_range,
+        )
+        self.assertAllClose(outputs, expected)


I think verifying the output values alone is not sufficient.
We can add an assertion for the gradient, similar to this one:
https://github.com/keras-team/keras/blob/master/keras/src/layers/core/dense_test.py#L584

james77777778 · 2024-12-19T05:52:55Z

keras/src/quantizers/quantizers_test.py

+        )
+        self.assertAllClose(outputs, expected)
+
+    def _TestGradOp(


For QAT purposes, I don't think we need *_gradient ops.

james77777778 · 2024-12-19T05:53:03Z

keras/src/quantizers/quantizers_test.py

+        )
+        self.assertAllClose(outputs, expected)
+
+    def _TestChannelsGradOp(


For QAT purposes, I don't think we need *_gradient ops.

doncarlos999 · 2024-12-19T08:28:37Z

@james77777778 thank you for the review. I'm working on revisions now.

Regarding the *_gradient functions, I added those as a way to test that the gradients that come from the main functions were being calculated correctly. Should we keep them just for testing purposes but not expose them in the public facing API? If not I will remove them.

james77777778 · 2024-12-19T08:38:36Z

Regarding the *_gradient functions, I added those as a way to test that the gradients that come from the main functions were being calculated correctly. Should we keep them just for testing purposes but not expose them in the public facing API? If not I will remove them.

We can test the gradients of fake_* functions using:

tensorflow: tf.GradientTape() + tape.gradient
torch: loss.backward() + variable.grad
jax: jax.grad

You can refer to this test for an example:
https://github.com/keras-team/keras/blob/master/keras/src/layers/core/dense_test.py#L584-L649

Using a different function, separate from the user-facing function, for testing purposes seems redundant and fragile to me. However, we should wait for calls from @fchollet

doncarlos999 · 2024-12-19T08:42:58Z

I agree that having two separate functions is fragile I simply kept the functions separate as that was how they were tested in the Tensorflow repo.
I will start adding tests based on your example in the meantime. Thank you.

…nt_with_min_max_vars function

… and torch

QAT (squashed this time) (#1)

fef281e

* adds fake_quant_with_min_max functions from TF to keras3

google-ml-butler bot added the size:XL label Dec 13, 2024

google-ml-butler bot assigned gbaned Dec 13, 2024

fchollet reviewed Dec 14, 2024

View reviewed changes

keras/src/quantizers/quantizers.py Outdated Show resolved Hide resolved

keras/src/quantizers/quantizers.py Outdated Show resolved Hide resolved

keras/src/quantizers/quantizers.py Outdated Show resolved Hide resolved

doncarlos999 and others added 3 commits December 16, 2024 11:19

Merge branch 'keras-team:master' into master

9043966

Addresses PR review comments

b66e9c3

drops another type hint

84338f2

doncarlos999 requested a review from fchollet December 18, 2024 08:34

google-ml-butler bot added the awaiting review label Dec 18, 2024

james77777778 reviewed Dec 19, 2024

View reviewed changes

doncarlos999 and others added 3 commits December 19, 2024 09:39

swaps out if statements, change float() to ops.cast and adds fake_qua…

170f31e

…nt_with_min_max_vars function

fix missed if statement, adds gradient tests via main function for tf…

c193d16

… and torch

Merge branch 'keras-team:master' into master

5c48be2

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Porting TF fake_quant_with_min_max functions #20641

Porting TF fake_quant_with_min_max functions #20641

doncarlos999 commented Dec 13, 2024 •

edited

Loading

google-cla bot commented Dec 13, 2024

codecov-commenter commented Dec 13, 2024 •

edited

Loading

fchollet left a comment

james77777778 left a comment

james77777778 Dec 19, 2024

james77777778 Dec 19, 2024

james77777778 Dec 19, 2024

james77777778 Dec 19, 2024

james77777778 Dec 19, 2024

james77777778 Dec 19, 2024

james77777778 Dec 19, 2024

james77777778 Dec 19, 2024

james77777778 Dec 19, 2024

james77777778 Dec 19, 2024

doncarlos999 commented Dec 19, 2024

james77777778 commented Dec 19, 2024

doncarlos999 commented Dec 19, 2024

Porting TF fake_quant_with_min_max functions #20641

Are you sure you want to change the base?

Porting TF fake_quant_with_min_max functions #20641

Conversation

doncarlos999 commented Dec 13, 2024 • edited Loading

google-cla bot commented Dec 13, 2024

codecov-commenter commented Dec 13, 2024 • edited Loading

Codecov Report

fchollet left a comment

Choose a reason for hiding this comment

james77777778 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

doncarlos999 commented Dec 19, 2024

james77777778 commented Dec 19, 2024

doncarlos999 commented Dec 19, 2024

doncarlos999 commented Dec 13, 2024 •

edited

Loading

codecov-commenter commented Dec 13, 2024 •

edited

Loading