CNNの計算量

そこら中でやられているので、やってみた記事。Tipsも含めて。
例えばConvolution層の計算量（ここでは計算量をFLOPsと定義し、biasに関する計算は除く）は、下記のように表されます。

$$
FLOps = HWC_{in}K^2C_{out}
$$

ここで、$H,W,C_{in}$はそれぞれ入力画像の縦横チャンネル数。$K$はカーネルサイズで、$C_{out}$は出力のチャンネル数です。
基本的な内容ですので、興味がある人はこちらの論文¹などを参照されると良いと思います。

これを全部の層について考えて、モデル全体でどのくらいあるのかを調べたいわけです。
演算の並列化や最適化があると厳密ではありませんが、だいたいの計算量の目安になります。

flops_counter

Computer Vision界隈で有名なオープンソースのmmcvですが、こちらのutilityの中に、flops_counterというものがあります。その名の通り、モデルを与えるとFLOPsをカウントしてくれる、超絶便利なモジュールです。

使い方は下記のように行います。

from mmcv.cnn.utils import flops_counter

model = <MODEL DEFINE>
input_shape = <INPUT SHAPE>
flops_counter.get_model_complexity_info(model, input_shape)

非常にシンプルです。

実際に動かしてみる

手元で適当な小さいPyTorchのモデルを作ってみます。
全結合層が二つ。途中に適当なReluとDropoutを入れています。
この二つは不要ですが、どのように表示されるか、見てみましょう。

import torch.nn as nn
from mmcv.cnn.utils import flops_counter
net_seq = nn.Sequential(
    nn.Linear(100, 10),  #  weight:100x10=1000, bias:10
    nn.ReLU(),
    nn.Dropout(0.2),
    nn.Linear(10, 2)  # weight:10x2=20, bias:2
)
flops_counter.get_model_complexity_info(net_seq, (100,), as_strings=False)

実行結果

Sequential(
  0.001 M, 100.000% Params, 0.0 GFLOPs, 100.000% FLOPs, 
  (0): Linear(0.001 M, 97.868% Params, 0.0 GFLOPs, 97.087% FLOPs, in_features=100, out_features=10, bias=True)
  (1): ReLU(0.0 M, 0.000% Params, 0.0 GFLOPs, 0.971% FLOPs, )
  (2): Dropout(0.0 M, 0.000% Params, 0.0 GFLOPs, 0.000% FLOPs, p=0.2, inplace=False)
  (3): Linear(0.0 M, 2.132% Params, 0.0 GFLOPs, 1.942% FLOPs, in_features=10, out_features=2, bias=True)
)
(1030.0, 1032)

期待する結果通り、約1000回の演算と、1032個のWeightが表示されました。
（演算も1032を期待していたのですが、1030となりました。誰か理由を教えて。どこかで丸められているだけ？）

また、どの層でどの程度のパラメータがあるのか、など詳細に見ることが可能です。
お手軽で素晴らしいですね。

flops_counter使用時のエラー

状況

pip install mmcv-fullなどをしたあとに、こんなエラーが出た人はこちらを参照してください。

RuntimeError: roi_width >= 0 && roi_height >= 0INTERNAL ASSERT FAILED at "/<PROJECT PATH>/mmcv/mmcv/ops/csrc/pytorch/cpu/roi_align.cpp":141, please report a bug to PyTorch. ROIs in ROIAlign cannot have non-negative size!

対策

ソースから直接インストールする。

まずはソースのダウンロード。

git clone https://github.com/open-mmlab/mmcv.git

下記ファイルの140行目付近に、２行を追加する。

mmcv/mmcv/ops/csrc/pytorch/cpu/roi_align.cpp

T roi_width = roi_end_w - roi_start_w;
T roi_height = roi_end_h - roi_start_h;
if (roi_width < 0) roi_width = 0;  // この行を追加
if (roi_height < 0) roi_height = 0;  // この行を追加
if (aligned) {
  AT_ASSERTM(roi_width >= 0 && roi_height >= 0,
            "ROIs in ROIAlign cannot have non-negative size!");
} else {  // for backward-compatibility only
  roi_width = std::max(roi_width, (T)1.);
  roi_height = std::max(roi_height, (T)1.);
}

インストールをする

command on linux

MMCV_WITH_OPS=1 pip install -e .

command on mac

CC=clang CXX=clang++ CFLAGS='-stdlib=libc++' MMCV_WITH_OPS=1 pip install -e .

これで無事に実行できました。

Pruning Convolutional Neural Networks for Resource Efficient Inference. P.Molchanov et al ↩