ColossalAI/colossalai/trainer/hooks/__init__.py

16 lines
731 B
Python

from ._base_hook import BaseHook
from ._checkpoint_hook import SaveCheckpointHook, LoadCheckpointHook
from ._metric_hook import (LossHook, Accuracy2DHook, AccuracyHook, MetricHook,
Accuracy1DHook, Accuracy2p5DHook, Accuracy3DHook)
from ._log_hook import LogMetricByEpochHook, TensorboardHook, LogTimingByEpochHook, LogMemoryByEpochHook
from ._lr_scheduler_hook import LRSchedulerHook
__all__ = [
'BaseHook', 'MetricHook',
'LoadCheckpointHook', 'SaveCheckpointHook',
'LossHook', 'AccuracyHook', 'Accuracy2DHook',
'Accuracy1DHook', 'Accuracy2p5DHook', 'Accuracy3DHook',
'LogMetricByEpochHook', 'TensorboardHook', 'LogTimingByEpochHook', 'LogMemoryByEpochHook',
'LRSchedulerHook'
]