Predictor.cc

Go to the documentation of this file.

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

#include "PhysicsTools/MXNet/interface/Predictor.h"

#include <cassert>
#include <memory>

#include "FWCore/Utilities/interface/Exception.h"

namespace mxnet {

  namespace cpp {

    Block::Block() {}

    Block::Block(const std::string& symbol_file, const std::string& param_file) {
      // load the symbol
      sym_ = Symbol::Load(symbol_file);
      // load the parameters
      load_parameters(param_file);
    }

    Block::~Block() {}

    void Block::load_parameters(const std::string& param_file) {
      std::map<std::string, NDArray> paramters;
      NDArray::Load(param_file, nullptr, &paramters);
      for (const auto& k : paramters) {
        if (k.first.substr(0, 4) == "aux:") {
          auto name = k.first.substr(4, k.first.size() - 4);
          aux_map_[name] = k.second;
        }
        if (k.first.substr(0, 4) == "arg:") {
          auto name = k.first.substr(4, k.first.size() - 4);
          arg_map_[name] = k.second;
        }
      }
    }

    std::mutex Predictor::mutex_;
    const Context Predictor::context_ = Context(DeviceType::kCPU, 0);

    Predictor::Predictor() {}

    Predictor::Predictor(const Block& block)
        : sym_(block.symbol()), arg_map_(block.arg_map()), aux_map_(block.aux_map()) {}

    Predictor::Predictor(const Block& block, const std::string& output_node)
        : sym_(block.symbol(output_node)), arg_map_(block.arg_map()), aux_map_(block.aux_map()) {}

    Predictor::~Predictor() {}

    void Predictor::set_input_shapes(const std::vector<std::string>& input_names,
                                     const std::vector<std::vector<mx_uint> >& input_shapes) {
      assert(input_names.size() == input_shapes.size());
      input_names_ = input_names;
      // init the input NDArrays and add them to the arg_map
      for (unsigned i = 0; i < input_names_.size(); ++i) {
        const auto& name = input_names_[i];
        arg_map_.emplace(name, NDArray(input_shapes[i], context_, false));
      }
    }

    const std::vector<float>& Predictor::predict(const std::vector<std::vector<mx_float> >& input_data) {
      assert(input_names_.size() == input_data.size());

      try {
        // create the executor (if not done yet)
        if (!exec_) {
          bind_executor();
        }
        assert(exec_);
        // set the inputs
        for (unsigned i = 0; i < input_names_.size(); ++i) {
          const auto& name = input_names_[i];
          arg_map_[name].SyncCopyFromCPU(input_data[i]);
        }
        // run forward
        exec_->Forward(false);
        // copy the output to pred_
        exec_->outputs[0].SyncCopyToCPU(&pred_);
        return pred_;
      } catch (const dmlc::Error& e) {
        throw cms::Exception("RuntimeError") << e.what() << MXGetLastError();
      }
    }

    void Predictor::bind_executor() {
      // acquire lock
      std::lock_guard<std::mutex> lock(mutex_);

      // infer shapes
      const auto arg_name_list = sym_.ListArguments();
      std::vector<std::vector<mx_uint> > in_shapes, aux_shapes, out_shapes;
      std::map<std::string, std::vector<mx_uint> > arg_shapes;

      for (const auto& arg_name : arg_name_list) {
        auto iter = arg_map_.find(arg_name);
        if (iter != arg_map_.end()) {
          arg_shapes[arg_name] = iter->second.GetShape();
        }
      }
      sym_.InferShape(arg_shapes, &in_shapes, &aux_shapes, &out_shapes);

      // init argument arrays
      std::vector<NDArray> arg_arrays;
      for (size_t i = 0; i < in_shapes.size(); ++i) {
        const auto& shape = in_shapes[i];
        const auto& arg_name = arg_name_list[i];
        auto iter_arg = arg_map_.find(arg_name);
        if (iter_arg != arg_map_.end()) {
          arg_arrays.push_back(iter_arg->second);
        } else {
          arg_arrays.push_back(NDArray(shape, context_, false));
        }
      }
      std::vector<NDArray> grad_arrays(arg_arrays.size());
      std::vector<OpReqType> grad_reqs(arg_arrays.size(), kNullOp);

      // init auxiliary array
      std::vector<NDArray> aux_arrays;
      const auto aux_name_list = sym_.ListAuxiliaryStates();
      for (size_t i = 0; i < aux_shapes.size(); ++i) {
        const auto& shape = aux_shapes[i];
        const auto& aux_name = aux_name_list[i];
        auto iter_aux = aux_map_.find(aux_name);
        if (iter_aux != aux_map_.end()) {
          aux_arrays.push_back(iter_aux->second);
        } else {
          aux_arrays.push_back(NDArray(shape, context_, false));
        }
      }

      // bind executor
      exec_ = std::make_unique<Executor>(sym_, context_, arg_arrays, grad_arrays, grad_reqs, aux_arrays);
    }

  }  // namespace cpp

} /* namespace mxnet */