MLエンジニアへの道 #16 - PyTorch vs TensorFlow

MLDL

Last Edited: 8/30/2024

このブログ記事では、ディープラーニングのライブラリであるPyTorchとTensorFlowを紹介します。

ML

初心者が深層学習を始める際に直面する大きな決断の一つは、PyTorchとTensorFlowという深層学習コミュニティで最も 人気のあるフレームワークのどちらを選ぶかです。ここでは、両方のフレームワークを使用して構築された例の機械学習 パイプラインを通じて、それらの違いを見ていきます。

ステップ 1 & 2. データ探索と前処理

ここで使用するデータセットは、これまでにも見たことがあるMNISTデータセットです。これは、60,000枚の手書き数字の 画像サンプルを含んでいます。MNISTデータセットのいくつかの画像サンプルを見てみましょう。

import keras
import matplotlib.pyplot as plt
 
# Download MNIST
(X_train, y_train), (X_test, y_test) = keras.datasets.mnist.load_data()
 
# Display 10 samples
plt.figure(figsize=(10, 4))
for i in range(10):
    plt.subplot(2, 5, i + 1)
    plt.imshow(X_train[i], cmap='gray')
    plt.axis('off')
plt.tight_layout()
plt.show()
MNIST samples

上記の画像はかなりぼやけているかもしれませんが、それでも28×28ピクセルのサイズで、合計784ピクセルを含んでいます。 前処理前のピクセル値は0から255の間であり、クラスラベルはワンホットエンコードされていません。したがって、 次のようにデータを正規化し、ワンホットエンコードを適用します。

X_train = X_train.reshape(X_train.shape[0], X_train.shape[1]*X_train.shape[2])
X_test = X_test.reshape(X_test.shape[0], X_test.shape[1]*X_test.shape[2])
 
def zscore(X, axis = None):
    X_mean = X.mean(axis=axis, keepdims=True)
    X_std  = np.std(X, axis=axis, keepdims=True)
    zscore = (X-X_mean)/X_std
    return zscore
 
X_train = zscore(X_train)
X_test = zscore(X_test)
 
y_train = keras.utils.to_categorical(y_train)
y_test = keras.utils.to_categorical(y_test)

過学習や未学習の可能性をチェックするために、通常は検証データセットを設定し、そこでのコストとメトリクスを追跡します。 次のようにして、訓練データセットの一部を使用して検証データセットを作成できます。

from sklearn.model_selection import train_test_split
 
X_train, X_val, y_train, y_val = train_test_split(X_train, y_train, test_size=10000, random_state=101)

TensorFlowの場合、モデルがNumPy配列と互換性があるため、この時点でデータ前処理はすでに完了しています。しかし、PyTorchの場合、 NumPy配列をtensor(テンソル)に変換する必要があります。

import torch.nn as nn
 
X_train, X_val, X_test = map(lambda X: torch.tensor(X, dtype=torch.float32), (X_train, X_val, X_test))
y_train, y_val, y_test = map(lambda y: torch.tensor(y, dtype=torch.int64), (y_train, y_val, y_test))

実際のところ、上記のコードはPyTorchにとってはすでに十分です。しかし、データに関するバッチサイズや他の ハイパーパラメータを設定するために、テンソルからDatasetとDataLoaderを構築することが強く推奨されます。

train_dataset = torch.utils.data.TensorDataset(X_train, y_train)
val_dataset = torch.utils.data.TensorDataset(X_val, y_val)
test_dataset = torch.utils.data.TensorDataset(X_test, y_test)
 
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=32, shuffle=True)
val_loader = torch.utils.data.DataLoader(dataset=val_dataset, batch_size=32, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=1, shuffle=True)

ここでは、PyTorchとTensorFlowの両方のモデルにバッチサイズ32を使用します。他のハイパーパラメータも両方のフレームワークで同じに設定されています。

ステップ 3. モデル

TensorFlowとPyTorchを使って分類モデルを作成し、それらの違いを確認しましょう。 以下をクリックすると、それぞれの実装を見ることができます。

ここでは、特に言及することもないのでトレーニングの結果やステップ4（モデル評価）については省略しますが、 練習としてご自身で試してみることを強くお勧めします。

結論

一般的に、TensorFlowは、定義済みの関数、オブジェクト、およびTensorBoardのようなトレーニングの監視機能や、 TensorFlow LiteやTensorFlow.jsなどの機能を使って複数のプラットフォームに簡単にデプロイできることから、 簡単なプロダクション向けと言われています。

一方で、PyTorchは、新しいカスタムモデルやレイヤーなど、様々なカスタムコンポーネントを構築できる高いカスタマイズ性のため、 研究向けとされています。しかし、どちらのフレームワークも機能を改善し続けており、どちらもコードを書くのが簡単で、 カスタマイズも容易です（少なくとも、私が両方を使った経験ではそう思います）。したがって、どちらを使用するかは 完全にあなた次第です。個人的には、両者は十分に似ているし、結局のところどちらも使う可能性が高いので、 どちらも交互に使いながら学ぶことをお勧めします。

TOPに戻る↑