deepx_core/src/graph/optimizer/sgd.cc at master · Tencent/deepx_core · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
// Copyright 2019 the deepx authors.
// Author: Yafei Zhang (kimmyzhang@tencent.com)
//

#include <deepx_core/graph/optimizer_impl.h>

namespace deepx_core {

class SGDOptimizer : public OptimizerBase0 {
 private:
  ll_optimizer_t::SGDConfig config_;

 public:
  DEFINE_OPTIMIZER_LIKE(SGDOptimizer);

 protected:
  bool InitConfigKV(const std::string& k, const std::string& v) override {
    if (k == "alpha") {
      config_.alpha = (float_t)std::stod(v);
      if (config_.alpha <= 0) {
        DXERROR("Invalid %s: %s.", k.c_str(), v.c_str());
        return false;
      }
    } else if (k == "min_alpha") {
      config_.min_alpha = (float_t)std::stod(v);
      if (config_.min_alpha <= 0) {
        DXERROR("Invalid %s: %s.", k.c_str(), v.c_str());
        return false;
      }
    } else if (k == "batch_decay") {
      config_.batch_decay = std::stoi(v);
      if (config_.batch_decay < 0) {
        DXERROR("Invalid %s: %s.", k.c_str(), v.c_str());
        return false;
      }
    } else if (k == "batch_decay_rate") {
      config_.batch_decay_rate = (float_t)std::stod(v);
      if (config_.batch_decay_rate <= 0 || config_.batch_decay_rate > 1) {
        DXERROR("Invalid %s: %s.", k.c_str(), v.c_str());
        return false;
      }
    } else {
      DXERROR("Unexpected config: %s=%s.", k.c_str(), v.c_str());
      return false;
    }
    return true;
  }

  bool PostInitConfig() override {
    ll_optimizer_t::Init(&config_);
    return true;
  }

  void WriteConfigLegacy(OutputStream& os) const override {
    os << config_.alpha << config_.min_alpha << config_.batch_decay
       << config_.batch_decay_rate << config_.real_batch << config_.real_alpha;
  }

  void ReadConfigLegacy(InputStream& is) override {
    is >> config_.alpha >> config_.min_alpha >> config_.batch_decay >>
        config_.batch_decay_rate >> config_.real_batch >> config_.real_alpha;
  }

  void CopyConfigLegacy(const Optimizer& other) override {
    config_ = ((const SGDOptimizer&)other).config_;
  }

  void PreUpdate() override { ll_optimizer_t::PreBatch(&config_); }

  void PostUpdate() override { ll_optimizer_t::PostBatch(&config_); }

  void UpdateTSR2TSR(const std::string& /*name*/, const tsr_t& G, tsr_t* W,
                     OptimizerTSRSlot* /*slot*/) const override {
    ll_optimizer_t::UpdateTSR2TSR(config_, G, W);
  }

  void UpdateSRM2TSR(const std::string& /*name*/, const srm_t& G, tsr_t* W,
                     OptimizerTSRSlot* /*slot*/) const override {
    ll_optimizer_t::UpdateSRM2TSR(config_, G, W);
  }

  void UpdateSRM2SRM(const std::string& /*name*/, const srm_t& G, srm_t* W,
                     OptimizerSRMSlot* slot) const override {
    if (use_lock_) {
      ll_optimizer_t::UpdateSRM2SRM(config_, G, W, slot->Wlock.get());
    } else {
      ll_optimizer_t::UpdateSRM2SRM(config_, G, W);
    }
  }
};

OPTIMIZER_REGISTER(SGDOptimizer, "SGDOptimizer");
OPTIMIZER_REGISTER(SGDOptimizer, "sgd");

}  // namespace deepx_core