| 269 | BENCHMARK(BM_Adam)->ArgPair(256 << 5, 1)->ArgPair(256 << 16, 1); |
| 270 | |
| 271 | static void RMSProp(int32 n, Graph** init_g, Graph** train_g) { |
| 272 | TensorShape shape({n}); |
| 273 | { |
| 274 | Graph* g = new Graph(OpRegistry::Global()); |
| 275 | auto var = Var(g, n); |
| 276 | auto ms = Var(g, n); |
| 277 | auto mom = Var(g, n); |
| 278 | auto zero = Zeros(g, n); |
| 279 | test::graph::Assign(g, var, zero); |
| 280 | test::graph::Assign(g, ms, zero); |
| 281 | test::graph::Assign(g, mom, zero); |
| 282 | *init_g = g; |
| 283 | } |
| 284 | { |
| 285 | Graph* g = new Graph(OpRegistry::Global()); |
| 286 | auto var = Var(g, n); |
| 287 | auto ms = Var(g, n); |
| 288 | auto mom = Var(g, n); |
| 289 | auto lr = Scalar(g, 0.01); |
| 290 | auto rho = Scalar(g, 0.9); |
| 291 | auto momentum = Scalar(g, 0.9); |
| 292 | auto epsilon = Scalar(g, 1e-8); |
| 293 | auto grad = Random(g, n); |
| 294 | test::graph::Multi(g, "ApplyRMSProp", |
| 295 | {var, ms, mom, lr, rho, momentum, epsilon, grad}); |
| 296 | *train_g = g; |
| 297 | } |
| 298 | } |
| 299 | |
| 300 | static void BM_RMSProp(int iters, int params) { |
| 301 | const int64 tot = static_cast<int64>(iters) * params; |