Performs a single optimization step. Args: param_name(String): the name of the param param_value(Tensor): param values to be update in-place grad(Tensor): param gradients; the values may be updated in this function; can
(self, param_name, param_value, param_grad)
| 478 | self.history = dict() |
| 479 | |
| 480 | def apply(self, param_name, param_value, param_grad): |
| 481 | """Performs a single optimization step. |
| 482 | |
| 483 | Args: |
| 484 | param_name(String): the name of the param |
| 485 | param_value(Tensor): param values to be update in-place |
| 486 | grad(Tensor): param gradients; the values may be updated |
| 487 | in this function; cannot use it anymore |
| 488 | """ |
| 489 | assert param_value.shape == param_grad.shape, ("shape mismatch", |
| 490 | param_value.shape, |
| 491 | param_grad.shape) |
| 492 | self.device_check(param_value, self.step_counter, self.lr_value, |
| 493 | self.epsilon_value, self.decay_value) |
| 494 | |
| 495 | # if self.decay_value != 0: |
| 496 | if self.weight_decay.init_value != 0: |
| 497 | singa.Axpy(self.decay_value.data, param_value.data, param_grad.data) |
| 498 | |
| 499 | if param_name not in self.history: |
| 500 | flag = param_value.device.graph_enabled() |
| 501 | param_value.device.EnableGraph(False) |
| 502 | self.history[param_name] = tensor.zeros_like(param_value) |
| 503 | param_value.device.EnableGraph(flag) |
| 504 | |
| 505 | # history = history + param_grad * param_grad |
| 506 | # param_value = param_value - lr * param_grad / sqrt(history + epsilon) |
| 507 | |
| 508 | tmp = self.history[param_name].data |
| 509 | tmp += singa.Square(param_grad.data) |
| 510 | |
| 511 | minus_lr = 0.0 - self.lr_value |
| 512 | tmp = self.history[param_name] + self.epsilon_value |
| 513 | tmp = singa.Sqrt(tmp.data) |
| 514 | tmp = singa.__div__(param_grad.data, tmp) |
| 515 | singa.Axpy(minus_lr.data, tmp, param_value.data) |
| 516 | |
| 517 | def step(self): |
| 518 | # increment step counter, lr and moment |