From b114ac9d4bb99284747ffbf71f58b230aec7f733 Mon Sep 17 00:00:00 2001
From: Tzu-Wei Huang <huang.dexter@gmail.com>
Date: Mon, 16 Oct 2017 09:13:19 +0800
Subject: [PATCH] chainer extension example

---
 examples/chainer/extension_logger/net.py      |  86 ++++++++++++
 .../chainer/extension_logger/train_dcgan.py   | 118 ++++++++++++++++
 examples/chainer/extension_logger/updater.py  |  48 +++++++
 .../chainer/extension_logger/visualize.py     |  24 ++++
 .../extension_logger/writetnesorboard.py      | 128 ++++++++++++++++++
 examples/chainer/{ => plain_logger}/data.py   |   0
 examples/chainer/{ => plain_logger}/net.py    |   0
 .../chainer/{ => plain_logger}/train_vae.py   |   0
 8 files changed, 404 insertions(+)
 create mode 100644 examples/chainer/extension_logger/net.py
 create mode 100644 examples/chainer/extension_logger/train_dcgan.py
 create mode 100644 examples/chainer/extension_logger/updater.py
 create mode 100644 examples/chainer/extension_logger/visualize.py
 create mode 100644 examples/chainer/extension_logger/writetnesorboard.py
 rename examples/chainer/{ => plain_logger}/data.py (100%)
 rename examples/chainer/{ => plain_logger}/net.py (100%)
 rename examples/chainer/{ => plain_logger}/train_vae.py (100%)

