model.h

Go to the documentation of this file.

/*
==============================================================================
MIT License
Copyright 2022 Institute for Automotive Engineering of RWTH Aachen University.
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
==============================================================================
*/
 
#pragma once
 
#include <stdexcept>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>
 
#include <tensorflow/core/platform/env.h>
#include <tensorflow/core/public/session.h>
#include <tensorflow_cpp/graph_utils.h>
#include <tensorflow_cpp/saved_model_utils.h>
#include <tensorflow_cpp/utils.h>
 
namespace tensorflow_cpp {
 
 
class Model {
 
 public:
  Model() {}
 
  Model(const std::string& model_path, const bool warmup = false,
        const bool allow_growth = true,
        const double per_process_gpu_memory_fraction = 0,
        const std::string& visible_device_list = "") {
 
    loadModel(model_path, warmup, allow_growth, per_process_gpu_memory_fraction,
              visible_device_list);
  }
 
  void loadModel(const std::string& model_path, const bool warmup = false,
                 const bool allow_growth = true,
                 const double per_process_gpu_memory_fraction = 0,
                 const std::string& visible_device_list = "") {
 
    is_frozen_graph_ = (model_path.substr(model_path.size() - 3) == ".pb");
    is_saved_model_ = !is_frozen_graph_;
 
    // load model
    if (is_frozen_graph_) {
      graph_def_ = loadFrozenGraph(model_path);
      session_ = createSession(allow_growth, per_process_gpu_memory_fraction,
                               visible_device_list);
      loadGraphIntoSession(session_, graph_def_);
    } else {
      saved_model_ =
        loadSavedModel(model_path, allow_growth,
                       per_process_gpu_memory_fraction, visible_device_list);
      session_ = saved_model_.GetSession();
    }
 
    // automatically find inputs and outputs
    if (is_frozen_graph_) {
      input_names_ = getGraphInputNames(graph_def_);
      output_names_ = getGraphOutputNames(graph_def_);
    } else {
      input_names_ = getSavedModelInputNames(saved_model_, true);
      output_names_ = getSavedModelOutputNames(saved_model_, true);
      const auto input_nodes_ = getSavedModelInputNames(saved_model_, false);
      const auto output_nodes_ = getSavedModelOutputNames(saved_model_, false);
      for (int k = 0; k < input_names_.size(); k++) {
        saved_model_node2layer_[input_nodes_[k]] = input_names_[k];
        saved_model_layer2node_[input_names_[k]] = input_nodes_[k];
      }
      for (int k = 0; k < output_names_.size(); k++) {
        saved_model_node2layer_[output_nodes_[k]] = output_names_[k];
        saved_model_layer2node_[output_names_[k]] = output_nodes_[k];
      }
    }
    n_inputs_ = input_names_.size();
    n_outputs_ = output_names_.size();
 
    // run dummy inference to warm-up
    if (warmup) dummyCall();
  }
 
  bool isLoaded() const {
 
    bool is_loaded = bool(session_);
 
    return is_loaded;
  }
 
  std::unordered_map<std::string, tf::Tensor> operator()(
    const std::vector<std::pair<std::string, tf::Tensor>>& inputs,
    const std::vector<std::string>& output_names) const{
 
    // properly set input/output names for session->Run()
    std::vector<std::pair<std::string, tf::Tensor>> input_nodes;
    std::vector<std::string> output_node_names;
    if (is_saved_model_) {
      for (const auto& input : inputs)
        input_nodes.push_back(
          {saved_model_layer2node_.find(input.first)->second, input.second});
      for (const auto& name : output_names)
        output_node_names.push_back(saved_model_layer2node_.find(name)->second);
    } else if (is_frozen_graph_) {
      input_nodes = inputs;
      output_node_names = output_names;
    } else {
      return {};
    }
 
    // run model
    tf::Status status;
    std::vector<tf::Tensor> output_tensors;
    status = session_->Run(input_nodes, output_node_names, {}, &output_tensors);
 
    // build outputs
    std::unordered_map<std::string, tf::Tensor> outputs;
    if (status.ok()) {
      for (int k = 0; k < output_tensors.size(); k++)
        outputs[output_names[k]] = output_tensors[k];
    } else {
      throw std::runtime_error("Failed to run model: " + status.ToString());
    }
 
    return outputs;
  }
 
