If no closure is supplied, :attr:`step` should be called after ``fp16_optimizer_obj.backward(loss)``. :attr:`step` updates the fp32 master copy of parameters using the optimizer supplied to :class:`FP16_Optimizer`'s constructor, then copies the updated fp32 params i
(self, closure=None)
| 409 | current.data.copy_(saved.data) |
| 410 | |
| 411 | def step(self, closure=None): # could add clip option. |
| 412 | """ |
| 413 | If no closure is supplied, :attr:`step` should be called after |
| 414 | ``fp16_optimizer_obj.backward(loss)``. |
| 415 | :attr:`step` updates the fp32 master copy of parameters using the optimizer supplied to |
| 416 | :class:`FP16_Optimizer`'s constructor, then copies the updated fp32 params into the fp16 params |
| 417 | originally referenced by :class:`FP16_Optimizer`'s constructor, so the user may immediately run |
| 418 | another forward pass using their model. |
| 419 | |
| 420 | If a closure is supplied, :attr:`step` may be called without a prior call to |
| 421 | :attr:`backward(loss)`. |
| 422 | This control flow is identical to `ordinary Pytorch optimizer use`_ with closures. |
| 423 | However, the user should take care that any ``loss.backward()`` call within the closure |
| 424 | has been replaced by ``fp16_optimizer_obj.backward(loss)``. |
| 425 | |
| 426 | Args: |
| 427 | closure (optional): Closure that will be supplied to the underlying optimizer originally passed to :class:`FP16_Optimizer`'s constructor. closure should call :attr:`zero_grad()` on the :class:`FP16_Optimizer` object, compute the loss, call :attr:`backward(loss)`, and return the loss. |
| 428 | |
| 429 | Example with closure:: |
| 430 | |
| 431 | # optimizer is assumed to be an FP16_Optimizer object, previously constructed from an |
| 432 | # existing pytorch optimizer. |
| 433 | for input, target in dataset: |
| 434 | def closure(): |
| 435 | optimizer.zero_grad() |
| 436 | output = model(input) |
| 437 | loss = loss_fn(output, target) |
| 438 | # loss.backward() becomes: |
| 439 | optimizer.backward(loss) |
| 440 | return loss |
| 441 | optimizer.step(closure) |
| 442 | |
| 443 | .. warning:: |
| 444 | Currently, calling :attr:`step` with a closure is not compatible with dynamic loss scaling. |
| 445 | |
| 446 | .. _`ordinary Pytorch optimizer use`: |
| 447 | http://pytorch.org/docs/master/optim.html#optimizer-step-closure |
| 448 | """ |
| 449 | |
| 450 | scale = self.loss_scaler.loss_scale |
| 451 | self._update_scale(self.overflow) |
| 452 | |
| 453 | if self.overflow: |
| 454 | self.maybe_print("OVERFLOW! Skipping step. Attempted loss scale: {}, reducing to {}" |
| 455 | .format(scale, self.loss_scale)) |
| 456 | return |
| 457 | |
| 458 | if closure is not None: |
| 459 | retval = self._step_with_closure(closure) |
| 460 | else: |
| 461 | retval = self.optimizer.step() |
| 462 | |
| 463 | self._master_params_to_model_params() |
| 464 | |
| 465 | return retval |
| 466 | |
| 467 | def _step_with_closure(self, closure): |
| 468 | def wrapped_closure(): |
no test coverage detected