PyTorchでモデルの演算量を確認する方法｜PyTorch TIPS

Aru

ディープラーニングのモデルを利用・作成するとき、「このモデルの演算量（計算量）はどれくらいだろう」と演算量が気になることがあるかと思います。実際にサービスに適用することを考えた場合、処理量と性能のバランスを適切に取ることを求められることがあります。この記事では、Pytorchのモデルの演算量を確認する方法を解説します。

Contents

はじめに
- モデルの演算量を調べたい理由
- 演算量の指標（MACs）
演算量（MACs）を求める方法
- ptflopsをインストールする
- 演算量を出力
まとめ

はじめに

モデルの演算量を調べたい理由

ディープラーニングのモデルを実際のサービスで利用する場合、演算量を把握することは非常に重要です。というのも、演算量の少ないモデルは処理を高速に行うことが可能なので、エッジデバイスなどのリソースが限られて環境で動作させたい場合には大切な判断基準となります。また、クラウドサービスで利用する場合も、処理時間＝コストなので、処理量が少ないモデルの方が有利です。

とはいえ、演算量が少なくても性能が低くてはどうしょうもありません。実務では、性能と処理用のバランスを考えてモデルを選定することが求められます。

では、演算量はどのようにして調査したらよいでしょうか？

この記事では、PyTorchで作成したモデルの演算量を調べる方法について解説します。

演算量の指標（MACs）

演算量の指標としてMACs（Multiply-Accumulate Operations）があります。

MACsとは、ディープラーニングにおけるモデルの計算量を評価するための指標の一つです。

MACsは乗算（Multiply）と加算（Accumulate）の組み合わせを意味しています。

この２つの操作（積和演算）は、ニューラルネットワークで頻繁に利用される操作で、畳み込み層や全結合層の処理の基本単位となります。

MACsは、性能評価の指標として利用されていて、MACsはモデルの演算量の単位となります。

FLOPs（Floating Point Operations Per Second）が演算量として用いられますが、多くの論文ではMACsが用いられています。

演算量（MACs）を求める方法

ptflopsをインストールする

MACsを求めるには、ptflopsパッケージを利用します。

github: Flops counting tool for neural networks in pytorch framework

以下の方法でインストール可能です。

pip install ptflops

演算量を出力

使い方は簡単です。以下のようにget_model_complexity_info関数を呼び出すことでモデルの計算量を出力できます。

from torchvision import models
from ptflops import get_model_complexity_info

model = models.resnet18()
result = get_model_complexity_info(model, (3, 224, 224), as_strings=False, print_per_layer_stat=False, verbose=False)
print(result)

引数の1つ目はモデルのオブジェクトです。

引数の2つ目はモデルに入力するデータの形式です。resnet18は(batch, 3, 224, 224)の入力を受け取るので(3, 224, 224)を設定します。

MACsは、設定した入力データ形式に対して計算されることに注意してください。上記の例では、(3, 224, 224)という入力サイズに対して計算されます。たとえば、入力を(3, 256, 256)とした場合は演算量は変化します。

結果は以下のように出力されます。出力の１つ目がGMACsで、２つ目はパラメータ数（M）になります。結果から、resnet18は、(3, 224, 224)の入力の場合、1.83GMACsで、パラメータ数が11.69Mということがわかります。

(1825313768, 11689512)

as_strings=Trueに設定すれば、1.8GMacなどの表記になり、見やすくなります。

get_model_complexity_info()の主なオプションは以下になります。

オプション名	説明
as_strings	True/Falseで設定。Trueを設定すると文字列で結果が出力される（‘1.83 GMac’, ‘11.69 M’のように読みやすくなります）
print_per_layer_stat	True/Falseで設定。Trueに設定するとレイヤごとの情報を表示
verbose	True/Falseで設定。Trueに設定するとWarningなどを出力

モデルの詳細出力

print_per_layer_statにTrueを設定すると以下のように各レイヤーの処理量を詳細に出力します。

from torchvision import models
from ptflops import get_model_complexity_info

model = models.resnet18()
result = get_model_complexity_info(model, (3, 224, 224), as_strings=True, print_per_layer_stat=True, verbose=False)
print(result)

