はじめに

ディープラーニングのモデルを実装するとき、大体の方がPytorchかTensorFlowを使われるかと思います。
私はPytorch派で、時々画像分類とかして遊んでいます。

お好きな方を使えば良いと思うのですが、初学者の方やTensorFlowをいつも使っている人と話すときによく聞くことがあります。

「Pytorch よくわからん」

学習のロジックを具体的に実装出来るPytorchは好きなのですが、TensorFlowをほとんど使ったことがないのでどれくらい違いがあるか把握してませんでした。

ですので、PytorchとTensorFlowを使ったコーディングってどれだけ違うのか確認しました。

2021/01/31 追記

TensorFlowはkerasを使った場合を対象としたので、正確には「Tensorflow + Keras」とPytorchを比較した内容の記事になります。
Tensorflow側は簡単にかけそうな印象を持つ内容になってしまっていますが、それはKerasの恩恵を受けている点をご留意ください。
元々のtensorflowの書き方としては以下の記事がわかりやすいと思います。
https://www.tdi.co.jp/miso/tensorflow-keras-pytorch#21_Tensorflow
※@pudding167様のコメントから追記

動作環境

• Google colaboratoryで実装、動作確認をしています。
• 記事投稿時点でのPytorch, tensorflowのバージョン
  • Pytorch：1.7.0+cu101
  • TensorFlow：2.4.0
  • それぞれ以下のコードを実行して確認できます

#Pytorch
import torch
print(torch.__version__)

# TensorFlow
import tensorflow as tf
print(tf.__version__)

コーディングを比較

PytorchとTensorFlowを使ったコーディングを比較するために、同じ問題（MNISTの画像分類問題）を同じやり方（モデル、誤差関数、最適化、、）で実装しました。
また、私自身のコーディングの癖などが出ないように以下で公開されているチュートリアルを参考にしました。

• Pytorch
  PyTorchチュートリアル（日本語翻訳版） 「ニューラルネットワーク入門」
• TensorFlow
  TensorFlowチュートリアル 畳み込みニューラルネットワーク

そのままのコードだと、畳み込み層でのパラメータの値や全結合層の数が微妙に違うのでTensorFlowのチュートリアル側に合わせてPytorchのコードを修正しました。

やること

みんな大好きMNIST手書き文字画像分類を行います。
ちょっと細かいかもしれませんが、学習処理だけではなく初めから比べていきます。

パッケージのインポート

• TensorFlow

import tensorflow as tf
from tensorflow import keras
from keras import datasets
from keras import layers
from keras import models

• Pytorch

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
from   torchvision import transforms

TensorFlowでは6個、Pytorchでは7個importしました。
「.」で繋げればTensorFlowは「tensorflow」から、Pytorchは「torch」から大体の機能を参照できるので上記のimport数の差に意味はありません。

Tensorflowはkerasが組み込まれていて、kerasからのインポートが多いですね。
layersは結合層や畳み込み層の機能（DenseとかConv2Dとか）を表現したモジュール、
modelsはモデルを表現したモジュールといった感じでしょうか。
Tensorflow + kerasで使われるケースが多いっぽいです。

Pytorchは主にtorchとtorchvisionにモジュールが分かれています
主要なものはtorchをインポートすれば使えると思います。
torchvisionはComputer Vision系のモデルやデータセットが使えるパッケージです。
tensorflowと比べると、torch.optimやtorch.nn.functionalが違うところかと思います。

MNISTデータ読み込み

まずはMNISTデータセットを取得します。

• TensorFlow

# データセット取得
(train_images, train_labels), (test_images, test_labels) = datasets.mnist.load_data()
# 28X28 の行列へ変換
train_images = train_images.reshape((60000, 28, 28, 1))
test_images = test_images.reshape((10000, 28, 28, 1))
# 値が[0 - 1]の値になるように変換
train_images, test_images = train_images / 255.0, test_images / 255.0

