| 309 | BENCHMARK(BM_RMSProp)->Arg(128 << 10)->Arg(256 << 10); |
| 310 | |
| 311 | static void AddSign(int32 n, Graph** init_g, Graph** train_g) { |
| 312 | TensorShape shape({n}); |
| 313 | { |
| 314 | Graph* g = new Graph(OpRegistry::Global()); |
| 315 | auto var = Var(g, n); |
| 316 | auto m = Var(g, n); |
| 317 | auto zero = Zeros(g, n); |
| 318 | test::graph::Assign(g, var, zero); |
| 319 | test::graph::Assign(g, m, zero); |
| 320 | *init_g = g; |
| 321 | } |
| 322 | { |
| 323 | Graph* g = new Graph(OpRegistry::Global()); |
| 324 | auto var = Var(g, n); |
| 325 | auto m = Var(g, n); |
| 326 | auto lr = Scalar(g, 0.01); |
| 327 | auto alpha = Scalar(g, 0.1); |
| 328 | auto sign_decay = Scalar(g, 0.9); |
| 329 | auto beta = Scalar(g, 0.8); |
| 330 | auto grad = Random(g, n); |
| 331 | test::graph::Multi(g, "ApplyAddSign", |
| 332 | {var, m, lr, alpha, sign_decay, beta, grad}); |
| 333 | *train_g = g; |
| 334 | } |
| 335 | } |
| 336 | |
| 337 | static void BM_AddSign(int iters, int params) { |
| 338 | const int64 tot = static_cast<int64>(iters) * params; |