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