fastreid.engine.train_loop.SimpleTrainer Class Reference

Inheritance diagram for fastreid.engine.train_loop.SimpleTrainer:

Public Member Functions
	__init__ (self, model, data_loader, optimizer, amp_enabled)
	run_step (self)
Public Member Functions inherited from fastreid.engine.train_loop.TrainerBase
	register_hooks (self, hooks)
	train (self, int start_iter, int max_iter)
	before_train (self)
	after_train (self)
	before_step (self)
	after_step (self)

Public Attributes
	model
	data_loader
	optimizer
	amp_enabled
	scaler
	iter
Public Attributes inherited from fastreid.engine.train_loop.TrainerBase
	iter
	max_iter

Protected Member Functions
	_detect_anomaly (self, losses, loss_dict)
	_write_metrics (self, dict metrics_dict)

Protected Attributes
	_data_loader_iter
Protected Attributes inherited from fastreid.engine.train_loop.TrainerBase
	_hooks

Detailed Description

A simple trainer for the most common type of task:
single-cost single-optimizer single-data-source iterative optimization.
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.
If you want to do anything fancier than this,
either subclass TrainerBase and implement your own `run_step`,
or write your own training loop.

Definition at line 154 of file train_loop.py.

Constructor & Destructor Documentation

__init__()

fastreid.engine.train_loop.SimpleTrainer.__init__	(		self,
			model,
			data_loader,
			optimizer,
			amp_enabled )

Args:
    model: a torch Module. Takes a data from data_loader and returns a
        dict of heads.
    data_loader: an iterable. Contains data to be used to call model.
    optimizer: a torch optimizer.

Reimplemented from fastreid.engine.train_loop.TrainerBase.

Reimplemented in fastreid.engine.defaults.DefaultTrainer.

Definition at line 167 of file train_loop.py.

    def __init__(self, model, data_loader, optimizer, amp_enabled):
        """
        Args:
            model: a torch Module. Takes a data from data_loader and returns a
                dict of heads.
            data_loader: an iterable. Contains data to be used to call model.
            optimizer: a torch optimizer.
        """
        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
        self._data_loader_iter = iter(data_loader)
        self.optimizer = optimizer
        self.amp_enabled = amp_enabled
 
        if amp_enabled:
            # Creates a GradScaler once at the beginning of training.
            self.scaler = amp.GradScaler()
 

Member Function Documentation

_detect_anomaly()

fastreid.engine.train_loop.SimpleTrainer._detect_anomaly ( self, losses, loss_dict )

protected

Definition at line 246 of file train_loop.py.

    def _detect_anomaly(self, losses, loss_dict):
        if not torch.isfinite(losses).all():
            raise FloatingPointError(
                "Loss became infinite or NaN at iteration={}!\nloss_dict = {}".format(
                    self.iter, loss_dict
                )
            )
 

_write_metrics()

fastreid.engine.train_loop.SimpleTrainer._write_metrics ( self, dict metrics_dict )

protected

Args:
    metrics_dict (dict): dict of scalar metrics

Definition at line 254 of file train_loop.py.

    def _write_metrics(self, metrics_dict: dict):
        """
        Args:
            metrics_dict (dict): dict of scalar metrics
        """
        metrics_dict = {
            k: v.detach().cpu().item() if isinstance(v, torch.Tensor) else float(v)
            for k, v in metrics_dict.items()
        }
        # gather metrics among all workers for logging
        # This assumes we do DDP-style training, which is currently the only
        # supported method in fastreid.
        all_metrics_dict = comm.gather(metrics_dict)
 
        if comm.is_main_process():
            if "data_time" in all_metrics_dict[0]:
                # 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])
                self.storage.put_scalar("data_time", data_time)
 
            # 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(loss for loss in metrics_dict.values())
 
            self.storage.put_scalar("total_loss", total_losses_reduced)
            if len(metrics_dict) > 1:
                self.storage.put_scalars(**metrics_dict)

run_step()

fastreid.engine.train_loop.SimpleTrainer.run_step

(

self

)

Implement the standard training logic described above.

Reimplemented from fastreid.engine.train_loop.TrainerBase.

Definition at line 195 of file train_loop.py.

    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 your 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 your want to do something with the heads, you can wrap the model.
        """
 
        with amp.autocast(enabled=self.amp_enabled):
            outs = self.model(data)
 
            # Compute loss
            if isinstance(self.model, DistributedDataParallel):
                loss_dict = self.model.module.losses(outs)
            else:
                loss_dict = self.model.losses(outs)
 
            losses = sum(loss_dict.values())
 
        with torch.cuda.stream(torch.cuda.Stream()):
            metrics_dict = loss_dict
            metrics_dict["data_time"] = data_time
            self._write_metrics(metrics_dict)
            self._detect_anomaly(losses, loss_dict)
 
        """
        If you need accumulate gradients or something similar, you can
        wrap the optimizer with your custom `zero_grad()` method.
        """
        self.optimizer.zero_grad()
 
        if self.amp_enabled:
            self.scaler.scale(losses).backward()
            self.scaler.step(self.optimizer)
            self.scaler.update()
        else:
            losses.backward()
            """
            If you need gradient clipping/scaling or other processing, you can
            wrap the optimizer with your custom `step()` method.
            """
            self.optimizer.step()
 

Member Data Documentation

_data_loader_iter

fastreid.engine.train_loop.SimpleTrainer._data_loader_iter

protected

Definition at line 187 of file train_loop.py.

amp_enabled

fastreid.engine.train_loop.SimpleTrainer.amp_enabled

Definition at line 189 of file train_loop.py.

data_loader

fastreid.engine.train_loop.SimpleTrainer.data_loader

Definition at line 186 of file train_loop.py.

iter

fastreid.engine.train_loop.SimpleTrainer.iter

Definition at line 250 of file train_loop.py.

model

fastreid.engine.train_loop.SimpleTrainer.model

Definition at line 185 of file train_loop.py.

optimizer

fastreid.engine.train_loop.SimpleTrainer.optimizer

Definition at line 188 of file train_loop.py.

scaler

fastreid.engine.train_loop.SimpleTrainer.scaler

Definition at line 193 of file train_loop.py.

---

The documentation for this class was generated from the following file:

• smreid/fastreid/engine/train_loop.py