はじめに
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 の間ぐらい。
