C++用機械学習ライブラリflashlightとpytorchの速度比較をしてみた(5/3追記)

はじめに

5/3　追記）学習だけでなく推論時の比較も追加しました

この前以下の記事を書きました。

話題のC++用Deep Learningライブラリ「Flashlight」をWSL2上のDockerで動かしてHelloWorldしてみた

この記事ではWSLにCPU-baseのDockerコンテナを立ててHelloWorldするだけでしたが、flashlightに期待することの一つとして実行速度があると思うため、今回はpytorchとflashloghtの速度比較を主に学習時間に絞って行いました。
今回は前回とは異なり、ubuntu機を使ってCUDA-baseのイメージを使っています。

ざっくりとした比較なのですが、結論から言うと

学習時

• flashlightの方が速い(1.4倍くらい)
• GPUのメモリ効率もflashlightの方がpytorchより20%ほど良い
• GPUの使用率はそこまで大差なし

推論時

• pytorchの方が速い(1.1倍くらい)
• GPUのメモリ効率はflashlightの方がpytorchより6%ほど良い
• GPUの使用率はpytorchのほうが効率が良い

といった具合でした(せっかくなのでメモリ等の情報も拾いました)。

実行環境

• OS:Ubuntu 18.04.4 LTS
• CPU:AMD Ryzen 9 3900X
• GPU:GeForce GTX 1070

実験条件

今回は以下の条件で比較しました。できるだけ元のプログラムをpytorchで再現するように書いています。

学習時

• データ：Mnist(DL先)
• モデル：flashlightのexampleにあるMnistを学習しているモデル
• 実行環境：どちらも同じPCでflashlightのイメージを使用(コンテナは分けている)
• 時間計測範囲：プログラム実行開始から終わりまで、具体的な測定方法は以下
  • pytorch：importの終わりから実行終了までをtimeモジュールで計測
  • flashlight：timeコマンドで計測
• その他条件：学習時間をメインで比較したいためepoch数は100に変更、その他パラメータは元のMnist.cppと同じ

推論時

• データ：Mnist(DL先)
• モデル：flashlightのexampleにあるMnistを学習しているモデル
• 実行環境：どちらも同じPCでflashlightのイメージを使用(コンテナは分けている)
• 時間計測範囲：どちらも推論部分のみ
• その他条件：推論時の時間計測は学習済みモデルとそうでないモデルで分けて時間を計測した

pytorchの学習コード

今回flashlightにコードを寄せたのでpytorchでコードを書きました。
notebookから直接起こしたので雑ですが・・・

import time
import torch
import torchvision.transforms as transforms
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import numpy as np
from PIL import Image

t1 = time.time()

TRAIN_SIZE = 60000
VAL_SIZE = 5000
TEST_SIZE = 10000
IM_DIM = 28
PIXEL_MAX = 255

learning_rate = 1e-2
epochs = 100
batch_size = 64

device = "cuda:0"

label_path = "../data/train-labels-idx1-ubyte"
image_path = "../data/train-images-idx3-ubyte"

# https://kyudy.hatenablog.com/entry/2018/11/05/202318 を参考
def read_labels_from_binary(filepath):
        with open(filepath, mode="rb") as f:
            data = f.read()
            array = np.fromstring(data, dtype='<u1')
            array = array[8:]

        return array

def read_images_from_binary(filepath):
    with open(filepath, mode="rb") as f:
        data = f.read()
        array = np.fromstring(data, dtype='<u1')
        array = array[16:]
        len_array = len(array) // IM_DIM // IM_DIM
        array = array.reshape((len_array, IM_DIM*IM_DIM))

    return array

class Mnistdatasets(torch.utils.data.Dataset):
    def __init__(self, images, labels, transform = None):
        self.images = images
        self.labels = labels
        self.transform = transform
        
    def __len__(self):
        return len(self.images)

    def __getitem__(self, idx):
        image = self.images[idx]
        label = self.labels[idx]

        if self.transform:
            batch_image = self.transform(image)

        return image, label

images = read_images_from_binary(image_path)/255
labels = read_labels_from_binary(label_path)

train_images = images[VAL_SIZE:]
val_images = images[:VAL_SIZE]
train_labels = labels[VAL_SIZE:]
val_lebels = labels[:VAL_SIZE]

transform = transforms.Compose([transforms.ToTensor()])

trainset = Mnistdatasets(train_images, train_labels)
trainloader = torch.utils.data.DataLoader(trainset,
                                            batch_size=batch_size)

valset = Mnistdatasets(train_images, train_labels)
valloader = torch.utils.data.DataLoader(trainset,
                                            batch_size=batch_size)

