Implement SampleRNN

1 files changed, 267 insertions, 0 deletions

diff --git a/trainer/plugins.py b/trainer/plugins.py
new file mode 100644
index 0000000..552804a
--- /dev/null
+++ b/trainer/plugins.py
@@ -0,0 +1,267 @@
+import matplotlib
+matplotlib.use('Agg')
+
+from model import Generator
+
+import torch
+from torch.autograd import Variable
+from torch.utils.trainer.plugins.plugin import Plugin
+from torch.utils.trainer.plugins.monitor import Monitor
+from torch.utils.trainer.plugins import LossMonitor
+
+from librosa.output import write_wav
+from matplotlib import pyplot
+
+from glob import glob
+import os
+import pickle
+import time
+
+
+class TrainingLossMonitor(LossMonitor):
+
+    stat_name = 'training_loss'
+
+
+class ValidationPlugin(Plugin):
+
+    def __init__(self, val_dataset, test_dataset):
+        super().__init__([(1, 'epoch')])
+        self.val_dataset = val_dataset
+        self.test_dataset = test_dataset
+
+    def register(self, trainer):
+        self.trainer = trainer
+        val_stats = self.trainer.stats.setdefault('validation_loss', {})
+        val_stats['log_epoch_fields'] = ['{last:.4f}']
+        test_stats = self.trainer.stats.setdefault('test_loss', {})
+        test_stats['log_epoch_fields'] = ['{last:.4f}']
+
+    def epoch(self, idx):
+        self.trainer.model.eval()
+
+        val_stats = self.trainer.stats.setdefault('validation_loss', {})
+        val_stats['last'] = self._evaluate(self.val_dataset)
+        test_stats = self.trainer.stats.setdefault('test_loss', {})
+        test_stats['last'] = self._evaluate(self.test_dataset)
+
+        self.trainer.model.train()
+
+    def _evaluate(self, dataset):
+        loss_sum = 0
+        n_examples = 0
+        for data in dataset:
+            batch_inputs = data[: -1]
+            batch_target = data[-1]
+            batch_size = batch_target.size()[0]
+
+            def wrap(input):
+                if torch.is_tensor(input):
+                    input = Variable(input, volatile=True)
+                    if self.trainer.cuda:
+                        input = input.cuda()
+                return input
+            batch_inputs = list(map(wrap, batch_inputs))
+
+            batch_target = Variable(batch_target, volatile=True)
+            if self.trainer.cuda:
+                batch_target = batch_target.cuda()
+
+            batch_output = self.trainer.model(*batch_inputs)
+            loss_sum += self.trainer.criterion(batch_output, batch_target) \
+                                    .data[0] * batch_size
+
+            n_examples += batch_size
+
+        return loss_sum / n_examples
+
+
+class AbsoluteTimeMonitor(Monitor):
+
+    stat_name = 'time'
+
+    def __init__(self, *args, **kwargs):
+        kwargs.setdefault('unit', 's')
+        kwargs.setdefault('precision', 0)
+        kwargs.setdefault('running_average', False)
+        kwargs.setdefault('epoch_average', False)
+        super(AbsoluteTimeMonitor, self).__init__(*args, **kwargs)
+        self.start_time = None
+
+    def _get_value(self, *args):
+        if self.start_time is None:
+            self.start_time = time.time()
+        return time.time() - self.start_time
+
+
+class SaverPlugin(Plugin):
+
+    last_pattern = 'ep{}-it{}'
+    best_pattern = 'best-ep{}-it{}'
+
+    def __init__(self, checkpoints_path, keep_old_checkpoints):
+        super().__init__([(1, 'epoch')])
+        self.checkpoints_path = checkpoints_path
+        self.keep_old_checkpoints = keep_old_checkpoints
+        self._best_val_loss = float('+inf')
+
+    def register(self, trainer):
+        self.trainer = trainer
+
+    def epoch(self, epoch_index):
+        if not self.keep_old_checkpoints:
+            self._clear(self.last_pattern.format('*', '*'))
+        torch.save(
+            self.trainer.model.state_dict(),
+            os.path.join(
+                self.checkpoints_path,
+                self.last_pattern.format(epoch_index, self.trainer.iterations)
+            )
+        )
+
+        cur_val_loss = self.trainer.stats['validation_loss']['last']
+        if cur_val_loss < self._best_val_loss:
+            self._clear(self.best_pattern.format('*', '*'))
+            torch.save(
+                self.trainer.model.state_dict(),
+                os.path.join(
+                    self.checkpoints_path,
+                    self.best_pattern.format(
+                        epoch_index, self.trainer.iterations
+                    )
+                )
+            )
+            self._best_val_loss = cur_val_loss
+
+    def _clear(self, pattern):
+        pattern = os.path.join(self.checkpoints_path, pattern)
+        for file_name in glob(pattern):
+            os.remove(file_name)
+
+
+class GeneratorPlugin(Plugin):
+
+    pattern = 'ep{}-s{}.wav'
+
+    def __init__(self, samples_path, n_samples, sample_length, sample_rate):
+        super().__init__([(1, 'epoch')])
+        self.samples_path = samples_path
+        self.n_samples = n_samples
+        self.sample_length = sample_length
+        self.sample_rate = sample_rate
+
+    def register(self, trainer):
+        self.generate = Generator(trainer.model.model, trainer.cuda)
+
+    def epoch(self, epoch_index):
+        samples = self.generate(self.n_samples, self.sample_length) \
+                      .cpu().float().numpy()
+        for i in range(self.n_samples):
+            write_wav(
+                os.path.join(
+                    self.samples_path, self.pattern.format(epoch_index, i + 1)
+                ),
+                samples[i, :], sr=self.sample_rate, norm=True
+            )
+
+
+class StatsPlugin(Plugin):
+
+    data_file_name = 'stats.pkl'
+    plot_pattern = '{}.svg'
+
+    def __init__(self, results_path, iteration_fields, epoch_fields, plots):
+        super().__init__([(1, 'iteration'), (1, 'epoch')])
+        self.results_path = results_path
+
+        self.iteration_fields = self._fields_to_pairs(iteration_fields)
+        self.epoch_fields = self._fields_to_pairs(epoch_fields)
+        self.plots = plots
+        self.data = {
+            'iterations': {
+                field: []
+                for field in self.iteration_fields + [('iteration', 'last')]
+            },
+            'epochs': {
+                field: []
+                for field in self.epoch_fields + [('iteration', 'last')]
+            }
+        }
+
+    def register(self, trainer):
+        self.trainer = trainer
+
+    def iteration(self, *args):
+        for (field, stat) in self.iteration_fields:
+            self.data['iterations'][field, stat].append(
+                self.trainer.stats[field][stat]
+            )
+
+        self.data['iterations']['iteration', 'last'].append(
+            self.trainer.iterations
+        )
+
+    def epoch(self, epoch_index):
+        for (field, stat) in self.epoch_fields:
+            self.data['epochs'][field, stat].append(
+                self.trainer.stats[field][stat]
+            )
+
+        self.data['epochs']['iteration', 'last'].append(
+            self.trainer.iterations
+        )
+
+        data_file_path = os.path.join(self.results_path, self.data_file_name)
+        with open(data_file_path, 'wb') as f:
+            pickle.dump(self.data, f)
+
+        for (name, info) in self.plots.items():
+            x_field = self._field_to_pair(info['x'])
+
+            try:
+                y_fields = info['ys']
+            except KeyError:
+                y_fields = [info['y']]
+
+            labels = list(map(
+                lambda x: ' '.join(x) if type(x) is tuple else x,
+                y_fields
+            ))
+            y_fields = self._fields_to_pairs(y_fields)
+
+            try:
+                formats = info['formats']
+            except KeyError:
+                formats = [''] * len(y_fields)
+
+            pyplot.gcf().clear()
+
+            for (y_field, format, label) in zip(y_fields, formats, labels):
+                if y_field in self.iteration_fields:
+                    part_name = 'iterations'
+                else:
+                    part_name = 'epochs'
+
+                xs = self.data[part_name][x_field]
+                ys = self.data[part_name][y_field]
+
+                pyplot.plot(xs, ys, format, label=label)
+
+            if 'log_y' in info and info['log_y']:
+                pyplot.yscale('log')
+
+            pyplot.legend()
+            pyplot.savefig(
+                os.path.join(self.results_path, self.plot_pattern.format(name))
+            )
+
+    @staticmethod
+    def _field_to_pair(field):
+        if type(field) is tuple:
+            return field
+        else:
+            return (field, 'last')
+
+    @classmethod
+    def _fields_to_pairs(cls, fields):
+        return list(map(cls._field_to_pair, fields))