agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 1 | /** |
| 2 | * Copyright (C) 2020 Savoir-faire Linux Inc. |
| 3 | * |
| 4 | * Author: Aline Gondim Santos <aline.gondimsantos@savoirfairelinux.com> |
| 5 | * |
| 6 | * This program is free software; you can redistribute it and/or modify |
| 7 | * it under the terms of the GNU General Public License as published by |
| 8 | * the Free Software Foundation; either version 3 of the License, or |
| 9 | * (at your option) any later version. |
| 10 | * |
| 11 | * This program is distributed in the hope that it will be useful, |
| 12 | * but WITHOUT ANY WARRANTY; without even the implied warranty of |
| 13 | * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the |
| 14 | * GNU General Public License for more details. |
| 15 | * |
| 16 | * You should have received a copy of the GNU General Public License |
| 17 | * along with this program; if not, write to the Free Software |
| 18 | * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA. |
| 19 | */ |
| 20 | |
| 21 | #include "TFInference.h" |
| 22 | // Std libraries |
| 23 | #include <fstream> |
| 24 | #include <numeric> |
| 25 | #include <iostream> |
| 26 | #include <stdlib.h> |
| 27 | |
| 28 | |
| 29 | #ifdef TFLITE |
| 30 | // Tensorflow headers |
| 31 | #include <tensorflow/lite/interpreter.h> |
| 32 | #include <tensorflow/lite/builtin_op_data.h> |
| 33 | #include <tensorflow/lite/interpreter.h> |
| 34 | #include <tensorflow/lite/kernels/register.h> |
| 35 | #include <tensorflow/lite/model.h> |
| 36 | #include <tensorflow/lite/optional_debug_tools.h> |
| 37 | #else |
| 38 | #ifdef WIN32 |
| 39 | #include <WinBase.h> |
| 40 | #endif |
| 41 | #include <tensorflow/core/graph/graph.h> |
| 42 | //#include <tensorflow/core/graph/default_device.h> |
| 43 | #include <tensorflow/core/platform/env.h> |
| 44 | #endif // TFLITE |
| 45 | |
| 46 | #include <pluglog.h> |
| 47 | |
| 48 | const char sep = separator(); |
| 49 | const std::string TAG = "FORESEG"; |
| 50 | |
| 51 | namespace jami { |
| 52 | TensorflowInference::TensorflowInference(TFModel tfModel) : tfModel(tfModel) {} |
| 53 | |
| 54 | TensorflowInference::~TensorflowInference() {} |
| 55 | |
| 56 | bool |
| 57 | TensorflowInference::isAllocated() const |
| 58 | { |
| 59 | return allocated; |
| 60 | } |
| 61 | |
| 62 | #ifdef TFLITE |
| 63 | |
| 64 | void |
| 65 | TensorflowInference::loadModel() |
| 66 | { |
| 67 | Plog::log(Plog::LogPriority::INFO, TAG, "inside loadModel()"); |
| 68 | flatbufferModel = tflite::FlatBufferModel::BuildFromFile(tfModel.modelPath.c_str()); |
| 69 | if (!flatbufferModel) { |
| 70 | std::runtime_error("Failed to load the model file"); |
| 71 | } |
| 72 | Plog::log(Plog::LogPriority::INFO, "TENSOR", "MODEL LOADED" ); |
| 73 | |
| 74 | } |
| 75 | void |
| 76 | TensorflowInference::buildInterpreter() |
| 77 | { |
| 78 | Plog::log(Plog::LogPriority::INFO, TAG, "inside buildInterpreter()"); |
| 79 | // Build the interpreter |
| 80 | tflite::ops::builtin::BuiltinOpResolver resolver; |
| 81 | tflite::InterpreterBuilder builder(*flatbufferModel, resolver); |
| 82 | builder(&interpreter); |
| 83 | if(interpreter) { |
| 84 | setInterpreterSettings(); |
| 85 | Plog::log(Plog::LogPriority::INFO, "TENSOR", "INTERPRETER BUILT" ); |
| 86 | |
| 87 | if (tfModel.useNNAPI) { |
| 88 | TfLiteDelegate* optionalNnApiDelegate = tflite::NnApiDelegate(); |
| 89 | |
| 90 | if (interpreter->ModifyGraphWithDelegate(optionalNnApiDelegate) != kTfLiteOk) { |
| 91 | Plog::log(Plog::LogPriority::INFO, "TENSOR", "INTERPRETER ERROR!!!" ); |
| 92 | } |
| 93 | else { |
| 94 | Plog::log(Plog::LogPriority::INFO, "TENSOR", "INTERPRETER SET" ); |
| 95 | allocateTensors(); |
| 96 | } |
| 97 | } |
| 98 | else { |
| 99 | allocateTensors(); |
| 100 | } |
| 101 | } |
| 102 | } |
| 103 | |
| 104 | void |
| 105 | TensorflowInference::allocateTensors() |
| 106 | { |
| 107 | if (interpreter->AllocateTensors() != kTfLiteOk) { |
| 108 | std::runtime_error("Failed to allocate tensors!"); |
| 109 | } else { |
| 110 | Plog::log(Plog::LogPriority::INFO, "TENSOR", "TENSORS ALLOCATED" ); |
| 111 | allocated = true; |
| 112 | } |
| 113 | } |
| 114 | |
| 115 | void |
| 116 | TensorflowInference::describeModelTensors() const |
| 117 | { |
| 118 | // PrintInterpreterState(interpreter.get()); |
| 119 | std::ostringstream oss; |
| 120 | oss << "=============== inputs/outputs dimensions ===============" |
| 121 | << "\n"; |
| 122 | const std::vector<int> inputs = interpreter->inputs(); |
| 123 | const std::vector<int> outputs = interpreter->outputs(); |
| 124 | oss << "number of inputs: " << inputs.size() << std::endl; |
| 125 | oss << "number of outputs: " << outputs.size() << std::endl; |
| 126 | |
| 127 | Plog::log(Plog::LogPriority::INFO, "TENSOR", oss.str() ); |
| 128 | int input = interpreter->inputs()[0]; |
| 129 | int output = interpreter->outputs()[0]; |
| 130 | oss << "input 0 index: " << input << std::endl; |
| 131 | oss << "output 0 index: " << output << std::endl; |
| 132 | oss << "=============== input dimensions ===============" |
| 133 | << std::endl; |
| 134 | Plog::log(Plog::LogPriority::INFO, "TENSOR", oss.str() ); |
| 135 | // get input dimension from the input tensor metadata |
| 136 | // assuming one input only |
| 137 | |
| 138 | for (size_t i = 0; i < inputs.size(); i++) { |
| 139 | std::stringstream ss; |
| 140 | ss << "Input " << i << " ➛ "; |
| 141 | describeTensor(ss.str(), interpreter->inputs()[i]); |
| 142 | } |
| 143 | oss.str(""); |
| 144 | oss << "=============== output dimensions ===============" |
| 145 | << "\n"; |
| 146 | Plog::log(Plog::LogPriority::INFO, "TENSOR", oss.str() ); |
| 147 | // get input dimension from the input tensor metadata |
| 148 | // assuming one input only |
| 149 | for (size_t i = 0; i < outputs.size(); i++) { |
| 150 | std::stringstream ss; |
| 151 | ss << "Output " << i << " ➛ "; |
| 152 | describeTensor(ss.str(), interpreter->outputs()[i]); |
| 153 | } |
| 154 | } |
| 155 | |
| 156 | void |
| 157 | TensorflowInference::setInterpreterSettings() |
| 158 | { |
| 159 | // interpreter->UseNNAPI(tfModel.useNNAPI); |
| 160 | interpreter->SetAllowFp16PrecisionForFp32(tfModel.allowFp16PrecisionForFp32); |
| 161 | |
| 162 | interpreter->SetNumThreads(static_cast<int>(tfModel.numberOfThreads)); |
| 163 | } |
| 164 | |
| 165 | void |
| 166 | TensorflowInference::describeTensor(std::string prefix, int index) const |
| 167 | { |
| 168 | std::vector<int> dimensions = getTensorDimensions(index); |
| 169 | size_t nbDimensions = dimensions.size(); |
| 170 | |
| 171 | std::ostringstream tensorDescription; |
| 172 | tensorDescription << prefix; |
| 173 | for (size_t i = 0; i < nbDimensions; i++) { |
| 174 | if (i == dimensions.size() - 1) { |
| 175 | tensorDescription << dimensions[i]; |
| 176 | } |
| 177 | else { |
| 178 | tensorDescription << dimensions[i] << " x "; |
| 179 | } |
| 180 | } |
| 181 | tensorDescription << std::endl; |
| 182 | Plog::log(Plog::LogPriority::INFO, "TENSOR", tensorDescription.str() ); |
| 183 | } |
| 184 | |
| 185 | std::vector<int> |
| 186 | TensorflowInference::getTensorDimensions(int index) const |
| 187 | { |
| 188 | TfLiteIntArray *dims = interpreter->tensor(index)->dims; |
| 189 | size_t size = static_cast<size_t>(interpreter->tensor(index)->dims->size); |
| 190 | std::vector<int> result; |
| 191 | result.reserve(size); |
| 192 | for (size_t i = 0; i != size; i++) { |
| 193 | result.push_back(dims->data[i]); |
| 194 | } |
| 195 | |
| 196 | dims = nullptr; |
| 197 | |
| 198 | return result; |
| 199 | } |
| 200 | |
| 201 | void |
| 202 | TensorflowInference::runGraph() |
| 203 | { |
| 204 | for (size_t i = 0; i < tfModel.numberOfRuns; i++) { |
| 205 | if (interpreter->Invoke() != kTfLiteOk) { |
| 206 | Plog::log(Plog::LogPriority::INFO, "RUN GRAPH", "A problem occured when running the graph"); |
| 207 | } |
| 208 | } |
| 209 | } |
| 210 | |
| 211 | void |
| 212 | TensorflowInference::init() |
| 213 | { |
| 214 | /// Loading the model |
| 215 | loadModel(); |
| 216 | buildInterpreter(); |
| 217 | describeModelTensors(); |
| 218 | } |
| 219 | |
| 220 | #else |
| 221 | // Reads a model graph definition from disk, and creates a session object you |
| 222 | // can use to run it. |
| 223 | void |
| 224 | TensorflowInference::LoadGraph() |
| 225 | { |
| 226 | tensorflow::GraphDef graph_def; |
| 227 | tensorflow::Status load_graph_status = tensorflow::ReadBinaryProto(tensorflow::Env::Default(), tfModel.modelPath, &graph_def); |
| 228 | if (!load_graph_status.ok()) { |
| 229 | allocated = false; |
| 230 | Plog::log(Plog::LogPriority::INFO, "LOAD GRAPH", "A problem occured when loading the graph"); |
| 231 | return ; |
| 232 | } |
| 233 | Plog::log(Plog::LogPriority::INFO, "LOAD GRAPH", "graph loaded"); |
| 234 | |
| 235 | // Plog::log(Plog::LogPriority::INFO, "GRAPH SIZE: ", std::to_string(graph_def.node_size())); |
| 236 | // for (auto& node : *graph_def.mutable_node()) |
| 237 | // { |
| 238 | // Plog::log(Plog::LogPriority::INFO, "GRAPH NODE: ", node.name().c_str()); |
| 239 | // // Plog::log(Plog::LogPriority::INFO, "\tNODE SIZE: ", node.().c_str()); |
| 240 | // } |
| 241 | |
| 242 | PluginParameters* parameters = getGlobalPluginParameters(); |
| 243 | |
| 244 | tensorflow::SessionOptions options; |
| 245 | if(parameters->useGPU) { |
| 246 | options.config.mutable_gpu_options()->set_allow_growth(true); |
| 247 | options.config.mutable_gpu_options()->set_per_process_gpu_memory_fraction(0.3); |
| 248 | } |
| 249 | else { |
| 250 | #ifdef WIN32 |
| 251 | options.config.mutable_device_count()->insert({ "CPU", 1 }); |
| 252 | options.config.mutable_device_count()->insert({ "GPU", 0 }); |
| 253 | #else |
| 254 | setenv("CUDA_VISIBLE_DEVICES", "", 1); |
| 255 | #endif |
| 256 | } |
| 257 | |
| 258 | (&session)->reset(tensorflow::NewSession(options)); |
| 259 | tensorflow::Status session_create_status = session->Create(graph_def); |
| 260 | if (!session_create_status.ok()) { |
| 261 | Plog::log(Plog::LogPriority::INFO, "INIT SESSION", "A problem occured when initializating session"); |
| 262 | allocated = true; |
| 263 | return ; |
| 264 | } |
| 265 | Plog::log(Plog::LogPriority::INFO, "INIT SESSION", "session initialized"); |
| 266 | |
| 267 | allocated = true; |
| 268 | } |
| 269 | |
| 270 | void |
| 271 | TensorflowInference::runGraph() |
| 272 | { |
| 273 | for (size_t i = 0; i < tfModel.numberOfRuns; i++) { |
| 274 | // Actually run the image through the model. |
| 275 | tensorflow::Status run_status = session->Run({{tfModel.inputLayer, imageTensor}}, {tfModel.outputLayer}, {}, &outputs); |
| 276 | if (!run_status.ok()) { |
| 277 | Plog::log(Plog::LogPriority::INFO, "RUN GRAPH", "A problem occured when running the graph"); |
| 278 | } |
| 279 | } |
| 280 | } |
| 281 | |
| 282 | void |
| 283 | TensorflowInference::init() |
| 284 | { |
| 285 | // Loading the model |
| 286 | LoadGraph(); |
| 287 | } |
| 288 | #endif |
| 289 | |
| 290 | } |