operator.h

Go to the documentation of this file.

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
 
#ifndef MXNET_CPP_OPERATOR_H_
#define MXNET_CPP_OPERATOR_H_
 
#include <map>
#include <string>
#include <vector>
#include "mxnet-cpp/base.h"
#include "mxnet-cpp/op_map.h"
#include "mxnet-cpp/symbol.h"
 
namespace mxnet {
namespace cpp {
class Mxnet;
class Operator {
 public:
  explicit Operator(const std::string& operator_name);
  Operator& operator=(const Operator& rhs);
  template <typename T>
  Operator& SetParam(const std::string& name, const T& value) {
    std::string value_str;
    std::stringstream ss;
    ss << value;
    ss >> value_str;
 
    params_[name] = value_str;
    return *this;
  }
  template <typename T>
  Operator& SetParam(int pos, const T& value) {
    std::string value_str;
    std::stringstream ss;
    ss << value;
    ss >> value_str;
 
    params_[arg_names_[pos]] = value_str;
    return *this;
  }
  Operator& SetInput(const std::string& name, const Symbol& symbol);
  template <int N = 0>
  void PushInput(const Symbol& symbol) {
    input_symbols_.push_back(symbol.GetHandle());
  }
  Operator& operator()() {
    return *this;
  }
  Operator& operator()(const Symbol& symbol) {
    input_symbols_.push_back(symbol.GetHandle());
    return *this;
  }
  Operator& operator()(const std::vector<Symbol>& symbols) {
    for (auto& s : symbols) {
      input_symbols_.push_back(s.GetHandle());
    }
    return *this;
  }
  Symbol CreateSymbol(const std::string& name = "");
 
  Operator& SetInput(const std::string& name, const NDArray& ndarray);
  template <int N = 0>
  Operator& PushInput(const NDArray& ndarray) {
    input_ndarrays_.push_back(ndarray.GetHandle());
    return *this;
  }
  template <class T, class... Args, int N = 0>
  Operator& PushInput(const T& t, Args... args) {
    SetParam(N, t);
    PushInput<Args..., N + 1>(args...);
    return *this;
  }
  template <class T, int N = 0>
  Operator& PushInput(const T& t) {
    SetParam(N, t);
    return *this;
  }
  Operator& operator()(const NDArray& ndarray) {
    input_ndarrays_.push_back(ndarray.GetHandle());
    return *this;
  }
  Operator& operator()(const std::vector<NDArray>& ndarrays) {
    for (auto& s : ndarrays) {
      input_ndarrays_.push_back(s.GetHandle());
    }
    return *this;
  }
  template <typename... Args>
  Operator& operator()(Args... args) {
    PushInput(args...);
    return *this;
  }
  std::vector<NDArray> Invoke();
  void Invoke(NDArray& output);
  void Invoke(std::vector<NDArray>& outputs);
 
 private:
  std::map<std::string, std::string> params_desc_;
  bool variable_params_ = false;
  std::map<std::string, std::string> params_;
  std::vector<SymbolHandle> input_symbols_;
  std::vector<NDArrayHandle> input_ndarrays_;
  std::vector<std::string> input_keys_;
  std::vector<std::string> arg_names_;
  AtomicSymbolCreator handle_;
  static OpMap*& op_map();
};
}  // namespace cpp
}  // namespace mxnet
 
#endif  // MXNET_CPP_OPERATOR_H_