class torch.cuda.amp.GradScaler(init_scale=65536.0, growth_factor=2.0, backoff_factor=0.5, growth_interval=2000, enabled=True) [source]
 
 
get_backoff_factor() [source]
 
Returns a Python float containing the scale backoff factor. 
  
get_growth_factor() [source]
 
Returns a Python float containing the scale growth factor. 
  
get_growth_interval() [source]
 
Returns a Python int containing the growth interval. 
  
get_scale() [source]
 
Returns a Python float containing the current scale, or 1.0 if scaling is disabled.  Warning get_scale() incurs a CPU-GPU sync.  
  
is_enabled() [source]
 
Returns a bool indicating whether this instance is enabled. 
  
load_state_dict(state_dict) [source]
 
Loads the scaler state. If this instance is disabled, load_state_dict() is a no-op.  Parameters 
state_dict (dict) – scaler state. Should be an object returned from a call to state_dict().   
  
scale(outputs) [source]
 
Multiplies (‘scales’) a tensor or list of tensors by the scale factor. Returns scaled outputs. If this instance of GradScaler is not enabled, outputs are returned unmodified.  Parameters 
outputs (Tensor or iterable of Tensors) – Outputs to scale.   
  
set_backoff_factor(new_factor) [source]
 
 Parameters 
new_scale (float) – Value to use as the new scale backoff factor.   
  
set_growth_factor(new_factor) [source]
 
 Parameters 
new_scale (float) – Value to use as the new scale growth factor.   
  
set_growth_interval(new_interval) [source]
 
 Parameters 
new_interval (int) – Value to use as the new growth interval.   
  
state_dict() [source]
 
Returns the state of the scaler as a dict. It contains five entries:  
"scale" - a Python float containing the current scale 
"growth_factor" - a Python float containing the current growth factor 
"backoff_factor" - a Python float containing the current backoff factor 
"growth_interval" - a Python int containing the current growth interval 
"_growth_tracker" - a Python int containing the number of recent consecutive unskipped steps.  If this instance is not enabled, returns an empty dict.  Note If you wish to checkpoint the scaler’s state after a particular iteration, state_dict() should be called after update().  
  
step(optimizer, *args, **kwargs) [source]
 
step() carries out the following two operations:  Internally invokes unscale_(optimizer) (unless unscale_() was explicitly called for optimizer earlier in the iteration). As part of the unscale_(), gradients are checked for infs/NaNs. If no inf/NaN gradients are found, invokes optimizer.step() using the unscaled gradients. Otherwise, optimizer.step() is skipped to avoid corrupting the params.  *args and **kwargs are forwarded to optimizer.step(). Returns the return value of optimizer.step(*args, **kwargs).  Parameters 
 
optimizer (torch.optim.Optimizer) – Optimizer that applies the gradients. 
args – Any arguments. 
kwargs – Any keyword arguments.     Warning Closure use is not currently supported.  
  
unscale_(optimizer) [source]
 
Divides (“unscales”) the optimizer’s gradient tensors by the scale factor. unscale_() is optional, serving cases where you need to modify or inspect gradients between the backward pass(es) and step(). If unscale_() is not called explicitly, gradients will be unscaled automatically during step(). Simple example, using unscale_() to enable clipping of unscaled gradients: ...
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
scaler.step(optimizer)
scaler.update()
  Parameters 
optimizer (torch.optim.Optimizer) – Optimizer that owns the gradients to be unscaled.    Note unscale_() does not incur a CPU-GPU sync.   Warning unscale_() should only be called once per optimizer per step() call, and only after all gradients for that optimizer’s assigned parameters have been accumulated. Calling unscale_() twice for a given optimizer between each step() triggers a RuntimeError.   Warning unscale_() may unscale sparse gradients out of place, replacing the .grad attribute.  
  
update(new_scale=None) [source]
 
Updates the scale factor. If any optimizer steps were skipped the scale is multiplied by backoff_factor to reduce it. If growth_interval unskipped iterations occurred consecutively, the scale is multiplied by growth_factor to increase it. Passing new_scale sets the scale directly.  Parameters 
new_scale (float or torch.cuda.FloatTensor, optional, default=None) – New scale factor.    Warning update() should only be called at the end of the iteration, after scaler.step(optimizer) has been invoked for all optimizers used this iteration.