TensorFlow 2.0 のコードの書き方基本集（PyTorchとの比較）

はじめに

最近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）なのか、「複数の層を格納した１つのリスト」（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の学習ラッパーをそのまま使えます。凄く便利です。

今回は疲れたのでこの辺で。

github.com