[PYTHON] Use tensorboard with Chainer

The tensorboard is quite convenient to use. Graphs, histograms, images, etc. can be drawn, and the interface is also good. Recently, I use chainer a lot, but I wanted to use this convenient tensorboard as a visualization tool, so I made a package that uses tensorboard with chainer.

https://github.com/neka-nat/tensorboard-chainer

There was People doing the same thing with pytorch, so I'm making it based on that.

How to use

Basically, you can see what it looks like by running demp.py or demo_graph.py. It supports drawing graphs, histograms, images, etc.

python demo.py
tensorboard --logdir runs

** SCARA **

histogram

image

About implementation

I was a little annoyed about the graph generation. Basically, I use chainer's computational_graph to calculate nodes and edges, and change it to the type of protobuf used in the tensorboard log. For the time being, since the graph of chainer is adopted as it is, VariableNode between functions that is not displayed in tensorflow is also drawn.

name_scope The rest is how to name the nodes, but we have prepared a chainer version of name_scope used in tensorflow to make the graph easier to see. It can be used in the following forms.

import chainer
import chainer.functions as F
import chainer.links as L
from tb_chainer import name_scope

class MLP(chainer.Chain):
    def __init__(self, n_units, n_out):
        super(MLP, self).__init__()
        with self.init_scope():
            self.l1 = L.Linear(None, n_units)  # n_in -> n_units
            self.l2 = L.Linear(None, n_units)  # n_units -> n_units
            self.l3 = L.Linear(None, n_out)  # n_units -> n_out

    def __call__(self, x):
        with name_scope('linear1', self.l1.params()):
            h1 = F.relu(self.l1(x))
        with name_scope('linear2', self.l2.params()):
            h2 = F.relu(self.l2(h1))
        with name_scope('linear3', self.l3.params()):
            o = self.l3(h2)
        return o

with name_scope ... is where the namespace is set. In this case, if you do not set the namespace, all nodes will be drawn as shown below, but if you set the namespace, it will be drawn as shown below.

** No namespace **

** With namespace **

Of course, it can be expanded in the namespace.

How to implement name_scope

It's a pretty brute force method, but it's done with with syntax + monkey patch. The constructors of VariableNode and Function have been rewritten to change to constructors that can hold the namespace stack.

from chainer import function
from chainer import variable
import functools
from types import MethodType

def copy_method(c):
    g = MethodType(c.__init__, None, c)
    return g

def _init_with_name_scope(self, *args, **kargs):
    self.name_scope = kargs['name_scope']
    org_init = kargs['org_init']
    del kargs['name_scope']
    del kargs['org_init']
    org_init(self, *args, **kargs)

#Increase the function class as needed.
_org_classes = [function.Function,
                variable.VariableNode]
_copy_org_inits = [copy_method(c) for c in _org_classes]

class name_scope(object):
    stack = []
    def __init__(self, name, values=[]):
        self.stack.append(name)
        self._org_inits = []
        for v in values:
            v.node.name_scope = '/'.join(self.stack)

    def __enter__(self):
        for idx, c in enumerate(_org_classes):
            self._org_inits.append(c.__init__)
            c.__init__ = MethodType(functools.partial(_init_with_name_scope,
                                                      name_scope='/'.join(self.stack),
                                                      org_init=_copy_org_inits[idx]),
                                    None, c)
        return self

    def __exit__(self, exec_type, exec_value, traceback):
        for idx, c in enumerate(_org_classes):
            c.__init__ = self._org_inits[idx]
        self.stack.pop(-1)

The above _init_with_name_scope normally calls the constructor and preserves the namespace, and only switches to this constructor in the with syntax.

Features you want to add

Tensor size notation on the edge

It seems that the size of the tensor can be indicated at the edge when drawing the network on the tensorboard, but I can't do it because I don't know how to do it. You may have to look at the contents of the tensorflow source to find out ...

reference

Blackmagic with Python Metaprogramming Python