Edited at

Chainerでtensorboardを使う

More than 1 year has passed since last update.

tensorboardは使ってみるとかなり便利です。グラフやヒストグラム、画像などが描画できて、しかもインターフェースも良く出来ています。

私は最近では、chainerをよく使っているのですが、可視化ツールとしてこの便利なtensorboardを使いたかったので、chainerでtensorboardを使うパッケージを作りました。

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

pytorchで同じことをやっている方がいましたので、そちらをもろにパクってベースにして作っています。


使い方

基本的にはdemp.pydemo_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...のところが名前空間を設定しているところです。

この場合、名前空間を設定しないと以下のように全ノードが描画されますが、名前空間を設定することでさらに下のように描画されます。

名前空間無し

名前空間有り

もちろん名前空間の中の展開が可能です。


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が通常コンストラクタの呼び出しと名前空間の保持を行うようになっており、with構文の中でだけこのコンストラクタに切り替わるようになっています。


追加したい機能


エッジにテンソルのサイズ表記

tensorboardにネットワーク描画させた際に、エッジのところにテンソルのサイズを表記できるっぽいのですが、どうやるかが良く分からなくてできてません。

tensorflowのソースの中身を見て調べないとダメかも…


参考

Pythonによる黒魔術

メタプログラミングPython