Skip to content

Commit f7a3369

Browse files
ut improve for objective (#1260)
Signed-off-by: Cheng, Zixuan <zixuan.cheng@intel.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
1 parent 2956db2 commit f7a3369

File tree

1 file changed

+5
-1
lines changed

1 file changed

+5
-1
lines changed

test/objective/test_objective.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -285,13 +285,14 @@ def test_performance(self):
285285
dataset = Datasets("tensorflow")["dummy"]((100, 256, 256, 1), label=True)
286286

287287
from neural_compressor.experimental import Quantization, common
288-
from neural_compressor.utils.utility import get_size
288+
from neural_compressor.model import tensorflow_model
289289

290290
quantizer = Quantization("fake_yaml.yaml")
291291
quantizer.calib_dataloader = common.DataLoader(dataset)
292292
quantizer.eval_dataloader = common.DataLoader(dataset)
293293
quantizer.model = self.constant_graph
294294
q_model = quantizer.fit()
295+
self.assertTrue(isinstance(q_model, tensorflow_model.TensorflowBaseModel))
295296

296297
from neural_compressor.experimental import Benchmark, common
297298

@@ -345,6 +346,7 @@ def eval(model):
345346

346347
from neural_compressor.conf.config import conf
347348
from neural_compressor.experimental import Quantization
349+
from neural_compressor.model import onnx_model
348350

349351
conf.model.framework = "onnxrt_integerops"
350352
conf.quantization.approach = "post_training_dynamic_quant"
@@ -357,6 +359,8 @@ def eval(model):
357359
quantize.model = model
358360
quantize.eval_func = eval
359361
q_model = quantize()
362+
self.assertTrue(isinstance(q_model, onnx_model.ONNXModel))
363+
self.assertTrue("quantize" in str(q_model.model.producer_name))
360364

361365
def test_tune_data(self):
362366
from neural_compressor.objective import MultiObjective

0 commit comments

Comments
 (0)