Line data Source code
1 : // SPDX-License-Identifier: Apache-2.0
2 : /**
3 : * Copyright (C) 2021 Jihoon Lee <jhoon.it.lee@samsung.com>
4 : *
5 : * @file tflite_opnode.cpp
6 : * @date 28 April 2021
7 : * @brief contains tflite opnode which has information to convert to tflite file
8 : * @see https://github.com/nnstreamer/nntrainer
9 : * @author Jihoon Lee <jhoon.it.lee@samsung.com>
10 : * @author Donghak Park <donghak.park@samsung.com>
11 : * @bug No known bugs except for NYI items
12 : */
13 :
14 : #include <tflite_opnode.h>
15 :
16 : #include <layer_context.h>
17 : #include <layer_node.h>
18 : #include <memory.h>
19 : namespace nntrainer {
20 :
21 22 : TfOpNode::TfOpNode() :
22 : inputs(),
23 : outputs(),
24 : weights(),
25 : weight_transform(nullptr),
26 22 : is_input(false),
27 22 : is_output(false),
28 22 : is_virtual(false),
29 22 : is_trainable(true),
30 22 : is_to_be_removed(false),
31 22 : need_reorder_weight(false),
32 : node_owned_variable(),
33 : /// @todo distinguish between uninitialized and ADD operator.
34 22 : op_type(tflite::BuiltinOperator_ADD),
35 : builtin_ops(),
36 22 : builtin_option_type(tflite::BuiltinOptions_NONE){};
37 :
38 22 : void TfOpNode::setLayerNode(const LayerNode &layer) {
39 22 : is_input = layer.getNumInputConnections() == 0;
40 22 : is_output = layer.getNumOutputConnections() == 0;
41 : /// @todo support more loss layers
42 22 : static const std::set<std::string> loss_type = {"mse", "cross"};
43 : /** set to graph output node if output connection of the node includes loss
44 : *layer string
45 : * @note this is workaround because it cannot be guaranteed that a loss layer
46 : *always has a loss type in its name.
47 : *
48 : * There are two ways to pass `GraphRepresentation` parameters to `serialize`
49 : *method.
50 : *
51 : * 1. with loss layer at the end of the graph
52 : * 2. without loss layer but last node has loss layer output connection
53 : *
54 : * Loss layer of the first case is removed by `LossRealizer` and the previous
55 : *layer of the loss layer is set as the output node. And, the below logic is
56 : *for the second case.
57 : **/
58 : /// assume that loss layers have single output
59 22 : if (layer.getNumOutputConnections() == 1) {
60 63 : for (auto &loss : loss_type) {
61 42 : if (layer.getOutputConnections()[0].find(loss) != std::string::npos) {
62 4 : is_output = true;
63 : }
64 : }
65 : }
66 : /// @todo support more virtual nodes
67 22 : is_virtual = layer.getType() == "multiout";
68 :
69 22 : auto &context = layer.getRunContext();
70 44 : auto create_variables = [](auto tensor_getter, unsigned size) {
71 : Variables v;
72 44 : v.reserve(size);
73 82 : for (unsigned i = 0; i < size; ++i) {
74 38 : v.push_back(tensor_getter(i));
75 : }
76 44 : return v;
77 0 : };
78 :
79 : /**
80 : * Q1) Why convert from NCHW to NHWC?
81 : * A1) the tflite uses NHWC format; nntrainer uses NCHW
82 : *
83 : * Q2) Why are only output tensors reshaped?
84 : * A2) the tflite needs only one tensor between nodes. Therefore,
85 : * basically, outputs are used for tflite tensors
86 : **/
87 22 : auto create_variables_with_NCHW_to_NHWC = [](auto tensor_getter,
88 : unsigned size) {
89 : Variables v;
90 22 : v.reserve(size);
91 44 : for (unsigned i = 0; i < size; ++i) {
92 22 : Tensor *tensor = const_cast<Tensor *>(tensor_getter(i));
93 44 : tensor->reshape(TensorDim{tensor->batch(), tensor->height(),
94 22 : tensor->width(), tensor->channel()});
95 22 : v.push_back(tensor);
96 : }
97 22 : return v;
98 0 : };
99 :
100 44 : inputs = create_variables(
101 22 : [&context](unsigned idx) { return &context.getInput(idx); },
102 22 : context.getNumInputs());
103 44 : outputs = create_variables_with_NCHW_to_NHWC(
104 22 : [&context](unsigned idx) { return &context.getOutput(idx); },
105 22 : context.getNumOutputs());
106 44 : weights = create_variables(
107 16 : [&context](unsigned idx) {
108 16 : auto &t = context.getWeight(idx);
109 16 : NNTR_THROW_IF(t.empty() || !t.isAllocated(), std::invalid_argument)
110 : << "every weight tensor must be allocated";
111 16 : return &t;
112 : },
113 22 : context.getNumWeights());
114 :
115 22 : if (context.getNumWeights() == 0) {
116 14 : is_trainable = false;
117 : }
118 22 : }
119 :
120 8 : void TfOpNode::setWeightTransformFn(TransformFn fn) { weight_transform = fn; }
121 :
122 5 : void TfOpNode::setInputTransformFn(TransformFn fn) { input_transform = fn; }
123 :
124 0 : void TfOpNode::setWeights(Variables weights_, bool weight_transpose) {
125 : unsigned int cnt = 0;
126 :
127 0 : for (auto &w : weights_) {
128 0 : if (cnt >= weights.size())
129 : break;
130 :
131 0 : const unsigned int unit = w->batch();
132 0 : const unsigned int channel = w->channel();
133 0 : const unsigned int height = w->height();
134 0 : const unsigned int width = w->width();
135 0 : auto weight_data = weights.at(cnt)->getData();
136 :
137 : auto *ptr = const_cast<float *>(weight_data);
138 0 : memcpy(&ptr[0], &w->getData()[0],
139 0 : sizeof(float) * (unit * channel * height * width));
140 0 : cnt++;
141 : }
142 :
143 0 : auto weight_transform_fn = [](std::vector<const Tensor *> &weights) {
144 : std::vector<Tensor> new_weights;
145 0 : new_weights.reserve(weights.size());
146 0 : if (weights.size() != 0) {
147 0 : new_weights.push_back(weights[0]->transpose("1:2:0"));
148 : }
149 0 : return new_weights;
150 0 : };
151 :
152 0 : auto transform_if = [this](TransformFn &fn, Variables &v) {
153 0 : if (fn) {
154 : auto result = fn(v);
155 0 : v.resize(result.size());
156 0 : node_owned_variable.insert(node_owned_variable.end(), result.begin(),
157 : result.end());
158 : std::transform(node_owned_variable.end() - result.size(),
159 0 : node_owned_variable.end(), v.begin(),
160 : [](Tensor &t) { return &t; });
161 0 : }
162 0 : };
163 :
164 0 : if (weight_transpose == true) {
165 0 : setWeightTransformFn(weight_transform_fn);
166 0 : transform_if(weight_transform, weights);
167 : }
168 0 : }
169 :
170 6 : void TfOpNode::weightReorder(unsigned int node_count) {
171 :
172 6 : if (need_reorder_weight == true) {
173 :
174 4 : auto previous_input_shape = input_nodes[0]->getInputs()[0];
175 :
176 4 : const unsigned int unit = outputs[0]->height();
177 4 : const unsigned int channel = previous_input_shape->channel();
178 4 : const unsigned int height = previous_input_shape->height();
179 4 : const unsigned int width = previous_input_shape->width();
180 :
181 4 : auto weight_data = weights[0]->getData();
182 : auto *ptr = const_cast<float *>(weight_data);
183 :
184 4 : std::vector<float> old_value_list(unit * channel * height * width);
185 4 : memcpy(&old_value_list[0], &ptr[0],
186 : sizeof(float) * (unit * channel * height * width));
187 :
188 14 : for (unsigned int h = 0; h < height; h++) {
189 62 : for (unsigned int w = 0; w < width; w++) {
190 643 : for (unsigned int c = 0; c < channel; c++) {
191 :
192 591 : unsigned int now_position = h * (width * channel) + w * channel + c;
193 591 : unsigned int next_position = c * (height * width) + h * width + w;
194 :
195 591 : memcpy(&ptr[now_position * unit],
196 591 : &old_value_list[next_position * unit], sizeof(float) * unit);
197 : }
198 : }
199 : }
200 4 : }
201 :
202 6 : auto weight_transform_fn = [](std::vector<const Tensor *> &weights) {
203 : std::vector<Tensor> new_weights;
204 6 : new_weights.reserve(weights.size());
205 18 : new_weights.push_back(weights[0]->transpose("0:2:1"));
206 6 : new_weights.push_back(*weights[1]);
207 6 : return new_weights;
208 0 : };
209 :
210 6 : setWeightTransformFn(weight_transform_fn);
211 :
212 6 : auto transform_if = [this](TransformFn &fn, Variables &v) {
213 6 : if (fn) {
214 : auto result = fn(v);
215 6 : v.resize(result.size());
216 6 : node_owned_variable.insert(node_owned_variable.end(), result.begin(),
217 : result.end());
218 : std::transform(node_owned_variable.end() - result.size(),
219 6 : node_owned_variable.end(), v.begin(),
220 : [](Tensor &t) { return &t; });
221 6 : }
222 6 : };
223 :
224 6 : transform_if(weight_transform, weights);
225 6 : }
226 :
227 22 : void TfOpNode::finalize() {
228 44 : auto transform_if = [this](TransformFn &fn, Variables &v) {
229 44 : if (fn) {
230 : auto result = fn(v);
231 7 : v.resize(result.size());
232 : /// basically, result.size() == v.size() except InputLayer because a
233 : /// Transpose operator is added for converting nchw to nhwc
234 : /// @todo comment out below codes. TfOpNode needs to have LayerNode
235 : /// pointer
236 : // NNTR_THROW_IF(dynamic_cast<InputLayer>(layer_ptr->getLayer()) ==
237 : // nullptr && result.size() != v.size(), std::invalid_argument)
238 : // << "result size must match with given variable size";
239 7 : node_owned_variable.insert(node_owned_variable.end(), result.begin(),
240 : result.end());
241 : std::transform(node_owned_variable.end() - result.size(),
242 7 : node_owned_variable.end(), v.begin(),
243 : [](Tensor &t) { return &t; });
244 7 : }
245 44 : };
246 :
247 22 : transform_if(weight_transform, weights);
248 22 : transform_if(input_transform, inputs);
249 22 : }
250 :
251 22 : flatbuffers::Offset<void> TfOpNode::getBuiltinOps() const {
252 22 : switch (op_type) {
253 22 : case tflite::BuiltinOperator_ADD:
254 : case tflite::BuiltinOperator_AVERAGE_POOL_2D:
255 : case tflite::BuiltinOperator_CONV_2D:
256 : case tflite::BuiltinOperator_FULLY_CONNECTED:
257 : case tflite::BuiltinOperator_RELU:
258 : case tflite::BuiltinOperator_RESHAPE:
259 : case tflite::BuiltinOperator_SOFTMAX:
260 : case tflite::BuiltinOperator_TRANSPOSE:
261 : case tflite::BuiltinOperator_MUL:
262 :
263 22 : return builtin_ops;
264 0 : default:
265 0 : throw std::runtime_error{"Unsupported operator"};
266 : }
267 : }
268 :
269 22 : void TfOpNode::setBuiltinOptions(
270 : tflite::BuiltinOptions builtin_option_type_,
271 : const flatbuffers::Offset<void> &builtin_ops_) {
272 22 : builtin_ops = builtin_ops_;
273 22 : builtin_option_type = builtin_option_type_;
274 22 : }
275 :
276 : } // namespace nntrainer
|