• pytorch

# データセット取得
train_dataset = torchvision.datasets.MNIST(root="./data", train=True, download=True, transform = transforms.ToTensor())
# データローダ定義
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=16, shuffle=True)

tensorflow、pytorchどちらもdatasetsモジュールで取得できますね。

tensorflowはtrain、testデータセットをまとめて取得できます。
取得時点の形は（データ数、高さ、幅）なので、
Conv2Dに入力するデータのため（データ数、高さ、幅、チャネル数）にreshapeします。

pytorchはtrain, testは別々で取得です。
テストデータの取得はtorchvision.datasets.MNISTメソッドのtrainパラメータをFALSEに設定します。
pytorchではDataLoaderを定義するところが異なる点ですね。
「batch_size」で1バッチに含めるデータ数を指定できます。
「shuffle」でデータをシャッフルするかを指定できます。
　　trainはTrueを指定します。
　　test用のデータセットではFalseを指定して再現性があるようにします。

Model定義

いよいよMNISTの画像分類モデル定義部分です。
とても簡単に構成を説明すると、
畳み込み層×3 + プーリング層×2 + 全結合層 + 出力層、となります

• tensorflow

# model
model = models.Sequential()
model.add(layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1) ))
model.add(layers.MaxPooling2D((2, 2)))
model.add(layers.Conv2D(64, (3, 3), activation='relu'))
model.add(layers.MaxPooling2D((2, 2)))
model.add(layers.Conv2D(64, (3, 3), activation='relu'))
model.add(layers.Flatten())
model.add(layers.Dense(64, activation="relu"))
model.add(layers.Dense(10, activation="softmax"))

tensorflow + kerasでは、models.Sequential()でインスタンスを作成して、そこにレイヤを追加していく作りです。
レイヤは、Conv2Dが畳み込み層、MaxPooling2Dがプーリング層、Denseが全結合層です。
Flatten()は行列（H × W × CH）のデータを1次元データに整列します
活性化はReLU関数を設定しているようですが、Conv2Dのactivationプロパティとして組み込まれてますね。
model.add(layers.Conv2D(... ))で畳み込み1層分の入出力データのシェイプ、活性化をまとめて追加しています。
それが2層目以降もプーリング層、全結合層と追加される。さらにFlatten()も追加します。
定義は1層毎に追加して各層の実行処理はmodelに任せる感じの実装ですね。

• Pytorch

Pytorchはまずモデルのクラス定義からです。

class Net(nn.Module):
  def __init__(self):
    super(Net, self).__init__()
    self.conv1 = nn.Conv2d(1, 32, (3, 3))
    self.conv2 = nn.Conv2d(32, 64, (3, 3))
    self.conv3 = nn.Conv2d(64, 64, (3, 3))
    self.fc1 = nn.Linear(576, 64)
    self.fc2 = nn.Linear(64, 10)

  def forward(self, x):
    x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
    x = F.max_pool2d(F.relu(self.conv2(x)), (2, 2))
    x = x.view(-1, self.num_flat_features(x))
    x = F.relu(self.fc1(x))
    x = self.fc2(x)
    return x

  def num_flat_features(self, x):
    size = x.size()[1:]  # all dimensions except the batch dimension
    num_features = 1
    for s in size:
      num_features *= s
    return num_features

net = Net()

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), lr=0.001)

torch.nn.Moduleを継承したモデル定義のクラスを作成し、各層の定義や順伝播の処理を作成していきます。この辺りが大きく違いますね。
Define by Run方式のPytorchとDefine and Run方式のTensorflowの違いとも言えるでしょうか。

クラスを定義した後は、Net()のところインスタンス化、誤差関数と最適化関数を定義します。
誤差関数にCrossEntropyLoss（交差エントロピー誤差）、最適化関数にAdamを使います。
Pytorchだとこの辺が個別に提供されてて、自分で選んで使うことになります。
後述しますが、tensorflow + kerasだと、誤差関数と最適化に何使うかはcompile()メソッドで指定してmodelに組み込まれます。
コンポーネントを自分で組み合わせて使うPytorch、なんでもmodelに集約して実行処理はお任せするtensorflow + kerasというイメージですね。

