はじめに
最近KerasからPyTorchに流れていく人たちが多く見受けられます。その中でて「Kerasで書いていたコードをPyTorchのコードで実装するにはどうすれば良いんだろう?」という声があります。要は、今まで使っていたフレームワークでやろうとしていたことを、別のフレームワークでやろうとしたときに、どんな対応になるのかが明確にわかっていればコードは圧倒的に書きやすくなるので、それを知りたいという話だと思われます。
ご要望に従って(?)、今後PyTorchからTF2.0に移りにあたり「PyTorchで書いていたコードに対応する、TF2.0でのコードはどんな感じになるんだ?」というところにお答えします。
線形回帰と学習のコード
データセット
下記のデータをPyTorchとTFで共通して使うことにします。
import numpy as np # Hyper-parameters input_size = 1 output_size = 1 num_epochs = 60 learning_rate = 0.001 # Toy dataset x_train = np.array([[3.3], [4.4], [5.5], [6.71], [6.93], [4.168], [9.779], [6.182], [7.59], [2.167], [7.042], [10.791], [5.313], [7.997], [3.1]], dtype=np.float32) y_train = np.array([[1.7], [2.76], [2.09], [3.19], [1.694], [1.573], [3.366], [2.596], [2.53], [1.221], [2.827], [3.465], [1.65], [2.904], [1.3]], dtype=np.float32)
PyTorch
import torch import torch.nn as nn import numpy as np # Linear regression model model = nn.Linear(input_size, output_size) # Loss and optimizer criterion = nn.MSELoss() optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate) # Train the model for epoch in range(num_epochs): # Convert numpy arrays to torch tensors inputs = torch.from_numpy(x_train) targets = torch.from_numpy(y_train) # Forward pass outputs = model(inputs) loss = criterion(outputs, targets) # Backward and optimize optimizer.zero_grad() loss.backward() optimizer.step()
TF2.0
import tensorflow as tf import numpy as np L = tf.keras.layers # Linear regression model model = L.Dense(output_size) # loss function def loss_fn(y_predict, y): return tf.keras.losses.mean_squared_error(y_predict, y) # optimizer optimizer = tf.keras.optimizers.SGD(learning_rate=learning_rate) for epoch in range(num_epochs): inputs = tf.convert_to_tensor(x_train) targets = tf.convert_to_tensor(y_train) with tf.GradientTape() as tape: y_predict = model(inputs) loss = loss_fn(y_predict, targets) grads = tape.gradient(loss, model.variables) # update prameters using grads optimizer.apply_gradients(zip(grads, model.variables))
違い
些細な違い:層の定義の仕方
PyTorchでは nn.Linear
のコンストラクタにインプットの次元とアウトプットの次元を両方を渡してやる必要があります。一方でTFではL.Dense
に対してアウトプットの次元のみを渡せばよく、インプットの次元は、最初にデータが入力されたときにそのデータの次元で自動的に決定されます。
この入力データによるインプットの次元を決定する恩恵は畳み込み層やLSTMなどのアウトプットを線形層にインプットするようなケースで感じるでしょう。
些細な違い:ロス関数の書き方
ここではロス関数の定義の仕方に違いがあります。PyTorchは色々なロス関数がnn.MSELoss()
のようにクラスとして定義されているので、そいつのインスタンスを拾ってくるということになります。当然インスタンスの__call__
メソッドに、具体的な計算が定義されているのですが、PyTorchは多くの場面で具体的な処理はメソッドとして隠蔽しつつ、内部でいろんな関数を呼び出すようになっています。
TF2.0ではtf.keras.losses
モジュールの関数を呼び出すという形式を取っています。今回、Python関数としてラッピングする意味は全く無いので、学習コードのところに直接 loss = tf.keras.losses.mean_squared_error(y_predict, y)
と書いてしまってもいいでしょう。
基本的には適宜自分で必要な関数を呼び出すという書き方になります。
大きな違い:勾配計算とパラメータ更新
学習コードでその違いは顕著になっており、まずPyTorchではoptimizer
に予め訓練するパラメータを渡してしまいます。そうすることで、optimizerは内部変数として訓練パラメータを保持するようになり更新時には更新メソッドであるupdate()
を呼ぶだけで更新が可能になっています。
一方でTFではoptimizer
は計算式を持っているだけで、どの変数をどの勾配を使って更新するのかは、後で更新メソッドapply_gradients
に対して引数として(訓練パラメータのリストとその勾配のリストをzip
して)渡します。
勾配の計算にも違いが出ており、PyTorchは基本的に常に計算グラフを保持しながら計算を進めるので、loss
は自分がどのように計算されたかを知ることができ、loss.backward()
のようにbackward()
メソッドを使うことで、loss
を計算するにあたって用いられてtorch.tensor
による勾配を
各々のtorch.tensor
の内部変数に格納 します。今回の場合は線形変換に使われるパラメータのtorch.tensor
の勾配がそれぞれのパラメータの内部変数に入ります。前述の通り、optimizer
は訓練するパラメータを保持しているので、それらの勾配にもアクセスができ、メソッドに引数を渡すことなく楽々と勾配計算から更新までができてしまうのです。
TFではwith tf.GradientTape() as tape
コンテキストで囲んだ部分で計算グラフの履歴を残す計算をすることになります。このとき履歴は tape
が保持するので、勾配計算などをする場合は tape
に処理をしてもらうことになります。このコンテキスト内には色々な計算式が含まれうるので、tape.gradient()
メソッドに微分される変数(今回はloss
)と、微分する変数(model.variables
)を明示的に渡してやります。逆にいろいろな計算が含まれていても、このように必要最低限の計算しか実施されません。
勾配を取り出したら前述の通り、optimizer
の更新メソッド apply_gradients
に訓練パラメータと勾配の zip
を渡してあげます。基本的に何を渡してどんな処理を行うのかをコードで明示して書いていくスタイルになります。
ニューラルネットワークの簡単な書き方
層をひたすら渡すだけで、簡単に多層のニューラルネットワークを書くことができるよう、それぞれSequential
というクラスが用意されています(おそらくkerasが最初に提供?)。
PyTorch
model = nn.Sequential(
nn.Linear(in_size, hidden_size),
nn.ReLU(),
nn.Linear(hidden_size, hidden_size),
nn.ReLU(),
nn.Linear(1),
)
nn.Sequential
のコンストラクタに層を順番に渡してあげます。
TF2.0
model = tf.keras.Sequential([ L.Dense(hidden_size, activation="relu"), L.Dense(hidden_size, activation="relu"), L.Dense(1), ])
tf.keras.Sequential
のコンストラクタに
層を順番に入れたリストを渡してあげます。
違い
見ての通り、似た書き方になりますが、渡すものが「複数の層」(PyTorch)なのか、「複数の層を格納した1つのリスト」(TF)なのかという違いがあり、地味ですがハマったりする人がたまにいます(これはKerasからPyTorchに行った場合にハマる可能性が在る部分かもしれません)。
また、TFでは層そのものが活性化関数を保持できるので行数を削減できます。ただし、バッチ正規化などを間に挟んだりドロップアウトを挟んだり、そういうケースにも対応できるように、層には活性化関数を持たせずに、活性化関数だけの層を挟む方法も用意されています。
model = tf.keras.Sequential([
L.Dense(hidden_size),
L.ReLU(),
L.Dense(hidden_size),
L.ReLU(),
L.Dense(1),
])
これはKerasがTFにマージされ、PyTorchを意識し始めた段階で準備されたように思います。
畳み込みニューラルネットワーク
PyTorch
class ConvNet(nn.Module): def __init__(self, num_classes=10): super(ConvNet, self).__init__() self.layer1 = nn.Sequential( nn.Conv2d(1, 16, kernel_size=5, stride=1, padding=2), nn.BatchNorm2d(16), nn.ReLU(), nn.MaxPool2d(kernel_size=2, stride=2)) self.layer2 = nn.Sequential( nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2), nn.BatchNorm2d(32), nn.ReLU(), nn.MaxPool2d(kernel_size=2, stride=2)) self.fc = nn.Linear(7*7*32, num_classes) def forward(self, x): out = self.layer1(x) out = self.layer2(out) out = out.reshape(out.size(0), -1) out = self.fc(out) return out
TF2.0
class Cifar10Model(tf.keras.Model): def __init__(self): super(Cifar10Model, self).__init__(name='cifar_cnn') self.conv_block1 = tf.keras.Sequential([ L.Conv2D( 8, 5, padding='same', activation=tf.nn.relu, kernel_regularizer=tf.keras.regularizers.l2(l=0.001), ), L.MaxPooling2D( (3, 3), (2, 2), padding='same' ), ]) self.conv_block2 = tf.keras.Sequential([ L.Conv2D( 16, 5, padding='same', activation=tf.nn.relu, kernel_regularizer=tf.keras.regularizers.l2(l=0.001), ), L.MaxPooling2D( (3, 3), (2, 2), padding='same', ), ]) self.conv_block3 = tf.keras.Sequential([ L.Conv2D( 32, 5, padding='same', activation=tf.nn.relu, kernel_regularizer=tf.keras.regularizers.l2(l=0.001), ), L.MaxPooling2D( (3, 3), (2, 2), padding='same', ), ]) self.flatten = L.Flatten() self.fc1 = L.Dense( 256, activation=tf.nn.relu, kernel_regularizer=tf.keras.regularizers.l2(l=0.001)) self.dropout = L.Dropout(0.8) self.fc2 = L.Dense(10) self.softmax = L.Softmax() def call(self, x, training=False): x = self.conv_block1(x, training=training) x = self.conv_block2(x, training=training) x = self.conv_block3(x, training=training) x = self.flatten(x) x = self.dropout(self.fc1(x), training=training) x = self.fc2(x) return self.softmax(x)
違い
パディング
まず、畳み込み層の大きな違いはpadding
の設定になると思われます。PyTorchではパディングを数値で渡しますが、TFの場合は'same'
か'valid'
で渡します。same
は縦横が変わらないようにパディングをするという意味で、valid
はパディングしないという意味になります。もしも縦横の大きさを変えたくない場合は same
は非常に便利です。PyTorchの場合畳み込みのカーネルに応じて自分で計算することになります。
一方で、意図した大きさに縦横を明示的に変えていきたい場合は、TFの畳み込み層では実施することすらできません。代わりに畳み込み層の方ではvalid
でパディングをしない設定にしつつ、L.ZeroPadding2D
という層を使うことで明示的にパディングを実施してやることができます。
畳み込み層→線形層
畳み込み層から線形層に変わる部分ではテンソルの形を調整する必要があります。
PyTorchでは自分で形を変えてあげる必要がありますが、TFには線形層に入れるためのL.Flatten()
があるため、そのようなコードは省くことができます。ただし、結局 tensor.reshape(batchsize, -1)
としてやればよく、batchsize=tensor.shape[0]
などで簡単にその場で取得できるので大したコードの節約にはなりません。
このL.Flatten()
が活躍するのは、実はtf.keras.Sequential
を利用するようなケースです。
traininig
フラグ
ここでついでに解説しますが、PyTorchのtorch.nn.Module
にあるmodel.train()
やmodel.eval()
のようなメソッドはtf.keras.Model
にはありません。call
メソッドで訓練時とテスト時の振る舞いを明示的に分ける必要があります。ちなみにtf.keras.Sequential
クラスはcall
メソッドを自分で書かないのですが、ちゃんと最初からtraining
引数を取れるようになっており、L.Dropout
などに自動で適用されるようになります。
RNN
今回はBidirectional RNNをLSTMを使って見ていきます。
PyTorch
class BiRNN(nn.Module): def __init__(self, input_size, hidden_size, num_layers, num_classes): super(BiRNN, self).__init__() self.hidden_size = hidden_size self.num_layers = num_layers self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True, bidirectional=True) self.fc = nn.Linear(hidden_size*2, num_classes) # 2 for bidirection def forward(self, x): # Set initial states h0 = torch.zeros(self.num_layers*2, x.size(0), self.hidden_size).to(device) # 2 for bidirection c0 = torch.zeros(self.num_layers*2, x.size(0), self.hidden_size).to(device) # Forward propagate LSTM out, _ = self.lstm(x, (h0, c0)) # out: tensor of shape (batch_size, seq_length, hidden_size*2) # Decode the hidden state of the last time step out = self.fc(out[:, -1, :]) return out
TF2.0
class BiRNN(tf.keras.Model): def __init__(self, hidden_size=10, num_layers=2, num_classes=10): super(BiRNN, self).__init__(name='mnist_rnn') self.hidden_size = hidden_size self.num_layers = num_layers self.lstm = self.get_lstm_layers(hidden_size, num_layers) self.fc = L.Dense(num_classes, activation="softmax") @staticmethod def get_lstm_layers(hidden_size, num_layers): lstm_layers = [] # we need all sequence data. write return_sequences=True! for i in range(num_layers-1): lstm_layers.append( L.Bidirectional( L.LSTM(units=hidden_size, return_sequences=True) ) ) # the final layer return only final sequence # if you need all sequences, you have to write return_sequences=True. lstm_layers.append( L.Bidirectional( L.LSTM(units=hidden_size) ) ) return tf.keras.Sequential(lstm_layers) def call(self, x): # Forward propagate LSTM out = self.lstm(x) out = self.fc(out) return out
違い
大きな違い:多層化
まず、PyTorchのnn.LSTM
は最初から多層にすることを想定しており、num_layers
なる引数を持っています。TFで多層化したい場合は、Sequential
にLSTM
がたくさん入ったリストを渡してあげる必要が在るため、この辺で面倒なコードを書く必要が出ます。
些細な違い:Bidirectional
PyTorchのLSTMではBidirectionalにしたければコンストラクタでその指定が可能です。TFではL.Bidirectoinal
でラッピングすることになります。
大きな違い:戻り値の並び
PyTorchのLSTMは (seq_len, batch, features)
という形でテンソルが返ってきます(それプラス隠れ状態とセル状態も帰ってきます)。一方でTFでは(batch, seq_len, features)
という形でテンソルが返ってきます。ただし、PyTorchは(batch, seq_len, features)
で返ってくるように指定することも可能です。上記のコードではPyTorch側をTF側に揃えています。
もしも最後のシーケンスのアウトプットのみが欲しい場合は、PyTorchの場合はスライシングで対応することになりますが、TFの場合LSTMのコンストラクタに全てのシーケンスを返させるか、最後のシーケンスのみを変えさせるかを指定できます。スライシングもできるので、特に迷ったのであれば全て返させても良いかもしれません。ただ、一旦全てのシーケンスを何らかの変数に保存して、再度取り出すのは無駄と言えば無駄かもしれません(そう影響は無いだろうが)。
学習
学習に関しては既に線形回帰の方で比較をしたのですが、実はTFの方はモデルを作った後 model.compile
してmodel.fit
をすればKerasの学習ラッパーをそのまま使えます。凄く便利です。
今回は疲れたのでこの辺で。