  tf::Tensor operator()(const tf::Tensor& input_tensor) const {
 
    if (n_inputs_ != 1 || n_outputs_ != 1) {
      throw std::runtime_error(
        "'tf::Tensor tensorflow_cpp::Model::operator()(const tf::Tensor&)' is "
        "only available for single-input/single-output models. Found " +
        std::to_string(n_inputs_) + " inputs and " +
        std::to_string(n_outputs_) + " outputs.");
    }
 
    // run model
    auto outputs =
      (*this)({{input_names_[0], input_tensor}}, {output_names_[0]});
 
    return outputs[output_names_[0]];
  }
 
  std::vector<tf::Tensor> operator()(
    const std::vector<tf::Tensor>& input_tensors) const {
 
    if (input_tensors.size() != n_inputs_) {
      throw std::runtime_error(
        "Model has " + std::to_string(n_inputs_) + " inputs, but " +
        std::to_string(input_tensors.size()) + " input tensors were given");
    }
 
    // assign inputs in default order
    std::vector<std::pair<std::string, tf::Tensor>> inputs;
    for (int k = 0; k < n_inputs_; k++)
      inputs.push_back({input_names_[k], input_tensors[k]});
 
    // run model
    auto outputs = (*this)(inputs, output_names_);
 
    // return output tensors in default order
    std::vector<tf::Tensor> output_tensors;
    for (const auto& name : output_names_)
      output_tensors.push_back(outputs[name]);
 
    return output_tensors;
  }
 
  std::vector<int> getNodeShape(const std::string& name) {
 
    if (is_saved_model_) {
      return getSavedModelNodeShape(saved_model_,
                                    saved_model_layer2node_[name]);
    } else if (is_frozen_graph_) {
      return getGraphNodeShape(graph_def_, name);
    } else {
      return {};
    }
  }
 
  std::vector<int> getInputShape() {
 
    if (n_inputs_ != 1) {
      throw std::runtime_error(
        "std::vector<int> tensorflow_cpp::Model::getInputShape()' is only "
        "available for single-input models. Found " +
        std::to_string(n_inputs_) + " inputs.");
    }
 
    return getNodeShape(input_names_[0]);
  }
 
  std::vector<int> getOutputShape() {
 
    if (n_outputs_ != 1) {
      throw std::runtime_error(
        "std::vector<int> tensorflow_cpp::Model::getOutputShape()' is only "
        "available for single-output models. Found " +
        std::to_string(n_outputs_) + " outputs.");
    }
 
    return getNodeShape(output_names_[0]);
  }
 
  std::vector<std::vector<int>> getInputShapes() {
 
    std::vector<std::vector<int>> shapes;
    for (const auto& name : input_names_) shapes.push_back(getNodeShape(name));
 
    return shapes;
  }
 
  std::vector<std::vector<int>> getOutputShapes() {
 
    std::vector<std::vector<int>> shapes;
    for (const auto& name : output_names_) shapes.push_back(getNodeShape(name));
 
    return shapes;
  }
 
  tf::DataType getNodeType(const std::string& name) {
 
    if (is_saved_model_) {
      return getSavedModelNodeType(saved_model_, saved_model_layer2node_[name]);
    } else if (is_frozen_graph_) {
      return getGraphNodeType(graph_def_, name);
    } else {
      return tf::DataType();
    }
  }
 
  tf::DataType getInputType() {
 
    if (n_inputs_ != 1) {
      throw std::runtime_error(
        "'tf::DataType tensorflow_cpp::Model::getInputType()' is only "
        "available for single-input models. Found " +
        std::to_string(n_inputs_) + " inputs.");
    }
 
    return getNodeType(input_names_[0]);
  }
 
  tf::DataType getOutputType() {
 
    if (n_outputs_ != 1) {
      throw std::runtime_error(
        "'tf::DataType tensorflow_cpp::Model::getOutputType()' is only "
        "available for single-output models. Found " +
        std::to_string(n_outputs_) + " outputs.");
    }
 
    return getNodeType(output_names_[0]);
  }
 
  std::vector<tf::DataType> getInputTypes() {
 
    std::vector<tf::DataType> types;
    for (const auto& name : input_names_) types.push_back(getNodeType(name));
 
    return types;
  }
 
  std::vector<tf::DataType> getOutputTypes() {
 
    std::vector<tf::DataType> types;
    for (const auto& name : output_names_) types.push_back(getNodeType(name));
 
    return types;
  }
 
