diff --git a/tensorflow/lite/tools/optimize/BUILD b/tensorflow/lite/tools/optimize/BUILD index c7c50437d8f..bf7e1baafd9 100644 --- a/tensorflow/lite/tools/optimize/BUILD +++ b/tensorflow/lite/tools/optimize/BUILD @@ -241,6 +241,7 @@ tf_cc_test( "//tensorflow/lite/tools/optimize:testdata/minimum.bin", "//tensorflow/lite/tools/optimize:testdata/mixed.bin", "//tensorflow/lite/tools/optimize:testdata/multi_input_add_reshape.bin", + "//tensorflow/lite/tools/optimize:testdata/pack.bin", "//tensorflow/lite/tools/optimize:testdata/single_avg_pool_min_minus_5_max_plus_5.bin", "//tensorflow/lite/tools/optimize:testdata/single_conv_no_bias.bin", "//tensorflow/lite/tools/optimize:testdata/single_conv_weights_min_0_max_plus_10.bin", diff --git a/tensorflow/lite/tools/optimize/quantize_model_test.cc b/tensorflow/lite/tools/optimize/quantize_model_test.cc index 344a6054bd6..da1b293c84b 100644 --- a/tensorflow/lite/tools/optimize/quantize_model_test.cc +++ b/tensorflow/lite/tools/optimize/quantize_model_test.cc @@ -1260,6 +1260,68 @@ TEST_F(QuantizeCustomOpTest, VerifyMixedQuantization) { } } +class QuantizePackTest : public QuantizeModelTest { + protected: + QuantizePackTest() { + input_model_ = ReadModel(internal::kModelPack); + readonly_model_ = input_model_->GetModel(); + readonly_model_->UnPackTo(&model_); + } +}; + +TEST_F(QuantizePackTest, VerifyPack) { + auto status = QuantizeModel(&builder_, &model_, &error_reporter_); + + ASSERT_EQ(kTfLiteOk, status); + + const auto subgraph = model_.subgraphs[0].get(); + + // The model should only have 3 inputs and 1 output. + EXPECT_EQ(subgraph->inputs.size(), 3); + EXPECT_EQ(subgraph->outputs.size(), 1); + + const auto& op1 = subgraph->operators[1].get(); + const auto& op2 = subgraph->operators[2].get(); + const auto& op3 = subgraph->operators[3].get(); + const auto& op4 = subgraph->operators[4].get(); + + ASSERT_EQ(model_.operator_codes[op1->opcode_index].get()->builtin_code, + BuiltinOperator_QUANTIZE); + ASSERT_EQ(model_.operator_codes[op2->opcode_index].get()->builtin_code, + BuiltinOperator_QUANTIZE); + ASSERT_EQ(model_.operator_codes[op3->opcode_index].get()->builtin_code, + BuiltinOperator_PACK); + ASSERT_EQ(model_.operator_codes[op4->opcode_index].get()->builtin_code, + BuiltinOperator_DEQUANTIZE); + + const auto& pack_input0 = subgraph->tensors[op3->inputs[0]].get(); + const auto& pack_input1 = subgraph->tensors[op3->inputs[1]].get(); + const auto& pack_input2 = subgraph->tensors[op3->inputs[2]].get(); + + const auto& pack_output = subgraph->tensors[op3->outputs[0]].get(); + + // Check quantization parameters for input and output. + EXPECT_FLOAT_EQ(pack_input0->quantization->scale[0], + pack_input1->quantization->scale[0]); + EXPECT_FLOAT_EQ(pack_input1->quantization->scale[0], + pack_input2->quantization->scale[0]); + EXPECT_FLOAT_EQ(pack_input0->quantization->zero_point[0], + pack_input1->quantization->zero_point[0]); + EXPECT_FLOAT_EQ(pack_input1->quantization->zero_point[0], + pack_input2->quantization->zero_point[0]); + + EXPECT_FLOAT_EQ(pack_input1->quantization->scale[0], + pack_output->quantization->scale[0]); + EXPECT_FLOAT_EQ(pack_input1->quantization->zero_point[0], + pack_output->quantization->zero_point[0]); + + // Check type of input and output. + EXPECT_EQ(pack_output->type, TensorType_INT8); + EXPECT_EQ(pack_input0->type, TensorType_INT8); + EXPECT_EQ(pack_input1->type, TensorType_INT8); + EXPECT_EQ(pack_input2->type, TensorType_INT8); +} + class QuantizeMinimumMaximumTest : public QuantizeModelTest, public testing::WithParamInterface { diff --git a/tensorflow/lite/tools/optimize/test_util.cc b/tensorflow/lite/tools/optimize/test_util.cc index 0d7cfd6622a..7d5e9d65f06 100644 --- a/tensorflow/lite/tools/optimize/test_util.cc +++ b/tensorflow/lite/tools/optimize/test_util.cc @@ -51,6 +51,8 @@ const char* kModelMixed = "mixed.bin"; const char* kModelSplit = "split.bin"; +const char* kModelPack = "pack.bin"; + const char* kLstmCalibrated = "lstm_calibrated.bin"; const char* kLstmQuantized = "lstm_quantized.bin"; diff --git a/tensorflow/lite/tools/optimize/test_util.h b/tensorflow/lite/tools/optimize/test_util.h index 525fbd0c573..abcdbc21d36 100644 --- a/tensorflow/lite/tools/optimize/test_util.h +++ b/tensorflow/lite/tools/optimize/test_util.h @@ -79,6 +79,9 @@ extern const char* kModelMixed; // Test model with split op. extern const char* kModelSplit; +// Test model with pack op. +extern const char* kModelPack; + // Test model with LSTM op that has layer norm, has projection, without // peephole, without cifg. extern const char* kLstmCalibrated; diff --git a/tensorflow/lite/tools/optimize/testdata/pack.bin b/tensorflow/lite/tools/optimize/testdata/pack.bin new file mode 100644 index 00000000000..c367eea06a5 Binary files /dev/null and b/tensorflow/lite/tools/optimize/testdata/pack.bin differ