github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-01 09:07:51 +00:00
Code Issues Packages Projects Releases Wiki Activity

[misc] update pre-commit and run all files (#4752)

Browse Source 
* [misc] update pre-commit

* [misc] run pre-commit

* [misc] remove useless configuration files

* [misc] ignore cuda for clang-format
This commit is contained in:
Hongxin Liu
2023-09-19 14:20:26 +08:00
committed by GitHub
parent 3c6b831c26
commit 079bf3cb26
1268 changed files with 50037 additions and 38444 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/legacy/trainer/__init__.py
View File

@@ -1,3 +1,3 @@
from ._trainer import Trainer
__all__ = ['Trainer']
__all__ = ["Trainer"]

5
colossalai/legacy/trainer/_trainer.py
View File

@@ -151,7 +151,7 @@ class Trainer:
@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())
return display_progress and is_dp_rank_0() and is_tp_rank_0() and is_no_pp_or_last_stage()
def _train_epoch(
self,
@@ -293,8 +293,7 @@ class Trainer:
assert isinstance(hooks, list), f"expected argument hooks be to list, but got {type(hooks)}"
for hook in hooks:
assert isinstance(hook, BaseHook), \
f'expected the hook to be of type BaseHook, but got {type(hook)}'
assert isinstance(hook, BaseHook), f"expected the hook to be of type BaseHook, but got {type(hook)}"
else:
hooks = []
self.hooks = hooks

15
colossalai/legacy/trainer/hooks/__init__.py
View File

@@ -11,7 +11,16 @@ from ._lr_scheduler_hook import LRSchedulerHook
from ._metric_hook import AccuracyHook, LossHook, MetricHook, ThroughputHook
__all__ = [
'BaseHook', 'MetricHook', 'LossHook', 'AccuracyHook', 'LogMetricByEpochHook', 'TensorboardHook',
'LogTimingByEpochHook', 'LogMemoryByEpochHook', 'LRSchedulerHook', 'ThroughputHook', 'LogMetricByStepHook',
'SaveCheckpointHook'
"BaseHook",
"MetricHook",
"LossHook",
"AccuracyHook",
"LogMetricByEpochHook",
"TensorboardHook",
"LogTimingByEpochHook",
"LogMemoryByEpochHook",
"LRSchedulerHook",
"ThroughputHook",
"LogMetricByStepHook",
"SaveCheckpointHook",
]

46
colossalai/legacy/trainer/hooks/_base_hook.py
View File

@@ -18,24 +18,16 @@ class BaseHook(ABC):
self.priority = priority
def after_hook_is_attached(self, trainer):
"""Actions after hooks are attached to trainer.
"""
pass
"""Actions after hooks are attached to trainer."""
def before_train(self, trainer):
"""Actions before training.
"""
pass
"""Actions before training."""
def after_train(self, trainer):
"""Actions after training.
"""
pass
"""Actions after training."""
def before_train_iter(self, trainer):
"""Actions before running a training iteration.
"""
pass
"""Actions before running a training iteration."""
def after_train_iter(self, trainer, output: Tensor, label: Tensor, loss: Tensor):
"""Actions after running a training iteration.
@@ -46,42 +38,27 @@ class BaseHook(ABC):
label (:class:`torch.Tensor`): Labels of the input data.
loss (:class:`torch.Tensor`): Loss between the output and input data.
"""
pass
def before_train_epoch(self, trainer):
"""Actions before starting a training epoch.
"""
pass
"""Actions before starting a training epoch."""
def after_train_epoch(self, trainer):
"""Actions after finishing a training epoch.
"""
pass
"""Actions after finishing a training epoch."""
def before_test(self, trainer):
"""Actions before evaluation.
"""
pass
"""Actions before evaluation."""
def after_test(self, trainer):
"""Actions after evaluation.
"""
pass
"""Actions after evaluation."""
def before_test_epoch(self, trainer):
"""Actions before starting a testing epoch.
"""
pass
"""Actions before starting a testing epoch."""
def after_test_epoch(self, trainer):
"""Actions after finishing a testing epoch.
"""
pass
"""Actions after finishing a testing epoch."""
def before_test_iter(self, trainer):
"""Actions before running a testing iteration.
"""
pass
"""Actions before running a testing iteration."""
def after_test_iter(self, trainer, output: Tensor, label: Tensor, loss: Tensor):
"""Actions after running a testing iteration.
@@ -92,7 +69,6 @@ class BaseHook(ABC):
label (:class:`torch.Tensor`): Labels of the input data
loss (:class:`torch.Tensor`): Loss between the output and input data
"""
pass
def init_runner_states(self, trainer, key, val):
"""Initializes trainer's state.

37
colossalai/legacy/trainer/hooks/_checkpoint_hook.py
View File

@@ -27,12 +27,14 @@ class SaveCheckpointHook(BaseHook):
depend on the hooks order in the hook list.
"""
def __init__(self,
interval: int = 1,
checkpoint_dir: str = None,
model: torch.nn.Module = None,
save_by_iter: bool = False,
priority: int = 10):
def __init__(
self,
interval: int = 1,
checkpoint_dir: str = None,
model: torch.nn.Module = None,
save_by_iter: bool = False,
priority: int = 10,
):
super().__init__(priority=priority)
self.interval = interval
self.checkpoint_dir = checkpoint_dir
@@ -52,22 +54,23 @@ class SaveCheckpointHook(BaseHook):
self.model = self.model if self.model is not None else trainer.engine.model
def after_train_iter(self, trainer, output, label, loss):
"""Saves the model after a training iter.
"""
"""Saves the model after a training iter."""
# save by interval
if self.save_by_iter and trainer.cur_step % self.interval == 0:
save_checkpoint(self.checkpoint_dir, trainer.cur_epoch, self.model, trainer.engine.optimizer,
self._lr_scheduler)
self.logger.info(f'checkpoint for iteration {trainer.cur_step} is saved to {self.checkpoint_dir}',
ranks=[0])
save_checkpoint(
self.checkpoint_dir, trainer.cur_epoch, self.model, trainer.engine.optimizer, self._lr_scheduler
)
self.logger.info(
f"checkpoint for iteration {trainer.cur_step} is saved to {self.checkpoint_dir}", ranks=[0]
)
else:
pass
def after_train_epoch(self, trainer):
"""Saves the model after a training epoch.
"""
"""Saves the model after a training epoch."""
# save by interval
if trainer.cur_epoch % self.interval == 0:
save_checkpoint(self.checkpoint_dir, trainer.cur_epoch, self.model, trainer.engine.optimizer,
self._lr_scheduler)
self.logger.info(f'checkpoint for epoch {trainer.cur_epoch} is saved to {self.checkpoint_dir}', ranks=[0])
save_checkpoint(
self.checkpoint_dir, trainer.cur_epoch, self.model, trainer.engine.optimizer, self._lr_scheduler
)
self.logger.info(f"checkpoint for epoch {trainer.cur_epoch} is saved to {self.checkpoint_dir}", ranks=[0])

4
colossalai/legacy/trainer/hooks/_commons_.py
View File

@@ -3,7 +3,7 @@ import torch
def _format_number(val, prec=5):
if isinstance(val, float):
return f'{val:.{prec}g}'
return f"{val:.{prec}g}"
elif torch.is_tensor(val) and torch.is_floating_point(val):
return f'{val.item():.{prec}g}'
return f"{val.item():.{prec}g}"
return val

126
colossalai/legacy/trainer/hooks/_log_hook.py
View File

@@ -51,20 +51,20 @@ class LogMetricByStepHook(BaseHook):
super().__init__(priority)
def after_train_iter(self, trainer, *args):
trainer.states['step_metrics'] = dict()
for metric_name, metric_calculator in trainer.states['metrics']['train'].items():
trainer.states["step_metrics"] = dict()
for metric_name, metric_calculator in trainer.states["metrics"]["train"].items():
if isinstance(metric_calculator, ThroughputMetric):
trainer.states['step_metrics'][metric_name.lower()] = metric_calculator.get_last_step_info()
trainer.states["step_metrics"][metric_name.lower()] = metric_calculator.get_last_step_info()
else:
trainer.states['step_metrics'][metric_name.lower()] = metric_calculator.get_last_step_value()
trainer.states["step_metrics"][metric_name.lower()] = metric_calculator.get_last_step_value()
def after_test_iter(self, trainer, *args):
trainer.states['step_metrics'] = dict()
for metric_name, metric_calculator in trainer.states['metrics']['test'].items():
trainer.states["step_metrics"] = dict()
for metric_name, metric_calculator in trainer.states["metrics"]["test"].items():
if isinstance(metric_calculator, ThroughputMetric):
trainer.states['step_metrics'][metric_name.lower()] = metric_calculator.get_last_step_info()
trainer.states["step_metrics"][metric_name.lower()] = metric_calculator.get_last_step_info()
else:
trainer.states['step_metrics'][metric_name.lower()] = metric_calculator.get_last_step_value()
trainer.states["step_metrics"][metric_name.lower()] = metric_calculator.get_last_step_value()
@HOOKS.register_module
@@ -85,24 +85,24 @@ class LogMetricByEpochHook(LogByEpochHook):
def _get_str(self, trainer, mode):
msg = []
for metric_name, metric_calculator in trainer.states['metrics'][mode].items():
msg.append(f'{metric_name} = {_format_number(metric_calculator.get_accumulated_value())}')
msg = ' | '.join(msg)
for metric_name, metric_calculator in trainer.states["metrics"][mode].items():
msg.append(f"{metric_name} = {_format_number(metric_calculator.get_accumulated_value())}")
msg = " | ".join(msg)
return msg
def after_train_epoch(self, trainer):
if self._is_epoch_to_log(trainer):
msg = self._get_str(trainer=trainer, mode='train')
msg = self._get_str(trainer=trainer, mode="train")
if self._is_rank_to_log:
self.logger.info(f'[Epoch {trainer.cur_epoch} / Train]: {msg}')
self.logger.info(f"[Epoch {trainer.cur_epoch} / Train]: {msg}")
# f'Training - Epoch {trainer.cur_epoch} - {self.__class__.__name__}: {msg}')
def after_test_epoch(self, trainer):
if self._is_epoch_to_log(trainer):
msg = self._get_str(trainer=trainer, mode='test')
msg = self._get_str(trainer=trainer, mode="test")
if self._is_rank_to_log:
self.logger.info(f'[Epoch {trainer.cur_epoch} / Test]: {msg}')
self.logger.info(f"[Epoch {trainer.cur_epoch} / Test]: {msg}")
# f'Testing - Epoch {trainer.cur_epoch} - {self.__class__.__name__}: {msg}')
@@ -145,8 +145,11 @@ class TensorboardHook(BaseHook):
self._is_valid_rank_to_log = True
# check for
if gpc.is_initialized(ParallelMode.PIPELINE) and \
not gpc.is_last_rank(ParallelMode.PIPELINE) and self._is_valid_rank_to_log:
if (
gpc.is_initialized(ParallelMode.PIPELINE)
and not gpc.is_last_rank(ParallelMode.PIPELINE)
and self._is_valid_rank_to_log
):
raise ValueError("Tensorboard hook can only log on the last rank of pipeline process group")
if self._is_valid_rank_to_log:
@@ -157,38 +160,38 @@ class TensorboardHook(BaseHook):
rank = 0
# create workspace
log_dir = osp.join(log_dir, f'{parallel_mode}_rank_{rank}')
log_dir = osp.join(log_dir, f"{parallel_mode}_rank_{rank}")
os.makedirs(log_dir, exist_ok=True)
self.writer = SummaryWriter(log_dir=log_dir, filename_suffix=f'_rank_{rank}')
self.writer = SummaryWriter(log_dir=log_dir, filename_suffix=f"_rank_{rank}")
def _log_by_iter(self, trainer, mode: str):
for metric_name, metric_calculator in trainer.states['metrics'][mode].items():
for metric_name, metric_calculator in trainer.states["metrics"][mode].items():
if metric_calculator.epoch_only:
continue
val = metric_calculator.get_last_step_value()
if self._is_valid_rank_to_log:
self.writer.add_scalar(f'{metric_name}/{mode}', val, trainer.cur_step)
self.writer.add_scalar(f"{metric_name}/{mode}", val, trainer.cur_step)
def _log_by_epoch(self, trainer, mode: str):
for metric_name, metric_calculator in trainer.states['metrics'][mode].items():
for metric_name, metric_calculator in trainer.states["metrics"][mode].items():
if metric_calculator.epoch_only:
val = metric_calculator.get_accumulated_value()
if self._is_valid_rank_to_log:
self.writer.add_scalar(f'{metric_name}/{mode}', val, trainer.cur_step)
self.writer.add_scalar(f"{metric_name}/{mode}", val, trainer.cur_step)
def after_test_iter(self, trainer, *args):
self._log_by_iter(trainer, mode='test')
self._log_by_iter(trainer, mode="test")
def after_test_epoch(self, trainer):
self._log_by_epoch(trainer, mode='test')
self._log_by_epoch(trainer, mode="test")
def after_train_iter(self, trainer, *args):
self._log_by_iter(trainer, mode='train')
self._log_by_iter(trainer, mode="train")
def after_train_epoch(self, trainer):
self._log_by_epoch(trainer, mode='train')
self._log_by_epoch(trainer, mode="train")
@HOOKS.register_module
@@ -206,13 +209,15 @@ class LogTimingByEpochHook(LogByEpochHook):
ignore_num_train_steps (int, optional): Number of training steps to ignore, defaults to 0.
"""
def __init__(self,
timer: MultiTimer,
logger: DistributedLogger,
interval: int = 1,
priority: int = 10,
log_eval: bool = True,
ignore_num_train_steps: int = 0) -> None:
def __init__(
self,
timer: MultiTimer,
logger: DistributedLogger,
interval: int = 1,
priority: int = 10,
log_eval: bool = True,
ignore_num_train_steps: int = 0,
) -> None:
super().__init__(logger=logger, interval=interval, priority=priority)
self._timer = timer
self._log_eval = log_eval
@@ -229,33 +234,31 @@ class LogTimingByEpochHook(LogByEpochHook):
if timer_name.startswith(mode):
last_elapsed_time = timer.get_elapsed_time()
if timer.has_history:
if timer_name == 'Train-step' and not self._is_train_step_history_trimmed:
timer._history = timer._history[self._ignore_num_train_steps:]
if timer_name == "Train-step" and not self._is_train_step_history_trimmed:
timer._history = timer._history[self._ignore_num_train_steps :]
self._is_train_step_history_trimmed = True
history_mean = timer.get_history_mean()
history_sum = timer.get_history_sum()
timer.get_history_sum()
msg.append(
f'{timer_name}: last = {_format_number(last_elapsed_time)} s, mean = {_format_number(history_mean)} s'
f"{timer_name}: last = {_format_number(last_elapsed_time)} s, mean = {_format_number(history_mean)} s"
)
else:
msg.append(f'{timer_name}: last = {_format_number(last_elapsed_time)} s')
msg.append(f"{timer_name}: last = {_format_number(last_elapsed_time)} s")
msg = ' | '.join(msg)
msg = " | ".join(msg)
return msg
def after_train_epoch(self, trainer):
"""Writes log after finishing a training epoch.
"""
"""Writes log after finishing a training epoch."""
if self._is_epoch_to_log(trainer) and self._is_rank_to_log:
msg = self._get_message('Train')
self.logger.info(f'[Epoch {trainer.cur_epoch} / Train]: {msg} | #steps/epoch = {trainer.steps_per_epoch}')
msg = self._get_message("Train")
self.logger.info(f"[Epoch {trainer.cur_epoch} / Train]: {msg} | #steps/epoch = {trainer.steps_per_epoch}")
def after_test_epoch(self, trainer):
"""Writes log after finishing a testing epoch.
"""
"""Writes log after finishing a testing epoch."""
if self._is_epoch_to_log(trainer) and self._is_rank_to_log and self._log_eval:
msg = self._get_message('Test')
self.logger.info(f'[Epoch {trainer.cur_epoch} / Test]: {msg}')
msg = self._get_message("Test")
self.logger.info(f"[Epoch {trainer.cur_epoch} / Test]: {msg}")
@HOOKS.register_module
@@ -272,31 +275,28 @@ class LogMemoryByEpochHook(LogByEpochHook):
"""
def __init__(
self,
logger: DistributedLogger,
interval: int = 1,
priority: int = 10,
log_eval: bool = True,
report_cpu: bool = False, # no reference
self,
logger: DistributedLogger,
interval: int = 1,
priority: int = 10,
log_eval: bool = True,
report_cpu: bool = False, # no reference
) -> None:
super().__init__(logger=logger, interval=interval, priority=priority)
self._log_eval = log_eval
self._is_rank_to_log = is_dp_rank_0() and is_tp_rank_0()
def before_train(self, trainer):
"""Resets before training.
"""
"""Resets before training."""
if self._is_epoch_to_log(trainer) and self._is_rank_to_log:
report_memory_usage('Before-train', self.logger)
report_memory_usage("Before-train", self.logger)
def after_train_epoch(self, trainer):
"""Writes log after finishing a training epoch.
"""
"""Writes log after finishing a training epoch."""
if self._is_epoch_to_log(trainer) and self._is_rank_to_log:
report_memory_usage(f'[Epoch {trainer.cur_epoch} / Train]', self.logger)
report_memory_usage(f"[Epoch {trainer.cur_epoch} / Train]", self.logger)
def after_test(self, trainer):
"""Reports after testing.
"""
"""Reports after testing."""
if self._is_epoch_to_log(trainer) and self._is_rank_to_log and self._log_eval:
report_memory_usage(f'[Epoch {trainer.cur_epoch} / Test]', self.logger)
report_memory_usage(f"[Epoch {trainer.cur_epoch} / Test]", self.logger)

9
colossalai/legacy/trainer/hooks/_lr_scheduler_hook.py
View File

@@ -34,15 +34,16 @@ class LRSchedulerHook(MetricHook):
def after_hook_is_attached(self, trainer):
self._check_metric_states_initialization(trainer)
trainer.states['metrics']['train']['LR'] = LearningRateMetric(epoch_only=self.by_epoch,
initial_lr=self.lr_scheduler.get_last_lr()[0])
trainer.states["metrics"]["train"]["LR"] = LearningRateMetric(
epoch_only=self.by_epoch, initial_lr=self.lr_scheduler.get_last_lr()[0]
)
def after_train_epoch(self, trainer):
if self.by_epoch:
self.lr_scheduler.step()
trainer.states['metrics']['train']['LR'].update(self.lr_scheduler.get_last_lr()[0])
trainer.states["metrics"]["train"]["LR"].update(self.lr_scheduler.get_last_lr()[0])
def after_train_iter(self, trainer, output: Tensor, label: Tensor, loss: Tensor):
if not self.by_epoch:
self.lr_scheduler.step()
trainer.states['metrics']['train']['LR'].update(self.lr_scheduler.get_last_lr()[0])
trainer.states["metrics"]["train"]["LR"].update(self.lr_scheduler.get_last_lr()[0])

71
colossalai/legacy/trainer/hooks/_metric_hook.py
View File

@@ -35,8 +35,7 @@ class Metric(ABC):
@property
def epoch_only(self):
"""Returns :attr:`epoch_only`.
"""
"""Returns :attr:`epoch_only`."""
return self._epoch_only
@abstractmethod
@@ -44,20 +43,16 @@ class Metric(ABC):
"""Resets the metric to it's initial state.
By default, this is called at the start of each epoch.
"""
pass
@abstractmethod
def update(self, *args, **kwargs) -> None:
"""Updates the metric's state using the passed batch output.
By default, this is called once for each batch.
"""
pass
@abstractmethod
def get_last_step_value(self) -> float:
"""Returns the metric value in the last iteration.
"""
pass
"""Returns the metric value in the last iteration."""
@abstractmethod
def get_accumulated_value(self):
@@ -67,7 +62,6 @@ class Metric(ABC):
:return: the actual quantity of interest
:rtype: Any
"""
pass
@staticmethod
@abstractmethod
@@ -77,7 +71,6 @@ class Metric(ABC):
:return: The result of comparison
:rtype: bool
"""
pass
class LossMetric(Metric):
@@ -94,8 +87,7 @@ class LossMetric(Metric):
self.count = 0
def reset(self) -> None:
"""Sets :attr:`last_step_loss` and :attr:`accum_loss` to zero.
"""
"""Sets :attr:`last_step_loss` and :attr:`accum_loss` to zero."""
self.last_step_loss.zero_()
self.accum_loss.zero_()
self.count = 0
@@ -114,8 +106,7 @@ class LossMetric(Metric):
self.count += 1
def get_accumulated_value(self):
"""Returns accumulated loss.
"""
"""Returns accumulated loss."""
if gpc.is_initialized(ParallelMode.DATA):
dist.all_reduce(self.accum_loss, op=dist.ReduceOp.SUM, group=gpc.get_group(ParallelMode.DATA))
self.accum_loss.div_(gpc.get_world_size(ParallelMode.DATA))
@@ -124,8 +115,7 @@ class LossMetric(Metric):
return self.accum_loss.item()
def get_last_step_value(self) -> float:
"""Returns :attr:`last_step_loss`.
"""
"""Returns :attr:`last_step_loss`."""
return self.last_step_loss.cpu().item()
@staticmethod
@@ -141,7 +131,7 @@ class LearningRateMetric(Metric):
initial_lr (float, optional): Initial learning rate, defaults to 0.0.
"""
def __init__(self, epoch_only: bool, initial_lr: float = 0.):
def __init__(self, epoch_only: bool, initial_lr: float = 0.0):
super().__init__(epoch_only=epoch_only)
self.lr = initial_lr
@@ -241,8 +231,8 @@ class MetricHook(BaseHook):
self._is_stage_to_compute = is_no_pp_or_last_stage()
def _check_metric_states_initialization(self, trainer):
if 'metrics' not in trainer.states:
self.init_runner_states(trainer, 'metrics', dict(train={}, test={}))
if "metrics" not in trainer.states:
self.init_runner_states(trainer, "metrics", dict(train={}, test={}))
@HOOKS.register_module
@@ -266,8 +256,8 @@ class LossHook(MetricHook):
self.test_loss = LossMetric(epoch_only=True)
# register the metric calculator
trainer.states['metrics']['train']['Loss'] = self.train_loss
trainer.states['metrics']['test']['Loss'] = self.test_loss
trainer.states["metrics"]["train"]["Loss"] = self.train_loss
trainer.states["metrics"]["test"]["Loss"] = self.test_loss
def before_train_epoch(self, trainer):
if self._is_stage_to_compute:
@@ -307,7 +297,7 @@ class AccuracyHook(MetricHook):
self.metric = AccuracyMetric(epoch_only=True, accuracy_func=self.accuracy_func)
# register the metric
trainer.states['metrics']['test']['Accuracy'] = self.metric
trainer.states["metrics"]["test"]["Accuracy"] = self.metric
def before_test(self, trainer):
if self._is_stage_to_compute:
@@ -356,8 +346,9 @@ class ThroughputMetric(Metric):
if self._use_local:
self.last_step_num_samples *= gpc.get_world_size(ParallelMode.DATA)
else:
self.last_step_used_time = all_reduce(self.last_step_used_time, ParallelMode.DATA) / \
gpc.get_world_size(ParallelMode.DATA)
self.last_step_used_time = all_reduce(self.last_step_used_time, ParallelMode.DATA) / gpc.get_world_size(
ParallelMode.DATA
)
self.last_step_num_samples = all_reduce(self.last_step_num_samples, ParallelMode.DATA)
sample_per_sec = _format_number(self.last_step_num_samples / (self.last_step_used_time + 1e-12).item())
@@ -367,8 +358,9 @@ class ThroughputMetric(Metric):
if self._use_local:
self.last_step_num_samples *= gpc.get_world_size(ParallelMode.DATA)
else:
self.last_step_used_time = all_reduce(self.last_step_used_time, ParallelMode.DATA) / \
gpc.get_world_size(ParallelMode.DATA)
self.last_step_used_time = all_reduce(self.last_step_used_time, ParallelMode.DATA) / gpc.get_world_size(
ParallelMode.DATA
)
self.last_step_num_samples = all_reduce(self.last_step_num_samples, ParallelMode.DATA)
sample_per_sec = _format_number(self.last_step_num_samples / (self.last_step_used_time + 1e-12).item())
@@ -379,8 +371,9 @@ class ThroughputMetric(Metric):
return f"{sample_per_sec} sample_per_sec"
def get_accumulated_value(self) -> float:
self.accumulated_used_time = all_reduce(self.accumulated_used_time, ParallelMode.DATA) / \
gpc.get_world_size(ParallelMode.DATA)
self.accumulated_used_time = all_reduce(self.accumulated_used_time, ParallelMode.DATA) / gpc.get_world_size(
ParallelMode.DATA
)
self.accumulated_num_samples = all_reduce(self.accumulated_num_samples, ParallelMode.DATA)
return (self.accumulated_num_samples / (self.accumulated_used_time + 1e-12)).item()
@@ -411,14 +404,16 @@ class ThroughputHook(MetricHook):
def after_hook_is_attached(self, trainer):
self._check_metric_states_initialization(trainer)
if self._is_stage_to_compute:
self.metric = ThroughputMetric(epoch_only=True,
ignored_steps=self.ignored_steps,
tflop_per_step=self._tflop_per_step,
use_local=self._use_local)
self.metric = ThroughputMetric(
epoch_only=True,
ignored_steps=self.ignored_steps,
tflop_per_step=self._tflop_per_step,
use_local=self._use_local,
)
# register the metric
trainer.states['metrics']['train']['Throughput'] = self.metric
trainer.states['metrics']['test']['Throughput'] = self.metric
trainer.states["metrics"]["train"]["Throughput"] = self.metric
trainer.states["metrics"]["test"]["Throughput"] = self.metric
def before_train_epoch(self, trainer):
if self._is_stage_to_compute:
@@ -426,8 +421,9 @@ class ThroughputHook(MetricHook):
def after_train_iter(self, trainer, *args):
if self._is_stage_to_compute:
self.metric.update(trainer.engine.schedule.batch_size,
trainer._timer.get_timer('Train-step').get_elapsed_time())
self.metric.update(
trainer.engine.schedule.batch_size, trainer._timer.get_timer("Train-step").get_elapsed_time()
)
def before_test(self, trainer):
if self._is_stage_to_compute:
@@ -435,5 +431,6 @@ class ThroughputHook(MetricHook):
def after_test_iter(self, trainer, *args):
if self._is_stage_to_compute:
self.metric.update(trainer.engine.schedule.batch_size,
trainer._timer.get_timer('Test-step').get_elapsed_time())
self.metric.update(
trainer.engine.schedule.batch_size, trainer._timer.get_timer("Test-step").get_elapsed_time()
)