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