text: # Copyright 2020-present, Pietro Buzzega, Matteo Boschini, Angelo Porrello, Davide Abati, Simone Calderara.
  # All rights reserved.
  # This source code is licensed under the license found in the
  # LICENSE file in the root directory of this source tree.
  
  from utils.buffer import Buffer
  from torch.nn import functional as F
  from utils.args import *
  from utils.distill import class_logits_from_subclass_logits,aux_loss
  from models.utils.continual_model import ContinualModel
  
  
  def get_parser() -> ArgumentParser:
      parser = ArgumentParser(description='Continual learning via'
                                          ' Dark Experience Replay.')
      add_management_args(parser)
      add_experiment_args(parser)
      add_rehearsal_args(parser)
      parser.add_argument('--alpha', type=float, required=True,
                          help='Penalty weight.')
      return parser
  
  
  class Der(ContinualModel):
      NAME = 'der'
      COMPATIBILITY = ['class-il', 'domain-il', 'task-il', 'general-continual']
  
      def __init__(self, backbone, loss, args, transform):
          super(Der, self).__init__(backbone, loss, args, transform)
          self.buffer = Buffer(self.args.buffer_size, self.device)
  
      def observe(self, inputs, labels, not_aug_inputs,num_classes):
          self.opt.zero_grad()
  
          outputs = self.net(inputs)
          class_outputs = class_logits_from_subclass_logits(outputs, num_classes)
          if self.args.aux:
              auxiliary = aux_loss(outputs, self.device)
              loss = self.loss(class_outputs, labels) + self.args.aux_weight * auxiliary.squeeze()
          else:
              loss = self.loss(class_outputs, labels)
  
          if not self.buffer.is_empty():
              buf_inputs, buf_logits = self.buffer.get_data(
                  self.args.minibatch_size, transform=self.transform)
              buf_outputs = self.net(buf_inputs)
              # class_buf_outputs = class_logits_from_subclass_logits(buf_outputs,20)
              loss += self.args.alpha * F.mse_loss(buf_outputs, buf_logits)
  
          loss.backward()
          self.opt.step()
          self.buffer.add_data(examples=not_aug_inputs, logits=outputs.data)
  
          return loss.item()
  