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_ */