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
(self, closure=None)
| 357 | current.data.copy_(saved.data) |
| 358 | |
| 359 | def step(self, closure=None): # could add clip option. |
| 360 | """ |
| 361 | If no closure is supplied, :attr:`step` should be called after |
| 362 | ``fp16_optimizer_obj.backward(loss)``. |
| 363 | :attr:`step` updates the fp32 master copy of parameters using the |
| 364 | optimizer supplied to |
| 365 | :class:`FP16_Optimizer`'s constructor, then copies the updated fp32 |
| 366 | params into the fp16 params originally referenced by |
| 367 | :class:`FP16_Optimizer`'s constructor, so the user may immediately run |
| 368 | another forward pass using their model. |
| 369 | |
| 370 | If a closure is supplied, :attr:`step` may be called without a prior |
| 371 | call to :attr:`backward(loss)`. |
| 372 | This control flow is identical to `ordinary Pytorch optimizer use`_ with |
| 373 | closures. However, the user should take care that any ``loss.backward()`` |
| 374 | call within the closure has been replaced by |
| 375 | ``fp16_optimizer_obj.backward(loss)``. |
| 376 | |
| 377 | Args: |
| 378 | closure (optional): Closure that will be supplied to the underlying |
| 379 | optimizer originally passed to :class:`FP16_Optimizer`'s |
| 380 | constructor. closure should call :attr:`zero_grad()` on the |
| 381 | :class:`FP16_Optimizer` object, compute the loss, call |
| 382 | :attr:`backward(loss)`, and return the loss. |
| 383 | |
| 384 | Example with closure:: |
| 385 | |
| 386 | # optimizer is assumed to be an FP16_Optimizer object, previously |
| 387 | # constructed from an existing pytorch optimizer. |
| 388 | for input, target in dataset: |
| 389 | def closure(): |
| 390 | optimizer.zero_grad() |
| 391 | output = model(input) |
| 392 | loss = loss_fn(output, target) |
| 393 | # loss.backward() becomes: |
| 394 | optimizer.backward(loss) |
| 395 | return loss |
| 396 | optimizer.step(closure) |
| 397 | |
| 398 | .. warning:: |
| 399 | Currently, calling :attr:`step` with a closure is not compatible |
| 400 | with dynamic loss scaling. |
| 401 | |
| 402 | .. _`ordinary Pytorch optimizer use`: |
| 403 | http://pytorch.org/docs/master/optim.html#optimizer-step-closure |
| 404 | """ |
| 405 | if self.first_step_call: |
| 406 | self._model_params_to_master_params() |
| 407 | self.first_step_call = True |
| 408 | try: |
| 409 | if self.overflow: |
| 410 | print("OVERFLOW! Skipping step. Reducing loss scale to {}" |
| 411 | .format(self.loss_scale)) |
| 412 | return None |
| 413 | |
| 414 | if closure is not None: |
| 415 | retval = self._step_with_closure(closure) |
| 416 | else: |
no test coverage detected