TensorFlow 2.0 の速度メモ 【vs PyTorch】

はじめに

TF2.0が出たので速度をメモしておきます。

TensorFlow 2.0

データ

とりあえずcifer の形式で適当にデータを作成。

import numpy as np
import time

train_data = np.random.randn(5000, 32, 32, 3).astype(np.float32)
label_data = np.random.randint(0, 10, 5000).astype(np.int32)

モジュール名

kerasを叩きまくることになるので省略名を付けておきます。

import tensorflow as tf

tfk = tf.keras
tfkl = tfk.layers

データセット

データをタプルで渡して、シャッフルのバッファーサイズとバッチサイズを指定します。

batch_size = 128
dataset = tf.data.Dataset.from_tensor_slices((train_data, label_data))
dataset = dataset.shuffle(buffer_size=50000)
dataset = dataset.batch(batch_size, drop_remainder=True)

モデル作成

tf.keras.Model を継承して下記で適当に作成。 注意点としては、call メソッドで training 引数を渡せるようにしておくこと。（実は tf.keras.Model が **kwargs で受け取れるようになっているので明示しなくても良さそう）

dropoutやバッチ正規化を呼び出すときに必ず渡すこと。

class Model(tf.keras.Model):
    def __init__(self, **kwargs):
        super(Model, self).__init__(**kwargs)

        self.activation = tf.keras.layers.ReLU()
        self.conv1 = tfkl.Conv2D(64, (3, 3), padding='VALID')
        self.bn1 = tfkl.BatchNormalization()
        self.conv2 = tfkl.Conv2D(128, (3, 3), padding='VALID')
        self.bn2 = tfkl.BatchNormalization()
        self.conv3 = tfkl.Conv2D(256, (3, 3), padding='VALID')
        self.bn3 = tfkl.BatchNormalization()
        self.flatten = tfkl.Flatten()
        self.dense = tfkl.Dense(10)

    def call(self, inputs, training=False):
        h = self.conv1(inputs)
        h = self.bn1(h, training=training)
        h = self.activation(h)
        h = self.conv2(h)
        h = self.bn2(h, training=training)
        h = self.activation(h)
        h = self.conv3(h)
        h = self.bn3(h, training=training)
        h = self.activation(h)
        h = self.flatten(h)
        return self.dense(h)

モデルのインスタンス化と訓練準備

tfk.metrics が地味に便利です。（running_lossとか必要なし）

train_loss = tfk.metrics.Mean() 
train_acc = tfk.metrics.SparseCategoricalAccuracy()
optimizer = tf.optimizers.Adam(1.0e-4)
model = Model()

訓練関数

学習の1iterationを書いておきます。tf.function でこの関数をTFのGraph op に変換してくれます。 なるべく中の処理はTFの関数で記述すること。あまりPythonネイティブの処理は書かないほうがバグを踏まずに済む。

@tf.function
def train_step(inputs):
    images, labels = inputs

    with tf.GradientTape() as tape:
        logits = model(images, training=True)
        loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels, logits)

    grad = tape.gradient(loss, sources=model.trainable_variables)
    optimizer.apply_gradients(zip(grad, model.trainable_variables))

    train_loss.update_state(loss)
    train_acc.update_state(labels, logits)

訓練

学習コードを書く。ここで、train_step をgpuに割り当てておけば、中のTensorをすべてgpuで処理してくれます。楽ちん。このコンテキストの外はすべてcpu処理なので、IOを書くところはコンテキストの外に。

epochs = 10
for epoch in range(epochs):
    time_start = time.time()

    for images, labels in dataset: 
        with tf.device("/gpu:0"):
            train_step((images, labels)) 

    epoch_loss = train_loss.result()
    epoch_acc = 100 * train_acc.result()

    time_epoch = time.time() - time_start
    print('epoch: {:} loss: {:.4f} acc: {:.2f}% time: {:.2f}s'.format(
        epoch + 1, epoch_loss, epoch_acc, time_epoch))

    train_loss.reset_states()
    train_acc.reset_states()

google colab で 1epoch 3.6 seconds 程でした。ただし最初のepochは tf.function の中のGraphを評価し計算グラフの構築に時間を要するため 6 seconds となりました。

ちなみに eagerのままだと 4.5seconds程。

PyTorch

import

おなじみ

import torch
import torch.nn as nn
import torch.nn.functional as F

データ準備

データローダーの作成周りは、色々カスタマイズが利きますが、とりあえずデフォルトで。

train_data = np.random.randn(5000, 3, 32, 32).astype(np.float32)
label_data = np.random.randint(0, 10, 5000).astype(np.int32)

trainset = torch.utils.data.TensorDataset(torch.tensor(train_data), 
                                          torch.tensor(label_data))
trainloader = torch.utils.data.DataLoader(trainset,
                                          batch_size=128,
                                          shuffle=True,
                                          num_workers=1)

モデル作成

畳み込みで縦横がどのように変化していくかは計算が必要。ここは地味に面倒だったりしますが、このレベルなら困りません。

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.activation = nn.ReLU()
        self.conv1 = nn.Conv2d(3, 64, 3)
        self.bn1 = nn.BatchNorm2d(64) 
        self.conv2 = nn.Conv2d(64, 128, 3) 
        self.bn2 = nn.BatchNorm2d(128) 
        self.conv3 = nn.Conv2d(128, 256, 3) 
        self.bn3 = nn.BatchNorm2d(256) 
        self.fc = nn.Linear(26*26*256, 10)
        
    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.activation(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.activation(x)
        x = self.conv3(x)
        x = self.bn3(x)
        x = self.activation(x)
        x = x.reshape(-1, x.size(1)*x.size(2)*x.size(3))
        return self.fc(x)

モデルのインスタンス化と訓練準備

このタイミングでモデルをGPUへ転送。 to(device)を利用すれば、device = "cuda" などで一括管理できます。 オプティマイザにこのタイミングで最適化すべきパラメータを渡します。 TFではアップデートをするときに渡す形式だったのですが、PyTorchではクラスインスタンスに直接パラメータがひも付きます。

net = Net().cuda()
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(net.parameters())

学習コード

running_loss などを書かなければならないのが地味に面倒。とは行っても dataloader は len() 関数でミニバッチの個数を返してくれるので、そこまで困りません。

for epoch in range(epochs):
    running_loss = 0.0
    start = time.time()
    for i, (inputs, labels) in enumerate(trainloader, 0):
        optimizer.zero_grad()

        outputs = net(inputs.cuda())
        loss = criterion(outputs, labels.long().cuda())
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
    print("time ", time.time() - start)
    print('epoch: {:d} loss: {:.3f}'.format(
        epoch + 1, running_loss / len(trainloader)))

1epoch 4.0 seconds 程。tf graph と tf eager の間ぐらい。