• Search by Module
• Search by Words
• Search Projects

• Java
• Python
• JavaScript
• TypeScript
• C++
• Scala
• Blog

• torchbearer-master
  • PULL_REQUEST_TEMPLATE.md
  • codecov.yml
  • LICENSE
  • torchbearer
    • callbacks
      • manifold_mixup.py
      • init.py
      • decorators.py
      • sample_pairing.py
      • checkpointers.py
      • torch_scheduler.py
      • cutout.py
      • csv_logger.py
      • pycm.py
      • printer.py
      • live_loss_plot.py
      • label_smoothing.py
      • __init__.py
      • between_class.py
      • callbacks.py
      • early_stopping.py
      • imaging
        • __init__.py
        • inside_cnns.py
        • imaging.py
      • gradient_clipping.py
      • tensor_board.py
      • unpack_state.py
      • weight_decay.py
      • mixup.py
      • aggregate_predictions.py
      • lsuv.py
      • terminate_on_nan.py
    • state.py
    • version.py
    • magics.py
    • cv_utils.py
    • trial.py
    • bases.py
    • __init__.py
    • metrics
      • decorators.py
      • primitives.py
      • default.py
      • wrappers.py
      • lr.py
      • metrics.py
      • roc_auc_score.py
      • aggregators.py
      • __init__.py
      • timer.py
  • CHANGELOG.md
  • setup.py
  • .travis.yml
  • README.md
  • tests
    • callbacks
      • test_manifold_mixup.py
      • test_early_stopping.py
      • test_weight_decay.py
      • test_cutout.py
      • test_callbacks.py
      • test_printer.py
      • test_checkpointers.py
      • test_decorators.py
      • test_live_loss_plot.py
      • test_sample_pairing.py
      • test_init.py
      • test_gradient_clipping.py
      • test_terminate_on_nan.py
      • test_pycm.py
      • test_label_smoothing.py
      • __init__.py
      • test_csv_logger.py
      • test_tensor_board.py
      • test_aggregate_predictions.py
      • imaging
        • test_imaging.py
        • test_inside_cnns.py
        • __init__.py
      • test_mixup.py
      • test_unpack_state.py
      • test_torch_scheduler.py
      • test_between_class.py
    • test_state.py
    • test_bases.py
    • test_cv_utils.py
    • test_end_to_end.py
    • test_trial.py
    • __init__.py
    • metrics
      • test_default.py
      • test_aggregators.py
      • test_roc_auc_score.py
      • test_decorators.py
      • test_wrappers.py
      • test_timer.py
      • test_primitives.py
      • test_lr.py
      • __init__.py
      • test_metrics.py
    • test_magics.py
  • requirements.txt
  • .gitignore
  • docs
    • code
      • callbacks.rst
      • metrics.rst
      • main.rst
    • Makefile
    • notes
      • callbacks.rst
      • metrics.rst
      • trial.rst
      • distributed.rst
      • visdom.rst
      • tensorboard.rst
    • requirements.txt
    • index.rst
    • conf.py
    • _static
      • img
        • logo_inkscape.svg
        • qr_code.svg
        • logo_dark.svg
        • logo_light.svg
        • nbviewer_logo.svg
        • logo_light_text.svg
        • logo_inkscape_text.svg
        • logo_dark_text.svg
      • examples
        • tensorboard.py
        • distributed_data_parallel.py
        • visdom_note.py
      • notebooks
        • svm_linear.ipynb
        • adversarial.ipynb
        • serialization.ipynb
        • callbacks.ipynb
        • gan.ipynb
        • pycm.ipynb
        • history.ipynb
        • manifold_mixup.ipynb

import unittest

import torch
from torch.nn import Module, Linear
import torch.nn.init as init

import torchbearer
from torchbearer import callbacks as c


class Net(Module):
    def __init__(self, x):
        super(Net, self).__init__()
        self.pars = torch.nn.Parameter(x)

    def f(self):
        """
        function to be minimised:
        f(x) = (x[0]-5)^2 + x[1]^2 + (x[2]-1)^2
        Solution:
        x = [5,0,1]
        """
        out = torch.zeros_like(self.pars)
        out[0] = self.pars[0]-5
        out[1] = self.pars[1]
        out[2] = self.pars[2]-1
        return torch.sum(out**2)

    def forward(self, _):
        return self.f()


class NetWithState(Net):
    def forward(self, _, state=None):
        if state is None:
            raise ValueError
        return super(NetWithState, self).forward(_)


def loss(y_pred, y_true):
    return y_pred


