| 186 | ->ArgPair(128, 128 << 10); |
| 187 | |
| 188 | static void Momentum(int32 n, Graph** init_g, Graph** train_g) { |
| 189 | TensorShape shape({n}); |
| 190 | { |
| 191 | Graph* g = new Graph(OpRegistry::Global()); |
| 192 | auto var = Var(g, n); |
| 193 | auto accum = Var(g, n); |
| 194 | auto zero = Zeros(g, n); |
| 195 | test::graph::Assign(g, var, zero); |
| 196 | test::graph::Assign(g, accum, zero); |
| 197 | *init_g = g; |
| 198 | } |
| 199 | { |
| 200 | Graph* g = new Graph(OpRegistry::Global()); |
| 201 | auto var = Var(g, n); |
| 202 | auto accum = Var(g, n); |
| 203 | auto lr = Scalar(g, 0.01); |
| 204 | auto grad = Random(g, n); |
| 205 | auto mom = Scalar(g, 0.01); |
| 206 | test::graph::Multi(g, "ApplyMomentum", {var, accum, lr, grad, mom}); |
| 207 | *train_g = g; |
| 208 | } |
| 209 | } |
| 210 | |
| 211 | static void BM_Momentum(int iters, int params) { |
| 212 | const int64 tot = static_cast<int64>(iters) * params; |