Class FP16_Optimizer

PATH/core/fp16/opt.py:22–585

:class:`FP16_Optimizer` is designed to wrap an existing PyTorch optimizer, and manage static or dynamic loss scaling and master weights in a manner transparent to the user. For standard use, only two lines must be changed: creating the :class:`FP16_Optimizer` instance, and chang

# TODO:  Update overflow check + downscale to use Carl's fused kernel.
class FP16_Optimizer(object):
    """
    :class:`FP16_Optimizer` is designed to wrap an existing PyTorch optimizer, 
    and manage static or dynamic loss scaling and master weights in a manner 
    transparent to the user. For standard use, only two lines must be changed:  
    creating the :class:`FP16_Optimizer` instance, and changing the call to 
    ``backward``.

    Example::

        model = torch.nn.Linear(D_in, D_out).cuda().half()
        optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
        # Name the FP16_Optimizer instance to replace the existing optimizer
        # (recommended but not required):
        optimizer = FP16_Optimizer(optimizer, static_loss_scale = 128.0)
        ...
        # loss.backward() becomes:
        optimizer.backward(loss)
        ...

    Example with dynamic loss scaling::

        ...
        optimizer = FP16_Optimizer(optimizer, dynamic_loss_scale=True)

    Args:
        init_optimizer (torch.optim.optimizer):  Existing optimizer created with
            the parameters to optimize.  Internally, :class:`FP16_Optimizer` 
            replaces the passed optimizer's fp16 parameters, if any, with fp32 
            master parameters copied from the original ones.  
            :class:`FP16_Optimizer` also stores references to the original fp16 
            parameters, and updates these fp16 parameters from the master fp32 
            copy at the end of each :attr:`step`.  
        static_loss_scale (float, optional, default=1.0):  Loss scale used 
            internally to scale gradients computed by the model.  Any fp16 
            gradients will be copied to fp32, then downscaled before being 
            applied to the fp32 master params, so ``static_loss_scale`` should 
            not affect learning rate.
        dynamic_loss_scale (bool, optional, default=False):  Use dynamic loss 
            scaling.  If True, this will override any ``static_loss_scale`` 
            option.
        verbose (bool, optional, default=False):  By default, FP16_Optimizer's 
            constructor prints out the parameters and parameter groups it is 
            ingesting, as a sanity check.  If this becomes annoying (e.g. for 
            large models), it can be disabled by passing ``verbose=False``.  
            ``verbose=False`` will not disable printing when the loss scale is 
            readjusted during dynamic loss scaling.

    ``init_optimizer`` is expected to have been constructed in the ordinary way.  
    It is recommended (although not required) that the newly constructed 
    :class:`FP16_Optimizer` instance be 
    named to replace ``init_optimizer``, for two reasons:  
    First, it means that references to the same name later in the file will not 
        have to change.  
    Second, :class:`FP16_Optimizer` reserves the right (as an implementation 
        detail) to modify ``init_optimizer``.  
    If you do choose a unique name for the new :class:`FP16_Optimizer` instance, 
    you should only work with this new instance, because the preexisting