  std::string getInfoString() {
 
    if (is_saved_model_) {
      return getSavedModelInfoString(saved_model_);
    } else if (is_frozen_graph_) {
      return getGraphInfoString(graph_def_);
    } else {
      return "";
    }
  }
 
  tf::Session* session() const {
    return session_;
  }
 
  const tf::SavedModelBundleLite& savedModel() const {
    return saved_model_;
  }
 
  const tf::GraphDef& frozenGraph() const {
    return graph_def_;
  }
 
  bool isSavedModel() const {
    return is_saved_model_;
  }
 
  bool isFrozenGraph() const {
    return is_frozen_graph_;
  }
 
  int nInputs() const {
    return n_inputs_;
  }
 
  int nOutputs() const {
    return n_outputs_;
  }
 
  std::vector<std::string> inputNames() const {
    return input_names_;
  }
 
  std::vector<std::string> outputNames() const {
    return output_names_;
  }
 
 protected:
  void dummyCall() {
 
    // infer input shapes/types to create dummy input tensors
    auto input_shapes = getInputShapes();
    auto input_types = getInputTypes();
    std::vector<tf::Tensor> input_dummies;
    for (int k = 0; k < n_inputs_; k++) {
      std::vector<long int> dummy_shape(input_shapes[k].begin(),
                                        input_shapes[k].end());
      // Replace -1 (batch size dimension, None in python) with 1
      std::replace(dummy_shape.begin(), dummy_shape.end(), -1l, 1l);
      auto dummy_tensor_shape =
        tf::TensorShape(tf::gtl::ArraySlice<long int>(dummy_shape));
      tf::Tensor dummy(input_types[k], dummy_tensor_shape);
      // init to zero, based on type
      switch (input_types[k]) {
        case tf::DT_FLOAT:
          dummy.flat<float>().setZero();
          break;
        case tf::DT_DOUBLE:
          dummy.flat<double>().setZero();
        case tf::DT_INT32:
          dummy.flat<tf::int32>().setZero();
          break;
        case tf::DT_UINT32:
          dummy.flat<tf::uint32>().setZero();
          break;
        case tf::DT_UINT8:
          dummy.flat<tf::uint8>().setZero();
          break;
        case tf::DT_UINT16:
          dummy.flat<tf::uint16>().setZero();
          break;
        case tf::DT_INT16:
          dummy.flat<tf::int16>().setZero();
          break;
        case tf::DT_INT8:
          dummy.flat<tf::int8>().setZero();
          break;
        case tf::DT_STRING:
          dummy.flat<tf::tstring>().setZero();
          break;
        case tf::DT_COMPLEX64:
          dummy.flat<tf::complex64>().setZero();
          break;
        case tf::DT_COMPLEX128:
          dummy.flat<tf::complex128>().setZero();
          break;
        case tf::DT_INT64:
          dummy.flat<tf::int64>().setZero();
          break;
        case tf::DT_UINT64:
          dummy.flat<tf::uint64>().setZero();
          break;
        case tf::DT_BOOL:
          dummy.flat<bool>().setZero();
          break;
        case tf::DT_QINT8:
          dummy.flat<tf::qint8>().setZero();
          break;
        case tf::DT_QUINT8:
          dummy.flat<tf::quint8>().setZero();
          break;
        case tf::DT_QUINT16:
          dummy.flat<tf::quint16>().setZero();
          break;
        case tf::DT_QINT16:
          dummy.flat<tf::qint16>().setZero();
          break;
        case tf::DT_QINT32:
          dummy.flat<tf::qint32>().setZero();
          break;
        case tf::DT_BFLOAT16:
          dummy.flat<tf::bfloat16>().setZero();
          break;
        case tf::DT_HALF:
          dummy.flat<Eigen::half>().setZero();
          break;
      }
      input_dummies.push_back(dummy);
    }
 
    // run dummy inference
    volatile auto output_dummies = (*this)(input_dummies);
  }
 
 protected:
  tf::Session* session_ = nullptr;
 
  tf::SavedModelBundleLite saved_model_;
 
  tf::GraphDef graph_def_;
 
  bool is_saved_model_ = false;
 
  bool is_frozen_graph_ = false;
 
  int n_inputs_;
 
  int n_outputs_;
 
  std::vector<std::string> input_names_;
 
  std::vector<std::string> output_names_;
 
  std::unordered_map<std::string, std::string> saved_model_node2layer_;
 
  std::unordered_map<std::string, std::string> saved_model_layer2node_;
};
 
 
}  // namespace tensorflow_cpp