mirror of https://github.com/hpcaitech/ColossalAI
44 lines
1.9 KiB
Python
44 lines
1.9 KiB
Python
import torch.nn as nn
|
|
from typing import List
|
|
from colossalai.engine import BaseGradientHandler
|
|
from typing import Iterable
|
|
from torch.optim import Optimizer
|
|
from torch.optim.lr_scheduler import _LRScheduler
|
|
from ._gradient_accumulation import GradAccumDataloader, GradAccumOptimizer, GradAccumLrSchedulerByStep, GradAccumGradientHandler
|
|
|
|
|
|
def accumulate_gradient(model: nn.Module,
|
|
optimizer: Optimizer,
|
|
dataloader: Iterable,
|
|
accumulate_size: int,
|
|
gradient_handlers: List[BaseGradientHandler] = None,
|
|
lr_scheduler: _LRScheduler = None):
|
|
"""
|
|
:param model: your model object
|
|
:type model: :class:`torch.nn.Module`
|
|
:param optimizer: your optimizer object
|
|
:type optimizer: :class:`torch.optim.Optimizer`
|
|
:param dataloader: your dataloader object
|
|
:type dataloader: Iterable
|
|
:param accumulate_size: the number of steps to accumulate gradients
|
|
:type accumulate_size: int
|
|
:param gradient_handlers: list of gradient handler objects. Default is None
|
|
:type gradient_handlers: List[:class:`colossalai.engine.BaseGradientHandler`]
|
|
:param lr_scheduler: your lr scheduler object. Default is None
|
|
:type lr_scheduler: `torch.optim.lr_scheduler._LRScheduler`
|
|
"""
|
|
optimizer = GradAccumOptimizer(optimizer, accumulate_size=accumulate_size, model=model)
|
|
dataloader = GradAccumDataloader(dataloader, accumulate_size=accumulate_size)
|
|
|
|
if gradient_handlers is not None:
|
|
gradient_handlers = [GradAccumGradientHandler(handler, accumulate_size) for handler in gradient_handlers]
|
|
|
|
if lr_scheduler is not None:
|
|
lr_scheduler = GradAccumLrSchedulerByStep(lr_scheduler, accumulate_size=accumulate_size)
|
|
|
|
return optimizer, dataloader, gradient_handlers, lr_scheduler
|
|
|
|
|
|
__all__ = ['accumulate_gradient', 'GradAccumDataloader', 'GradAccumOptimizer',
|
|
'GradAccumLrSchedulerByStep', 'GradAccumGradientHandler']
|