はじめに
下記記事に影響を受けてPyTorchとTensorFlowの速度比較をしました。
結論から言えば、PyTorchはPythonicに書いても速く、現状TensorFlow Eagerで書いたコードをgraphへ変換した場合と同等以上かなという印象です（上記の記事ではEagerをGraphに変換したコードのほうが速い）。
PyTorchもgraphモードに変換するtorch.jitが1.0から搭載され、更に高速化が可能になっていますので、これを使いこなすと更にPyTorchは速くなるでしょう。ただし、PyTorchでtorch.jitを利用するにはコードの書き換えが必要っぽいです（専用のクラスがあったり、変換される処理内でのデータ型は静的で無ければいけなかったり、まあまあ制約があります）。一方でTensorFlowはGraphへの変換が極めて手軽で、ひとまず@tf.contrib.eager.defunでデコレートしてやれば変換できてしまいます。ここは元々graphモードで何でも記述できる構えを持っていたTFの強みなのかもしれません。
TensorFlow
設定
計算させるのは下記の数式で表されるモデルの順伝播と逆伝播です。ボールド体にして無くてもベクトルあるいは行列です。
$$ f( x) = W_3 W_2 W_1 x $$
ここで、$x \in \mathbb R^{1000}$ で $W_1, W_2 \in \mathbb R^{1000 \times 1000}$ とし、出力層だけ $W_3 \in \mathbb R^{1\times 1000}$です。損失関数はmean squared errorとして逆伝播を計算します。
問題設定の共通コード
まずは必要なモジュールをインポートします。
import tensorflow as tf import numpy as np import time tf.enable_eager_execution()
設定したモデルはTensorFlowでは下記のように書けます。
model = tf.keras.Sequential([
tf.keras.layers.Dense(1000),
tf.keras.layers.Dense(1000),
tf.keras.layers.Dense(1),
])
続いてデータはバッチサイズを1024として与えておきます。 損失関数と最適化手法も指定しておきましょう。
# batch_size is 1024. x = tf.random_normal([1024, 1000]) y = tf.random_normal([1024, 1]) def loss(y, y_pre): return tf.losses.mean_squared_error(y, y_pre) optimizer = tf.train.GradientDescentOptimizer(1e-4)
Eager Execution
下記の関数を実行し終えるのに掛かる時間を計測します。
def measurement_forloop(gpu=False): if gpu: device = "/gpu:0" else: device = "/cpu:0" for _ in range(10): with tf.device(device): with tf.GradientTape() as tape: y_pre = model(x) loss_value = loss(y, y_pre) grads = tape.gradient(loss_value, model.variables) optimizer.apply_gradients(zip(grads, model.variables))
cpuの実行結果は下記
%%timeit measurement_forloop(False) # > 1 loop, best of 3: 1.47 s per loop
gpuの実行結果は下記
%%timeit measurement_forloop(True) # > 1 loop, best of 3: 260 ms per loop
graph mode
評価する関数は下記。ただ単にデコレートしてやるだけ。
@tf.contrib.eager.defun def graph_measurement_forloop(gpu=False): if gpu: device = "/gpu:0" else: device = "/cpu:0" with tf.device(device): for _ in range(10): with tf.GradientTape() as tape: y_pre = model(x) loss_value = loss(y, y_pre) grads = tape.gradient(loss_value, model.variables) optimizer.apply_gradients(zip(grads, model.variables))
cpuの結果
%%timeit graph_measurement_forloop(False) # > 1 loop, best of 3: 1.62 s per loop
gpuの結果
%%timeit graph_measurement_forloop(True) # > 1 loop, best of 3: 110 ms per loop
PyTorch
問題設定共通のコード
モジュールをインポートします。
import torch import torch.nn as nn
次にモデルと損失関数、最適化手法、データを取り揃えてしまいましょう。
model = nn.Sequential(
nn.Linear(1000, 1000),
nn.Linear(1000, 1000),
nn.Linear(1000, 1),
)
x = torch.randn(1024, 1000)
y = torch.randn(1024, 1)
loss = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-4)
Eager
評価する関数は下記
def measurement_forloop(gpu=False): if gpu: device = "cuda" else: device = "cpu" model.to(device) for _ in range(10): optimizer.zero_grad() y_pre = model(x.to(device)) loss_value = loss(y_pre, y.to(device)) loss_value.backward() optimizer.step()
cpuは以下の結果
%%timeit measurement_forloop(False) # > 1 loop, best of 3: 1.53 s per loop
gpuは以下の結果（速い！）
%%timeit measurement_forloop(True) # > 10 loops, best of 3: 72.2 ms per loop
graph mode
正直torch.jitがよく分からず、どう考えてもイレギュラーな書き方になりました（素直に@torch.jit.scriptでデコレートするだけでは動かなかったので、仕方なく、最低限動くであろう形式で試しました）。まず、モデルのクラスを作成しますがこのときにはtorch.jit.ScriptModuleを継承します。また、メソッドを作るときに@torch.jit.script_methodでデコレートするようです。ちなみにコンストラクタが動く段階でパラメータたちをgpuに転送しています。これ、インスタンスを作ってから転送が何故かできなかったのでこうしていますが、どなたがご存知であれば教えてください（インスタンスを転送できないって、めちゃくちゃ不便な気がしますが…）。
class Model(torch.jit.ScriptModule): def __init__(self): super(Model, self).__init__() self.model = nn.Sequential( nn.Linear(1000, 1000), nn.Linear(1000, 1000), nn.Linear(1000, 1), ).to('cuda') @torch.jit.script_method def forward(self, x): return self.model(x) model = Model()
続いて評価する関数は下記のようになっています。予めデータをgpuに送っているので、これまで違いgpuへの転送時間が評価されなくなってしまうのですが、動くコードを突貫で書いたので申し訳ありません。事実上コンパイルされているのはモデルの順伝播だけということになるのでしょうか…？
x = x.to('cuda') y = y.to('cuda') def measurement_forloop_script(): for _ in range(10): optimizer.zero_grad() y_pre = model(x) loss_value = loss(y_pre, y) loss_value.backward() optimizer.step()
結果としては下記
%%timeit
measurement_forloop_script()
# > 10 loops, best of 3: 68.8 ms per loop
ちょっとだけ速くなっているっぽい。PyTorchはよく調べず書いているのでフェアではありませんが、コンパイルしなくてもそれなりに速いのは驚きでした。
