mirror of https://github.com/hpcaitech/ColossalAI
457 lines
16 KiB
Python
457 lines
16 KiB
Python
from typing import Union, List
|
|
from colossalai.context.parallel_mode import ParallelMode
|
|
|
|
import torch
|
|
from torch import Tensor
|
|
from torch.utils.data import DataLoader
|
|
from tqdm import tqdm
|
|
|
|
from colossalai.core import global_context as gpc
|
|
|
|
from colossalai.engine import Engine
|
|
from colossalai.engine.schedule import NonPipelineSchedule, BaseSchedule
|
|
from colossalai.logging import DistributedLogger
|
|
from colossalai.utils import MultiTimer
|
|
from colossalai.utils import is_dp_rank_0, is_tp_rank_0, is_no_pp_or_last_stage
|
|
from colossalai.trainer.hooks import BaseHook
|
|
|
|
|
|
class Trainer:
|
|
r"""This is a class tending for easy deployments of users' training and evaluation instead of
|
|
writing their own scripts. It is similar with ``ignite.engine`` and ``keras.engine``, but is
|
|
called `Trainer`.
|
|
|
|
Args:
|
|
engine (:class:`Engine`): Engine responsible for the process function.
|
|
schedule (:class:`BaseSchedule`, optional): Schedule responsible for forward and backward steps.
|
|
timer (:class:`MultiTimer`, optional): Timer used to monitor the whole training.
|
|
logger (:class:`colossalai.logging.DistributedLogger`, optional): Logger used to record the whole training log.
|
|
|
|
Note:
|
|
when `schedule` is None, the ``NonPipelineSchedule`` would be used. If you would like to use pipeline,
|
|
you should choose ``PipelineSchedule`` or ``InterleavedPipelineSchedule`` for the `schedule`
|
|
|
|
Examples:
|
|
>>> # define model, criterion, optimizer, lr_scheduler, train_dataloader for your training
|
|
>>> model = ...
|
|
>>> criterion = ...
|
|
>>> optimizer = ...
|
|
>>> train_dataloader = ...
|
|
>>> # Initialize your engine, train_dataloader, test_dataloader, lr_scheduler
|
|
>>> engine, train_dataloader, _, _ = colossalai.initialize(model, optimizer, criterion)
|
|
>>> # Beginning training progress
|
|
>>> timier = ...
|
|
>>> logger = ...
|
|
>>> trainer = Trainer(engine=engine, logger=logger, schedule=schedule, timer=timier)
|
|
>>> # add hooks you would like to use here.
|
|
>>> hook_list = []
|
|
>>> trainer.fit(
|
|
>>> train_dataloader=train_dataloader,
|
|
>>> epochs=gpc.config.NUM_EPOCHS,
|
|
>>> test_interval=1,
|
|
>>> hooks=hook_list,
|
|
>>> display_progress=True,
|
|
>>> return_output_label=False
|
|
>>> )
|
|
|
|
More examples and details could be found in
|
|
`Training with engine and trainer <https://www.colossalai.org/docs/basics/engine_trainer>`_
|
|
and `ColossalAI-Examples <https://github.com/hpcaitech/ColossalAI-Examples/tree/main>`_.
|
|
"""
|
|
def __init__(
|
|
self,
|
|
engine: Engine,
|
|
schedule: BaseSchedule = None,
|
|
timer: MultiTimer = None,
|
|
logger: DistributedLogger = None,
|
|
):
|
|
# training-ralated params
|
|
self._engine = engine
|
|
self._max_epochs = 0
|
|
self._cur_epoch = 0
|
|
self._max_steps = 0
|
|
self._cur_step = 0
|
|
self._steps_per_epoch = 0
|
|
|
|
# misc params
|
|
self._logger = logger
|
|
self._verbose = logger is not None
|
|
|
|
# hooks can store states in this dict, and could be consumed by other hooks
|
|
self.states = dict()
|
|
|
|
# build hooks
|
|
self.hooks = list()
|
|
|
|
# multi-timer for time benchmarking
|
|
self._timer = timer
|
|
|
|
# set schedule which specifies the training iteration for the engine
|
|
if schedule is None:
|
|
schedule = NonPipelineSchedule()
|
|
if (gpc.is_initialized(ParallelMode.PIPELINE)
|
|
and gpc.get_world_size(ParallelMode.PIPELINE) > 1):
|
|
assert not isinstance(
|
|
schedule, NonPipelineSchedule
|
|
), "NonPipelineSchedule cannot be used for pipeline parallel training, please use PipelineSchedule instead."
|
|
self._schedule = schedule
|
|
self._schedule.pre_processing(engine)
|
|
|
|
@property
|
|
def cur_epoch(self):
|
|
"""Returns the index of the current epoch."""
|
|
return self._cur_epoch
|
|
|
|
@cur_epoch.setter
|
|
def cur_epoch(self, epoch: int):
|
|
"""Set how many epochs have been processed."""
|
|
# allow setter for training resumption
|
|
self._cur_epoch = epoch
|
|
|
|
@property
|
|
def cur_step(self):
|
|
"""Returns how many iteration steps have been processed."""
|
|
return self._cur_step
|
|
|
|
@property
|
|
def max_epochs(self):
|
|
return self._max_epochs
|
|
|
|
@property
|
|
def max_steps(self):
|
|
return self._max_steps
|
|
|
|
@property
|
|
def steps_per_epoch(self):
|
|
return self._steps_per_epoch
|
|
|
|
@property
|
|
def engine(self):
|
|
return self._engine
|
|
|
|
@property
|
|
def schedule(self):
|
|
return self._schedule
|
|
|
|
def _set_current_step(self, epoch: int):
|
|
"""Sets current step number.
|
|
|
|
Args:
|
|
epoch (int): Step number to be set.
|
|
"""
|
|
self._cur_step = epoch * self._steps_per_epoch
|
|
|
|
def _call_timer(self, action: str, item: str, *args, **kwargs) -> None:
|
|
"""Call timer funciton with a given timer name.
|
|
|
|
Args:
|
|
action (str): Function to be called on timer.
|
|
item (str): Name of the timer.
|
|
args (list): args used for action function.
|
|
kwargs (dict): kwargs used for action function.
|
|
"""
|
|
|
|
if self._timer is not None:
|
|
getattr(self._timer, action)(item, *args, **kwargs)
|
|
|
|
def _reset_states(self) -> None:
|
|
"""Clear trainer states"""
|
|
self.states = dict()
|
|
|
|
def _call_hooks(self, func, output=None):
|
|
"""Calls specific hooks in the current time point.
|
|
|
|
Args:
|
|
func (str): A string represents the time point.
|
|
output (Any, optional): Output of the model after running an iteration or None in any other time points.
|
|
"""
|
|
# Only after iter hook will receive output
|
|
for hook in self.hooks:
|
|
if output is None:
|
|
getattr(hook, func)(self)
|
|
else:
|
|
getattr(hook, func)(self, *output)
|
|
|
|
@staticmethod
|
|
def _should_display_progress(display_progress: bool):
|
|
"""Only display progress on DP rank 0, TP rank 0 and PP last rank"""
|
|
return (display_progress and is_dp_rank_0() and is_tp_rank_0()
|
|
and is_no_pp_or_last_stage())
|
|
|
|
def _train_epoch(
|
|
self,
|
|
train_dataloader: DataLoader,
|
|
epoch: int = None,
|
|
display_progress: bool = False,
|
|
return_output_label: bool = True,
|
|
):
|
|
# set training state
|
|
self._engine.train()
|
|
data_iter = iter(train_dataloader)
|
|
progress = range(self._steps_per_epoch)
|
|
if display_progress:
|
|
if epoch is None:
|
|
progress = tqdm(progress, desc="[Train]")
|
|
else:
|
|
progress = tqdm(progress, desc=f"[Epoch {epoch} / Train]")
|
|
|
|
self._call_hooks("before_train_epoch")
|
|
self._call_timer(action="start", item="Train-epoch")
|
|
for i in progress:
|
|
self._call_hooks("before_train_iter")
|
|
self._call_timer(action="start", item="Train-step")
|
|
|
|
# run 1 training step
|
|
self.engine.zero_grad()
|
|
logits, label, loss = self.schedule.forward_backward_step(
|
|
self.engine,
|
|
data_iter,
|
|
forward_only=False,
|
|
return_loss=True,
|
|
return_output_label=return_output_label,
|
|
)
|
|
self.engine.step()
|
|
self._call_timer(action="stop",
|
|
item="Train-step",
|
|
keep_in_history=True)
|
|
self._call_hooks("after_train_iter", output=(logits, label, loss))
|
|
|
|
self._cur_step += 1
|
|
|
|
if display_progress:
|
|
if "step_metrics" in self.states:
|
|
progress.set_postfix(**self.states["step_metrics"])
|
|
|
|
# stop when max iter is reached
|
|
if self._exceed_max_step():
|
|
break
|
|
|
|
self._call_timer(action="stop",
|
|
item="Train-epoch",
|
|
keep_in_history=True)
|
|
self._call_hooks("after_train_epoch")
|
|
self._call_timer(action="reset", item="Train-epoch")
|
|
|
|
def _eval(
|
|
self,
|
|
test_dataloader: DataLoader,
|
|
epoch: int = None,
|
|
display_progress: bool = False,
|
|
return_output_label: bool = True,
|
|
):
|
|
# switch engine status
|
|
self._engine.eval()
|
|
|
|
data_iter = iter(test_dataloader)
|
|
num_steps = len(test_dataloader)
|
|
|
|
self._call_hooks("before_test")
|
|
# prepare progress bar
|
|
progress = range(num_steps)
|
|
if display_progress:
|
|
desc = "Evaluation"
|
|
if epoch is not None:
|
|
desc = "[Epoch %d / Test]" % epoch
|
|
progress = tqdm(progress, desc=desc)
|
|
|
|
self._call_hooks("before_test_epoch")
|
|
self._call_timer(action="start", item="Test-epoch")
|
|
with torch.no_grad():
|
|
for _ in progress:
|
|
self._call_hooks("before_test_iter")
|
|
self._call_timer(action="start", item="Test-step")
|
|
logits, label, loss = self.schedule.forward_backward_step(
|
|
self.engine,
|
|
data_iter,
|
|
forward_only=True,
|
|
return_loss=True,
|
|
return_output_label=return_output_label,
|
|
)
|
|
self._call_timer(action="stop",
|
|
item="Test-step",
|
|
keep_in_history=True)
|
|
self._call_hooks("after_test_iter",
|
|
output=(logits, label, loss))
|
|
|
|
if display_progress:
|
|
if "step_metrics" in self.states:
|
|
progress.set_postfix(**self.states["step_metrics"])
|
|
|
|
self._call_timer(action="stop",
|
|
item="Test-epoch",
|
|
keep_in_history=True)
|
|
self._call_hooks("after_test_epoch")
|
|
self._call_hooks("after_test")
|
|
self._call_timer(action="reset", item="Test-step")
|
|
self._call_timer(action="reset", item="Test-epoch")
|
|
|
|
def _exceed_max_step(self):
|
|
return self._max_steps is not None and self._cur_step >= self._max_steps
|
|
|
|
def fit(
|
|
self,
|
|
train_dataloader: DataLoader,
|
|
epochs: int,
|
|
max_steps: int = None,
|
|
test_dataloader: DataLoader = None,
|
|
test_interval: int = 1,
|
|
hooks: List[BaseHook] = None,
|
|
display_progress: bool = False,
|
|
return_output_label: bool = True,
|
|
):
|
|
r"""Trains the model to fit training data.
|
|
|
|
Args:
|
|
train_dataloader (:class:`torch.utils.data.DataLoader`): DataLoader for training.
|
|
epochs (int): Maximum number of epochs.
|
|
max_steps (int, optional): Maximum number of running iterations.
|
|
test_dataloader (:class:`torch.utils.data.DataLoader`, optional): DataLoader for validation.
|
|
test_interval (int, optional): Interval of validation
|
|
hooks (list[BaseHook], optional): A list of hooks used in training.
|
|
display_progress (bool, optional): If True, a progress bar will be displayed.
|
|
"""
|
|
|
|
# set epochs and steps, consider gradient accumulation
|
|
self._steps_per_epoch = len(train_dataloader)
|
|
self._max_steps = max_steps
|
|
self._max_epochs = epochs
|
|
|
|
# check if testing is required
|
|
should_test = False
|
|
if test_dataloader is not None:
|
|
should_test = True
|
|
|
|
display_progress = self._should_display_progress(display_progress)
|
|
|
|
# reset hooks
|
|
self._reset_states()
|
|
if hooks is not None:
|
|
assert isinstance(
|
|
hooks, list
|
|
), f"expected argument hooks be to list, but got {type(hooks)}"
|
|
else:
|
|
hooks = []
|
|
self.hooks = hooks
|
|
self.hooks.sort(key=lambda hook: hook.priority)
|
|
if self._verbose:
|
|
for hook in self.hooks:
|
|
self._logger.info(
|
|
f"Using {hook.__class__.__name__} for training, priority = {hook.priority}",
|
|
ranks=[0],
|
|
)
|
|
self._logger.info(
|
|
"Lower value means higher priority for calling hook function",
|
|
ranks=[0])
|
|
self._call_hooks("after_hook_is_attached")
|
|
|
|
self._engine.train()
|
|
self._call_hooks("before_train")
|
|
|
|
# recover step value if resuming training
|
|
last_epoch = self._cur_epoch
|
|
if self.cur_epoch != 0:
|
|
self._set_current_step(last_epoch)
|
|
|
|
for epoch in range(last_epoch, epochs):
|
|
# train for one epoch
|
|
self._train_epoch(
|
|
train_dataloader=train_dataloader,
|
|
epoch=epoch,
|
|
display_progress=display_progress,
|
|
return_output_label=return_output_label,
|
|
)
|
|
|
|
# start eval
|
|
if should_test and epoch % test_interval == 0:
|
|
self._eval(
|
|
test_dataloader=test_dataloader,
|
|
display_progress=display_progress,
|
|
epoch=epoch,
|
|
return_output_label=return_output_label,
|
|
)
|
|
|
|
self._cur_epoch += 1
|
|
|
|
# check for termination
|
|
if self._exceed_max_step():
|
|
self._logger.info(
|
|
f"Max number of steps {max_steps} has been reached, training is stopped automatically",
|
|
ranks=[0],
|
|
)
|
|
break
|
|
self._call_hooks("after_train")
|
|
self._call_timer("reset", "Train-epoch")
|
|
|
|
def evaluate(
|
|
self,
|
|
test_dataloader: DataLoader,
|
|
hooks: List[BaseHook] = None,
|
|
display_progress: bool = False,
|
|
return_output_label: bool = True,
|
|
):
|
|
"""Evaluates the model with testing data.
|
|
|
|
Args:
|
|
test_dataloader (:class:`torch.utils.data.DataLoader`, optional): Dataloader for testing.
|
|
hooks (list, optional): A list of hooks used in evaluation. Defaults to None.
|
|
display_progress (bool, optional): If True, the evaluation progress will be printed. Defaults to False.
|
|
return_output_label (bool, optional): If True, the output of model and the label
|
|
will be returned. Defaults to True.
|
|
"""
|
|
# set display
|
|
display_progress = self._should_display_progress(display_progress)
|
|
|
|
# reset hooks
|
|
self._reset_states()
|
|
if hooks is not None:
|
|
assert isinstance(
|
|
hooks, list
|
|
), f"expected argument hooks be to list, but got {type(hooks)}"
|
|
else:
|
|
hooks = []
|
|
self.hooks = hooks
|
|
self.hooks.sort(key=lambda hook: hook.priority)
|
|
if self._verbose:
|
|
for hook in self.hooks:
|
|
self._logger.info(
|
|
f"Using {hook.__class__.__name__} for training, priority = {hook.priority}",
|
|
ranks=[0],
|
|
)
|
|
self._logger.info(
|
|
"Lower value means higher priority for calling hook function",
|
|
ranks=[0])
|
|
self._call_hooks("after_hook_is_attached")
|
|
|
|
# eval
|
|
self._eval(
|
|
test_dataloader=test_dataloader,
|
|
display_progress=display_progress,
|
|
return_output_label=return_output_label,
|
|
)
|
|
|
|
def predict(self, data: Union[Tensor, List[Tensor]]):
|
|
"""Uses trained model to make a prediction for a tensor or a tensor list.
|
|
|
|
Args:
|
|
data (Union[:class:`torch.tensor`, List[:class:`torch.tensor`]]): Data as the input.
|
|
|
|
Returns:
|
|
:class:`torch.tensor`: The output of model as the prediction
|
|
"""
|
|
# predict without labels
|
|
if isinstance(data, (list, tuple)):
|
|
assert isinstance(data[0], Tensor)
|
|
else:
|
|
assert isinstance(data, Tensor)
|
|
self._engine.eval()
|
|
|
|
# prepare a list of (data, label) to make it iterable
|
|
# for compatibility with schedule
|
|
simple_dataloader = [(data, None)]
|
|
data_iter = iter(simple_dataloader)
|
|
output, _, _ = self.schedule.forward_backward_step(self.engine,
|
|
data_iter,
|
|
forward_only=True,
|
|
return_loss=False)
|
|
return output
|