updater.h

#ifndef SINGA_MODEL_UPDATER_H_
#define SINGA_MODEL_UPDATER_H_

#include "singa/model/optimizer.h"
#include "singa/core/device.h"
#include "singa/core/tensor.h"
#include "singa/utils/logging.h"

#include <memory>
#include <vector>
#include <mutex>
#include <condition_variable>
#include <string>
#include <utility>
#include <unordered_map>
#include <atomic>

namespace singa {
class Updater {
 public:
  explicit Updater(Optimizer* opt) : opt_{opt} {}
  virtual ~Updater() {}
  virtual void Setup(const OptimizerConf& conf);
  virtual void Register(const string& name, const ParamSpec& specs);
  virtual void Apply(int step, const string& name, Tensor& grad, Tensor& value);
  Optimizer* GetOptimizer() { return opt_; }

  // No copy allowed.
  Updater(const Updater&) = delete;
  void operator=(const Updater&) = delete;

 protected:
  Optimizer* opt_;
};

class LocalUpdater : public Updater {
 public:
  LocalUpdater(int total_num, Optimizer* opt,
               std::shared_ptr<Device> dev = defaultDevice)
      : Updater(opt), total_num_{total_num}, dev_(dev) {}
  virtual ~LocalUpdater() override {}
  virtual void Register(const string& name, const ParamSpec& specs) override;
  virtual void Apply(int step, const string& name, Tensor& grad,
                     Tensor& value) override;
 private:
  template <typename T1, typename T2>
  struct key_hasher {
    size_t operator() (const std::pair<T1, T2>& p) const {
      auto h1 = std::hash<T1>{}(p.first);
      auto h2 = std::hash<T2>{}(p.second);
      return h1 ^ h2;
    }
  };

  int total_num_;
  std::shared_ptr<Device> dev_;
  std::unordered_map<std::string, std::atomic<int>> dev_index_;
  std::unordered_map<std::string, int> to_updater_finished_;
  std::unordered_map<std::pair<int, std::string>, Tensor,
    key_hasher<int, std::string>> grad_buffer_;
  std::unordered_map<std::string, Tensor> sum_, param_buffer_;
  std::unordered_map<std::string, std::mutex> mtx_;
  std::unordered_map<std::string, std::condition_variable>
    to_updater_all_finished_;
};
}  //  namespace singa

#endif  //  SINGA_MODEL_UPDATER_H_