Skip to content

Commit ac1fae3

Browse files
rkuesterrkuester
authored
refactor(compression): replace test_models with model_editor in compress_test (#3586)
Replace dictionary-based test_models.build() with model_editor's declarative API for building test models. BUG=part of #3256
1 parent 730449f commit ac1fae3

2 files changed

Lines changed: 66 additions & 153 deletions

File tree

‎tensorflow/lite/micro/compression/BUILD‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -162,7 +162,6 @@ py_test(
162162
":metadata_py",
163163
":model_editor",
164164
":spec",
165-
":test_models",
166165
"//tensorflow/lite/python:schema_py",
167166
requirement("bitarray"),
168167
requirement("numpy"),

‎tensorflow/lite/micro/compression/compress_test.py‎

Lines changed: 66 additions & 152 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,6 @@
2121
from tflite_micro.tensorflow.lite.micro.compression import metadata_py_generated as schema
2222
from tflite_micro.tensorflow.lite.micro.compression import model_editor
2323
from tflite_micro.tensorflow.lite.micro.compression import spec
24-
from tflite_micro.tensorflow.lite.micro.compression import test_models
2524
from tflite_micro.tensorflow.lite.python import schema_py_generated as tflite
2625

2726

@@ -170,153 +169,70 @@ def test_multiple_tables_with_padding(self):
170169
self.assertEqual(result, expected_output)
171170

172171

173-
# yapf: disable
174-
TEST_MODEL = {
175-
"operator_codes": {
176-
0: {
177-
"builtin_code": tflite.BuiltinOperator.ADD,
178-
},
179-
},
180-
"metadata": {
181-
0: {
182-
"name": "metadata0",
183-
"buffer": 0
184-
},
185-
},
186-
"subgraphs": {
187-
0: {
188-
"operators": {
189-
0: {
190-
"opcode_index": 0,
191-
"inputs": (
192-
0,
193-
1,
194-
),
195-
"outputs": (2, ),
196-
},
197-
},
198-
"tensors": {
199-
0: {
200-
"shape": (16, 1),
201-
"type": tflite.TensorType.UINT8,
202-
"buffer": 1,
203-
"quantization": {
204-
"quantized_dimension": 1,
205-
"scale": (1,),
206-
"zero_point": (0,),
207-
},
208-
},
209-
1: {
210-
"shape": (16, 1),
211-
"type": tflite.TensorType.INT8,
212-
"buffer": 2,
213-
"quantization": {
214-
"quantized_dimension": 1,
215-
"scale": (1,),
216-
"zero_point": (0,),
217-
},
218-
},
219-
2: {
220-
"shape": (16, 1),
221-
"type": tflite.TensorType.INT16,
222-
"buffer": 3,
223-
"quantization": {
224-
"quantized_dimension": 1,
225-
"scale": (1,),
226-
"zero_point": (0,),
227-
},
228-
},
229-
3: {
230-
"shape": (16, 1),
231-
"type": tflite.TensorType.INT32,
232-
"buffer": 4,
233-
"quantization": {
234-
"quantized_dimension": 1,
235-
"scale": (1,),
236-
"zero_point": (0,),
237-
},
238-
},
239-
4: {
240-
"shape": (16, 1),
241-
"type": tflite.TensorType.INT32,
242-
"buffer": 5,
243-
"quantization": {
244-
"quantized_dimension": 1,
245-
"scale": (1,),
246-
"zero_point": (0,),
247-
},
248-
},
249-
5: {
250-
"shape": (4, 5),
251-
"type": tflite.TensorType.INT16,
252-
"buffer": 6,
253-
"quantization": {
254-
"quantized_dimension": 1,
255-
"scale": (1, 1, 1, 1, 1),
256-
"zero_point": (0, 0, 0, 0, 0),
257-
},
258-
},
259-
6: {
260-
"shape": (5, 4),
261-
"type": tflite.TensorType.INT16,
262-
"buffer": 7,
263-
"quantization": {
264-
"quantized_dimension": 0,
265-
"scale": (1, 1, 1, 1, 1),
266-
"zero_point": (0, 0, 0, 0, 0),
267-
},
268-
},
269-
7: {
270-
"shape": (5, 4),
271-
"type": tflite.TensorType.INT16,
272-
"buffer": 8,
273-
"quantization": {
274-
"quantized_dimension": 0,
275-
"scale": (1,),
276-
"zero_point": (0,),
277-
},
278-
},
279-
8: {
280-
"shape": (16, 1),
281-
"type": tflite.TensorType.UINT8,
282-
"buffer": 9,
283-
},
284-
},
285-
},
286-
},
287-
"buffers": {
288-
0: None,
289-
290-
1: np.array(range(16), dtype=np.dtype("<u1")),
291-
292-
2: np.array(range(-16, 0), dtype=np.dtype("<i1")),
293-
294-
3: np.array(range(-1616, -1600), dtype=np.dtype("<i2")),
295-
296-
4: np.array(range(-160_016, -160_000), dtype=np.dtype("<i4")),
297-
298-
5: np.array(range(16), dtype=np.dtype("<i4")),
299-
300-
6: np.array(((1, 5, 9, 13, 17),
301-
(2, 6, 10, 14, 18),
302-
(3, 7, 11, 15, 19),
303-
(4, 8, 12, 16, 20)), dtype=np.dtype("<i2")),
304-
305-
7: np.array(((1, 2, 3, 4),
306-
(5, 6, 7, 8),
307-
(9, 10, 11, 12),
308-
(13, 14, 15, 16),
309-
(17, 18, 19, 20)), dtype=np.dtype("<i2")),
310-
311-
8: np.array(((1, 2, 3, 4),
312-
(1, 2, 3, 4),
313-
(1, 2, 3, 4),
314-
(1, 2, 3, 4),
315-
(1, 2, 3, 4)), dtype=np.dtype("<i2")),
316-
317-
9: np.array(range(16), dtype=np.dtype("<u1")),
318-
},
319-
}
172+
def _build_test_model():
173+
"""Build test model using model_editor API."""
174+
from tflite_micro.tensorflow.lite.micro.compression.model_editor import (
175+
Model, Subgraph, Tensor, Operator, Quantization)
176+
177+
# Pre-declare tensors with stable indices for compression specs
178+
t0 = Tensor(shape=(16, 1),
179+
dtype=tflite.TensorType.UINT8,
180+
data=np.array(range(16), dtype="<u1"),
181+
quantization=Quantization(scales=1, zero_points=0))
182+
t1 = Tensor(shape=(16, 1),
183+
dtype=tflite.TensorType.INT8,
184+
data=np.array(range(-16, 0), dtype="<i1"),
185+
quantization=Quantization(scales=1, zero_points=0))
186+
t2 = Tensor(shape=(16, 1),
187+
dtype=tflite.TensorType.INT16,
188+
data=np.array(range(-1616, -1600), dtype="<i2"),
189+
quantization=Quantization(scales=1, zero_points=0))
190+
t3 = Tensor(shape=(16, 1),
191+
dtype=tflite.TensorType.INT32,
192+
data=np.array(range(-160_016, -160_000), dtype="<i4"),
193+
quantization=Quantization(scales=1, zero_points=0))
194+
t4 = Tensor(shape=(16, 1),
195+
dtype=tflite.TensorType.INT32,
196+
data=np.array(range(16), dtype="<i4"),
197+
quantization=Quantization(scales=1, zero_points=0))
198+
t5 = Tensor(shape=(4, 5),
199+
dtype=tflite.TensorType.INT16,
200+
data=np.array(((1, 5, 9, 13, 17), (2, 6, 10, 14, 18),
201+
(3, 7, 11, 15, 19), (4, 8, 12, 16, 20)),
202+
dtype="<i2"),
203+
quantization=Quantization(scales=[1, 1, 1, 1, 1],
204+
zero_points=[0, 0, 0, 0, 0],
205+
axis=1))
206+
t6 = Tensor(shape=(5, 4),
207+
dtype=tflite.TensorType.INT16,
208+
data=np.array(((1, 2, 3, 4), (5, 6, 7, 8), (9, 10, 11, 12),
209+
(13, 14, 15, 16), (17, 18, 19, 20)),
210+
dtype="<i2"),
211+
quantization=Quantization(scales=[1, 1, 1, 1, 1],
212+
zero_points=[0, 0, 0, 0, 0],
213+
axis=0))
214+
t7 = Tensor(shape=(5, 4),
215+
dtype=tflite.TensorType.INT16,
216+
data=np.array(((1, 2, 3, 4), (1, 2, 3, 4), (1, 2, 3, 4),
217+
(1, 2, 3, 4), (1, 2, 3, 4)),
218+
dtype="<i2"),
219+
quantization=Quantization(scales=1, zero_points=0))
220+
t8 = Tensor(shape=(16, 1),
221+
dtype=tflite.TensorType.UINT8,
222+
data=np.array(range(16), dtype="<u1"))
223+
224+
model = Model(metadata={"metadata0": b""},
225+
subgraphs=[
226+
Subgraph(tensors=[t0, t1, t2, t3, t4, t5, t6, t7, t8],
227+
operators=[
228+
Operator(opcode=tflite.BuiltinOperator.ADD,
229+
inputs=[t0, t1],
230+
outputs=[t2])
231+
])
232+
])
233+
234+
return model.build()
235+
320236

321237
TEST_COMPRESSION_SPEC = [
322238
spec.Tensor( # spec 0
@@ -341,7 +257,6 @@ def test_multiple_tables_with_padding(self):
341257
),
342258

343259
# Tensor 4 intentionally left uncompressed
344-
345260
spec.Tensor( # spec 4
346261
subgraph=0,
347262
tensor=5,
@@ -358,7 +273,6 @@ def test_multiple_tables_with_padding(self):
358273
compression=[spec.LookUpTableCompression(index_bitwidth=2)],
359274
),
360275
]
361-
# yapf: enable
362276

363277

364278
class TestsCompression(unittest.TestCase):
@@ -367,7 +281,7 @@ class TestsCompression(unittest.TestCase):
367281
@classmethod
368282
def setUpClass(cls):
369283
super().setUpClass()
370-
cls.flatbuffer = test_models.build(TEST_MODEL)
284+
cls.flatbuffer = _build_test_model()
371285
cls.uncompressed = model_editor.read(cls.flatbuffer)
372286

373287
def test_compression_metadata(self):
@@ -460,7 +374,7 @@ class TestCompressedModel(unittest.TestCase):
460374
def setUpClass(cls):
461375
super().setUpClass()
462376
# Create a model
463-
uncompressed_fb = test_models.build(TEST_MODEL)
377+
uncompressed_fb = _build_test_model()
464378
cls.uncompressed = model_editor.read(uncompressed_fb)
465379

466380
# Compress the model

0 commit comments

Comments
 (0)