GCC Code Coverage Report

Directory: src/
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 100.0% 13 / 0 / 13
Functions: 100.0% 4 / 0 / 4
Branches: 75.0% 3 / 0 / 4
optimizer/optimizer.h
Line Branch Exec Source
1 #pragma once
2
3 #include <string>
4 #include <vector>
5 #include <unordered_map>
6 #include <cmath>
7 #include <memory>
8 #include <stdexcept>
9 #include "sparse_tensor.h"
10 #include "ps/base/base_client.h"
11 #include "ps/base/parameters.h"
12
13 using ::ParameterCompressReader;
14 using recstore::EmbeddingTableConfig;
15
16 class Optimizer {
17 protected:
18 std::unordered_map<std::string, SparseTensor*> tensor_map_;
19
20 public:
21 40 virtual ~Optimizer() {
22
2/2
✓ Branch 5 taken 44 times.
✓ Branch 6 taken 40 times.
 84 for (auto& pair : tensor_map_) {
23
1/2
✓ Branch 0 taken 44 times.
✗ Branch 1 not taken.
 44 delete pair.second;
24 }
25 40 }
26
27 virtual void Init(const std::vector<std::string> table_name,
28 const EmbeddingTableConfig& config,
29 BaseKV* base_kv) = 0;
30
31 virtual void Update(std::string table,
32 const ParameterCompressReader* reader,
33 unsigned tid) = 0;
34 virtual void UpdateFlat(
35 std::string table,
36 const base::ConstArray<uint64_t>& keys,
37 const float* grads,
38 int64_t num_rows,
39 int64_t embedding_dim,
40 unsigned tid) = 0;
41 };
42
43 class SGD : public Optimizer {
44 private:
45 float learning_rate_;
46
47 public:
48 30 explicit SGD(float lr = 0.01) : learning_rate_(lr) {}
49
50 void Init(const std::vector<std::string> table_name,
51 const EmbeddingTableConfig& config,
52 BaseKV* base_kv) override;
53 void Update(std::string table,
54 const ParameterCompressReader* reader,
55 unsigned tid) override;
56 void UpdateFlat(std::string table,
57 const base::ConstArray<uint64_t>& keys,
58 const float* grads,
59 int64_t num_rows,
60 int64_t embedding_dim,
61 unsigned tid) override;
62 };
63
64 class AdaGrad : public Optimizer {
65 private:
66 float learning_rate_;
67 float epsilon_;
68
69 public:
70 explicit AdaGrad(float lr = 0.01, float epsilon = 1e-10)
71 : learning_rate_(lr), epsilon_(epsilon) {}
72
73 void Init(const std::vector<std::string> table_name,
74 const EmbeddingTableConfig& config,
75 BaseKV* base_kv) override;
76 void Update(std::string table,
77 const ParameterCompressReader* reader,
78 unsigned tid) override;
79 void UpdateFlat(std::string table,
80 const base::ConstArray<uint64_t>& keys,
81 const float* grads,
82 int64_t num_rows,
83 int64_t embedding_dim,
84 unsigned tid) override;
85 };
86
87 class RowWiseAdaGrad : public Optimizer {
88 private:
89 float learning_rate_;
90 float epsilon_;
91
92 public:
93 4 explicit RowWiseAdaGrad(float lr = 0.01, float epsilon = 1e-10)
94 4 : learning_rate_(lr), epsilon_(epsilon) {}
95
96 void Init(const std::vector<std::string> table_name,
97 const EmbeddingTableConfig& config,
98 BaseKV* base_kv) override;
99 void Update(std::string table,
100 const ParameterCompressReader* reader,
101 unsigned tid) override;
102 void UpdateFlat(std::string table,
103 const base::ConstArray<uint64_t>& keys,
104 const float* grads,
105 int64_t num_rows,
106 int64_t embedding_dim,
107 unsigned tid) override;
108 };
109
110 // Sparse AdamW keeps first/second moments and a persisted step counter in
111 // RecStore. Updates are applied to rows present in the submitted sparse
112 // gradient (the same sparse visibility contract as the existing optimizers).
113 class AdamW : public Optimizer {
114 private:
115 float learning_rate_;
116 float beta1_;
117 float beta2_;
118 float epsilon_;
119 float weight_decay_;
120
121 void UpdateRows(const std::string& table,
122 const uint64_t* keys,
123 const float* grads,
124 int64_t num_rows,
125 int64_t embedding_dim,
126 unsigned tid);
127
128 public:
129 6 explicit AdamW(float lr = 0.001,
130 float beta1 = 0.9,
131 float beta2 = 0.98,
132 float epsilon = 1e-8,
133 float weight_decay = 0.0)
134 12 : learning_rate_(lr),
135 6 beta1_(beta1),
136 6 beta2_(beta2),
137 6 epsilon_(epsilon),
138 6 weight_decay_(weight_decay) {}
139
140 void Init(const std::vector<std::string> table_name,
141 const EmbeddingTableConfig& config,
142 BaseKV* base_kv) override;
143 void Update(std::string table,
144 const ParameterCompressReader* reader,
145 unsigned tid) override;
146 void UpdateFlat(std::string table,
147 const base::ConstArray<uint64_t>& keys,
148 const float* grads,
149 int64_t num_rows,
150 int64_t embedding_dim,
151 unsigned tid) override;
152 };
153
154 std::unique_ptr<Optimizer> CreateOptimizer(const json& config);
155