Predictor.h

Go to the documentation of this file.

/*
 * MXNetCppPredictor.h
 *
 *  Created on: Jul 19, 2018
 *      Author: hqu
 */

#ifndef PHYSICSTOOLS_MXNET_MXNETCPPPREDICTOR_H_
#define PHYSICSTOOLS_MXNET_MXNETCPPPREDICTOR_H_

#include <map>
#include <vector>
#include <memory>
#include <mutex>

#include "mxnet-cpp/MxNetCpp.h"

namespace mxnet {

namespace cpp {

// note: Most of the objects in mxnet::cpp are effective just shared_ptr's

// Simple class to hold MXNet model (symbol + params)
// designed to be sharable by multiple threads
class Block {
public:
  Block();
  Block(const std::string &symbol_file, const std::string &param_file);
  virtual ~Block();

  const Symbol& symbol() const { return sym_; }
  Symbol symbol(const std::string &output_node) const { return sym_.GetInternals()[output_node]; }
  const std::map<std::string, NDArray>& arg_map() const { return arg_map_; }
  const std::map<std::string, NDArray>& aux_map() const { return aux_map_; }

private:
  void load_parameters(const std::string& param_file);

  // symbol
  Symbol sym_;
  // argument arrays
  std::map<std::string, NDArray> arg_map_;
  // auxiliary arrays
  std::map<std::string, NDArray> aux_map_;
};

// Simple helper class to run prediction
// this cannot be shared between threads
class Predictor {
public:
  Predictor();
  Predictor(const Block &block);
  Predictor(const Block &block, const std::string &output_node);
  virtual ~Predictor();

  // set input array shapes
  void set_input_shapes(const std::vector<std::string>& input_names, const std::vector<std::vector<mx_uint>>& input_shapes);

  // run prediction
  const std::vector<float>& predict(const std::vector<std::vector<mx_float>>& input_data);

private:
  static std::mutex mutex_;

  void bind_executor();

  // context
  static const Context context_;
  // executor
  std::unique_ptr<Executor> exec_;
  // symbol
  Symbol sym_;
  // argument arrays
  std::map<std::string, NDArray> arg_map_;
  // auxiliary arrays
  std::map<std::string, NDArray> aux_map_;
  // output of the prediction
  std::vector<float> pred_;
  // names of the input nodes
  std::vector<std::string> input_names_;

};

} /* namespace cpp */
} /* namespace mxnet */

#endif /* PHYSICSTOOLS_MXNET_MXNETCPPPREDICTOR_H_ */