ResNet(
  11.69 M, 100.000% Params, 1.82 GMac, 99.828% MACs, 
  (conv1): Conv2d(9.41 k, 0.080% Params, 118.01 MMac, 6.465% MACs, 3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
  (bn1): BatchNorm2d(128, 0.001% Params, 1.61 MMac, 0.088% MACs, 64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (relu): ReLU(0, 0.000% Params, 802.82 KMac, 0.044% MACs, inplace=True)
  (maxpool): MaxPool2d(0, 0.000% Params, 802.82 KMac, 0.044% MACs, kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
  (layer1): Sequential(
    147.97 k, 1.266% Params, 464.83 MMac, 25.466% MACs, 
    (0): BasicBlock(
      73.98 k, 0.633% Params, 232.42 MMac, 12.733% MACs, 
      (conv1): Conv2d(36.86 k, 0.315% Params, 115.61 MMac, 6.333% MACs, 64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn1): BatchNorm2d(128, 0.001% Params, 401.41 KMac, 0.022% MACs, 64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(0, 0.000% Params, 401.41 KMac, 0.022% MACs, inplace=True)
      (conv2): Conv2d(36.86 k, 0.315% Params, 115.61 MMac, 6.333% MACs, 64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): BatchNorm2d(128, 0.001% Params, 401.41 KMac, 0.022% MACs, 64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    )
    (1): BasicBlock(
      73.98 k, 0.633% Params, 232.42 MMac, 12.733% MACs, 
      (conv1): Conv2d(36.86 k, 0.315% Params, 115.61 MMac, 6.333% MACs, 64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn1): BatchNorm2d(128, 0.001% Params, 401.41 KMac, 0.022% MACs, 64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(0, 0.000% Params, 401.41 KMac, 0.022% MACs, inplace=True)
      (conv2): Conv2d(36.86 k, 0.315% Params, 115.61 MMac, 6.333% MACs, 64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): BatchNorm2d(128, 0.001% Params, 401.41 KMac, 0.022% MACs, 64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    )
  )
  (layer2): Sequential(
    525.57 k, 4.496% Params, 412.45 MMac, 22.596% MACs, 
    (0): BasicBlock(
      230.14 k, 1.969% Params, 180.63 MMac, 9.896% MACs, 
      (conv1): Conv2d(73.73 k, 0.631% Params, 57.8 MMac, 3.167% MACs, 64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
      (bn1): BatchNorm2d(256, 0.002% Params, 200.7 KMac, 0.011% MACs, 128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(0, 0.000% Params, 200.7 KMac, 0.011% MACs, inplace=True)
      (conv2): Conv2d(147.46 k, 1.261% Params, 115.61 MMac, 6.333% MACs, 128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): BatchNorm2d(256, 0.002% Params, 200.7 KMac, 0.011% MACs, 128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (downsample): Sequential(
        8.45 k, 0.072% Params, 6.62 MMac, 0.363% MACs, 
        (0): Conv2d(8.19 k, 0.070% Params, 6.42 MMac, 0.352% MACs, 64, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)
        (1): BatchNorm2d(256, 0.002% Params, 200.7 KMac, 0.011% MACs, 128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      )
    )
    (1): BasicBlock(
      295.42 k, 2.527% Params, 231.81 MMac, 12.700% MACs, 
      (conv1): Conv2d(147.46 k, 1.261% Params, 115.61 MMac, 6.333% MACs, 128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn1): BatchNorm2d(256, 0.002% Params, 200.7 KMac, 0.011% MACs, 128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(0, 0.000% Params, 200.7 KMac, 0.011% MACs, inplace=True)
      (conv2): Conv2d(147.46 k, 1.261% Params, 115.61 MMac, 6.333% MACs, 128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): BatchNorm2d(256, 0.002% Params, 200.7 KMac, 0.011% MACs, 128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    )
  )
  (layer3): Sequential(
    2.1 M, 17.962% Params, 411.74 MMac, 22.557% MACs, 
    (0): BasicBlock(
      919.04 k, 7.862% Params, 180.23 MMac, 9.874% MACs, 
      (conv1): Conv2d(294.91 k, 2.523% Params, 57.8 MMac, 3.167% MACs, 128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
      (bn1): BatchNorm2d(512, 0.004% Params, 100.35 KMac, 0.005% MACs, 256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(0, 0.000% Params, 100.35 KMac, 0.005% MACs, inplace=True)
      (conv2): Conv2d(589.82 k, 5.046% Params, 115.61 MMac, 6.333% MACs, 256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): BatchNorm2d(512, 0.004% Params, 100.35 KMac, 0.005% MACs, 256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (downsample): Sequential(
        33.28 k, 0.285% Params, 6.52 MMac, 0.357% MACs, 
        (0): Conv2d(32.77 k, 0.280% Params, 6.42 MMac, 0.352% MACs, 128, 256, kernel_size=(1, 1), stride=(2, 2), bias=False)
        (1): BatchNorm2d(512, 0.004% Params, 100.35 KMac, 0.005% MACs, 256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      )
    )
    (1): BasicBlock(
      1.18 M, 10.100% Params, 231.51 MMac, 12.683% MACs, 
      (conv1): Conv2d(589.82 k, 5.046% Params, 115.61 MMac, 6.333% MACs, 256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn1): BatchNorm2d(512, 0.004% Params, 100.35 KMac, 0.005% MACs, 256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(0, 0.000% Params, 100.35 KMac, 0.005% MACs, inplace=True)
      (conv2): Conv2d(589.82 k, 5.046% Params, 115.61 MMac, 6.333% MACs, 256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): BatchNorm2d(512, 0.004% Params, 100.35 KMac, 0.005% MACs, 256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    )
  )
  (layer4): Sequential(
    8.39 M, 71.806% Params, 411.39 MMac, 22.538% MACs, 
    (0): BasicBlock(
      3.67 M, 31.422% Params, 180.03 MMac, 9.863% MACs, 
      (conv1): Conv2d(1.18 M, 10.092% Params, 57.8 MMac, 3.167% MACs, 256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
      (bn1): BatchNorm2d(1.02 k, 0.009% Params, 50.18 KMac, 0.003% MACs, 512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(0, 0.000% Params, 50.18 KMac, 0.003% MACs, inplace=True)
      (conv2): Conv2d(2.36 M, 20.183% Params, 115.61 MMac, 6.333% MACs, 512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): BatchNorm2d(1.02 k, 0.009% Params, 50.18 KMac, 0.003% MACs, 512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (downsample): Sequential(
        132.1 k, 1.130% Params, 6.47 MMac, 0.355% MACs, 
        (0): Conv2d(131.07 k, 1.121% Params, 6.42 MMac, 0.352% MACs, 256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)
        (1): BatchNorm2d(1.02 k, 0.009% Params, 50.18 KMac, 0.003% MACs, 512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      )
    )
    (1): BasicBlock(
      4.72 M, 40.384% Params, 231.36 MMac, 12.675% MACs, 
      (conv1): Conv2d(2.36 M, 20.183% Params, 115.61 MMac, 6.333% MACs, 512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn1): BatchNorm2d(1.02 k, 0.009% Params, 50.18 KMac, 0.003% MACs, 512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(0, 0.000% Params, 50.18 KMac, 0.003% MACs, inplace=True)
      (conv2): Conv2d(2.36 M, 20.183% Params, 115.61 MMac, 6.333% MACs, 512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): BatchNorm2d(1.02 k, 0.009% Params, 50.18 KMac, 0.003% MACs, 512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    )
  )
  (avgpool): AdaptiveAvgPool2d(0, 0.000% Params, 25.09 KMac, 0.001% MACs, output_size=(1, 1))
  (fc): Linear(513.0 k, 4.389% Params, 513.0 KMac, 0.028% MACs, in_features=512, out_features=1000, bias=True)
)
('1.83 GMac', '11.69 M')

まとめ

timm（PyTorch Image Models）などをよく利用するのですが、演算量を知りたい場合に重宝します。

参考までに、公式ページにあるtorchvisionのモデルの処理量一覧を貼っておきます

Model	Input Resolution	Params(M)	MACs(G) (`pytorch`)	MACs(G) (`aten`)
alexnet	224×224	61.10	0.72	0.71
convnext_base	224×224	88.59	15.43	15.38
densenet121	224×224	7.98	2.90
efficientnet_b0	224×224	5.29	0.41
efficientnet_v2_m	224×224	54.14	5.43
googlenet	224×224	13.00	1.51
inception_v3	224×224	27.16	5.75	5.71
maxvit_t	224×224	30.92	5.48
mnasnet1_0	224×224	4.38	0.33
mobilenet_v2	224×224	3.50	0.32
mobilenet_v3_large	224×224	5.48	0.23
regnet_y_1_6gf	224×224	11.20	1.65
resnet18	224×224	11.69	1.83	1.81
resnet50	224×224	25.56	4.13	4.09
resnext50_32x4d	224×224	25.03	4.29
shufflenet_v2_x1_0	224×224	2.28	0.15
squeezenet1_0	224×224	1.25	0.84	0.82
vgg16	224×224	138.36	15.52	15.48
vit_b_16	224×224	86.57	17.61 (wrong)	16.86
wide_resnet50_2	224×224	68.88	11.45