code: import torchbearer
  
  from torchbearer.callbacks import Callback
  
  import torch
  
  
  class WeightDecay(Callback):
      """Create a WeightDecay callback which uses the given norm on the given parameters and with the given decay rate.
      If params is None (default) then the parameters will be retrieved from the model.
  
      Example: ::
  
          >>> from torchbearer import Trial
          >>> from torchbearer.callbacks import WeightDecay
  
          # Example Trial which runs a trial with weight decay on the model
          >>> decay = WeightDecay()
          >>> trial = Trial(None, callbacks=[decay], metrics=['loss'], verbose=2).for_steps(10).run(1)
  
      Args:
          rate (float): The decay rate or lambda
          p (int): The norm level
          params (Iterable[Tensor] or Tensor, optional): an iterable of Tensors or a
              single Tensor that will have gradients normalized, otherwise this is retrieved from state
  
      State Requirements:
          - :attr:`torchbearer.state.MODEL`: Model should have the `parameters` method
          - :attr:`torchbearer.state.LOSS`: Loss should be a tensor that can be incremented
      """
      def __init__(self, rate=5e-4, p=2, params=None):
          super(WeightDecay, self).__init__()
  
          self.p = p
          self.params = params
          self.rate = rate
  
      def on_start(self, state):
          """Retrieve params from state['model'] if required.
  
          Args:
              state (dict): The :class:`.Trial` state
          """
          if self.params is None:
              self.params = state[torchbearer.MODEL].parameters()
  
      def on_criterion(self, state):
          """Calculate the decay term and add to state['loss'].
  
          Args:
              state (dict): The :class:`.Trial` state
          """
          for param in self.params:
              state[torchbearer.LOSS] += self.rate * torch.norm(param, self.p)
  
  
  class L1WeightDecay(WeightDecay):
      """WeightDecay callback which uses an L1 norm with the given rate and parameters. If params is None (default) then
      the parameters will be retrieved from the model.
  
      Example: ::
  
          >>> from torchbearer import Trial
          >>> from torchbearer.callbacks import L1WeightDecay
  
          # Example Trial which runs a trial with weight decay on the model using an L1 norm
          >>> decay = L1WeightDecay()
          >>> trial = Trial(None, callbacks=[decay], metrics=['loss'], verbose=2).for_steps(10).run(1)
  
      Args:
          rate (float): The decay rate or lambda
          params (Iterable[Tensor] or Tensor, optional): an iterable of Tensors or a
              single Tensor that will have gradients normalized, otherwise this is retrieved from state
  
      State Requirements:
          - :attr:`torchbearer.state.MODEL`: Model should have the `parameters` method
          - :attr:`torchbearer.state.LOSS`: Loss should be a tensor that can be incremented
      """
      def __init__(self, rate=5e-4, params=None):
          super(L1WeightDecay, self).__init__(rate=rate, p=1, params=params)
  
  
  class L2WeightDecay(WeightDecay):
      """WeightDecay callback which uses an L2 norm with the given rate and parameters. If params is None (default) then
      the parameters will be retrieved from the model.
  
      Example: ::
  
          >>> from torchbearer import Trial
          >>> from torchbearer.callbacks import L2WeightDecay
  
          # Example Trial which runs a trial with weight decay on the model using an L2 norm
          >>> decay = L2WeightDecay()
          >>> trial = Trial(None, callbacks=[decay], metrics=['loss'], verbose=2).for_steps(10).run(1)
  
      Args:
          rate (float): The decay rate or lambda
          params (Iterable[Tensor] or Tensor, optional): an iterable of Tensors or a
              single Tensor that will have gradients normalized, otherwise this is retrieved from state
  
      State Requirements:
          - :attr:`torchbearer.state.MODEL`: Model should have the `parameters` method
          - :attr:`torchbearer.state.LOSS`: Loss should be a tensor that can be incremented
      """
      def __init__(self, rate=5e-4, params=None):
          super(L2WeightDecay, self).__init__(rate=rate, p=2, params=params)
   import torchbearer
  from torchbearer import Callback
  import torch
  import torch.nn.functional as F
  from torch.distributions import Beta
  
  from torchbearer.bases import cite
  
  bc = """
  @inproceedings{tokozume2018between,
    title={Between-class learning for image classification},
    author={Tokozume, Yuji and Ushiku, Yoshitaka and Harada, Tatsuya},
    booktitle={Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition},
    pages={5486--5494},
    year={2018}
  }
  """
  
  
  @cite(bc)
  class BCPlus(Callback):
      """BC+ callback which mixes images by treating them as waveforms. For standard BC, see :class:`.Mixup`.
      This callback can optionally convert labels to one hot before combining them according to the lambda parameters,
      sampled from a beta distribution, use alpha=1 to replicate the paper. Use with :meth:`BCPlus.bc_loss` or set
      `mixup_loss = True` and use :meth:`.Mixup.mixup_loss`.
  
      .. note::
  
         This callback first sets all images to have zero mean. Consider adding an offset (e.g. 0.5) back before
         visualising.
  
      Example: ::
  
          >>> from torchbearer import Trial
          >>> from torchbearer.callbacks import BCPlus
  
          # Example Trial which does BCPlus regularisation
          >>> bcplus = BCPlus(classes=10)
          >>> trial = Trial(None, criterion=BCPlus.bc_loss, callbacks=[bcplus], metrics=['acc'])
  
      Args:
          mixup_loss (bool): If True, the lambda and targets will be stored for use with the mixup loss function.
          alpha (float): The alpha value for the beta distribution.
          classes (int): The number of classes for conversion to one hot.
  
      State Requirements:
          - :attr:`torchbearer.state.X`: State should have the current data stored and correctly normalised
          - :attr:`torchbearer.state.Y_TRUE`: State should have the current data stored
      """
  
      def __init__(self, mixup_loss=False, alpha=1, classes=-1):
          super(BCPlus, self).__init__()
          self.mixup_loss = mixup_loss
          self.classes = classes
          self.dist = Beta(torch.tensor([float(alpha)]), torch.tensor([float(alpha)]))
  
      @staticmethod
      def bc_loss(state):
          """The KL divergence between the outputs of the model and the ratio labels. Model ouputs should be un-normalised
          logits as this function performs a log_softmax.
  
          Args:
              state: The current :class:`Trial` state.
          """
          prediction, target = state[torchbearer.Y_PRED], state[torchbearer.Y_TRUE]
  
          entropy = - (target[target.nonzero().split(1, dim=1)] * target[target.nonzero().split(1, dim=1)].log()).sum()
          cross = - (target * F.log_softmax(prediction, dim=1)).sum()
  
          return (cross - entropy) / prediction.size(0)
  
      def _to_one_hot(self, target):
          if target.dim() == 1:
              target = target.unsqueeze(1)
              one_hot = torch.zeros_like(target).repeat(1, self.classes)
              one_hot.scatter_(1, target, 1)
              return one_hot
          return target.float()
  
      def on_sample(self, state):
          super(BCPlus, self).on_sample(state)
  
          lam = self.dist.sample().to(state[torchbearer.DEVICE])
  
          permutation = torch.randperm(state[torchbearer.X].size(0))
  
          batch1 = state[torchbearer.X]
          batch1 = batch1 - batch1.view(batch1.size(0), -1).mean(1, keepdim=True).view(*tuple([batch1.size(0)] + [1] * (batch1.dim() - 1)))
          g1 = batch1.view(batch1.size(0), -1).std(1, keepdim=True).view(*tuple([batch1.size(0)] + [1] * (batch1.dim() - 1)))
  
          batch2 = batch1[permutation]
          g2 = g1[permutation]
  
          p = 1. / (1 + ((g1 / g2) * ((1 - lam) / lam)))
  
          state[torchbearer.X] = (batch1 * p + batch2 * (1 - p)) / (p.pow(2) + (1 - p).pow(2)).sqrt()
  
          if not self.mixup_loss:
              target = self._to_one_hot(state[torchbearer.TARGET]).float()
              state[torchbearer.Y_TRUE] = lam * target + (1 - lam) * target[permutation]
          else:
              state[torchbearer.MIXUP_LAMBDA] = lam
              state[torchbearer.MIXUP_PERMUTATION] = permutation
              state[torchbearer.Y_TRUE] = (state[torchbearer.Y_TRUE], state[torchbearer.Y_TRUE][state[torchbearer.MIXUP_PERMUTATION]])
  
      def on_sample_validation(self, state):
          super(BCPlus, self).on_sample_validation(state)
          if not self.mixup_loss:
              state[torchbearer.TARGET] = self._to_one_hot(state[torchbearer.TARGET]).float()
   import unittest
  
  from mock import Mock, call
  
  from torchbearer.metrics import RunningMean, Metric, RunningMetric, Mean, Std, Var
  
  import torch
  
  
  class TestVar(unittest.TestCase):
      def test_variance_dim(self):
          var = Var('test', dim=0)
          var.process(torch.Tensor([[1., 2.], [3., 4.]]))
          var.process(torch.Tensor([[4., 3.], [2., 1.]]))
          var.process(torch.Tensor([[1., 1.], [1., 1.]]))
  
          res = var.process_final()
          self.assertTrue(len(res) == 2)
          for m in res:
              self.assertTrue(abs(m - 1.6000) < 0.0001)
  
  
  class TestStd(unittest.TestCase):
      def setUp(self):
          self._metric = Metric('test')
          self._metric.process = Mock()
          self._metric.process.side_effect = [torch.zeros(torch.Size([])),
                                              torch.FloatTensor([0.1, 0.2, 0.3]),
                                              torch.FloatTensor([0.4, 0.5, 0.6]),
                                              torch.FloatTensor([0.7, 0.8, 0.9]),
                                              torch.ones(torch.Size([]))]
  
          self._std = Std('test', unbiased=False)
          self._std.reset({})
          self._target = 0.31622776601684
  
      def test_train(self):
          self.setUp()
          self._std.train()
          for i in range(5):
              self._std.process(self._metric.process())
          result = self._std.process_final({})
          self.assertAlmostEqual(self._target, result, places=5)
  
      def test_validate(self):
          self.setUp()
          self._std.eval()
          for i in range(5):
              self._std.process(self._metric.process())
          result = self._std.process_final({})
          self.assertAlmostEqual(self._target, result, places=5)
  
      def test_precision_error(self):
          self.setUp()
          self._std.train()
          val = torch.tensor([0.55])
          for i in range(2):
              self._std.process(val)
  
          result = self._std.process_final({})
          self.assertEqual(0, result)
  
      def setUpMoreDims(self):
          self._metric = Metric('test')
          self._metric.process = Mock()
          self._metric.process.side_effect = [torch.zeros(torch.Size([])),
                                              torch.FloatTensor([[0.1, 0.2, 0.3], [1.1, 1.2, 1.3]]),
                                              torch.FloatTensor([[0.4, 0.5, 0.6], [1.4, 1.5, 1.6]]),
                                              torch.FloatTensor([[0.7, 0.8, 0.9], [1.7, 1.8, 1.9]]),
                                              torch.ones(torch.Size([]))]
          self._std = Std('test', unbiased=False)
          self._std.reset({})
          self._target = 0.57662804083742
  
      def test_more_dims(self):
          self.setUpMoreDims()
          for i in range(5):
              self._std.process(self._metric.process())
          result = self._std.process_final({})
          self.assertAlmostEqual(self._target, result, places=5)
  
      def test_std_dim(self):
          std = Std('test', dim=0)
          std.process(torch.Tensor([[1., 2.], [3., 4.]]))
          std.process(torch.Tensor([[4., 3.], [2., 1.]]))
          std.process(torch.Tensor([[1., 1.], [1., 1.]]))
  
          res = std.process_final()
          self.assertTrue(len(res) == 2)
          for m in res:
              self.assertTrue(abs(m - 1.2649) < 0.0001)
  
  
  class TestMean(unittest.TestCase):
      def setUp(self):
          self._metric = Metric('test')
          self._metric.process = Mock()
          self._metric.process.side_effect = [torch.zeros(torch.Size([])),
                                              torch.FloatTensor([0.1, 0.2, 0.3]),
                                              torch.FloatTensor([0.4, 0.5, 0.6]),
                                              torch.FloatTensor([0.7, 0.8, 0.9]),
                                              torch.ones(torch.Size([]))]
  
          self._mean = Mean('test')
          self._mean.reset({})
          self._target = 0.5
  
      def test_train_dict(self):
          self.setUp()
          self._mean.train()
          for i in range(5):
              self._mean.process(self._metric.process())
          result = self._mean.process_final({})
          self.assertAlmostEqual(self._target, result, places=5)
  
      def test_validate_dict(self):
          self.setUp()
          self._mean.eval()
          for i in range(5):
              self._mean.process(self._metric.process())
          result = self._mean.process_final({})
          self.assertAlmostEqual(self._target, result, places=5)
  
      def setUpMoreDims(self):
          self._metric = Metric('test')
          self._metric.process = Mock()
          self._metric.process.side_effect = [torch.zeros(torch.Size([])),
                                              torch.FloatTensor([[0.1, 0.2, 0.3], [1.1, 1.2, 1.3]]),
                                              torch.FloatTensor([[0.4, 0.5, 0.6], [1.4, 1.5, 1.6]]),
                                              torch.FloatTensor([[0.7, 0.8, 0.9], [1.7, 1.8, 1.9]]),
                                              torch.ones(torch.Size([]))]
          self._mean = Mean('test')
          self._mean.reset({})
          self._target = 0.95
  
      def test_more_dims(self):
          self.setUpMoreDims()
          for i in range(5):
              self._mean.process(self._metric.process())
          result = self._mean.process_final({})
          self.assertAlmostEqual(self._target, result, places=5)
  
      def test_mean_dim(self):
          mean = Mean('test', dim=0)
          mean.process(torch.Tensor([[1., 2.], [3., 4.]]))
          mean.process(torch.Tensor([[4., 3.], [2., 1.]]))
          mean.process(torch.Tensor([[1., 1.], [1., 1.]]))
  
          res = mean.process_final()
          self.assertTrue(len(res) == 2)
          for m in res:
              self.assertTrue(abs(m - 2.0) < 0.0001)
  
  
  class TestRunningMetric(unittest.TestCase):
      def setUp(self):
          self._metric = RunningMetric('test', batch_size=5, step_size=5)
          self._metric.reset({})
          self._metric._process_train = Mock(return_value=3)
          self._metric._step = Mock(return_value='output')
  
      def test_train_called_with_state(self):
          self._metric.train()
          self._metric.process({'test': -1})
          self._metric._process_train.assert_called_with({'test': -1})
  
      def test_cache_one_step(self):
          self._metric.train()
          for i in range(6):
              self._metric.process({})
          self._metric._step.assert_has_calls([call([3]), call([3, 3, 3, 3, 3])])
  
      def test_empty_methods(self):
          metric = RunningMetric('test')
          self.assertRaises(NotImplementedError, lambda: metric._step(['test']) is None)
          self.assertRaises(NotImplementedError, lambda: metric._process_train(['test']) is None)
  
  
  class TestRunningMean(unittest.TestCase):
      def setUp(self):
          self._metric = Metric('test')
          self._mean = RunningMean('test')
          self._cache = [torch.Tensor([1.0]), torch.Tensor([1.5]), torch.Tensor([2.0])]
          self._target = 1.5
  
      def test_train(self):
          result = self._mean._process_train(torch.FloatTensor([1.0, 1.5, 2.0]))
          self.assertAlmostEqual(self._target, result, 3, 0.002)
  
      def test_step(self):
          result = self._mean._step(self._cache)
          self.assertEqual(self._target, result)
  
      def test_dims(self):
          mean = RunningMean('test', dim=0)
          cache = [mean._process_train(torch.Tensor([[1., 2.], [3., 4.]])),
                   mean._process_train(torch.Tensor([[4., 3.], [2., 1.]])),
                   mean._process_train(torch.Tensor([[1., 1.], [1., 1.]]))]
  
          res = mean._step(cache)
          self.assertTrue(len(res) == 2)
          for m in res:
              self.assertTrue(abs(m - 2.0) < 0.0001)
  
apis: torch.norm
apis: torch.zeros_like torch.nn.functional.log_softmax
apis: torch.Size torch.FloatTensor torch.tensor torch.Tensor