| 367 | return lr_min, mu |
| 368 | |
| 369 | def caffe2_yellowfin(self, zero_debias, grad_coef, n_dim, n_iter, gpu): |
| 370 | caffe2_res = {} |
| 371 | |
| 372 | alpha = 1.0 |
| 373 | mu = 0.0 |
| 374 | beta = 0.999 |
| 375 | curv_win_width = 20 |
| 376 | epsilon = 1e-6 |
| 377 | |
| 378 | net = core.Net("net") |
| 379 | param_init_net = core.Net("param_init_net") |
| 380 | workspace.ResetWorkspace() |
| 381 | |
| 382 | with core.DeviceScope(core.DeviceOption(caffe2_pb2.CPU)): |
| 383 | iteration = param_init_net.ConstantFill( |
| 384 | [], |
| 385 | "iteration", |
| 386 | shape=[1], |
| 387 | value=0, |
| 388 | dtype=core.DataType.INT64) |
| 389 | iter_mutex = param_init_net.CreateMutex([], ["iteration_mutex"]) |
| 390 | net.AtomicIter([iter_mutex, iteration], [iteration]) |
| 391 | pre_grad = param_init_net.ConstantFill( |
| 392 | [], |
| 393 | "pre_grad", |
| 394 | shape=[n_dim], |
| 395 | value=grad_coef |
| 396 | ) |
| 397 | if gpu: |
| 398 | iteration = net.CopyCPUToGPU( |
| 399 | [iteration], |
| 400 | "iteration_cpu" |
| 401 | ) |
| 402 | iteration_float = net.Cast([iteration], "iteration_float") |
| 403 | grad = net.Mul([pre_grad, iteration_float], "grad", broadcast=True) |
| 404 | w = param_init_net.ConstantFill([], "w", shape=[n_dim], value=0.0) |
| 405 | |
| 406 | # a hack to create an object with __dict__ |
| 407 | param_info = lambda: None |
| 408 | param_info.blob = w |
| 409 | param_info.grad = grad |
| 410 | |
| 411 | optimizer.YellowFinOptimizer( |
| 412 | alpha=alpha, |
| 413 | mu=mu, |
| 414 | beta=beta, |
| 415 | curv_win_width=curv_win_width, |
| 416 | epsilon=epsilon, |
| 417 | zero_debias=zero_debias |
| 418 | )._run( |
| 419 | net, |
| 420 | param_init_net, |
| 421 | param_info |
| 422 | ) |
| 423 | |
| 424 | workspace.RunNetOnce(param_init_net) |
| 425 | workspace.CreateNet(net, overwrite=True) |
| 426 | for i in range(n_iter): |