Skip to content

Commit eb38fe1

Browse files
xhaetensorflower-gardener
xhae
authored and
tensorflower-gardener
committed
Basic implementation for ternarization
PiperOrigin-RevId: 426282499
1 parent 604f679 commit eb38fe1

14 files changed

+2935
-6
lines changed

‎tensorflow_model_optimization/python/core/quantization/keras/experimental/ternarization/BUILD‎

Lines changed: 170 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,170 @@
1+
load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "py_strict_library", "py_strict_test")
2+
3+
package(default_visibility = [
4+
"//tensorflow_model_optimization:__subpackages__",
5+
])
6+
7+
licenses(["notice"])
8+
9+
py_strict_library(
10+
name = "ternarization",
11+
srcs = [
12+
"__init__.py",
13+
],
14+
srcs_version = "PY3",
15+
deps = [],
16+
)
17+
18+
py_strict_library(
19+
name = "ternarization_quantizers",
20+
srcs = [
21+
"ternarization_quantizers.py",
22+
],
23+
srcs_version = "PY3",
24+
deps = [
25+
# tensorflow dep1,
26+
"//tensorflow_model_optimization/python/core/quantization/keras:quantizers",
27+
],
28+
)
29+
30+
# TODO(b/185543626): Fix keras dependency
31+
py_test(
32+
name = "ternarization_quantizers_test",
33+
srcs = [
34+
"ternarization_quantizers_test.py",
35+
],
36+
python_version = "PY3",
37+
deps = [
38+
":ternarization_quantizers",
39+
# absl/testing:parameterized dep1,
40+
# tensorflow dep1,
41+
],
42+
)
43+
44+
py_strict_library(
45+
name = "ternarization_quantize_configs",
46+
srcs = [
47+
"ternarization_quantize_configs.py",
48+
],
49+
srcs_version = "PY3",
50+
deps = [
51+
":ternarization_quantizers",
52+
# tensorflow:tensorflow_no_contrib dep1,
53+
"//tensorflow_model_optimization/python/core/quantization/keras:quantize_config",
54+
],
55+
)
56+
57+
py_strict_library(
58+
name = "ternarization_quantize_registry",
59+
srcs = [
60+
"ternarization_quantize_registry.py",
61+
],
62+
srcs_version = "PY3",
63+
deps = [
64+
":ternarization_quantize_configs",
65+
":ternarization_quantizers",
66+
# tensorflow dep1,
67+
"//tensorflow_model_optimization/python/core/quantization/keras:quantize_config",
68+
"//tensorflow_model_optimization/python/core/quantization/keras:quantize_registry",
69+
"//tensorflow_model_optimization/python/core/quantization/keras:quantizers",
70+
],
71+
)
72+
73+
# TODO(b/185543626): Fix keras dependency
74+
py_test(
75+
name = "ternarization_quantize_registry_test",
76+
srcs = [
77+
"ternarization_quantize_registry_test.py",
78+
],
79+
python_version = "PY3",
80+
deps = [
81+
":ternarization_quantize_registry",
82+
# absl/testing:parameterized dep1,
83+
# numpy dep1,
84+
# tensorflow dep1,
85+
"//tensorflow_model_optimization/python/core/quantization/keras:quantizers",
86+
],
87+
)
88+
89+
# TODO(b/185543626): Fix keras dependency
90+
py_library(
91+
name = "ternarization_transforms",
92+
srcs = [
93+
"ternarization_transforms.py",
94+
],
95+
srcs_version = "PY3",
96+
deps = [
97+
":ternarization_quantize_configs",
98+
":ternarization_quantize_registry",
99+
# numpy dep1,
100+
# tensorflow dep1,
101+
"//tensorflow_model_optimization/python/core/quantization/keras:quantize_aware_activation",
102+
"//tensorflow_model_optimization/python/core/quantization/keras:quantize_layer",
103+
"//tensorflow_model_optimization/python/core/quantization/keras:quantizers",
104+
"//tensorflow_model_optimization/python/core/quantization/keras/graph_transformations:transforms",
105+
],
106+
)
107+
108+
py_strict_test(
109+
name = "ternarization_transforms_test",
110+
size = "large",
111+
srcs = [
112+
"ternarization_transforms_test.py",
113+
],
114+
python_version = "PY3",
115+
deps = [
116+
":ternarization_quantize_configs",
117+
":ternarization_transforms",
118+
# absl/testing:parameterized dep1,
119+
# numpy dep1,
120+
# tensorflow dep1,
121+
"//tensorflow_model_optimization/python/core/quantization/keras:quantize_aware_activation",
122+
"//tensorflow_model_optimization/python/core/quantization/keras:quantize_layer",
123+
"//tensorflow_model_optimization/python/core/quantization/keras:quantizers",
124+
"//tensorflow_model_optimization/python/core/quantization/keras/graph_transformations:model_transformer",
125+
"//tensorflow_model_optimization/python/core/quantization/keras/layers:conv_batchnorm_test_utils",
126+
"//tensorflow_model_optimization/python/core/quantization/keras/layers:dense_batchnorm_test_utils",
127+
],
128+
)
129+
130+
py_strict_library(
131+
name = "ternarization_quantize_layout_transform",
132+
srcs = [
133+
"ternarization_quantize_layout_transform.py",
134+
],
135+
srcs_version = "PY3",
136+
deps = [
137+
":ternarization_transforms",
138+
# tensorflow dep1,
139+
"//tensorflow_model_optimization/python/core/quantization/keras:quantize_layout_transform",
140+
"//tensorflow_model_optimization/python/core/quantization/keras/graph_transformations:model_transformer",
141+
],
142+
)
143+
144+
# TODO(b/185543626): Fix keras dependency
145+
py_test(
146+
name = "quantize_numerical_test",
147+
srcs = ["quantize_numerical_test.py"],
148+
python_version = "PY3",
149+
deps = [
150+
":ternarization_quantize_scheme",
151+
# absl/testing:parameterized dep1,
152+
# numpy dep1,
153+
# tensorflow dep1,
154+
"//tensorflow_model_optimization/python/core/quantization/keras:quantize",
155+
"//tensorflow_model_optimization/python/core/quantization/keras:utils",
156+
],
157+
)
158+
159+
py_strict_library(
160+
name = "ternarization_quantize_scheme",
161+
srcs = [
162+
"ternarization_quantize_scheme.py",
163+
],
164+
srcs_version = "PY3",
165+
deps = [
166+
":ternarization_quantize_layout_transform",
167+
":ternarization_quantize_registry",
168+
"//tensorflow_model_optimization/python/core/quantization/keras:quantize_scheme",
169+
],
170+
)

‎tensorflow_model_optimization/python/core/quantization/keras/experimental/ternarization/README.md‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
This directory contains code specific to an experimental ternary quantization
2+
scheme implemented using the QAT API.
3+
4+
This quantization scheme is experimental and subject to change.

‎tensorflow_model_optimization/python/core/quantization/keras/experimental/ternarization/__init__.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
# ==============================================================================

0 commit comments

Comments
 (0)