chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,511 @@
|
||||
import signal
|
||||
import typing
|
||||
from typing import Dict, List, Callable
|
||||
from typing import Optional, Tuple, Any, Collection
|
||||
|
||||
import torch.optim
|
||||
import torch.optim
|
||||
import torch.utils.data
|
||||
import torch.utils.data
|
||||
from labml import tracker, logger, monit
|
||||
from labml.configs import BaseConfigs, meta_config, option
|
||||
from labml.internal.monitor import Loop
|
||||
from labml.logger import Text
|
||||
from torch import nn
|
||||
from .device import DeviceConfigs
|
||||
from .metrics import StateModule
|
||||
|
||||
|
||||
class TrainingLoopIterator(Collection):
|
||||
def __init__(self, start: int, total: int, step: Optional[int]):
|
||||
self.step = step
|
||||
self.total = total
|
||||
self.start = start
|
||||
self.i = None
|
||||
|
||||
def __iter__(self):
|
||||
self.i = None
|
||||
return self
|
||||
|
||||
def __next__(self):
|
||||
if self.step is not None:
|
||||
if self.i is None:
|
||||
self.i = self.start
|
||||
else:
|
||||
self.i += self.step
|
||||
else:
|
||||
if self.i is None:
|
||||
self.i = 0
|
||||
else:
|
||||
self.i += 1
|
||||
|
||||
if self.i >= self.total:
|
||||
raise StopIteration()
|
||||
|
||||
if self.step is None:
|
||||
return tracker.get_global_step()
|
||||
else:
|
||||
return self.i
|
||||
|
||||
def __len__(self) -> int:
|
||||
if self.step is not None:
|
||||
return (self.total - self.start) // self.step
|
||||
else:
|
||||
return self.total
|
||||
|
||||
def __contains__(self, x: object) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
class TrainingLoop:
|
||||
_iter: Optional[TrainingLoopIterator]
|
||||
__loop: Loop
|
||||
__signal_received: Optional[Tuple[Any, Any]]
|
||||
|
||||
def __init__(self, *,
|
||||
loop_count: int,
|
||||
loop_step: Optional[int],
|
||||
log_new_line_interval: int,
|
||||
log_write_interval: int,
|
||||
is_loop_on_interrupt: bool):
|
||||
self.__loop_count = loop_count
|
||||
self.__loop_step = loop_step
|
||||
self.__log_new_line_interval = log_new_line_interval
|
||||
self.__log_write_interval = log_write_interval
|
||||
self.__last_write_step = 0
|
||||
self.__last_new_line_step = 0
|
||||
self.__last_save_step = 0
|
||||
self.__signal_received = None
|
||||
self.__is_loop_on_interrupt = is_loop_on_interrupt
|
||||
self._iter = None
|
||||
|
||||
def __iter__(self):
|
||||
self._iter = TrainingLoopIterator(tracker.get_global_step(),
|
||||
self.__loop_count,
|
||||
self.__loop_step)
|
||||
|
||||
self.__loop = monit.loop(typing.cast(Collection, self._iter))
|
||||
|
||||
iter(self.__loop)
|
||||
try:
|
||||
self.old_handler = signal.signal(signal.SIGINT, self.__handler)
|
||||
except ValueError:
|
||||
pass
|
||||
return self
|
||||
|
||||
@property
|
||||
def idx(self):
|
||||
if not self._iter:
|
||||
return 0
|
||||
if not self._iter.i:
|
||||
return 0
|
||||
if self.__loop_step is None:
|
||||
return self._iter.i
|
||||
return self._iter.i / self.__loop_step
|
||||
|
||||
def __finish(self):
|
||||
try:
|
||||
signal.signal(signal.SIGINT, self.old_handler)
|
||||
except ValueError:
|
||||
pass
|
||||
tracker.save()
|
||||
tracker.new_line()
|
||||
|
||||
def __next__(self):
|
||||
if self.__signal_received is not None:
|
||||
logger.log('\nKilling Loop.', Text.danger)
|
||||
monit.finish_loop()
|
||||
self.__finish()
|
||||
raise StopIteration("SIGINT")
|
||||
|
||||
try:
|
||||
global_step = next(self.__loop)
|
||||
except StopIteration as e:
|
||||
self.__finish()
|
||||
raise e
|
||||
|
||||
tracker.set_global_step(global_step)
|
||||
|
||||
if global_step - self.__last_write_step >= self.__log_write_interval:
|
||||
tracker.save()
|
||||
self.__last_write_step = global_step
|
||||
if global_step - self.__last_new_line_step >= self.__log_new_line_interval:
|
||||
tracker.new_line()
|
||||
self.__last_new_line_step = global_step
|
||||
|
||||
return global_step
|
||||
|
||||
def __handler(self, sig, frame):
|
||||
# Pass second interrupt without delaying
|
||||
if self.__signal_received is not None:
|
||||
logger.log('\nSIGINT received twice. Stopping...', Text.danger)
|
||||
self.old_handler(*self.__signal_received)
|
||||
return
|
||||
|
||||
if self.__is_loop_on_interrupt:
|
||||
# Store the interrupt signal for later
|
||||
self.__signal_received = (sig, frame)
|
||||
logger.log('\nSIGINT received. Delaying KeyboardInterrupt.', Text.danger)
|
||||
else:
|
||||
self.__finish()
|
||||
logger.log('Killing loop...', Text.danger)
|
||||
self.old_handler(sig, frame)
|
||||
|
||||
def __str__(self):
|
||||
return "LabTrainingLoop"
|
||||
|
||||
|
||||
class TrainingLoopConfigs(BaseConfigs):
|
||||
r"""
|
||||
This is a configurable training loop. You can extend this class for your configurations
|
||||
if it involves a training loop.
|
||||
|
||||
>>> for step in conf.training_loop:
|
||||
>>> ...
|
||||
|
||||
Arguments:
|
||||
loop_count (int): Total number of steps. Defaults to ``10``.
|
||||
loop_step (int): Number of steps to increment per iteration. Defaults to ``1``.
|
||||
log_new_line_interval (int): The interval (in steps) to print a new line to the screen.
|
||||
Defaults to ``1``.
|
||||
log_write_interval (int): The interval (in steps) to call :func:`labml.tracker.save`.
|
||||
Defaults to ``1``.
|
||||
is_loop_on_interrupt (bool): Whether to handle keyboard interrupts and wait until a iteration is complete.
|
||||
Defaults to ``False``.
|
||||
"""
|
||||
loop_count: int = 10
|
||||
loop_step: int = 1
|
||||
log_new_line_interval: int = 1
|
||||
log_write_interval: int = 1
|
||||
is_loop_on_interrupt: bool = False
|
||||
|
||||
training_loop: TrainingLoop
|
||||
|
||||
|
||||
@option(TrainingLoopConfigs.training_loop)
|
||||
def _loop_configs(c: TrainingLoopConfigs):
|
||||
return TrainingLoop(loop_count=c.loop_count,
|
||||
loop_step=c.loop_step,
|
||||
log_new_line_interval=c.log_new_line_interval,
|
||||
log_write_interval=c.log_write_interval,
|
||||
is_loop_on_interrupt=c.is_loop_on_interrupt)
|
||||
|
||||
|
||||
meta_config(TrainingLoopConfigs.loop_step,
|
||||
TrainingLoopConfigs.loop_count,
|
||||
TrainingLoopConfigs.log_new_line_interval,
|
||||
TrainingLoopConfigs.log_write_interval,
|
||||
TrainingLoopConfigs.is_loop_on_interrupt)
|
||||
|
||||
|
||||
class ModeState:
|
||||
def __init__(self):
|
||||
self._rollback_stack = []
|
||||
|
||||
self.is_train = False
|
||||
self.is_optimize = False
|
||||
|
||||
def _enter(self, mode: Dict[str, any]):
|
||||
rollback = {}
|
||||
for k, v in mode.items():
|
||||
if v is None:
|
||||
continue
|
||||
rollback[k] = getattr(self, k)
|
||||
setattr(self, k, v)
|
||||
|
||||
self._rollback_stack.append(rollback)
|
||||
|
||||
return len(self._rollback_stack)
|
||||
|
||||
def _exit(self, n: int):
|
||||
assert n == len(self._rollback_stack)
|
||||
|
||||
rollback = self._rollback_stack[-1]
|
||||
self._rollback_stack.pop(-1)
|
||||
|
||||
for k, v in rollback.items():
|
||||
setattr(self, k, v)
|
||||
|
||||
def update(self, *,
|
||||
is_train: Optional[bool] = None,
|
||||
is_optimize: Optional[bool] = None):
|
||||
return Mode(self,
|
||||
is_train=is_train,
|
||||
is_optimize=is_optimize)
|
||||
|
||||
|
||||
class Mode:
|
||||
def __init__(self, mode: ModeState, **kwargs: any):
|
||||
self.mode = mode
|
||||
self.update = {}
|
||||
for k, v in kwargs.items():
|
||||
if v is not None:
|
||||
self.update[k] = v
|
||||
|
||||
self.idx = -1
|
||||
|
||||
def __enter__(self):
|
||||
self.idx = self.mode._enter(self.update)
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
self.mode._exit(self.idx)
|
||||
|
||||
|
||||
class Trainer:
|
||||
def __init__(self, *,
|
||||
name: str,
|
||||
mode: ModeState,
|
||||
data_loader: torch.utils.data.DataLoader,
|
||||
inner_iterations: int,
|
||||
state_modules: List[StateModule],
|
||||
is_track_time: bool,
|
||||
step: Callable[[any, 'BatchIndex'], None]):
|
||||
self.is_track_time = is_track_time
|
||||
self.mode = mode
|
||||
self.name = name
|
||||
self.step = step
|
||||
self.state_modules = state_modules
|
||||
self.__iterable = None
|
||||
self.__states = [sm.create_state() for sm in self.state_modules]
|
||||
self.inner_iterations = inner_iterations
|
||||
self.data_loader = data_loader
|
||||
self._batch_index = BatchIndex(len(self.data_loader), self.inner_iterations)
|
||||
|
||||
def set_data_loader(self, data_loader: torch.utils.data.DataLoader):
|
||||
self.data_loader = data_loader
|
||||
self._batch_index = BatchIndex(len(data_loader), self.inner_iterations)
|
||||
self.__iterable = None
|
||||
|
||||
def __call__(self):
|
||||
for sm, s in zip(self.state_modules, self.__states):
|
||||
sm.set_state(s)
|
||||
|
||||
if self.__iterable is None or self._batch_index.completed:
|
||||
self.__iterable = iter(self.data_loader)
|
||||
self._batch_index.reset(len(self.data_loader), self.inner_iterations)
|
||||
for sm in self.state_modules:
|
||||
sm.on_epoch_start()
|
||||
with torch.set_grad_enabled(self.mode.is_train):
|
||||
self.__iterate()
|
||||
|
||||
if self._batch_index.completed:
|
||||
for sm in self.state_modules:
|
||||
sm.on_epoch_end()
|
||||
|
||||
def __iterate(self):
|
||||
with monit.section(self.name, is_partial=True, is_track=self.is_track_time):
|
||||
if self._batch_index.idx == 0:
|
||||
monit.progress(0)
|
||||
while not self._batch_index.iteration_completed:
|
||||
batch = next(self.__iterable)
|
||||
|
||||
self.step(batch, self._batch_index)
|
||||
|
||||
self._batch_index.step()
|
||||
monit.progress(self._batch_index.epoch_progress)
|
||||
|
||||
self._batch_index.step_inner()
|
||||
|
||||
|
||||
class BatchIndex:
|
||||
idx: int
|
||||
total: int
|
||||
iteration: int
|
||||
total_iterations: int
|
||||
|
||||
def __init__(self, total: int, total_iterations: int):
|
||||
self.total_iterations = total_iterations
|
||||
self.total = total
|
||||
|
||||
def is_interval(self, interval: int):
|
||||
if interval <= 0:
|
||||
return False
|
||||
if self.idx + 1 == self.total:
|
||||
return True
|
||||
else:
|
||||
return (self.idx + 1) % interval == 0
|
||||
|
||||
@property
|
||||
def is_last(self):
|
||||
return self.idx + 1 == self.total
|
||||
|
||||
@property
|
||||
def completed(self):
|
||||
return self.iteration >= self.total_iterations
|
||||
|
||||
@property
|
||||
def iteration_completed(self):
|
||||
# // is important so that the last step happens on the last iteration
|
||||
return self.idx >= (self.iteration + 1) * self.total // self.total_iterations
|
||||
|
||||
@property
|
||||
def epoch_progress(self):
|
||||
return self.idx / self.total
|
||||
|
||||
def step(self):
|
||||
self.idx += 1
|
||||
|
||||
def step_inner(self):
|
||||
self.iteration += 1
|
||||
|
||||
def reset(self, total: int, total_iterations: int):
|
||||
self.total = total
|
||||
self.total_iterations = total_iterations
|
||||
self.idx = 0
|
||||
self.iteration = 0
|
||||
|
||||
|
||||
class TrainValidConfigs(TrainingLoopConfigs):
|
||||
r"""
|
||||
This is a configurable module that you can extend for experiments that involve a
|
||||
training and validation datasets (i.e. most DL experiments).
|
||||
|
||||
Arguments:
|
||||
epochs (int): Number of epochs to train on. Defaults to ``10``.
|
||||
train_loader (torch.utils.data.DataLoader): Training data loader.
|
||||
valid_loader (torch.utils.data.DataLoader): Training data loader.
|
||||
inner_iterations (int): Number of times to switch between training and validation
|
||||
within an epoch. Defaults to ``1``.
|
||||
|
||||
You can override ``init``, ``step`` functions. There is also a ``sample`` function
|
||||
that you can override to generate samples ever time it switches between training and validation.
|
||||
"""
|
||||
state_modules: List[StateModule]
|
||||
|
||||
mode: ModeState
|
||||
|
||||
epochs: int = 10
|
||||
|
||||
trainer: Trainer
|
||||
validator: Trainer
|
||||
train_loader: torch.utils.data.DataLoader
|
||||
valid_loader: torch.utils.data.DataLoader
|
||||
|
||||
loop_count = '_data_loop_count'
|
||||
loop_step = None
|
||||
|
||||
inner_iterations: int = 1
|
||||
|
||||
is_track_time: bool = False
|
||||
|
||||
def init(self):
|
||||
pass
|
||||
|
||||
def step(self, batch: Any, batch_idx: BatchIndex):
|
||||
raise NotImplementedError
|
||||
|
||||
def run_step(self):
|
||||
for i in range(self.inner_iterations):
|
||||
with tracker.namespace('sample'):
|
||||
self.sample()
|
||||
with self.mode.update(is_train=True):
|
||||
with tracker.namespace('train'):
|
||||
self.trainer()
|
||||
if self.validator:
|
||||
with tracker.namespace('valid'):
|
||||
self.validator()
|
||||
tracker.save()
|
||||
|
||||
def run(self):
|
||||
with monit.section("Initialize"):
|
||||
self.init()
|
||||
_ = self.validator
|
||||
_ = self.trainer
|
||||
for _ in self.training_loop:
|
||||
self.run_step()
|
||||
|
||||
def sample(self):
|
||||
pass
|
||||
|
||||
|
||||
@option(TrainValidConfigs.trainer)
|
||||
def _default_trainer(c: TrainValidConfigs):
|
||||
return Trainer(name='Train',
|
||||
mode=c.mode,
|
||||
data_loader=c.train_loader,
|
||||
inner_iterations=c.inner_iterations,
|
||||
state_modules=c.state_modules,
|
||||
is_track_time=c.is_track_time,
|
||||
step=c.step)
|
||||
|
||||
|
||||
@option(TrainValidConfigs.validator)
|
||||
def _default_validator(c: TrainValidConfigs):
|
||||
return Trainer(name='Valid',
|
||||
mode=c.mode,
|
||||
data_loader=c.valid_loader,
|
||||
inner_iterations=c.inner_iterations,
|
||||
state_modules=c.state_modules,
|
||||
is_track_time=c.is_track_time,
|
||||
step=c.step)
|
||||
|
||||
|
||||
@option(TrainValidConfigs.loop_count)
|
||||
def _data_loop_count(c: TrainValidConfigs):
|
||||
return c.epochs
|
||||
|
||||
|
||||
class SimpleTrainValidConfigs(TrainValidConfigs):
|
||||
r"""
|
||||
This is a configurable module that works for many standard DL experiments.
|
||||
|
||||
Arguments:
|
||||
model: A PyTorch model.
|
||||
optimizer: A PyTorch optimizer to update model.
|
||||
device: The device to train the model on. This defaults to a configurable device
|
||||
loss_function: A function to calculate the loss. This should accept ``model_output, target`` as
|
||||
arguments.
|
||||
update_batches (int): Number of batches to accumulate before taking an optimizer step.
|
||||
Defaults to ``1``.
|
||||
log_save_batches (int): How often to call :func:`labml.tracker.save`.
|
||||
"""
|
||||
optimizer: torch.optim.Adam
|
||||
model: nn.Module
|
||||
device: torch.device = DeviceConfigs()
|
||||
|
||||
loss_func: nn.Module
|
||||
|
||||
update_batches: int = 1
|
||||
log_save_batches: int = 1
|
||||
|
||||
state_modules: List[StateModule] = []
|
||||
|
||||
def init(self):
|
||||
pass
|
||||
|
||||
def step(self, batch: Any, batch_idx: BatchIndex):
|
||||
self.model.train(self.mode.is_train)
|
||||
data, target = batch[0].to(self.device), batch[1].to(self.device)
|
||||
|
||||
if self.mode.is_train:
|
||||
tracker.add_global_step(len(data))
|
||||
|
||||
with monit.section("model"):
|
||||
output = self.model(data)
|
||||
|
||||
loss = self.loss_func(output, target)
|
||||
tracker.add("loss.", loss)
|
||||
|
||||
if self.mode.is_train:
|
||||
with monit.section('backward'):
|
||||
loss.backward()
|
||||
|
||||
if batch_idx.is_interval(self.update_batches):
|
||||
with monit.section('optimize'):
|
||||
self.optimizer.step()
|
||||
self.optimizer.zero_grad()
|
||||
|
||||
if batch_idx.is_interval(self.log_save_batches):
|
||||
tracker.save()
|
||||
|
||||
|
||||
meta_config(SimpleTrainValidConfigs.update_batches,
|
||||
)
|
||||
|
||||
|
||||
@option(SimpleTrainValidConfigs.optimizer)
|
||||
def _default_optimizer(c: SimpleTrainValidConfigs):
|
||||
from .optimizer import OptimizerConfigs
|
||||
opt_conf = OptimizerConfigs()
|
||||
opt_conf.parameters = c.model.parameters()
|
||||
return opt_conf
|
||||
Reference in New Issue
Block a user