def zero_grad(self, set_to_none: bool = False):
r"""Sets the gradients of all optimized :class:`torch.Tensor` s to zero.
Arguments:
set_to_none (bool): instead of setting to zero, set the grads to None.
This is will in general have lower memory footprint, and can modestly improve performance.
However, it changes certain behaviors. For example:
1. When the user tries to access a gradient and perform manual ops on it,
a None attribute or a Tensor full of 0s will behave differently.
2. If the user requests ``zero_grad(set_to_none=True)`` followed by a backward pass, ``.grad``\ s
are guaranteed to be None for params that did not receive a gradient.
3. ``torch.optim`` optimizers have a different behavior if the gradient is 0 or None
(in one case it does the step with a gradient of 0 and in the other it skips
the step altogether).
"""
per_device_and_dtype_grads = defaultdict(lambda: defaultdict(list))
for group in self.param_groups:
for p in group['params']:
if p.grad is not None:
if set_to_none:
p.grad = None
else:
if p.grad.grad_fn is not None:
p.grad.detach_()
else:
p.grad.requires_grad_(False)
if p.grad.is_sparse:
p.grad.zero_()
else:
per_device_and_dtype_grads[p.grad.device][p.grad.dtype].append(p.grad)
for _, per_dtype_grads in per_device_and_dtype_grads.items():
for grads in per_dtype_grads.values():
torch._foreach_zero_(grads)
def multi_tensor_scale(
self,
src: Sequence[torch.Tensor],
dst: Sequence[torch.Tensor],
scale: float,
) -> None:
with torch.no_grad(): # type: ignore[no-untyped-call]
# _foreach_zero for long type is not supported in CUDA
if self._enable_foreach and src[0].is_floating_point():
# scale
val = torch._foreach_mul(tuple(src), scale)
# copy tensor
torch._foreach_zero_(tuple(dst))
torch._foreach_add_(tuple(dst), val)
else:
for s, d in zip(src, dst):
d.copy_(s * scale)
def zero_grad(self, set_to_none: bool = False):
per_device_and_dtype_grads = defaultdict(lambda: defaultdict(list))
for group in self.param_groups:
for p in group['params']:
if p.grad is not None:
if set_to_none:
p.grad = None
else:
if p.grad.grad_fn is not None:
p.grad.detach_()
else:
p.grad.requires_grad_(False)
if p.grad.is_sparse:
p.grad.zero_()
else:
per_device_and_dtype_grads[p.grad.device][p.grad.dtype].append(p.grad)
for _, per_dtype_grads in per_device_and_dtype_grads.items():
for grads in per_dtype_grads.values():
torch._foreach_zero_(grads)
