MobileNet, SqueezeNet, ResNet等の処理時間の比較

はじめに

PyTorchのMobileNet実装のリポジトリに、SqueezeNet等の推論時の処理時間を比較しているコードがあったので、ちょっと改変してCPUも含めて処理時間の比較を行った。
環境はUbuntu 16.04, CPU: i7-7700 3.60GHz、GPU: GeForce GTX1080。
PyTorchのバージョンは0.1.11.

結果

use_gpu: True, nb_batches: 1
  resnet18 : 0.001915 (sd 0.000057)
   alexnet : 0.000691 (sd 0.000005)
     vgg16 : 0.002390 (sd 0.002091)
squeezenet : 0.002086 (sd 0.000104)
 mobilenet : 0.048602 (sd 0.000380)
use_gpu: True, nb_batches: 16
  resnet18 : 0.006055 (sd 0.005111)
   alexnet : 0.000744 (sd 0.000014)
     vgg16 : 0.025156 (sd 0.029848)
squeezenet : 0.012983 (sd 0.000024)
 mobilenet : 0.064022 (sd 0.000411)
use_gpu: False, nb_batches: 1
  resnet18 : 0.218282 (sd 0.002961)
   alexnet : 0.081834 (sd 0.000445)
     vgg16 : 1.484166 (sd 0.001384)
squeezenet : 0.102657 (sd 0.002118)
 mobilenet : 0.141093 (sd 0.005197)
use_gpu: False, nb_batches: 16
  resnet18 : 0.896854 (sd 0.004594)
   alexnet : 0.283497 (sd 0.003010)
     vgg16 : 5.622119 (sd 0.020102)
squeezenet : 0.514910 (sd 0.004134)
 mobilenet : 0.892604 (sd 0.017502)

元々はCPUでの推論時の処理時間に興味があったのだが、やはりパラメータ数と処理時間は結構切り離して考える必要があり、パラメータが少ないからと言って早いとは限らない。
面白いのは、VGGとMobileNetが対照的で、VGGはガッツリ重い畳み込みが少数あるのに対して、MobileNetは細かい畳込みがいっぱいあり、前者はGPUでかなり効率的に計算ができる一方、後者はGPUの恩恵をあまり受けられない結果となっているところ。CPUでVGGは動かすなって感じだけど、他のフレームワークでも同様の傾向なのか気になる。
逆にCPUだと理論的な計算量に近い直感的な結果となっている印象がある。

ちなみにcudnn.fastest = True, cudnn.benchmark = Trueとした場合。早くなるモデルと遅くなるモデルがある。

use_gpu: True, nb_batches: 1
  resnet18 : 0.001858 (sd 0.000051)
   alexnet : 0.000694 (sd 0.000017)
     vgg16 : 0.002456 (sd 0.001442)
squeezenet : 0.002058 (sd 0.000012)
 mobilenet : 0.040444 (sd 0.000894)
use_gpu: True, nb_batches: 16
  resnet18 : 0.005278 (sd 0.004701)
   alexnet : 0.000699 (sd 0.000026)
     vgg16 : 0.026446 (sd 0.026279)
squeezenet : 0.013775 (sd 0.000225)
 mobilenet : 0.052172 (sd 0.000935)

コード

import time
import numpy as np
import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn
import torchvision.models as models
from torch.autograd import Variable

# cudnn.fastest = True
# cudnn.benchmark = True
nb_itr = 20


class MobileNet(nn.Module):
    def __init__(self):
        super(MobileNet, self).__init__()

        def conv_bn(inp, oup, stride):
            return nn.Sequential(
                nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
                nn.BatchNorm2d(oup),
                nn.ReLU(inplace=True)
            )

        def conv_dw(inp, oup, stride):
            return nn.Sequential(
                nn.Conv2d(inp, inp, 3, stride, 1, groups=inp, bias=False),
                nn.BatchNorm2d(inp),
                nn.ReLU(inplace=True),

                nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
                nn.BatchNorm2d(oup),
                nn.ReLU(inplace=True),
            )

        self.model = nn.Sequential(
            conv_bn(  3,  32, 2), 
            conv_dw( 32,  64, 1),
            conv_dw( 64, 128, 2),
            conv_dw(128, 128, 1),
            conv_dw(128, 256, 2),
            conv_dw(256, 256, 1),
            conv_dw(256, 512, 2),
            conv_dw(512, 512, 1),
            conv_dw(512, 512, 1),
            conv_dw(512, 512, 1),
            conv_dw(512, 512, 1),
            conv_dw(512, 512, 1),
            conv_dw(512, 1024, 2),
            conv_dw(1024, 1024, 1),
            nn.AvgPool2d(7),
        )
        self.fc = nn.Linear(1024, 1000)

    def forward(self, x):
        x = self.model(x)
        x = x.view(-1, 1024)
        x = self.fc(x)
        return x


def speed(model, name, nb_batches):
    t0 = time.time()
    input = torch.rand(nb_batches,3,224,224)

    if use_gpu:
        input = Variable(input.cuda(), volatile = True)
    else:
        input = Variable(input, volatile = True)

    model(input)

    timings = []

    for _ in range(nb_itr):
        start_time = time.time()
        model(input)
        timings.append(time.time() - start_time)

    print('%10s : %f (sd %f)'% (name, np.array(timings).mean(), np.array(timings).std()))


if __name__ == '__main__':
    for use_gpu in [True, False]:
        for nb_batches in [1, 16]:
            if use_gpu:
                resnet18 = models.resnet18().cuda()
                alexnet = models.alexnet().cuda()
                vgg16 = models.vgg16().cuda()
                squeezenet = models.squeezenet1_0().cuda()
                mobilenet = MobileNet().cuda()
            else:
                resnet18 = models.resnet18()
                alexnet = models.alexnet()
                vgg16 = models.vgg16()
                squeezenet = models.squeezenet1_0()
                mobilenet = MobileNet()

            print("use_gpu: {}, nb_batches: {}".format(use_gpu, nb_batches))
            speed(resnet18, 'resnet18', nb_batches)
            speed(alexnet, 'alexnet', nb_batches)
            speed(vgg16, 'vgg16', nb_batches)
            speed(squeezenet, 'squeezenet', nb_batches)
            speed(mobilenet, 'mobilenet', nb_batches)