aboutsummaryrefslogtreecommitdiffhomepage
path: root/tensorflow/contrib/lite/toco/tflite/export.h
blob: 915d5dd3d67d17b6051ccc2cd483883637d84323 (plain)
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
/* 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 {

// Transform the given tf.mini model into a TF Lite flatbuffer and deposit the
// result in the given string.
void Export(const Model& model, bool allow_custom_ops, bool quantize_weights,
            string* output_file_contents);

// This if backward-compatibility.
// TODO(ycling): Remove the deprecated entry functions.
inline void Export(const Model& model, string* output_file_contents) {
  Export(model, true, false, output_file_contents);
}

// Export API with custom TFLite operator mapping.
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);

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(OperatorType type, const std::string& custom_code, int version)
      : type(type), custom_code(custom_code), version(version) {}
  const OperatorType type;
  const std::string custom_code;
  const int version;

  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)});
    }
  };
};

// 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);

}  // namespace details
}  // namespace tflite
}  // namespace toco

#endif  // TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_EXPORT_H_