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)
| 380 | self.running_average = dict() |
| 381 | |
| 382 | def apply(self, param_name, param_value, param_grad): |
| 383 | """Performs a single optimization step. |
| 384 | |
| 385 | Args: |
| 386 | param_name(String): the name of the param |
| 387 | param_value(Tensor): param values to be update in-place |
| 388 | grad(Tensor): param gradients; the values may be updated |
| 389 | in this function; cannot use it anymore |
| 390 | """ |
| 391 | assert param_value.shape == param_grad.shape, ("shape mismatch", |
| 392 | param_value.shape, |
| 393 | param_grad.shape) |
| 394 | self.device_check(param_value, self.step_counter, self.lr_value, |
| 395 | self.rho_value, self.epsilon_value, self.decay_value) |
| 396 | |
| 397 | # if self.decay_value != 0: |
| 398 | if self.weight_decay.init_value != 0: |
| 399 | singa.Axpy(self.decay_value.data, param_value.data, param_grad.data) |
| 400 | |
| 401 | if param_name not in self.running_average: |
| 402 | flag = param_value.device.graph_enabled() |
| 403 | param_value.device.EnableGraph(False) |
| 404 | self.running_average[param_name] = tensor.zeros_like(param_value) |
| 405 | param_value.device.EnableGraph(flag) |
| 406 | |
| 407 | # running_average = running_average * rho + param_grad * param_grad * (1 - rho) |
| 408 | # param_value = param_value - lr * param_grad / sqrt(running_average + epsilon) |
| 409 | |
| 410 | self.running_average[param_name] *= self.rho_value |
| 411 | |
| 412 | tmp1 = singa.Square(param_grad.data) |
| 413 | tmp2 = 1.0 - self.rho_value |
| 414 | singa.Axpy(tmp2.data, tmp1, self.running_average[param_name].data) |
| 415 | |
| 416 | minus_lr = 0.0 - self.lr_value |
| 417 | tmp3 = self.running_average[param_name] + self.epsilon_value |
| 418 | tmp3 = singa.Sqrt(tmp3.data) |
| 419 | tmp3 = singa.__div__(param_grad.data, tmp3) |
| 420 | |
| 421 | singa.Axpy(minus_lr.data, tmp3, param_value.data) |
| 422 | |
| 423 | def step(self): |
| 424 | # increment step counter, lr and moment |