xref: /aosp_15_r20/external/tensorflow/tensorflow/lite/toco/logging/conversion_log_util_test.cc (revision b6fb3261f9314811a0f4371741dbb8839866f948)
1 /* Copyright 2019 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 ==============================================================================*/
15 #include "tensorflow/lite/toco/logging/conversion_log_util.h"
16 
17 #include <memory>
18 #include <string>
19 #include <utility>
20 #include <vector>
21 
22 #include <gmock/gmock.h>
23 #include <gtest/gtest.h>
24 #include "absl/memory/memory.h"
25 #include "tensorflow/core/framework/node_def.pb.h"
26 #include "tensorflow/lite/toco/model.h"
27 #include "tensorflow/lite/toco/model_flags.pb.h"
28 
29 namespace toco {
30 namespace {
31 
32 using ::testing::ElementsAre;
33 using ::testing::UnorderedElementsAre;
34 
TEST(ConversionLogUtilTest,TestGetOperatorNames)35 TEST(ConversionLogUtilTest, TestGetOperatorNames) {
36   Model model;
37   // Built-in ops.
38   model.operators.push_back(std::make_unique<ConvOperator>());
39   model.operators.push_back(std::make_unique<MeanOperator>());
40   model.operators.push_back(std::make_unique<NegOperator>());
41   // Flex ops.
42   auto avg_pool_3d = std::make_unique<TensorFlowUnsupportedOperator>();
43   avg_pool_3d->tensorflow_op = "AvgPool3D";
44   tensorflow::NodeDef node_def;
45   node_def.set_op("AvgPool3D");
46   node_def.SerializeToString(&avg_pool_3d->tensorflow_node_def);
47   model.operators.push_back(std::move(avg_pool_3d));
48   // Custom ops.
49   auto my_custom_op = std::make_unique<TensorFlowUnsupportedOperator>();
50   my_custom_op->tensorflow_op = "MyAwesomeCustomOp";
51   model.operators.push_back(std::move(my_custom_op));
52 
53   const auto& output = GetOperatorNames(model);
54   EXPECT_THAT(output, ElementsAre("Conv", "Mean", "Neg", "AvgPool3D",
55                                   "MyAwesomeCustomOp"));
56 }
57 
TEST(ConversionLogUtilTest,TestCountOperatorsByType)58 TEST(ConversionLogUtilTest, TestCountOperatorsByType) {
59   Model model;
60   // 1st Conv operator.
61   std::unique_ptr<ConvOperator> conv1(new ConvOperator());
62   const std::string conv1_input_name = "conv_input1";
63   const std::string conv1_filter_name = "conv_filter1";
64   const std::string conv1_output_name = "conv_output1";
65   conv1->inputs.push_back(conv1_input_name);
66   conv1->inputs.push_back(conv1_filter_name);
67   conv1->outputs.push_back(conv1_output_name);
68   auto& array_map = model.GetMutableArrayMap();
69   array_map[conv1_input_name] = std::make_unique<Array>();
70   array_map[conv1_filter_name] = std::make_unique<Array>();
71   array_map[conv1_output_name] = std::make_unique<Array>();
72 
73   // 2nd Conv operator.
74   std::unique_ptr<ConvOperator> conv2(new ConvOperator());
75   const std::string conv2_input_name = "conv_input2";
76   const std::string conv2_filter_name = "conv_filter2";
77   const std::string conv2_output_name = "conv_output2";
78   conv2->inputs.push_back(conv2_input_name);
79   conv2->inputs.push_back(conv2_filter_name);
80   conv2->outputs.push_back(conv2_output_name);
81   array_map[conv2_input_name] = std::make_unique<Array>();
82   array_map[conv2_filter_name] = std::make_unique<Array>();
83   array_map[conv2_output_name] = std::make_unique<Array>();
84 
85   // Mean operator.
86   std::unique_ptr<MeanOperator> mean(new MeanOperator());
87   const std::string mean_input_name = "mean_input";
88   mean->inputs.push_back(mean_input_name);
89   array_map[mean_input_name] = std::make_unique<Array>();
90 
91   // 1st flex operator 'AvgPool3D'.
92   auto avg_pool_3d = std::make_unique<TensorFlowUnsupportedOperator>();
93   avg_pool_3d->tensorflow_op = "AvgPool3D";
94   tensorflow::NodeDef node_def;
95   node_def.set_op("AvgPool3D");
96   node_def.SerializeToString(&avg_pool_3d->tensorflow_node_def);
97 
98   // 2nd flex operator 'EluGrad'.
99   auto elu_grad = std::make_unique<TensorFlowUnsupportedOperator>();
100   elu_grad->tensorflow_op = "EluGrad";
101   node_def.set_op("EluGrad");
102   node_def.SerializeToString(&elu_grad->tensorflow_node_def);
103 
104   // 1st custom operator 'MyAwesomeCustomOp'.
105   auto my_custom_op = std::make_unique<TensorFlowUnsupportedOperator>();
106   my_custom_op->tensorflow_op = "MyAwesomeCustomOp";
107 
108   model.operators.push_back(std::move(conv1));
109   model.operators.push_back(std::move(conv2));
110   model.operators.push_back(std::move(mean));
111   model.operators.push_back(std::move(avg_pool_3d));
112   model.operators.push_back(std::move(elu_grad));
113   model.operators.push_back(std::move(my_custom_op));
114 
115   std::map<std::string, int> built_in_ops, select_ops, custom_ops;
116   CountOperatorsByType(model, &built_in_ops, &custom_ops, &select_ops);
117 
118   EXPECT_THAT(built_in_ops,
119               UnorderedElementsAre(std::pair<std::string, int>("Conv", 2),
120                                    std::pair<std::string, int>("Mean", 1)));
121   EXPECT_THAT(select_ops,
122               UnorderedElementsAre(std::pair<std::string, int>("AvgPool3D", 1),
123                                    std::pair<std::string, int>("EluGrad", 1)));
124   EXPECT_THAT(custom_ops, UnorderedElementsAre(std::pair<std::string, int>(
125                               "MyAwesomeCustomOp", 1)));
126 }
127 
TEST(ConversionLogUtilTest,TestGetInputAndOutputTypes)128 TEST(ConversionLogUtilTest, TestGetInputAndOutputTypes) {
129   Model model;
130   auto& array_map = model.GetMutableArrayMap();
131   const std::string input1 = "conv_input";
132   const std::string input2 = "conv_filter";
133   const std::string input3 = "feature";
134   const std::string output = "softmax";
135   array_map[input1] = std::make_unique<Array>();
136   array_map[input1]->data_type = ArrayDataType::kFloat;
137   array_map[input2] = std::make_unique<Array>();
138   array_map[input2]->data_type = ArrayDataType::kFloat;
139   array_map[input3] = std::make_unique<Array>();
140   array_map[input3]->data_type = ArrayDataType::kInt16;
141   array_map[output] = std::make_unique<Array>();
142   array_map[output]->data_type = ArrayDataType::kFloat;
143 
144   InputArray input_arrays[3];
145   input_arrays[0].set_name(input1);
146   input_arrays[1].set_name(input2);
147   input_arrays[2].set_name(input3);
148   *model.flags.add_input_arrays() = input_arrays[0];
149   *model.flags.add_input_arrays() = input_arrays[1];
150   *model.flags.add_input_arrays() = input_arrays[2];
151   model.flags.add_output_arrays(output);
152 
153   TFLITE_PROTO_NS::RepeatedPtrField<std::string> input_types, output_types;
154   GetInputAndOutputTypes(model, &input_types, &output_types);
155 
156   EXPECT_THAT(input_types, ElementsAre("float", "float", "int16"));
157   EXPECT_THAT(output_types, ElementsAre("float"));
158 }
159 
TEST(ConversionLogUtilTest,TestGetOpSignatures)160 TEST(ConversionLogUtilTest, TestGetOpSignatures) {
161   Model model;
162   auto& array_map = model.GetMutableArrayMap();
163 
164   std::unique_ptr<ConvOperator> conv(new ConvOperator());
165   const std::string conv_input_name = "conv_input";
166   const std::string conv_filter_name = "conv_filter";
167   const std::string conv_output_name = "conv_output";
168   conv->inputs.push_back(conv_input_name);
169   conv->inputs.push_back(conv_filter_name);
170   conv->outputs.push_back(conv_output_name);
171   array_map[conv_input_name] = std::make_unique<Array>();
172   array_map[conv_input_name]->data_type = ArrayDataType::kFloat;
173   array_map[conv_input_name]->copy_shape({4, 4, 3});
174   array_map[conv_filter_name] = std::make_unique<Array>();
175   array_map[conv_filter_name]->data_type = ArrayDataType::kFloat;
176   array_map[conv_filter_name]->copy_shape({2, 2});
177   array_map[conv_output_name] = std::make_unique<Array>();
178   array_map[conv_output_name]->data_type = ArrayDataType::kFloat;
179   array_map[conv_output_name]->copy_shape({4, 4, 2});
180 
181   const std::string mean_input_name = "mean_input";
182   const std::string mean_output_name = "mean_output";
183   std::unique_ptr<MeanOperator> mean(new MeanOperator());
184   mean->inputs.push_back(mean_input_name);
185   mean->outputs.push_back(mean_output_name);
186   array_map[mean_input_name] = std::make_unique<Array>();
187   array_map[mean_output_name] = std::make_unique<Array>();
188 
189   const std::string avg_pool_3d_output_name = "avg_pool_output";
190   auto avg_pool_3d = std::make_unique<TensorFlowUnsupportedOperator>();
191   avg_pool_3d->tensorflow_op = "AvgPool3D";
192   tensorflow::NodeDef node_def;
193   node_def.set_op("AvgPool3D");
194   node_def.SerializeToString(&avg_pool_3d->tensorflow_node_def);
195   avg_pool_3d->inputs.push_back(conv_output_name);
196   avg_pool_3d->outputs.push_back(avg_pool_3d_output_name);
197   array_map[avg_pool_3d_output_name] = std::make_unique<Array>();
198   array_map[avg_pool_3d_output_name]->data_type = ArrayDataType::kInt32;
199   array_map[avg_pool_3d_output_name]->copy_shape({2, 2});
200 
201   const std::string custom_op_output_name = "custom_op_output";
202   auto my_custom_op = std::make_unique<TensorFlowUnsupportedOperator>();
203   my_custom_op->tensorflow_op = "MyAwesomeCustomOp";
204   my_custom_op->inputs.push_back(avg_pool_3d_output_name);
205   my_custom_op->outputs.push_back(custom_op_output_name);
206   array_map[custom_op_output_name] = std::make_unique<Array>();
207   array_map[custom_op_output_name]->data_type = ArrayDataType::kFloat;
208   array_map[custom_op_output_name]->copy_shape({3});
209 
210   model.operators.push_back(std::move(conv));
211   model.operators.push_back(std::move(mean));
212   model.operators.push_back(std::move(avg_pool_3d));
213   model.operators.push_back(std::move(my_custom_op));
214 
215   TFLITE_PROTO_NS::RepeatedPtrField<std::string> op_signatures;
216   GetOpSignatures(model, &op_signatures);
217   EXPECT_THAT(op_signatures,
218               UnorderedElementsAre(
219                   "INPUT:[4,4,3]::float::[2,2]::float::OUTPUT:[4,4,2]::float::"
220                   "NAME:Conv::VERSION:1",
221                   "INPUT:None::None::OUTPUT:None::None::NAME:Mean::VERSION:1",
222                   "INPUT:[4,4,2]::float::OUTPUT:[2,2]::int32::NAME:AvgPool3D::"
223                   "VERSION:1",
224                   "INPUT:[2,2]::int32::OUTPUT:[3]::float::NAME:"
225                   "MyAwesomeCustomOp::VERSION:1"));
226 }
227 
TEST(ConversionLogUtilTest,TestSanitizeErrorMessage)228 TEST(ConversionLogUtilTest, TestSanitizeErrorMessage) {
229   const std::string error =
230       "error: failed while converting: 'main': Ops that can be supported by "
231       "the flex runtime (enabled via setting the -emit-select-tf-ops flag): "
232       "ResizeNearestNeighbor,ResizeNearestNeighbor. Ops that need custom "
233       "implementation (enabled via setting the -emit-custom-ops flag): "
234       "CombinedNonMaxSuppression.\nTraceback (most recent call last): File "
235       "/usr/local/bin/toco_from_protos, line 8, in <module>";
236   const std::string pruned_error =
237       "Ops that can be supported by "
238       "the flex runtime (enabled via setting the -emit-select-tf-ops flag): "
239       "ResizeNearestNeighbor,ResizeNearestNeighbor.Ops that need custom "
240       "implementation (enabled via setting the -emit-custom-ops flag): "
241       "CombinedNonMaxSuppression.";
242   EXPECT_EQ(SanitizeErrorMessage(error), pruned_error);
243 }
244 
TEST(ConversionLogUtilTest,TestSanitizeErrorMessageNoMatching)245 TEST(ConversionLogUtilTest, TestSanitizeErrorMessageNoMatching) {
246   const std::string error =
247       "error: failed while converting: 'main': Traceback (most recent call "
248       "last): File "
249       "/usr/local/bin/toco_from_protos, line 8, in <module>";
250   EXPECT_EQ(SanitizeErrorMessage(error), "");
251 }
252 
253 }  // namespace
254 }  // namespace toco
255