<gh_stars>0
import time
import torch
from torch.distributions.dirichlet import Dirichlet

from detectron2.engine.train_loop import SimpleTrainer


class LinearScalarizationTrainer(SimpleTrainer):
    """
    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, preference_vector: float = torch.ones(2) / 2):
        """
        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.
            preference_vector: Vector detailing the weight between losses
        """
        super().__init__(model, data_loader, optimizer)

        self.preference_vector = preference_vector

    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 you want to do something with the losses, you can wrap the model.
        """
        loss_dict = self.model(data)
        losses = torch.matmul(torch.stack(list(loss_dict.values())), self.preference_vector)
        loss_dict = self.model(data)
        loss_dict['total_loss'] = losses

        """
        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._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()


import time
import torch
from torch.distributions.dirichlet import Dirichlet

from detectron2.engine.train_loop import SimpleTrainer


class LinearScalarizationTrainer(SimpleTrainer):
    """
    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, preference_vector: float = torch.ones(2) / 2):
        """
        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.
            preference_vector: Vector detailing the weight between losses
        """
        super().__init__(model, data_loader, optimizer)

        self.preference_vector = preference_vector

    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 you want to do something with the losses, you can wrap the model.
        """
        loss_dict = self.model(data)
        losses = torch.matmul(torch.stack(list(loss_dict.values())), self.preference_vector)
        loss_dict = self.model(data)
        loss_dict['total_loss'] = losses

        """
        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._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()
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. 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.             preference_vector: Vector detailing the weight between losses Implement the standard training logic described above. If you want to do something with the data, you can wrap the dataloader. If you want to do something with the losses, you can wrap the model. If you need to accumulate gradients or do something similar, you can         wrap the optimizer with your custom `zero_grad()` method. 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