mirror of
https://github.com/huggingface/pytorch-image-models.git
synced 2025-06-03 15:01:08 +08:00
7f29a46d44
option: iters-to-accum(iterations to accmulate) Gradient accumulation improves training performance(samples/s). It can reduce the number of parameter sharing between each node. This option can be helpful when network is bottleneck. Signed-off-by: Taeksang Kim <voidbag@puzzle-ai.com>
58 lines
1.7 KiB
Python
58 lines
1.7 KiB
Python
""" CUDA / AMP utils
|
|
|
|
Hacked together by / Copyright 2020 Ross Wightman
|
|
"""
|
|
import torch
|
|
|
|
try:
|
|
from apex import amp
|
|
has_apex = True
|
|
except ImportError:
|
|
amp = None
|
|
has_apex = False
|
|
|
|
from .clip_grad import dispatch_clip_grad
|
|
|
|
|
|
class ApexScaler:
|
|
state_dict_key = "amp"
|
|
|
|
def __call__(self, loss, optimizer, clip_grad=None, clip_mode='norm', parameters=None, create_graph=False, need_step=True):
|
|
with amp.scale_loss(loss, optimizer) as scaled_loss:
|
|
scaled_loss.backward(create_graph=create_graph)
|
|
if clip_grad is not None:
|
|
dispatch_clip_grad(amp.master_params(optimizer), clip_grad, mode=clip_mode)
|
|
if need_step:
|
|
optimizer.step()
|
|
|
|
def state_dict(self):
|
|
if 'state_dict' in amp.__dict__:
|
|
return amp.state_dict()
|
|
|
|
def load_state_dict(self, state_dict):
|
|
if 'load_state_dict' in amp.__dict__:
|
|
amp.load_state_dict(state_dict)
|
|
|
|
|
|
class NativeScaler:
|
|
state_dict_key = "amp_scaler"
|
|
|
|
def __init__(self):
|
|
self._scaler = torch.cuda.amp.GradScaler()
|
|
|
|
def __call__(self, loss, optimizer, clip_grad=None, clip_mode='norm', parameters=None, create_graph=False, need_step=True):
|
|
self._scaler.scale(loss).backward(create_graph=create_graph)
|
|
if clip_grad is not None:
|
|
assert parameters is not None
|
|
self._scaler.unscale_(optimizer) # unscale the gradients of optimizer's assigned params in-place
|
|
dispatch_clip_grad(parameters, clip_grad, mode=clip_mode)
|
|
if need_step:
|
|
self._scaler.step(optimizer)
|
|
self._scaler.update()
|
|
|
|
def state_dict(self):
|
|
return self._scaler.state_dict()
|
|
|
|
def load_state_dict(self, state_dict):
|
|
self._scaler.load_state_dict(state_dict)
|