class TestEndToEnd(unittest.TestCase):
    def test_basic_opt(self):
        p = torch.tensor([2.0, 1.0, 10.0])
        training_steps = 1000

        model = NetWithState(p)
        optim = torch.optim.SGD(model.parameters(), lr=0.01)

        trial = torchbearer.Trial(model, optim, loss).for_train_steps(training_steps).for_val_steps(1).for_test_steps(1)
        trial.run()
        trial.predict()
        trial.evaluate()

        self.assertAlmostEqual(model.pars[0].item(), 5.0, places=4)
        self.assertAlmostEqual(model.pars[1].item(), 0.0, places=4)
        self.assertAlmostEqual(model.pars[2].item(), 1.0, places=4)

    def test_callbacks(self):
        from torch.utils.data import TensorDataset
        traingen = TensorDataset(torch.rand(10, 1, 3), torch.rand(10, 1))
        valgen = TensorDataset(torch.rand(10, 1, 3), torch.rand(10, 1))
        testgen = TensorDataset(torch.rand(10, 1, 3), torch.rand(10, 1))

        model = torch.nn.Linear(3, 1)
        optim = torch.optim.SGD(model.parameters(), lr=0.01)
        cbs = []
        cbs.extend([c.EarlyStopping(), c.GradientClipping(10, model.parameters()), c.Best('test.pt'),
                    c.MostRecent('test.pt'), c.ReduceLROnPlateau(), c.CosineAnnealingLR(0.1, 0.01),
                    c.ExponentialLR(1), c.Interval('test.pt'), c.CSVLogger('test_csv.pt'),
                    c.L1WeightDecay(), c.L2WeightDecay(), c.TerminateOnNaN(monitor='fail_metric')])

        trial = torchbearer.Trial(model, optim, torch.nn.MSELoss(), metrics=['loss'], callbacks=cbs)
        trial = trial.with_generators(traingen, valgen, testgen)
        trial.run(2)
        trial.predict()
        trial.evaluate(data_key=torchbearer.TEST_DATA)
        trial.evaluate()

        import os
        os.remove('test.pt')
        os.remove('test_csv.pt')

    def test_zero_model(self):
        model = Linear(3, 1)
        init.constant_(model.weight, 0)
        init.constant_(model.bias, 0)
        optim = torch.optim.SGD(model.parameters(), lr=0.01)

        trial = torchbearer.Trial(model, optim, loss)
        trial.with_test_data(torch.rand(10, 3), batch_size=3)
        preds = trial.predict()

        for i in range(len(preds)):
            self.assertAlmostEqual(preds[i], 0)

    def test_basic_checkpoint(self):
        p = torch.tensor([2.0, 1.0, 10.0])
        training_steps = 500

        model = Net(p)
        optim = torch.optim.SGD(model.parameters(), lr=0.01)

        trial = torchbearer.Trial(model, optim, loss, callbacks=[torchbearer.callbacks.MostRecent(filepath='test.pt')]).for_train_steps(training_steps).for_val_steps(1)
        trial.run(2)  # Simulate 2 'epochs'

        # Reload
        p = torch.tensor([2.0, 1.0, 10.0])
        model = Net(p)
        optim = torch.optim.SGD(model.parameters(), lr=0.01)

        trial = torchbearer.Trial(model, optim, loss, callbacks=[torchbearer.callbacks.MostRecent(filepath='test.pt')]).for_train_steps(training_steps)
        trial.load_state_dict(torch.load('test.pt'))
        self.assertEqual(len(trial.state[torchbearer.HISTORY]), 2)
        self.assertAlmostEqual(model.pars[0].item(), 5.0, places=4)
        self.assertAlmostEqual(model.pars[1].item(), 0.0, places=4)
        self.assertAlmostEqual(model.pars[2].item(), 1.0, places=4)

        import os
        os.remove('test.pt')

    def test_with_loader(self):
        p = torch.tensor([2.0, 1.0, 10.0])
        training_steps = 2

        model = Net(p)
        optim = torch.optim.SGD(model.parameters(), lr=0.01)
        test_var = {'loaded': False}

        def custom_loader(state):
            state[torchbearer.X], state[torchbearer.Y_TRUE] = None, None
            test_var['loaded'] = True

        trial = torchbearer.Trial(model, optim, loss, callbacks=[torchbearer.callbacks.MostRecent(filepath='test.pt')]).for_train_steps(training_steps).for_val_steps(1)
        trial.with_loader(custom_loader)
        self.assertTrue(not test_var['loaded'])
        trial.run(1)
        self.assertTrue(test_var['loaded'])

        import os
        os.remove('test.pt')

    def test_only_model(self):
        p = torch.tensor([2.0, 1.0, 10.0])
        model = Net(p)
        trial = torchbearer.Trial(model)
        self.assertListEqual(trial.run(), [])

    def test_no_model(self):
        trial = torchbearer.Trial(None)
        trial.run()
        self.assertTrue(torchbearer.trial.MockModel()(torch.rand(1)) is None)

    def test_no_train_steps(self):
        trial = torchbearer.Trial(None)
        trial.for_val_steps(10)
        trial.run()