| .. _amp-examples: |
| |
| CUDA Automatic Mixed Precision examples |
| ======================================= |
| |
| .. currentmodule:: torch.cuda.amp |
| |
| Ordinarily, "automatic mixed precision training" means training with |
| :class:`torch.autocast` and :class:`torch.cuda.amp.GradScaler` together. |
| |
| Instances of :class:`torch.autocast` enable autocasting for chosen regions. |
| Autocasting automatically chooses the precision for GPU operations to improve performance |
| while maintaining accuracy. |
| |
| Instances of :class:`torch.cuda.amp.GradScaler` help perform the steps of |
| gradient scaling conveniently. Gradient scaling improves convergence for networks with ``float16`` |
| gradients by minimizing gradient underflow, as explained :ref:`here<gradient-scaling>`. |
| |
| :class:`torch.autocast` and :class:`torch.cuda.amp.GradScaler` are modular. |
| In the samples below, each is used as its individual documentation suggests. |
| |
| (Samples here are illustrative. See the |
| `Automatic Mixed Precision recipe <https://pytorch.org/tutorials/recipes/recipes/amp_recipe.html>`_ |
| for a runnable walkthrough.) |
| |
| .. contents:: :local: |
| |
| Typical Mixed Precision Training |
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ |
| |
| :: |
| |
| # Creates model and optimizer in default precision |
| model = Net().cuda() |
| optimizer = optim.SGD(model.parameters(), ...) |
| |
| # Creates a GradScaler once at the beginning of training. |
| scaler = GradScaler() |
| |
| for epoch in epochs: |
| for input, target in data: |
| optimizer.zero_grad() |
| |
| # Runs the forward pass with autocasting. |
| with autocast(device_type='cuda', dtype=torch.float16): |
| output = model(input) |
| loss = loss_fn(output, target) |
| |
| # Scales loss. Calls backward() on scaled loss to create scaled gradients. |
| # Backward passes under autocast are not recommended. |
| # Backward ops run in the same dtype autocast chose for corresponding forward ops. |
| scaler.scale(loss).backward() |
| |
| # scaler.step() first unscales the gradients of the optimizer's assigned params. |
| # If these gradients do not contain infs or NaNs, optimizer.step() is then called, |
| # otherwise, optimizer.step() is skipped. |
| scaler.step(optimizer) |
| |
| # Updates the scale for next iteration. |
| scaler.update() |
| |
| .. _working-with-unscaled-gradients: |
| |
| Working with Unscaled Gradients |
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ |
| |
| All gradients produced by ``scaler.scale(loss).backward()`` are scaled. If you wish to modify or inspect |
| the parameters' ``.grad`` attributes between ``backward()`` and ``scaler.step(optimizer)``, you should |
| unscale them first. For example, gradient clipping manipulates a set of gradients such that their global norm |
| (see :func:`torch.nn.utils.clip_grad_norm_`) or maximum magnitude (see :func:`torch.nn.utils.clip_grad_value_`) |
| is :math:`<=` some user-imposed threshold. If you attempted to clip *without* unscaling, the gradients' norm/maximum |
| magnitude would also be scaled, so your requested threshold (which was meant to be the threshold for *unscaled* |
| gradients) would be invalid. |
| |
| ``scaler.unscale_(optimizer)`` unscales gradients held by ``optimizer``'s assigned parameters. |
| If your model or models contain other parameters that were assigned to another optimizer |
| (say ``optimizer2``), you may call ``scaler.unscale_(optimizer2)`` separately to unscale those |
| parameters' gradients as well. |
| |
| Gradient clipping |
| ----------------- |
| |
| Calling ``scaler.unscale_(optimizer)`` before clipping enables you to clip unscaled gradients as usual:: |
| |
| scaler = GradScaler() |
| |
| for epoch in epochs: |
| for input, target in data: |
| optimizer.zero_grad() |
| with autocast(device_type='cuda', dtype=torch.float16): |
| output = model(input) |
| loss = loss_fn(output, target) |
| scaler.scale(loss).backward() |
| |
| # Unscales the gradients of optimizer's assigned params in-place |
| scaler.unscale_(optimizer) |
| |
| # Since the gradients of optimizer's assigned params are unscaled, clips as usual: |
| torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm) |
| |
| # optimizer's gradients are already unscaled, so scaler.step does not unscale them, |
| # although it still skips optimizer.step() if the gradients contain infs or NaNs. |
| scaler.step(optimizer) |
| |
| # Updates the scale for next iteration. |
| scaler.update() |
| |
| ``scaler`` records that ``scaler.unscale_(optimizer)`` was already called for this optimizer |
| this iteration, so ``scaler.step(optimizer)`` knows not to redundantly unscale gradients before |
| (internally) calling ``optimizer.step()``. |
| |
| .. currentmodule:: torch.cuda.amp.GradScaler |
| |
| .. warning:: |
| :meth:`unscale_<unscale_>` should only be called once per optimizer per :meth:`step<step>` call, |
| and only after all gradients for that optimizer's assigned parameters have been accumulated. |
| Calling :meth:`unscale_<unscale_>` twice for a given optimizer between each :meth:`step<step>` triggers a RuntimeError. |
| |
| |
| Working with Scaled Gradients |
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ |
| |
| Gradient accumulation |
| --------------------- |
| |
| Gradient accumulation adds gradients over an effective batch of size ``batch_per_iter * iters_to_accumulate`` |
| (``* num_procs`` if distributed). The scale should be calibrated for the effective batch, which means inf/NaN checking, |
| step skipping if inf/NaN grads are found, and scale updates should occur at effective-batch granularity. |
| Also, grads should remain scaled, and the scale factor should remain constant, while grads for a given effective |
| batch are accumulated. If grads are unscaled (or the scale factor changes) before accumulation is complete, |
| the next backward pass will add scaled grads to unscaled grads (or grads scaled by a different factor) |
| after which it's impossible to recover the accumulated unscaled grads :meth:`step<step>` must apply. |
| |
| Therefore, if you want to :meth:`unscale_<unscale_>` grads (e.g., to allow clipping unscaled grads), |
| call :meth:`unscale_<unscale_>` just before :meth:`step<step>`, after all (scaled) grads for the upcoming |
| :meth:`step<step>` have been accumulated. Also, only call :meth:`update<update>` at the end of iterations |
| where you called :meth:`step<step>` for a full effective batch:: |
| |
| scaler = GradScaler() |
| |
| for epoch in epochs: |
| for i, (input, target) in enumerate(data): |
| with autocast(device_type='cuda', dtype=torch.float16): |
| output = model(input) |
| loss = loss_fn(output, target) |
| loss = loss / iters_to_accumulate |
| |
| # Accumulates scaled gradients. |
| scaler.scale(loss).backward() |
| |
| if (i + 1) % iters_to_accumulate == 0: |
| # may unscale_ here if desired (e.g., to allow clipping unscaled gradients) |
| |
| scaler.step(optimizer) |
| scaler.update() |
| optimizer.zero_grad() |
| |
| .. currentmodule:: torch.cuda.amp |
| |
| Gradient penalty |
| ---------------- |
| |
| A gradient penalty implementation commonly creates gradients using |
| :func:`torch.autograd.grad`, combines them to create the penalty value, |
| and adds the penalty value to the loss. |
| |
| Here's an ordinary example of an L2 penalty without gradient scaling or autocasting:: |
| |
| for epoch in epochs: |
| for input, target in data: |
| optimizer.zero_grad() |
| output = model(input) |
| loss = loss_fn(output, target) |
| |
| # Creates gradients |
| grad_params = torch.autograd.grad(outputs=loss, |
| inputs=model.parameters(), |
| create_graph=True) |
| |
| # Computes the penalty term and adds it to the loss |
| grad_norm = 0 |
| for grad in grad_params: |
| grad_norm += grad.pow(2).sum() |
| grad_norm = grad_norm.sqrt() |
| loss = loss + grad_norm |
| |
| loss.backward() |
| |
| # clip gradients here, if desired |
| |
| optimizer.step() |
| |
| To implement a gradient penalty *with* gradient scaling, the ``outputs`` Tensor(s) |
| passed to :func:`torch.autograd.grad` should be scaled. The resulting gradients |
| will therefore be scaled, and should be unscaled before being combined to create the |
| penalty value. |
| |
| Also, the penalty term computation is part of the forward pass, and therefore should be |
| inside an :class:`autocast` context. |
| |
| Here's how that looks for the same L2 penalty:: |
| |
| scaler = GradScaler() |
| |
| for epoch in epochs: |
| for input, target in data: |
| optimizer.zero_grad() |
| with autocast(device_type='cuda', dtype=torch.float16): |
| output = model(input) |
| loss = loss_fn(output, target) |
| |
| # Scales the loss for autograd.grad's backward pass, producing scaled_grad_params |
| scaled_grad_params = torch.autograd.grad(outputs=scaler.scale(loss), |
| inputs=model.parameters(), |
| create_graph=True) |
| |
| # Creates unscaled grad_params before computing the penalty. scaled_grad_params are |
| # not owned by any optimizer, so ordinary division is used instead of scaler.unscale_: |
| inv_scale = 1./scaler.get_scale() |
| grad_params = [p * inv_scale for p in scaled_grad_params] |
| |
| # Computes the penalty term and adds it to the loss |
| with autocast(device_type='cuda', dtype=torch.float16): |
| grad_norm = 0 |
| for grad in grad_params: |
| grad_norm += grad.pow(2).sum() |
| grad_norm = grad_norm.sqrt() |
| loss = loss + grad_norm |
| |
| # Applies scaling to the backward call as usual. |
| # Accumulates leaf gradients that are correctly scaled. |
| scaler.scale(loss).backward() |
| |
| # may unscale_ here if desired (e.g., to allow clipping unscaled gradients) |
| |
| # step() and update() proceed as usual. |
| scaler.step(optimizer) |
| scaler.update() |
| |
| |
| Working with Multiple Models, Losses, and Optimizers |
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ |
| |
| .. currentmodule:: torch.cuda.amp.GradScaler |
| |
| If your network has multiple losses, you must call :meth:`scaler.scale<scale>` on each of them individually. |
| If your network has multiple optimizers, you may call :meth:`scaler.unscale_<unscale_>` on any of them individually, |
| and you must call :meth:`scaler.step<step>` on each of them individually. |
| |
| However, :meth:`scaler.update<update>` should only be called once, |
| after all optimizers used this iteration have been stepped:: |
| |
| scaler = torch.cuda.amp.GradScaler() |
| |
| for epoch in epochs: |
| for input, target in data: |
| optimizer0.zero_grad() |
| optimizer1.zero_grad() |
| with autocast(device_type='cuda', dtype=torch.float16): |
| output0 = model0(input) |
| output1 = model1(input) |
| loss0 = loss_fn(2 * output0 + 3 * output1, target) |
| loss1 = loss_fn(3 * output0 - 5 * output1, target) |
| |
| # (retain_graph here is unrelated to amp, it's present because in this |
| # example, both backward() calls share some sections of graph.) |
| scaler.scale(loss0).backward(retain_graph=True) |
| scaler.scale(loss1).backward() |
| |
| # You can choose which optimizers receive explicit unscaling, if you |
| # want to inspect or modify the gradients of the params they own. |
| scaler.unscale_(optimizer0) |
| |
| scaler.step(optimizer0) |
| scaler.step(optimizer1) |
| |
| scaler.update() |
| |
| Each optimizer checks its gradients for infs/NaNs and makes an independent decision |
| whether or not to skip the step. This may result in one optimizer skipping the step |
| while the other one does not. Since step skipping occurs rarely (every several hundred iterations) |
| this should not impede convergence. If you observe poor convergence after adding gradient scaling |
| to a multiple-optimizer model, please report a bug. |
| |
| .. currentmodule:: torch.cuda.amp |
| |
| .. _amp-multigpu: |
| |
| Working with Multiple GPUs |
| ^^^^^^^^^^^^^^^^^^^^^^^^^^ |
| |
| The issues described here only affect :class:`autocast`. :class:`GradScaler`\ 's usage is unchanged. |
| |
| .. _amp-dataparallel: |
| |
| DataParallel in a single process |
| -------------------------------- |
| |
| Even if :class:`torch.nn.DataParallel` spawns threads to run the forward pass on each device. |
| The autocast state is propagated in each one and the following will work:: |
| |
| model = MyModel() |
| dp_model = nn.DataParallel(model) |
| |
| # Sets autocast in the main thread |
| with autocast(device_type='cuda', dtype=torch.float16): |
| # dp_model's internal threads will autocast. |
| output = dp_model(input) |
| # loss_fn also autocast |
| loss = loss_fn(output) |
| |
| DistributedDataParallel, one GPU per process |
| -------------------------------------------- |
| |
| :class:`torch.nn.parallel.DistributedDataParallel`'s documentation recommends one GPU per process for best |
| performance. In this case, ``DistributedDataParallel`` does not spawn threads internally, |
| so usages of :class:`autocast` and :class:`GradScaler` are not affected. |
| |
| DistributedDataParallel, multiple GPUs per process |
| -------------------------------------------------- |
| |
| Here :class:`torch.nn.parallel.DistributedDataParallel` may spawn a side thread to run the forward pass on each |
| device, like :class:`torch.nn.DataParallel`. :ref:`The fix is the same<amp-dataparallel>`: |
| apply autocast as part of your model's ``forward`` method to ensure it's enabled in side threads. |
| |
| .. _amp-custom-examples: |
| |
| Autocast and Custom Autograd Functions |
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ |
| |
| If your network uses :ref:`custom autograd functions<extending-autograd>` |
| (subclasses of :class:`torch.autograd.Function`), changes are required for |
| autocast compatibility if any function |
| |
| * takes multiple floating-point Tensor inputs, |
| * wraps any autocastable op (see the :ref:`Autocast Op Reference<autocast-op-reference>`), or |
| * requires a particular ``dtype`` (for example, if it wraps |
| `CUDA extensions <https://pytorch.org/tutorials/advanced/cpp_extension.html>`_ |
| that were only compiled for ``dtype``). |
| |
| In all cases, if you're importing the function and can't alter its definition, a safe fallback |
| is to disable autocast and force execution in ``float32`` ( or ``dtype``) at any points of use where errors occur:: |
| |
| with autocast(device_type='cuda', dtype=torch.float16): |
| ... |
| with autocast(device_type='cuda', dtype=torch.float16, enabled=False): |
| output = imported_function(input1.float(), input2.float()) |
| |
| If you're the function's author (or can alter its definition) a better solution is to use the |
| :func:`torch.cuda.amp.custom_fwd` and :func:`torch.cuda.amp.custom_bwd` decorators as shown in |
| the relevant case below. |
| |
| Functions with multiple inputs or autocastable ops |
| -------------------------------------------------- |
| |
| Apply :func:`custom_fwd<custom_fwd>` and :func:`custom_bwd<custom_bwd>` (with no arguments) to ``forward`` and |
| ``backward`` respectively. These ensure ``forward`` executes with the current autocast state and ``backward`` |
| executes with the same autocast state as ``forward`` (which can prevent type mismatch errors):: |
| |
| class MyMM(torch.autograd.Function): |
| @staticmethod |
| @custom_fwd |
| def forward(ctx, a, b): |
| ctx.save_for_backward(a, b) |
| return a.mm(b) |
| @staticmethod |
| @custom_bwd |
| def backward(ctx, grad): |
| a, b = ctx.saved_tensors |
| return grad.mm(b.t()), a.t().mm(grad) |
| |
| Now ``MyMM`` can be invoked anywhere, without disabling autocast or manually casting inputs:: |
| |
| mymm = MyMM.apply |
| |
| with autocast(device_type='cuda', dtype=torch.float16): |
| output = mymm(input1, input2) |
| |
| Functions that need a particular ``dtype`` |
| ------------------------------------------ |
| |
| Consider a custom function that requires ``torch.float32`` inputs. |
| Apply :func:`custom_fwd(cast_inputs=torch.float32)<custom_fwd>` to ``forward`` |
| and :func:`custom_bwd<custom_bwd>` (with no arguments) to ``backward``. |
| If ``forward`` runs in an autocast-enabled region, the decorators cast floating-point CUDA Tensor |
| inputs to ``float32``, and locally disable autocast during ``forward`` and ``backward``:: |
| |
| class MyFloat32Func(torch.autograd.Function): |
| @staticmethod |
| @custom_fwd(cast_inputs=torch.float32) |
| def forward(ctx, input): |
| ctx.save_for_backward(input) |
| ... |
| return fwd_output |
| @staticmethod |
| @custom_bwd |
| def backward(ctx, grad): |
| ... |
| |
| Now ``MyFloat32Func`` can be invoked anywhere, without manually disabling autocast or casting inputs:: |
| |
| func = MyFloat32Func.apply |
| |
| with autocast(device_type='cuda', dtype=torch.float16): |
| # func will run in float32, regardless of the surrounding autocast state |
| output = func(input) |
