Chainer에서 tensorboard 사용

tensorboard는 사용해 보면 꽤 편리합니다. 그래프나 히스토그램, 이미지 등을 그릴 수 있어 게다가 인터페이스도 잘 되어 있습니다.
최근에는 chainer를 자주 사용하고 있지만 시각화 도구로이 편리한 tensorboard를 사용하고 싶었기 때문에 chainer에서 tensorboard를 사용하는 패키지를 만들었습니다.

pytorch에서 같은 일을하는 사람 가 있었으므로, 그쪽을 물론 팩하고 베이스로 해 만들고 있습니다.

사용법

기본적으로는 demp.py 이나 demo_graph.py 를 실행해 주면 어떤 느낌인지 알 수 있을까 생각합니다. 그래프 드로잉, 히스토그램, 이미지 등의 드로잉을 지원합니다.

python demo.py
tensorboard --logdir runs

스칼라


히스토그램



이미지

구현 정보

그래프 생성의 곳이 조금 고민되었습니다.
기본적으로 chainer computational_graph를 사용하여 노드와 에지를 계산하고 tensorboard 로그에서 사용되는 protobuf 유형으로 변경합니다. 우선 chainer의 그래프를 그대로 채용하고 있으므로 tensorflow에서는 표시되지 않는 함수와 함수 사이의 VariableNode 도 그려집니다.

name_scope

그리고는 노드의 이름을 붙이는 방법입니다만, tensorflow로 사용되고 있는 name_scope 의 chainer판을 준비해, 그래프를 보기 쉽게 할 수 있도록 하고 있습니다.
다음과 같은 형태로 사용할 수 있습니다.

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... 가 네임 스페이스를 설정하고 있습니다.
이 경우, 네임스페이스를 설정하지 않으면 다음과 같이 모든 노드가 그려집니다만, 네임스페이스를 설정하는 것으로 한층 더 아래와 같이 paint 됩니다.

네임스페이스 없음


네임스페이스 있음


물론 네임스페이스 안의 전개가 가능합니다.

name_scope 구현 방법

상당히 강인한 방법입니다만, with 구문 + 원숭이 패치로 실시하고 있습니다.
VariableNode와 Function의 생성자를 다시 작성하여 네임 스페이스의 스택을 유지할 수있는 생성자로 변경했습니다.

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)

# 必要に応じて関数クラスを増やす。
_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)

위의 _init_with_name_scope 가 통상 constructor 의 호출과 네임스페이스의 보관 유지를 실시하게 되어 있어, with 구문 안에서만 이 constructor 으로 전환하게 되어 있습니다.

추가하려는 기능

가장자리에 텐서 크기 표기

tensorboard에 네트워크 묘화시켰을 때에, 엣지의 곳에 텐서의 사이즈를 표기할 수 있는 것 같습니다만, 어떻게 할까를 잘 모르고 할 수 없습니다.
tensorflow 소스의 내용을 보고 조사하지 않으면 안 될지도…

참고

파이썬으로 블랙 마술
메타 프로그래밍 Python