diff --git a/examples/chainer/extension_logger/net.py b/examples/chainer/extension_logger/net.py
new file mode 100644
index 00000000..300c3d33
--- /dev/null
+++ b/examples/chainer/extension_logger/net.py
@@ -0,0 +1,86 @@
+#!/usr/bin/env python
+
+from __future__ import print_function
+
+import numpy
+
+import chainer
+from chainer import cuda
+import chainer.functions as F
+import chainer.links as L
+
+
+def add_noise(h, sigma=0.2):
+    xp = cuda.get_array_module(h.data)
+    if chainer.config.train:
+        return h + sigma * xp.random.randn(*h.shape)
+    else:
+        return h
+
+
+class Generator(chainer.Chain):
+
+    def __init__(self, n_hidden, bottom_width=4, ch=512, wscale=0.02):
+        super(Generator, self).__init__()
+        self.n_hidden = n_hidden
+        self.ch = ch
+        self.bottom_width = bottom_width
+
+        with self.init_scope():
+            w = chainer.initializers.Normal(wscale)
+            self.l0 = L.Linear(self.n_hidden, bottom_width * bottom_width * ch,
+                               initialW=w)
+            self.dc1 = L.Deconvolution2D(ch, ch // 2, 4, 2, 1, initialW=w)
+            self.dc2 = L.Deconvolution2D(ch // 2, ch // 4, 4, 2, 1, initialW=w)
+            self.dc3 = L.Deconvolution2D(ch // 4, ch // 8, 4, 2, 1, initialW=w)
+            self.dc4 = L.Deconvolution2D(ch // 8, 3, 3, 1, 1, initialW=w)
+            self.bn0 = L.BatchNormalization(bottom_width * bottom_width * ch)
+            self.bn1 = L.BatchNormalization(ch // 2)
+            self.bn2 = L.BatchNormalization(ch // 4)
+            self.bn3 = L.BatchNormalization(ch // 8)
+
+    def make_hidden(self, batchsize):
+        return numpy.random.uniform(-1, 1, (batchsize, self.n_hidden, 1, 1))\
+            .astype(numpy.float32)
+
+    def __call__(self, z):
+        h = F.reshape(F.relu(self.bn0(self.l0(z))),
+                      (len(z), self.ch, self.bottom_width, self.bottom_width))
+        h = F.relu(self.bn1(self.dc1(h)))
+        h = F.relu(self.bn2(self.dc2(h)))
+        h = F.relu(self.bn3(self.dc3(h)))
+        x = F.sigmoid(self.dc4(h))
+        return x
+
+
+class Discriminator(chainer.Chain):
+
+    def __init__(self, bottom_width=4, ch=512, wscale=0.02):
+        w = chainer.initializers.Normal(wscale)
+        super(Discriminator, self).__init__()
+        with self.init_scope():
+            self.c0_0 = L.Convolution2D(3, ch // 8, 3, 1, 1, initialW=w)
+            self.c0_1 = L.Convolution2D(ch // 8, ch // 4, 4, 2, 1, initialW=w)
+            self.c1_0 = L.Convolution2D(ch // 4, ch // 4, 3, 1, 1, initialW=w)
+            self.c1_1 = L.Convolution2D(ch // 4, ch // 2, 4, 2, 1, initialW=w)
+            self.c2_0 = L.Convolution2D(ch // 2, ch // 2, 3, 1, 1, initialW=w)
+            self.c2_1 = L.Convolution2D(ch // 2, ch // 1, 4, 2, 1, initialW=w)
+            self.c3_0 = L.Convolution2D(ch // 1, ch // 1, 3, 1, 1, initialW=w)
+            self.l4 = L.Linear(bottom_width * bottom_width * ch, 1, initialW=w)
+            self.bn0_1 = L.BatchNormalization(ch // 4, use_gamma=False)
+            self.bn1_0 = L.BatchNormalization(ch // 4, use_gamma=False)
+            self.bn1_1 = L.BatchNormalization(ch // 2, use_gamma=False)
+            self.bn2_0 = L.BatchNormalization(ch // 2, use_gamma=False)
+            self.bn2_1 = L.BatchNormalization(ch // 1, use_gamma=False)
+            self.bn3_0 = L.BatchNormalization(ch // 1, use_gamma=False)
+
+    def __call__(self, x):
+        h = add_noise(x)
+        h = F.leaky_relu(add_noise(self.c0_0(h)))
+        h = F.leaky_relu(add_noise(self.bn0_1(self.c0_1(h))))
+        h = F.leaky_relu(add_noise(self.bn1_0(self.c1_0(h))))
+        h = F.leaky_relu(add_noise(self.bn1_1(self.c1_1(h))))
+        h = F.leaky_relu(add_noise(self.bn2_0(self.c2_0(h))))
+        h = F.leaky_relu(add_noise(self.bn2_1(self.c2_1(h))))
+        h = F.leaky_relu(add_noise(self.bn3_0(self.c3_0(h))))
+        return self.l4(h)
diff --git a/examples/chainer/extension_logger/train_dcgan.py b/examples/chainer/extension_logger/train_dcgan.py
new file mode 100644
index 00000000..06fce1f1
--- /dev/null
+++ b/examples/chainer/extension_logger/train_dcgan.py
@@ -0,0 +1,118 @@
+#!/usr/bin/env python
+
+from __future__ import print_function
+import argparse
+import os
+
+import chainer
+from chainer import training
+from chainer.training import extensions
+
+from net import Discriminator
+from net import Generator
+from updater import DCGANUpdater
+from visualize import out_generated_image
+from tensorboardX import SummaryWriter
+from writetnesorboard import LogTensorboard
+def main():
+    parser = argparse.ArgumentParser(description='Chainer example: DCGAN')
+    parser.add_argument('--batchsize', '-b', type=int, default=50,
+                        help='Number of images in each mini-batch')
+    parser.add_argument('--epoch', '-e', type=int, default=1000,
+                        help='Number of sweeps over the dataset to train')
+    parser.add_argument('--gpu', '-g', type=int, default=-1,
+                        help='GPU ID (negative value indicates CPU)')
+    parser.add_argument('--dataset', '-i', default='',
+                        help='Directory of image files.  Default is cifar-10.')
+    parser.add_argument('--out', '-o', default='result',
+                        help='Directory to output the result')
+    parser.add_argument('--resume', '-r', default='',
+                        help='Resume the training from snapshot')
+    parser.add_argument('--n_hidden', '-n', type=int, default=100,
+                        help='Number of hidden units (z)')
+    parser.add_argument('--seed', type=int, default=0,
+                        help='Random seed of z at visualization stage')
+    parser.add_argument('--snapshot_interval', type=int, default=1000,
+                        help='Interval of snapshot')
+    parser.add_argument('--display_interval', type=int, default=100,
+                        help='Interval of displaying log to console')
+    args = parser.parse_args()
+
+    print('GPU: {}'.format(args.gpu))
+    print('# Minibatch-size: {}'.format(args.batchsize))
+    print('# n_hidden: {}'.format(args.n_hidden))
+    print('# epoch: {}'.format(args.epoch))
+    print('')
+    writer = SummaryWriter()
+    # Set up a neural network to train
+    gen = Generator(n_hidden=args.n_hidden)
+    dis = Discriminator()
+
+    if args.gpu >= 0:
+        # Make a specified GPU current
+        chainer.cuda.get_device_from_id(args.gpu).use()
+        gen.to_gpu()  # Copy the model to the GPU
+        dis.to_gpu()
+
+    # Setup an optimizer
+    def make_optimizer(model, alpha=0.0002, beta1=0.5):
+        optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1)
+        optimizer.setup(model)
+        optimizer.add_hook(chainer.optimizer.WeightDecay(0.0001), 'hook_dec')
+        return optimizer
+    opt_gen = make_optimizer(gen)
+    opt_dis = make_optimizer(dis)
+
+    if args.dataset == '':
+        # Load the CIFAR10 dataset if args.dataset is not specified
+        train, _ = chainer.datasets.get_cifar10(withlabel=False, scale=255.)
+    else:
+        all_files = os.listdir(args.dataset)
+        image_files = [f for f in all_files if ('png' in f or 'jpg' in f)]
+        print('{} contains {} image files'
+              .format(args.dataset, len(image_files)))
+        train = chainer.datasets\
+            .ImageDataset(paths=image_files, root=args.dataset)
+
+    train_iter = chainer.iterators.SerialIterator(train, args.batchsize)
+
+    # Set up a trainer
+    updater = DCGANUpdater(
+        models=(gen, dis),
+        iterator=train_iter,
+        optimizer={
+            'gen': opt_gen, 'dis': opt_dis},
+        device=args.gpu)
+    trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out)
+
+    snapshot_interval = (args.snapshot_interval, 'iteration')
+    display_interval = (args.display_interval, 'iteration')
+    trainer.extend(
+        extensions.snapshot(filename='snapshot_iter_{.updater.iteration}.npz'),
+        trigger=snapshot_interval)
+    trainer.extend(extensions.snapshot_object(
+        gen, 'gen_iter_{.updater.iteration}.npz'), trigger=snapshot_interval)
+    trainer.extend(extensions.snapshot_object(
+        dis, 'dis_iter_{.updater.iteration}.npz'), trigger=snapshot_interval)
+    trainer.extend(extensions.LogReport(trigger=display_interval))
+    trainer.extend(LogTensorboard(trigger=display_interval, logger=writer))
+    trainer.extend(extensions.PrintReport([
+        'epoch', 'iteration', 'gen/loss', 'dis/loss',
+    ]), trigger=display_interval)
+    trainer.extend(extensions.ProgressBar(update_interval=10))
+    trainer.extend(
+        out_generated_image(
+            gen, dis,
+            10, 10, args.seed, args.out, writer),
+        trigger=snapshot_interval)
+
+    if args.resume:
+        # Resume from a snapshot
+        chainer.serializers.load_npz(args.resume, trainer)
+
+    # Run the training
+    trainer.run()
+
+
+if __name__ == '__main__':
+    main()
diff --git a/examples/chainer/extension_logger/updater.py b/examples/chainer/extension_logger/updater.py
new file mode 100644
index 00000000..b9fda48a
--- /dev/null
+++ b/examples/chainer/extension_logger/updater.py
@@ -0,0 +1,48 @@
+#!/usr/bin/env python
+
+from __future__ import print_function
+
+import chainer
+import chainer.functions as F
+from chainer import Variable
+
+
+class DCGANUpdater(chainer.training.StandardUpdater):
+
+    def __init__(self, *args, **kwargs):
+        self.gen, self.dis = kwargs.pop('models')
+        super(DCGANUpdater, self).__init__(*args, **kwargs)
+
+    def loss_dis(self, dis, y_fake, y_real):
+        batchsize = len(y_fake)
+        L1 = F.sum(F.softplus(-y_real)) / batchsize
+        L2 = F.sum(F.softplus(y_fake)) / batchsize
+        loss = L1 + L2
+        chainer.report({'loss': loss}, dis)
+        return loss
+
+    def loss_gen(self, gen, y_fake):
+        batchsize = len(y_fake)
+        loss = F.sum(F.softplus(-y_fake)) / batchsize
+        chainer.report({'loss': loss}, gen)
+        return loss
+
+    def update_core(self):
+        gen_optimizer = self.get_optimizer('gen')
+        dis_optimizer = self.get_optimizer('dis')
+
+        batch = self.get_iterator('main').next()
+        x_real = Variable(self.converter(batch, self.device)) / 255.
+        xp = chainer.cuda.get_array_module(x_real.data)
+
+        gen, dis = self.gen, self.dis
+        batchsize = len(batch)
+
+        y_real = dis(x_real)
+
+        z = Variable(xp.asarray(gen.make_hidden(batchsize)))
+        x_fake = gen(z)
+        y_fake = dis(x_fake)
+
+        dis_optimizer.update(self.loss_dis, dis, y_fake, y_real)
+        gen_optimizer.update(self.loss_gen, gen, y_fake)
diff --git a/examples/chainer/extension_logger/visualize.py b/examples/chainer/extension_logger/visualize.py
new file mode 100644
index 00000000..9b8421c4
--- /dev/null
+++ b/examples/chainer/extension_logger/visualize.py
@@ -0,0 +1,24 @@
+#!/usr/bin/env python
+
+import os
+
+import numpy as np
+from PIL import Image
+
+import chainer
+import chainer.cuda
+from chainer import Variable
+
+
+def out_generated_image(gen, dis, rows, cols, seed, dst, writer):
+    @chainer.training.make_extension()
+    def make_image(trainer):
+        np.random.seed(seed)
+        n_images = rows * cols
+        xp = gen.xp
+        z = Variable(xp.asarray(gen.make_hidden(n_images)))
+        with chainer.using_config('train', False):
+            x = gen(z)
+        writer.add_image('img', x, trainer.updater.iteration)
+
+    return make_image
diff --git a/examples/chainer/extension_logger/writetnesorboard.py b/examples/chainer/extension_logger/writetnesorboard.py
new file mode 100644
index 00000000..f3e68399
--- /dev/null
+++ b/examples/chainer/extension_logger/writetnesorboard.py
@@ -0,0 +1,128 @@
+import json
+import os
+import shutil
+import tempfile
+
+import six
+from chainer import reporter
+from chainer import serializer as serializer_module
+from chainer.training import extension
+from chainer.training import trigger as trigger_module
+
+
+class LogTensorboard(extension.Extension):
+
+    """Trainer extension to output the accumulated results to a log file.
+
+    This extension accumulates the observations of the trainer to
+    :class:`~chainer.DictSummary` at a regular interval specified by a supplied
+    trigger, and writes them into a log file in JSON format.
+
+    There are two triggers to handle this extension. One is the trigger to
+    invoke this extension, which is used to handle the timing of accumulating
+    the results. It is set to ``1, 'iteration'`` by default. The other is the
+    trigger to determine when to emit the result. When this trigger returns
+    True, this extension appends the summary of accumulated values to the list
+    of past summaries, and writes the list to the log file. Then, this
+    extension makes a new fresh summary object which is used until the next
+    time that the trigger fires.
+
+    It also adds some entries to each result dictionary.
+
+    - ``'epoch'`` and ``'iteration'`` are the epoch and iteration counts at the
+      output, respectively.
+    - ``'elapsed_time'`` is the elapsed time in seconds since the training
+      begins. The value is taken from :attr:`Trainer.elapsed_time`.
+
+    Args:
+        keys (iterable of strs): Keys of values to accumulate. If this is None,
+            all the values are accumulated and output to the log file.
+        trigger: Trigger that decides when to aggregate the result and output
+            the values. This is distinct from the trigger of this extension
+            itself. If it is a tuple in the form ``<int>, 'epoch'`` or
+            ``<int>, 'iteration'``, it is passed to :class:`IntervalTrigger`.
+        postprocess: Callback to postprocess the result dictionaries. Each
+            result dictionary is passed to this callback on the output. This
+            callback can modify the result dictionaries, which are used to
+            output to the log file.
+        log_name (str): Name of the log file under the output directory. It can
+            be a format string: the last result dictionary is passed for the
+            formatting. For example, users can use '{iteration}' to separate
+            the log files for different iterations. If the log name is None, it
+            does not output the log to any file.
+
+    """
+
+    def __init__(self, keys=None, trigger=(1, 'epoch'), postprocess=None,
+                 log_name='log', logger=None):
+        self._keys = keys
+        self._trigger = trigger_module.get_trigger(trigger)
+        self._postprocess = postprocess
+        self._log_name = log_name
+        self._log = []
+        self._logger = logger
+        self._init_summary()
+
+    def __call__(self, trainer):
+        # accumulate the observations
+        keys = self._keys
+        observation = trainer.observation
+        summary = self._summary
+
+        if keys is None:
+            summary.add(observation)
+        else:
+            summary.add({k: observation[k] for k in keys if k in observation})
+        for k, v in observation.items():
+            #self._logger.add_scalar(k, chainer.cuda.to_cpu(observation[k].data), trainer.updater.iteration)
+            self._logger.add_scalar(k, observation[k], trainer.updater.iteration)
+        if self._trigger(trainer):
+            # output the result
+            stats = self._summary.compute_mean()
+            stats_cpu = {}
+            for name, value in six.iteritems(stats):
+                stats_cpu[name] = float(value)  # copy to CPU
+
+            updater = trainer.updater
+            stats_cpu['epoch'] = updater.epoch
+            stats_cpu['iteration'] = updater.iteration
+            stats_cpu['elapsed_time'] = trainer.elapsed_time
+
+            if self._postprocess is not None:
+                self._postprocess(stats_cpu)
+
+            self._log.append(stats_cpu)
+
+            # write to the log file
+            if self._log_name is not None:
+                log_name = self._log_name.format(**stats_cpu)
+                fd, path = tempfile.mkstemp(prefix=log_name, dir=trainer.out)
+                with os.fdopen(fd, 'w') as f:
+                    json.dump(self._log, f, indent=4)
+
+                new_path = os.path.join(trainer.out, log_name)
+                shutil.move(path, new_path)
+
+            # reset the summary for the next output
+            self._init_summary()
+
+    @property
+    def log(self):
+        """The current list of observation dictionaries."""
+        return self._log
+
+    def serialize(self, serializer):
+        if hasattr(self._trigger, 'serialize'):
+            self._trigger.serialize(serializer['_trigger'])
+
+        # Note that this serialization may lose some information of small
+        # numerical differences.
+        if isinstance(serializer, serializer_module.Serializer):
+            log = json.dumps(self._log)
+            serializer('_log', log)
+        else:
+            log = serializer('_log', '')
+            self._log = json.loads(log)
+
+    def _init_summary(self):
+        self._summary = reporter.DictSummary()
diff --git a/examples/chainer/data.py b/examples/chainer/plain_logger/data.py
similarity index 100%
rename from examples/chainer/data.py
rename to examples/chainer/plain_logger/data.py
diff --git a/examples/chainer/net.py b/examples/chainer/plain_logger/net.py
similarity index 100%
rename from examples/chainer/net.py
rename to examples/chainer/plain_logger/net.py
diff --git a/examples/chainer/train_vae.py b/examples/chainer/plain_logger/train_vae.py
similarity index 100%
rename from examples/chainer/train_vae.py
rename to examples/chainer/plain_logger/train_vae.py