Spaces:
Sleeping
Sleeping
# -*- coding: utf-8 -*- | |
# Copyright (c) Facebook, Inc. and its affiliates. | |
import concurrent.futures | |
import logging | |
import numpy as np | |
import time | |
import weakref | |
from typing import List, Mapping, Optional | |
import torch | |
from torch.nn.parallel import DataParallel, DistributedDataParallel | |
import detectron2.utils.comm as comm | |
from detectron2.utils.events import EventStorage, get_event_storage | |
from detectron2.utils.logger import _log_api_usage | |
__all__ = ["HookBase", "TrainerBase", "SimpleTrainer", "AMPTrainer"] | |
class HookBase: | |
""" | |
Base class for hooks that can be registered with :class:`TrainerBase`. | |
Each hook can implement 4 methods. The way they are called is demonstrated | |
in the following snippet: | |
:: | |
hook.before_train() | |
for iter in range(start_iter, max_iter): | |
hook.before_step() | |
trainer.run_step() | |
hook.after_step() | |
iter += 1 | |
hook.after_train() | |
Notes: | |
1. In the hook method, users can access ``self.trainer`` to access more | |
properties about the context (e.g., model, current iteration, or config | |
if using :class:`DefaultTrainer`). | |
2. A hook that does something in :meth:`before_step` can often be | |
implemented equivalently in :meth:`after_step`. | |
If the hook takes non-trivial time, it is strongly recommended to | |
implement the hook in :meth:`after_step` instead of :meth:`before_step`. | |
The convention is that :meth:`before_step` should only take negligible time. | |
Following this convention will allow hooks that do care about the difference | |
between :meth:`before_step` and :meth:`after_step` (e.g., timer) to | |
function properly. | |
""" | |
trainer: "TrainerBase" = None | |
""" | |
A weak reference to the trainer object. Set by the trainer when the hook is registered. | |
""" | |
def before_train(self): | |
""" | |
Called before the first iteration. | |
""" | |
pass | |
def after_train(self): | |
""" | |
Called after the last iteration. | |
""" | |
pass | |
def before_step(self): | |
""" | |
Called before each iteration. | |
""" | |
pass | |
def after_backward(self): | |
""" | |
Called after the backward pass of each iteration. | |
""" | |
pass | |
def after_step(self): | |
""" | |
Called after each iteration. | |
""" | |
pass | |
def state_dict(self): | |
""" | |
Hooks are stateless by default, but can be made checkpointable by | |
implementing `state_dict` and `load_state_dict`. | |
""" | |
return {} | |
class TrainerBase: | |
""" | |
Base class for iterative trainer with hooks. | |
The only assumption we made here is: the training runs in a loop. | |
A subclass can implement what the loop is. | |
We made no assumptions about the existence of dataloader, optimizer, model, etc. | |
Attributes: | |
iter(int): the current iteration. | |
start_iter(int): The iteration to start with. | |
By convention the minimum possible value is 0. | |
max_iter(int): The iteration to end training. | |
storage(EventStorage): An EventStorage that's opened during the course of training. | |
""" | |
def __init__(self) -> None: | |
self._hooks: List[HookBase] = [] | |
self.iter: int = 0 | |
self.start_iter: int = 0 | |
self.max_iter: int | |
self.storage: EventStorage | |
_log_api_usage("trainer." + self.__class__.__name__) | |
def register_hooks(self, hooks: List[Optional[HookBase]]) -> None: | |
""" | |
Register hooks to the trainer. The hooks are executed in the order | |
they are registered. | |
Args: | |
hooks (list[Optional[HookBase]]): list of hooks | |
""" | |
hooks = [h for h in hooks if h is not None] | |
for h in hooks: | |
assert isinstance(h, HookBase) | |
# To avoid circular reference, hooks and trainer cannot own each other. | |
# This normally does not matter, but will cause memory leak if the | |
# involved objects contain __del__: | |
# See http://engineering.hearsaysocial.com/2013/06/16/circular-references-in-python/ | |
h.trainer = weakref.proxy(self) | |
self._hooks.extend(hooks) | |
def train(self, start_iter: int, max_iter: int): | |
""" | |
Args: | |
start_iter, max_iter (int): See docs above | |
""" | |
logger = logging.getLogger(__name__) | |
logger.info("Starting training from iteration {}".format(start_iter)) | |
self.iter = self.start_iter = start_iter | |
self.max_iter = max_iter | |
with EventStorage(start_iter) as self.storage: | |
try: | |
self.before_train() | |
for self.iter in range(start_iter, max_iter): | |
self.before_step() | |
self.run_step() | |
self.after_step() | |
# self.iter == max_iter can be used by `after_train` to | |
# tell whether the training successfully finished or failed | |
# due to exceptions. | |
self.iter += 1 | |
except Exception: | |
logger.exception("Exception during training:") | |
raise | |
finally: | |
self.after_train() | |
def before_train(self): | |
for h in self._hooks: | |
h.before_train() | |
def after_train(self): | |
self.storage.iter = self.iter | |
for h in self._hooks: | |
h.after_train() | |
def before_step(self): | |
# Maintain the invariant that storage.iter == trainer.iter | |
# for the entire execution of each step | |
self.storage.iter = self.iter | |
for h in self._hooks: | |
h.before_step() | |
def after_backward(self): | |
for h in self._hooks: | |
h.after_backward() | |
def after_step(self): | |
for h in self._hooks: | |
h.after_step() | |
def run_step(self): | |
raise NotImplementedError | |
def state_dict(self): | |
ret = {"iteration": self.iter} | |
hooks_state = {} | |
for h in self._hooks: | |
sd = h.state_dict() | |
if sd: | |
name = type(h).__qualname__ | |
if name in hooks_state: | |
# TODO handle repetitive stateful hooks | |
continue | |
hooks_state[name] = sd | |
if hooks_state: | |
ret["hooks"] = hooks_state | |
return ret | |
def load_state_dict(self, state_dict): | |
logger = logging.getLogger(__name__) | |
self.iter = state_dict["iteration"] | |
for key, value in state_dict.get("hooks", {}).items(): | |
for h in self._hooks: | |
try: | |
name = type(h).__qualname__ | |
except AttributeError: | |
continue | |
if name == key: | |
h.load_state_dict(value) | |
break | |
else: | |
logger.warning(f"Cannot find the hook '{key}', its state_dict is ignored.") | |
class SimpleTrainer(TrainerBase): | |
""" | |
A simple trainer for the most common type of task: | |
single-cost single-optimizer single-data-source iterative optimization, | |
optionally using data-parallelism. | |
It assumes that every step, you: | |
1. Compute the loss with a data from the data_loader. | |
2. Compute the gradients with the above loss. | |
3. Update the model with the optimizer. | |
All other tasks during training (checkpointing, logging, evaluation, LR schedule) | |
are maintained by hooks, which can be registered by :meth:`TrainerBase.register_hooks`. | |
If you want to do anything fancier than this, | |
either subclass TrainerBase and implement your own `run_step`, | |
or write your own training loop. | |
""" | |
def __init__( | |
self, | |
model, | |
data_loader, | |
optimizer, | |
gather_metric_period=1, | |
zero_grad_before_forward=False, | |
async_write_metrics=False, | |
): | |
""" | |
Args: | |
model: a torch Module. Takes a data from data_loader and returns a | |
dict of losses. | |
data_loader: an iterable. Contains data to be used to call model. | |
optimizer: a torch optimizer. | |
gather_metric_period: an int. Every gather_metric_period iterations | |
the metrics are gathered from all the ranks to rank 0 and logged. | |
zero_grad_before_forward: whether to zero the gradients before the forward. | |
async_write_metrics: bool. If True, then write metrics asynchronously to improve | |
training speed | |
""" | |
super().__init__() | |
""" | |
We set the model to training mode in the trainer. | |
However it's valid to train a model that's in eval mode. | |
If you want your model (or a submodule of it) to behave | |
like evaluation during training, you can overwrite its train() method. | |
""" | |
model.train() | |
self.model = model | |
self.data_loader = data_loader | |
# to access the data loader iterator, call `self._data_loader_iter` | |
self._data_loader_iter_obj = None | |
self.optimizer = optimizer | |
self.gather_metric_period = gather_metric_period | |
self.zero_grad_before_forward = zero_grad_before_forward | |
self.async_write_metrics = async_write_metrics | |
# create a thread pool that can execute non critical logic in run_step asynchronically | |
# use only 1 worker so tasks will be executred in order of submitting. | |
self.concurrent_executor = concurrent.futures.ThreadPoolExecutor(max_workers=1) | |
def run_step(self): | |
""" | |
Implement the standard training logic described above. | |
""" | |
assert self.model.training, "[SimpleTrainer] model was changed to eval mode!" | |
start = time.perf_counter() | |
""" | |
If you want to do something with the data, you can wrap the dataloader. | |
""" | |
data = next(self._data_loader_iter) | |
data_time = time.perf_counter() - start | |
if self.zero_grad_before_forward: | |
""" | |
If you need to accumulate gradients or do something similar, you can | |
wrap the optimizer with your custom `zero_grad()` method. | |
""" | |
self.optimizer.zero_grad() | |
""" | |
If you want to do something with the losses, you can wrap the model. | |
""" | |
loss_dict = self.model(data) | |
if isinstance(loss_dict, torch.Tensor): | |
losses = loss_dict | |
loss_dict = {"total_loss": loss_dict} | |
else: | |
losses = sum(loss_dict.values()) | |
if not self.zero_grad_before_forward: | |
""" | |
If you need to accumulate gradients or do something similar, you can | |
wrap the optimizer with your custom `zero_grad()` method. | |
""" | |
self.optimizer.zero_grad() | |
losses.backward() | |
self.after_backward() | |
if self.async_write_metrics: | |
# write metrics asynchronically | |
self.concurrent_executor.submit( | |
self._write_metrics, loss_dict, data_time, iter=self.iter | |
) | |
else: | |
self._write_metrics(loss_dict, data_time) | |
""" | |
If you need gradient clipping/scaling or other processing, you can | |
wrap the optimizer with your custom `step()` method. But it is | |
suboptimal as explained in https://arxiv.org/abs/2006.15704 Sec 3.2.4 | |
""" | |
self.optimizer.step() | |
def _data_loader_iter(self): | |
# only create the data loader iterator when it is used | |
if self._data_loader_iter_obj is None: | |
self._data_loader_iter_obj = iter(self.data_loader) | |
return self._data_loader_iter_obj | |
def reset_data_loader(self, data_loader_builder): | |
""" | |
Delete and replace the current data loader with a new one, which will be created | |
by calling `data_loader_builder` (without argument). | |
""" | |
del self.data_loader | |
data_loader = data_loader_builder() | |
self.data_loader = data_loader | |
self._data_loader_iter_obj = None | |
def _write_metrics( | |
self, | |
loss_dict: Mapping[str, torch.Tensor], | |
data_time: float, | |
prefix: str = "", | |
iter: Optional[int] = None, | |
) -> None: | |
logger = logging.getLogger(__name__) | |
iter = self.iter if iter is None else iter | |
if (iter + 1) % self.gather_metric_period == 0: | |
try: | |
SimpleTrainer.write_metrics(loss_dict, data_time, iter, prefix) | |
except Exception: | |
logger.exception("Exception in writing metrics: ") | |
raise | |
def write_metrics( | |
loss_dict: Mapping[str, torch.Tensor], | |
data_time: float, | |
cur_iter: int, | |
prefix: str = "", | |
) -> None: | |
""" | |
Args: | |
loss_dict (dict): dict of scalar losses | |
data_time (float): time taken by the dataloader iteration | |
prefix (str): prefix for logging keys | |
""" | |
metrics_dict = {k: v.detach().cpu().item() for k, v in loss_dict.items()} | |
metrics_dict["data_time"] = data_time | |
storage = get_event_storage() | |
# Keep track of data time per rank | |
storage.put_scalar("rank_data_time", data_time, cur_iter=cur_iter) | |
# Gather metrics among all workers for logging | |
# This assumes we do DDP-style training, which is currently the only | |
# supported method in detectron2. | |
all_metrics_dict = comm.gather(metrics_dict) | |
if comm.is_main_process(): | |
# data_time among workers can have high variance. The actual latency | |
# caused by data_time is the maximum among workers. | |
data_time = np.max([x.pop("data_time") for x in all_metrics_dict]) | |
storage.put_scalar("data_time", data_time, cur_iter=cur_iter) | |
# average the rest metrics | |
metrics_dict = { | |
k: np.mean([x[k] for x in all_metrics_dict]) for k in all_metrics_dict[0].keys() | |
} | |
total_losses_reduced = sum(metrics_dict.values()) | |
if not np.isfinite(total_losses_reduced): | |
raise FloatingPointError( | |
f"Loss became infinite or NaN at iteration={cur_iter}!\n" | |
f"loss_dict = {metrics_dict}" | |
) | |
storage.put_scalar( | |
"{}total_loss".format(prefix), total_losses_reduced, cur_iter=cur_iter | |
) | |
if len(metrics_dict) > 1: | |
storage.put_scalars(cur_iter=cur_iter, **metrics_dict) | |
def state_dict(self): | |
ret = super().state_dict() | |
ret["optimizer"] = self.optimizer.state_dict() | |
return ret | |
def load_state_dict(self, state_dict): | |
super().load_state_dict(state_dict) | |
self.optimizer.load_state_dict(state_dict["optimizer"]) | |
def after_train(self): | |
super().after_train() | |
self.concurrent_executor.shutdown(wait=True) | |
class AMPTrainer(SimpleTrainer): | |
""" | |
Like :class:`SimpleTrainer`, but uses PyTorch's native automatic mixed precision | |
in the training loop. | |
""" | |
def __init__( | |
self, | |
model, | |
data_loader, | |
optimizer, | |
gather_metric_period=1, | |
zero_grad_before_forward=False, | |
grad_scaler=None, | |
precision: torch.dtype = torch.float16, | |
log_grad_scaler: bool = False, | |
async_write_metrics=False, | |
): | |
""" | |
Args: | |
model, data_loader, optimizer, gather_metric_period, zero_grad_before_forward, | |
async_write_metrics: same as in :class:`SimpleTrainer`. | |
grad_scaler: torch GradScaler to automatically scale gradients. | |
precision: torch.dtype as the target precision to cast to in computations | |
""" | |
unsupported = "AMPTrainer does not support single-process multi-device training!" | |
if isinstance(model, DistributedDataParallel): | |
assert not (model.device_ids and len(model.device_ids) > 1), unsupported | |
assert not isinstance(model, DataParallel), unsupported | |
super().__init__( | |
model, data_loader, optimizer, gather_metric_period, zero_grad_before_forward | |
) | |
if grad_scaler is None: | |
from torch.cuda.amp import GradScaler | |
grad_scaler = GradScaler() | |
self.grad_scaler = grad_scaler | |
self.precision = precision | |
self.log_grad_scaler = log_grad_scaler | |
def run_step(self): | |
""" | |
Implement the AMP training logic. | |
""" | |
assert self.model.training, "[AMPTrainer] model was changed to eval mode!" | |
assert torch.cuda.is_available(), "[AMPTrainer] CUDA is required for AMP training!" | |
from torch.cuda.amp import autocast | |
start = time.perf_counter() | |
data = next(self._data_loader_iter) | |
data_time = time.perf_counter() - start | |
if self.zero_grad_before_forward: | |
self.optimizer.zero_grad() | |
with autocast(dtype=self.precision): | |
loss_dict = self.model(data) | |
if isinstance(loss_dict, torch.Tensor): | |
losses = loss_dict | |
loss_dict = {"total_loss": loss_dict} | |
else: | |
losses = sum(loss_dict.values()) | |
if not self.zero_grad_before_forward: | |
self.optimizer.zero_grad() | |
self.grad_scaler.scale(losses).backward() | |
if self.log_grad_scaler: | |
storage = get_event_storage() | |
storage.put_scalar("[metric]grad_scaler", self.grad_scaler.get_scale()) | |
self.after_backward() | |
if self.async_write_metrics: | |
# write metrics asynchronically | |
self.concurrent_executor.submit( | |
self._write_metrics, loss_dict, data_time, iter=self.iter | |
) | |
else: | |
self._write_metrics(loss_dict, data_time) | |
self.grad_scaler.step(self.optimizer) | |
self.grad_scaler.update() | |
def state_dict(self): | |
ret = super().state_dict() | |
ret["grad_scaler"] = self.grad_scaler.state_dict() | |
return ret | |
def load_state_dict(self, state_dict): | |
super().load_state_dict(state_dict) | |
self.grad_scaler.load_state_dict(state_dict["grad_scaler"]) | |