1 /* Copyright 2017 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 <algorithm>
16 #include <memory>
17 #include <string>
18 #include <unordered_map>
19 #include <vector>
20
21 #include "tensorflow/lite/toco/graph_transformations/graph_transformations.h"
22 #include "tensorflow/lite/toco/model.h"
23 #include "tensorflow/lite/toco/tooling_util.h"
24 #include "tensorflow/core/platform/logging.h"
25
26 namespace toco {
27
28 namespace {
29
RenameArray(Model * model,const string & oldname,const string & desired_newname)30 void RenameArray(Model* model, const string& oldname,
31 const string& desired_newname) {
32 const string& newname = AvailableArrayName(*model, desired_newname);
33 auto& arrays = model->GetMutableArrayMap();
34 arrays[newname] = std::move(arrays[oldname]);
35 arrays.erase(oldname);
36 for (const auto& op : model->operators) {
37 for (string& input : op->inputs) {
38 if (input == oldname) {
39 input = newname;
40 }
41 }
42 for (string& output : op->outputs) {
43 if (output == oldname) {
44 output = newname;
45 }
46 }
47 }
48 }
49
50 } // namespace
51
52 // Reorder the elements of an input_array according to the input_axes_order and
53 // output_axes_order. Then adjust the shapes of the input and output arrays
54 // accordingly. Note that input_array must have a buffer (that is, it is a
55 // constant array).
56 template <typename T, ArrayDataType DataType>
ReorderAxes(AxesOrder input_axes_order,AxesOrder output_axes_order,const Array & input_array,Array * output_array)57 void ReorderAxes(AxesOrder input_axes_order, AxesOrder output_axes_order,
58 const Array& input_array, Array* output_array) {
59 DCHECK(input_array.buffer->type == DataType);
60 DCHECK(!output_array->buffer);
61 const auto& input_data = input_array.GetBuffer<DataType>().data;
62 auto& output_data = output_array->GetMutableBuffer<DataType>().data;
63 output_data.resize(RequiredBufferSizeForShape(output_array->shape()));
64 // TODO(b/62904716) Shapes should be used directly.
65 Shape input_shape = input_array.shape();
66 Shape output_shape = output_array->shape();
67 if (AxesCount(input_axes_order) == 2) {
68 UnextendShape(&input_shape, 2);
69 UnextendShape(&output_shape, 2);
70 }
71 ShuffleArray(input_shape, input_axes_order, output_axes_order, output_shape,
72 input_data.data(), output_data.data());
73 if (input_array.minmax) {
74 output_array->GetOrCreateMinMax() = input_array.GetMinMax();
75 }
76 if (input_array.narrow_range) {
77 output_array->narrow_range = true;
78 }
79 }
80
Run(Model * model,std::size_t op_index,bool * modified)81 ::tensorflow::Status ResolveReorderAxes::Run(Model* model, std::size_t op_index,
82 bool* modified) {
83 *modified = false;
84 auto it = model->operators.begin() + op_index;
85 auto* op = it->get();
86 if (op->type != OperatorType::kReorderAxes) {
87 return ::tensorflow::Status::OK();
88 }
89 auto* reorder_op = static_cast<ReorderAxesOperator*>(op);
90
91 // Intentionally copies, not references.
92 const string input_array_name = reorder_op->inputs[0];
93 const string output_array_name = reorder_op->outputs[0];
94
95 auto& input_array = model->GetArray(input_array_name);
96 auto& output_array = model->GetArray(output_array_name);
97 if (!input_array.buffer) {
98 return ::tensorflow::Status::OK();
99 }
100 // Yield until output dims have been resolved.
101 if (!output_array.has_shape()) {
102 return ::tensorflow::Status::OK();
103 }
104 // Reorder the input array dims and buffer data
105 if (input_array.buffer->type == ArrayDataType::kFloat) {
106 ReorderAxes<float, ArrayDataType::kFloat>(reorder_op->input_axes_order,
107 reorder_op->output_axes_order,
108 input_array, &output_array);
109 } else if (input_array.buffer->type == ArrayDataType::kUint8) {
110 // TODO(benoitjacob): This path seems unused.
111 // ReorderAxes is only used when importing from
112 // TensorFlow GraphDef, which does not support quantized nodes.
113 ReorderAxes<uint8, ArrayDataType::kUint8>(reorder_op->input_axes_order,
114 reorder_op->output_axes_order,
115 input_array, &output_array);
116 } else {
117 LOG(FATAL) << "Cannot ReorderAxes unless input buffer is float or uint8.";
118 }
119
120 AddMessageF("Reordered axes for array %s", input_array_name);
121
122 DeleteOpAndArraysIfUnused(model, op);
123 RenameArray(model, output_array_name, input_array_name);
124
125 *modified = true;
126 return ::tensorflow::Status::OK();
127 }
128
129 } // namespace toco
130