initializer.h

#ifndef SINGA_MODEL_INITIALIZER_H_
#define SINGA_MODEL_INITIALIZER_H_
#include <string>
#include "singa/core/tensor.h"
#include "singa/proto/model.pb.h"
#include "singa/utils/string.h"
namespace singa {
using InitializerConf = FillerConf;
class Initializer {
 public:
  Initializer() = default;
  void Setup(const std::string& str) {
    InitializerConf conf;
    conf.ParseFromString(str);
    Setup(conf);
  }

  virtual void Setup(const InitializerConf& conf) {}

  virtual void Fill(Tensor& t) = 0;
};

namespace init {
class Constant : public Initializer {
public:
  Constant() = default;
  Constant(const float x) : v_(x) {}
  void Setup(const InitializerConf& conf) override { v_ = conf.value(); }
  void Fill(Tensor& t) override { t.SetValue(v_); }

 private:
  float v_ = 0;
};

class Uniform : public Initializer {
public:
  Uniform() = default;
  Uniform(const float low, const float high) : min_(low), max_(high) {}
  void Setup(const InitializerConf& conf) override {
    min_ = conf.min();
    max_ = conf.max();
  }
  void Fill(Tensor& t) override { singa::Uniform(min_, max_, &t); }

 private:
  float min_ = 0, max_ = 1;
};

class Gaussian : public Initializer {
public:
  Gaussian() = default;
  Gaussian(const float m, const float s): mean_(m), std_(s) {}
  void Setup(const InitializerConf& conf) override {
    mean_ = conf.mean();
    std_ = conf.std();
  }
  void Fill(Tensor& t) override { singa::Gaussian(mean_, std_, &t); }

 private:
  float mean_ = 0, std_ = 1;
};

class Xavier : public Initializer {
public:
  void Fill(Tensor& t) override {
    CHECK_EQ(t.nDim(), 2u);
    float scale = sqrt(6.0f / (t.shape(0) + t.shape(1)));
    LOG(INFO) << "xavier scale " << scale;
    singa::Uniform(-scale, scale, &t);
  }
};

class MSRA : public Initializer {
 public:
  void Fill(Tensor& t) override {
    CHECK_EQ(t.nDim(), 2u);
    float std = sqrt(2.0f / t.shape(0));
    singa::Gaussian(0.0f, std, &t);
  }
};

}  // namespace init

std::shared_ptr<Initializer> CreateInitializer(const InitializerConf& conf) {
  std::shared_ptr<Initializer> init;
  if (ToLowerCase(conf.type()) == "constant") {
    init = std::make_shared<init::Constant>();
  } else if (ToLowerCase(conf.type()) == "uniform") {
    init = std::make_shared<init::Uniform>();
  } else if (ToLowerCase(conf.type()) == "gaussian") {
    init = std::make_shared<init::Gaussian>();
  } else if (ToLowerCase(conf.type()) == "xavier") {
    init = std::make_shared<init::Xavier>();
  } else if (ToLowerCase(conf.type()) == "msra") {
    init = std::make_shared<init::MSRA>();
  } else {
    LOG(FATAL) << "Unknown initialization type: " << conf.type();
  }
  init->Setup(conf);
  return init;
}
}  // namespace singa
#endif  // SINGA_MODEL_INITIALIZER_H_