Chainerで書いたニューラルネットの理論計算量を推定するchainer_computational_cost

海洋の生物量の推定をするときには、実際にある空間の種ごとの個体数を1匹1匹数えるわけにはいきません。そこで、水中に残った生物の細胞のDNAすなわち環境DNAの密度を手がかりにして推計をするのだそうです(バケツ一杯の水で海洋生物の量や種類を知る)。

Chainerで書いたNNとダミー入力があるとき、そのNNのforward passの理論的な計算量・メモリ転送量を計算するchainer_computational_costを作りました。
ChainerのFunction Hookをベースにしているため、NNの定義コードに手を入れる必要は一切ありません。

import chainer
import chainer.links as L
import numpy as np
import chainer_computational_cost as c3
net = L.VGG16Layers()
x = np.random.random((1, 3, 224, 224)).astype(np.float32) # dummy input
with chainer.no_backprop_mode(), chainer.using_config('train', False):
with c3.ComputationalCostHook(fma_1flop=True) as cch:
y = net(x)
cch.show_summary_report(mode='table') # ←
+-----------------------+----------+--------+---------+----------+--------+
| Layer type | # Layers | GFLOPs | MemRead | MemWrite | MemR+W |
| | | | GiB | GiB | GiB |
+=======================+==========+========+=========+==========+========+
| Convolution2DFunction | 13 | 15.347 | 0.089 | 0.05 | 0.139 |
+-----------------------+----------+--------+---------+----------+--------+
| ReLU | 15 | 0.014 | 0.05 | 0.05 | 0.101 |
+-----------------------+----------+--------+---------+----------+--------+
| MaxPooling2D | 5 | 0.005 | 0.023 | 0.006 | 0.029 |
+-----------------------+----------+--------+---------+----------+--------+
| Reshape | 1 | 0.0 | 0.0 | 0.0 | 0.0 |
+-----------------------+----------+--------+---------+----------+--------+
| LinearFunction | 3 | 0.124 | 0.461 | 0.0 | 0.461 |
+-----------------------+----------+--------+---------+----------+--------+
| Softmax | 1 | 0.0 | 0.0 | 0.0 | 0.0 |
+-----------------------+----------+--------+---------+----------+--------+
| total | 38 | 15.488 | 0.623 | 0.107 | 0.729 |
+-----------------------+----------+--------+---------+----------+--------+

Continue reading