1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
|
/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_EXPORT_H_
#define TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_EXPORT_H_
#include "tensorflow/contrib/lite/toco/model.h"
#include "tensorflow/contrib/lite/toco/tflite/operator.h"
#include "tensorflow/contrib/lite/util.h"
namespace toco {
namespace tflite {
// The parameters for exporting a TFLite model.
struct ExportParams {
bool allow_custom_ops = false;
bool allow_flex_ops = false;
bool quantize_weights = false;
};
// Transform the given tf.mini model into a TF Lite flatbuffer and deposit the
// result in the given string.
void Export(const Model& model, string* output_file_contents,
const ExportParams& params);
// Export API with custom TFLite operator mapping.
void Export(
const Model& model, string* output_file_contents,
const ExportParams& params,
const std::map<OperatorType, std::unique_ptr<BaseOperator>>& ops_by_type);
// This is for backward-compatibility.
// TODO(ycling): Remove the deprecated entry functions.
inline void Export(const Model& model, bool allow_custom_ops,
bool quantize_weights, string* output_file_contents) {
ExportParams params;
params.allow_custom_ops = allow_custom_ops;
params.quantize_weights = quantize_weights;
Export(model, output_file_contents, params);
}
// This is for backward-compatibility.
// TODO(ycling): Remove the deprecated entry functions.
inline void Export(
const Model& model, bool allow_custom_ops, bool quantize_weights,
string* output_file_contents,
const std::map<OperatorType, std::unique_ptr<BaseOperator>>& ops_by_type) {
ExportParams params;
params.allow_custom_ops = allow_custom_ops;
params.quantize_weights = quantize_weights;
Export(model, output_file_contents, params, ops_by_type);
}
// This is for backward-compatibility.
// TODO(ycling): Remove the deprecated entry functions.
inline void Export(const Model& model, string* output_file_contents) {
ExportParams params;
params.allow_custom_ops = true;
Export(model, output_file_contents, params);
Export(model, true, false, output_file_contents);
}
namespace details {
// A maps from tensor name to its final position in the TF Lite buffer.
using TensorsMap = std::unordered_map<string, int>;
// A key to identify an operator.
// Only when `type` is `kUnsupported`, `custom_code` is filled to
// identify which operation is used.
struct OperatorKey {
OperatorKey() {}
OperatorKey(::tflite::BuiltinOperator type, const std::string& custom_code,
int version)
: type(type), custom_code(custom_code), version(version) {}
// Only `type`, `custom_code` and `version` is used to compute hash and
// identity.
::tflite::BuiltinOperator type = ::tflite::BuiltinOperator_CUSTOM;
std::string custom_code;
int version = 1;
// The fields below are not used to compute hash and identity.
// TODO(ycling): Consider to change these fields to accessor functions.
bool is_custom_op = false;
bool is_flex_op = false;
bool is_unsupported_flex_op = false;
// The original TensorFlow op name for the flex op. Filled only when
// `is_flex_op` is true.
std::string flex_tensorflow_op;
bool operator<(const OperatorKey& other) const {
if (type < other.type) return true;
else if (type > other.type)
return false;
else if (custom_code < other.custom_code)
return true;
else if (custom_code > other.custom_code)
return false;
else
return version < other.version;
}
bool operator==(const OperatorKey& other) const {
return type == other.type && custom_code == other.custom_code &&
version == other.version;
}
struct Hash {
size_t operator()(const OperatorKey& key) const {
return ::tflite::CombineHashes(
{std::hash<size_t>()(static_cast<size_t>(key.type)),
std::hash<std::string>()(key.custom_code),
std::hash<int>()(key.version)});
}
};
};
OperatorKey GetOperatorKey(
const ::toco::Operator& op,
const std::map<OperatorType, std::unique_ptr<BaseOperator>>& ops_by_type,
bool allow_flex_ops);
// A maps from operator type to its final position in the TF Lite buffer.
using OperatorsMap = std::unordered_map<OperatorKey, int, OperatorKey::Hash>;
void LoadTensorsMap(const Model& model, TensorsMap* tensors_map);
void LoadOperatorsMap(
const Model& model, OperatorsMap* operators_map,
const std::map<OperatorType, std::unique_ptr<BaseOperator>>& ops_by_type,
bool allow_flex_ops);
} // namespace details
} // namespace tflite
} // namespace toco
#endif // TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_EXPORT_H_
|