# モデル定義
class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 5, 1, padding=[5 // 2, 5 // 2], padding_mode='replicate')
        self.relu1 = nn.ReLU()
        self.pool1 = nn.MaxPool2d(2,2)
        self.conv2 = nn.Conv2d(32, 64, 5, 1, padding=[5 // 2, 5 // 2], padding_mode='replicate')
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(2,2)
        self.linear1 = nn.Linear(7 * 7 * 64, 1024)
        self.relu3 = nn.ReLU()
        self.dropout1 = nn.Dropout(0.5)
        self.linear2 = nn.Linear(1024, 10)
        
    def forward(self, x):
        x = x.view(-1, 1, IM_DIM, IM_DIM)
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(-1, 7 * 7 * 64)
        x = self.linear1(x)
        x = self.relu3(x)
        x = self.dropout1(x)
        x = self.linear2(x)
        
        return x

model = Model().to(device)

# 学習
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(),lr=learning_rate, momentum=0, dampening=0, weight_decay=0, nesterov=False)

def learning_step(model, batch_x, batch_y, optimizer, criterion, train = True):
    if train:
        optimizer.zero_grad()
        out = model(batch_x)
        loss = criterion(out, batch_y)
        loss.backward()
        optimizer.step()
    else:
        out = model(batch_x)
        loss = criterion(out, batch_y)
    
    return loss.item()

for e in range(epochs):
    print(f"now {e}/{epochs}")
    train_loss = 0
    # train
    model.train()
    for batch_image ,batch_label in trainloader:
        batch_image = batch_image.float().to(device)
        batch_label = batch_label.long().to(device)
        loss = learning_step(model, batch_image, batch_label, optimizer, criterion)
        train_loss += loss
        
    train_loss /= len(trainloader)
        
    # val
    model.eval()
    with torch.no_grad():
        val_loss = 0
        for batch_image ,batch_label in valloader:
            batch_image = batch_image.float().to(device)
            batch_label = batch_label.long().to(device)
            loss = learning_step(model, batch_image, batch_label, optimizer, criterion, False)
            val_loss += loss
        val_loss /= len(valloader)
    print(f"train_loss : {train_loss} , val_loss : {val_loss}")

t2 = time.time()
print(t2-t1)

結果

上記の条件とコードで比較した結果は以下の様になりました。
意外と有意に差が出たため、全て1発実行の結果です。

速度

学習時

学習速度に関してはflashlightの勝利となりました。
そんなに変わらないと思っていたのですが、1.4倍ほども差があってびっくりです。

• pytroch：512.4(s)
• flashlight：365.4(s)

各ライブラリとも学習時間以外のオーバーヘッドを無視しても、大半が学習時間であるため有意に差が出ていると思います。
pytorchも裏で色々うまくやっているはずですが、python自体がインタプリンタということもありそもそも遅いことがでた結果でしょうか？

推論時

学習せずに時間を計測した場合

• pytroch：172.8(s)
• flashlight：189.0(s)

学習して時間を計測した場合

• pytroch：172.8(s)
• flashlight：177.5(s)

こちらはなんとpytorchのほうが速い結果に・・・。
また、学習時と非学習時でflashlightの計測に差が出たのも意外でした。
重みが疎になってる場合に効率良くなるとかあるんですかね。

使用GPUメモリ&GPU使用率

付録程度に使用GPUメモリとGPU使用率の比較も載せておきます。
実行コマンドは以下で、計測範囲はプログラムの開始から終わりまでです。

nvidia-smi   --query-gpu=timestamp,name,utilization.gpu,utilization.memory,memory.used,memory.free,memory.used  --format=csv -l 1 -f gpu.log

使用GPUメモリ

学習時

こちらもflashlightの勝利となりました。※下は最大値比較

• pytorch：968(Mb)
• flashlight：777(Mb)

実行時間とGPUメモリ使用率のグラフも載せておきます。
基本平坦ですがpytorchだけ途中で変わるのはなんでなんでしょうかね。詳しい方いたら教えてください。

推論時

flashlightの勝利となりました。※下は最大値比較

• pytorch：804(Mb)
• flashlight：761(Mb)

同様なグラフも載せておきます。

pytorchにおいて推論時は逆伝播させる必要がないためメモリが効率化されますが、flashlightにそのような機能があるかどうかは知らないです(おそらくあると思いますが)。

GPU使用率

学習時

こちらは大差ありませんでした。
時間平均とかすると違いがわかるかもですがめんどくさいのでパスです

推論時

推論時はなぜかpytorchのほうが効率が良かったです。

おわりに

今回はflashlightとpytorchの速度等の比較を実施し、速度とメモリ効率の点でflashlightのほうが良いだろうという結論を得ました。
単純に速度が必要な場合はflashlightを使うと良いと思いますが、結局データの扱いやすさや実装の手間を考えるとC++が得意でない限りは実務上pytorchが有利だと思っています。
推論時にflashlightとpytorchに差があまり見られなかったため、pytorchでいいじゃんという気持ちになりました。

また今回は色々面倒なのでモデルはflashlightのサンプルに寄せましたが、よく使われるresnetなどを組んでみて比較すると実際発生しそうな差がわかるかもしれません。
多分やりませんが、気が向いて測定まで行けたらこの記事に追記したいと思います(多分flashlightが速いことには変わりないと思いますが)。