Take on an explicit dependency on pytorch

.travis.yml

@@ -21,7 +21,7 @@ install:
   #- conda update -q conda
   # Useful for debugging any issues with conda
   - conda info -a
-  - conda create -q -n test-environment python=$TRAVIS_PYTHON_VERSION numpy=1.12 scipy=0.18 cython
+  - conda create -q -n test-environment python=$TRAVIS_PYTHON_VERSION numpy=1.12 scipy=0.18 cython pytorch=0.3.1 torchvision
   - source activate test-environment
   - python setup.py install
   - pip install coveralls

requirements.txt

@@ -6,4 +6,7 @@ requests
 tornado==4.5.3
 pysoundfile
 matplotlib
-intervaltree
+intervaltree
+numpy==1.12.1
+scipy==0.19
+torch==0.3.1

setup.py

@@ -84,7 +84,10 @@
         'matplotlib',
         'argparse',
         'intervaltree',
-        'ujson'
+        'ujson',
+        'numpy==1.12.1',
+        'scipy==0.19',
+        'torch==0.3.1'
     ],
     package_data={
         'nputil': ['*.pyx', '*.pyxbld'],

zounds/learn/__init__.py

@@ -32,20 +32,18 @@
 
 from functional import hyperplanes, simhash, example_wise_unit_norm
 
-try:
-    from util import \
-        Conv1d, ConvTranspose1d, Conv2d, ConvTranspose2d, to_var, from_var, \
-        try_network, apply_network, feature_map_size, sample_norm, \
-        FrequencyDecompositionAnalyzer, FrequencyDecompositionGenerator, gradients
-    from gan_experiment import GanExperiment
-    from sample_embedding import RawSampleEmbedding
-    from dct_transform import DctTransform
-    from gated import GatedConvTransposeLayer, GatedConvLayer, GatedLinearLayer
-    from multiresolution import MultiResolutionConvLayer
-    from loss import PerceptualLoss, BandLoss, CategoricalLoss, \
-        WassersteinCriticLoss, WassersteinGradientPenaltyLoss, \
-        LearnedWassersteinLoss
-except ImportError:
-    pass
+
+from util import \
+    Conv1d, ConvTranspose1d, Conv2d, ConvTranspose2d, to_var, from_var, \
+    try_network, apply_network, feature_map_size, sample_norm, gradients
+from gan_experiment import GanExperiment
+from sample_embedding import RawSampleEmbedding
+from dct_transform import DctTransform
+from gated import GatedConvTransposeLayer, GatedConvLayer, GatedLinearLayer
+from multiresolution import MultiResolutionConvLayer
+from loss import PerceptualLoss, BandLoss, CategoricalLoss, \
+    WassersteinCriticLoss, WassersteinGradientPenaltyLoss, \
+    LearnedWassersteinLoss
+
 
 

zounds/learn/embedding.py

@@ -1,6 +1,9 @@
 from trainer import Trainer
 from random import choice
 import numpy as np
+from torch import nn
+from torch.optim import Adam
+import torch
 
 
 class TripletEmbeddingTrainer(Trainer):
@@ -46,6 +49,11 @@ def __init__(
         # of the paper https://arxiv.org/pdf/1711.02209.pdf
         self.margin = 0.1
         self.register_batch_complete_callback(self._log)
+        self.loss = nn.TripletMarginLoss(margin=self.margin)
+
+    def _cuda(self, device=None):
+        self.loss = self.loss.cuda()
+        self.network = self.network.cuda()
 
     def _driver(self, data):
         batches_in_epoch = len(data) // self.batch_size
@@ -65,8 +73,6 @@ def _apply_network_and_normalize(self, x):
         Pass x through the network, and give the output unit norm, as specified
         by section 4.2 of https://arxiv.org/pdf/1711.02209.pdf
         """
-
-        import torch
         x = self.network(x)
         return x / torch.norm(x, dim=1).view(-1, 1)
 
@@ -76,40 +82,35 @@ def _select_batch(self, training_set):
         return indices, batch.astype(np.float32)
 
     def train(self, data):
-        from torch import nn
-        from torch.optim import Adam
-        from util import to_var
 
         data = data['data']
 
-        self.network.cuda()
         self.network.train()
 
         optimizer = Adam(self.network.parameters(), lr=1e-5)
-        loss = nn.TripletMarginLoss(margin=self.margin).cuda()
 
         for epoch, batch in self._driver(data):
             self.network.zero_grad()
 
             # choose a batch of anchors
             indices, anchor = self._select_batch(data)
-            anchor_v = to_var(anchor)
+            anchor_v = self._variable(anchor)
             a = self._apply_network_and_normalize(anchor_v)
 
             # choose negative examples
             negative_indices, negative = self._select_batch(data)
-            negative_v = to_var(negative)
+            negative_v = self._variable(negative)
             n = self._apply_network_and_normalize(negative_v)
 
             # choose a deformation for this batch and apply it to produce the
             # positive examples
             deformation = choice(self.deformations)
             positive = deformation(anchor, data[indices, ...]) \
                 .astype(np.float32)
-            positive_v = to_var(positive)
+            positive_v = self._variable(positive)
             p = self._apply_network_and_normalize(positive_v)
 
-            error = loss.forward(a, p, n)
+            error = self.loss.forward(a, p, n)
             error.backward()
             optimizer.step()
 

zounds/learn/loss.py

@@ -56,6 +56,7 @@ def cuda(self, device=None):
         return super(PerceptualLoss, self).cuda(device=device)
 
     def _transform(self, x):
+        x = x.view(x.shape[0], 1, -1)
         features = F.conv1d(
             x, self.weights, stride=self.lap, padding=self.basis_size)
 

zounds/learn/pytorch_model.py

@@ -1,8 +1,8 @@
-
 import warnings
 import featureflow as ff
 from preprocess import Preprocessor, PreprocessResult, Op
 from zounds.persistence.util import extract_init_args
+import torch
 
 
 class PyTorchPreprocessResult(PreprocessResult):
@@ -41,21 +41,28 @@ def __setstate__(self, state):
         """
         Re-hydrate an instance from serialized state
         """
-        import torch
 
         restored_weights = dict(
-            ((k, torch.from_numpy(v).cuda())
+            ((k, torch.from_numpy(v))
              for k, v in state['weights'].iteritems()))
         init_args = state['init_args']
         network = state['cls'](*init_args)
         network.load_state_dict(restored_weights)
-        network.cuda()
+
+        # KLUDGE: Should we *ever* implicitly move things to the GPU?  If not,
+        # this would need to be done explicitly by the user when re-hydrating
+        # the learning pipeline
+        if torch.cuda.is_available():
+            network = network.cuda()
+
         network.eval()
 
         self.op = Op(
-            state['forward_func'], network=network, **state['op_kwargs'])
-        self.inversion_data = Op(state['inv_data_func'],
-                                 network=network)
+            state['forward_func'],
+            network=network,
+            **state['op_kwargs'])
+
+        self.inversion_data = Op(state['inv_data_func'], network=network)
         self.inverse = Op(state['backward_func'])
         self.name = state['name']
 
@@ -102,6 +109,7 @@ def x(d, network=None, chunk_size=None):
     def _backward_func(self):
         def x(_):
             raise NotImplementedError()
+
         return x
 
     def _enqueue(self, data, pusher):
@@ -174,27 +182,16 @@ def __init__(self, apply_network='generator', trainer=None, needs=None):
 
     def _forward_func(self):
         def x(d, network=None, apply_network=None):
-            import torch
-            from torch.autograd import Variable
-            import numpy as np
             from zounds.core import ArrayWithUnits, IdentityDimension
+            from zounds.learn import apply_network as apply
+            import numpy as np
 
             if apply_network == 'generator':
                 n = network.generator
             else:
                 n = network.discriminator
 
-            chunks = []
-            batch_size = 128
-
-            for i in xrange(0, len(d), batch_size):
-                tensor = torch.from_numpy(
-                    d[i:i + batch_size].astype(np.float32))
-                gpu = tensor.cuda()
-                v = Variable(gpu)
-                chunks.append(n(v).data.cpu().numpy())
-
-            result = np.concatenate(chunks)
+            result = apply(n, d.astype(np.float32), chunksize=128)
 
             try:
                 return ArrayWithUnits(
@@ -251,21 +248,12 @@ def __init__(self, trainer=None, needs=None):
 
     def _forward_func(self):
         def x(d, network=None):
-            import torch
-            from torch.autograd import Variable
-            import numpy as np
             from zounds.core import ArrayWithUnits, IdentityDimension
+            from zounds.learn import apply_network
+            import numpy as np
 
-            chunks = []
-            batch_size = 128
-
-            for i in xrange(0, len(d), batch_size):
-                tensor = torch.from_numpy(d[i:i + batch_size])
-                gpu = tensor.cuda()
-                v = Variable(gpu)
-                chunks.append(network.encoder(v).data.cpu().numpy())
-
-            encoded = np.concatenate(chunks)
+            encoded = apply_network(
+                network.encoder, d.astype(np.float32), chunksize=128)
 
             try:
                 extra_dims = (IdentityDimension(),) * (encoded.ndim - 1)
@@ -278,21 +266,10 @@ def x(d, network=None):
 
     def _backward_func(self):
         def x(d, network=None):
-            import torch
-            from torch.autograd import Variable
+            from zounds.learn import apply_network
             import numpy as np
-
-            chunks = []
-            batch_size = 128
-
-            for i in xrange(0, len(d), batch_size):
-                tensor = torch.from_numpy(d.astype(np.float32))
-                gpu = tensor.cuda()
-                v = Variable(gpu)
-                chunks.append(network.decoder(v).data.cpu().numpy())
-
-            decoded = np.concatenate(chunks)
-            return decoded
+            return apply_network(
+                network.decoder, d.astype(np.float32), chunksize=128)
 
         return x
 

zounds/learn/random_samples.py

@@ -136,9 +136,11 @@ def _dequeue(self):
 
 
 class InfiniteSampler(Node):
-    def __init__(self, nsamples=None, dtype=None, needs=None):
+    def __init__(self, nsamples=None, multiplexed=True, dtype=None, needs=None):
         super(InfiniteSampler, self).__init__(needs=needs)
-        self.reservoir = Reservoir(nsamples, dtype)
+        self.multiplexed = multiplexed
+        self.reservoir = MultiplexedReservoir(nsamples, dtype=dtype) \
+            if multiplexed else Reservoir(nsamples, dtype=dtype)
 
     def _total_samples(self, cls, feature, _ids):
         pool = ThreadPool(cpu_count())

zounds/learn/supervised.py

@@ -1,6 +1,8 @@
 from trainer import Trainer
 import numpy as np
 import warnings
+import torch
+from torch.autograd import Variable
 
 
 class SupervisedTrainer(Trainer):
@@ -26,7 +28,7 @@ def __init__(
         self.holdout_percent = holdout_percent
         self.optimizer = optimizer(model)
         self.loss = loss
-        self.model = model
+        self.network = model
         self.register_batch_complete_callback(self._log)
         self.samples = None
 
@@ -46,13 +48,11 @@ def random_sample(self):
         inp, label = self.samples[index]
         return inp, label
 
-    def train(self, data):
-        import torch
-        from torch.autograd import Variable
-
-        model = self.model.cuda()
-        loss = self.loss.cuda()
+    def _cuda(self, device=None):
+        self.network = self.network.cuda()
+        self.loss = self.loss.cuda()
 
+    def train(self, data):
         data, labels = data['data'], data['labels']
 
         test_size = int(self.holdout_percent * len(data))
@@ -62,16 +62,12 @@ def train(self, data):
         def batch(d, l, test=False):
             d = self.data_preprocessor(d)
             l = self.label_preprocessor(l)
-            inp = torch.from_numpy(d)
-            inp = inp.cuda()
-            inp_v = Variable(inp, volatile=test)
-            output = model(inp_v)
+            inp_v = self._variable(d, volatile=test)
+            output = self.network(inp_v)
 
-            labels_t = torch.from_numpy(l)
-            labels_t = labels_t.cuda()
-            labels_v = Variable(labels_t)
+            labels_v = self._variable(l)
 
-            error = loss(output, labels_v)
+            error = self.loss(output, labels_v)
 
             if not test:
                 error.backward()
@@ -90,7 +86,7 @@ def batch(d, l, test=False):
 
             for i in xrange(0, len(data), self.batch_size):
 
-                model.zero_grad()
+                self.network.zero_grad()
 
                 # training batch
                 minibatch_slice = slice(i, i + self.batch_size)
@@ -127,4 +123,4 @@ def batch(d, l, test=False):
 
             self._current_epoch += 1
 
-        return model
+        return self.network

zounds/learn/test_dct.py

@@ -1,21 +1,14 @@
 import unittest2
-
-try:
-    from dct_transform import DctTransform
-    import torch
-    from torch.autograd import Variable
-except ImportError:
-    DctTransform = None
+from dct_transform import DctTransform
+import torch
+from torch.autograd import Variable
+from zounds.spectral import HanningWindowingFunc
 
 
 class DctTransformTests(unittest2.TestCase):
-    def setUp(self):
-        if DctTransform is None:
-            self.skipTest('PyTorch is not available')
-
     def test_can_do_short_time_dct_transform(self):
         t = torch.FloatTensor(3, 1, 512)
         v = Variable(t)
         dct_trasform = DctTransform()
-        stdct = dct_trasform.short_time_dct(v, 128, 64)
+        stdct = dct_trasform.short_time_dct(v, 128, 64, HanningWindowingFunc())
         self.assertEqual((3, 128, 7), tuple(stdct.shape))