では、各メソッドの定義内容を確認します。

• コンストラクタ
  • 各層のインスタンスをメンバ変数へ設定します。
    この段階では各層の定義だけなので各層がどんな順で重なっているかはまだ分かりません。

def __init__(self):
  super(Net, self).__init__()
  self.conv1 = nn.Conv2d(1, 32, (3, 3))  #畳み込み層1
  self.conv2 = nn.Conv2d(32, 64, (3, 3)) #畳み込み層2
  self.conv3 = nn.Conv2d(64, 64, (3, 3)) #畳み込み層3 
  self.fc1 = nn.Linear(576, 64) #全結合
  self.fc2 = nn.Linear(64, 10)   #出力層

• forwardメソッド
  • 順伝播の処理を定義してあります。
    
  • self.conv1, 2, 3はコンストラクタで定義した畳み込み層のインスタンスが入ったメンバ
  • Fから活性化関数が使えます。
  • ここで、各層の繋がりが定義されます

def forward(self, x):
  x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
  x = F.max_pool2d(F.relu(self.conv2(x)), (2, 2))
  x = F.relu(self.conv3(x))
  x = x.view(-1, self.num_flat_features(x))
  x = F.relu(self.fc1(x))
  x = self.fc2(x)
  return x

tensorflow + kerasではmodelに追加しておけば後はFWにお任せでしたが、Pytorchでは上記のように実装します。
この辺がめんどくさいと感じるか、流れがわかって良いと感じるかは分かれるかなと思います。
あとは、Pytorchの方がデバッグしやすいかなーと思います。自分でバグ混入させることも多いですが、、、

学習処理

• tensorflow

model.compile(optimizer="adam", loss="sparse_categorical_crossentropy", metrics=['accuracy'])
model.fit(train_images, train_labels, epochs=5)

compile、model.fitメソッドで学習処理を実行してくれて簡単ですね。
誤差計算とか、順伝播処理実行とか基本的なことはtensorflowがやってくれます。
使う側は入力データの用意やエポック数などのパラメータを設定するだけ。

• Pytorch

for epoch in range(5):
  print(f"Epoch {epoch+1}/5")
  for idx, data in enumerate(train_loader, 0):
    inputs, labels = data
    optimizer.zero_grad() # 重み初期化
    output = net(inputs)  # 順伝播
    loss = criterion(output, labels) # 誤差計算
    loss.backward() # 誤差伝播
    optimizer.step()

Pytorchは色々実装しないといけません。
これまで書いてきたように、Pytorchでは各機能がコンポーネントで提供されている状態なので、それらを繋げた一連の学習ループは実装しないといけません。
tensorflowだとfit()メソッドでやってくれてる部分です。
ただ、書籍などで学んだ学習の流れ（データ入力→順伝播→誤差計測→誤差逆伝播）をコードに反映できるので個人的にはこっちの方が好きです。
fitでよくわからんことあるけど学習できる、より、バグコード書くこともあるけど学習の流れを理解して実装出来るという差がありますね。
どちらも一長一短なところはあります。

まとめ

簡単にまとめると、細かく実装するPytorchとなんとなくでも学習が出来るTensorflowという感じでしたね。
個人的にはやはりPytorchが好みですが、状況に応じてお好きな方を選べばよいと思います。
あと、久しぶりにTensorflow + kerasの実装を読みましたが、とても楽ですね。
compile → fitの流れには惹かれるものがあり、ちょっと使ってみようかな。。なんて気も出てきました。

Pytorch、Tensorflow + kerasを使った実装を比べてみましたがいかがでしたでしょうか。
何か参考になることがあればうれしいです。