| 220 | BENCHMARK(BM_Momentum)->Arg(128 << 10)->Arg(256 << 10); |
| 221 | |
| 222 | static void Adam(int32 n, Graph** init_g, Graph** train_g) { |
| 223 | TensorShape shape({n}); |
| 224 | { |
| 225 | Graph* g = new Graph(OpRegistry::Global()); |
| 226 | auto var = Var(g, n); |
| 227 | auto m = Var(g, n); |
| 228 | auto v = Var(g, n); |
| 229 | auto zero = Zeros(g, n); |
| 230 | test::graph::Assign(g, var, zero); |
| 231 | test::graph::Assign(g, m, zero); |
| 232 | test::graph::Assign(g, v, zero); |
| 233 | *init_g = g; |
| 234 | } |
| 235 | { |
| 236 | Graph* g = new Graph(OpRegistry::Global()); |
| 237 | auto var = Var(g, n); |
| 238 | auto m = Var(g, n); |
| 239 | auto v = Var(g, n); |
| 240 | auto beta1_power = Scalar(g, 0.9); |
| 241 | auto beta2_power = Scalar(g, 0.99); |
| 242 | auto lr = Scalar(g, 0.01); |
| 243 | auto beta1 = Scalar(g, 0.9); |
| 244 | auto beta2 = Scalar(g, 0.99); |
| 245 | auto epsilon = Scalar(g, 1e-8); |
| 246 | auto grad = Random(g, n); |
| 247 | test::graph::Multi( |
| 248 | g, "ApplyAdam", |
| 249 | {var, m, v, beta1_power, beta2_power, lr, beta1, beta2, epsilon, grad}); |
| 250 | *train_g = g; |
| 251 | } |
| 252 | } |
| 253 | |
| 254 | static void BM_Adam(int iters, int params, int is_multi_threaded) { |
| 255 | const int64 tot = static_cast<int64>(iters) * params; |