はじめに
普段，機械学習を研究で利用しているのですが，Tensorflowをメインに利用していました．しかし，研究の内容が高度化するに従ってTensorflowでは無理が出てくる場面が増えてきました．（私のコーディング能力の問題かもしれませんが．）
また，勉強会やコミュニティに参加していても，多くの場面でPyTorchを利用する場面が多いように感じました．
こうした現状からPyTorchを真面目に勉強する決断をしたのですが，それに当たって意外に大変だったのは，初学者向けの教材は多く見られるのですが，機械学習をある程度やっている人がPyTorchを利用する場合の解説記事があまり見られなかったことでした．初学者向けの教材は丁寧な説明が非常に魅力的ですが，その分詳細な説明はぼかしたり，冗長になりがちです．
本記事は，CNNを題材にTensorflowと比較しながらPyTorchへ移行する際の足掛かりとなるような記事を目指します．TensorflowとPyTorchの違いを本記事で確認した後，詳細を最後に挙げる参考文献で確認して頂くような利用を想定しています．
本記事は以下に当てはまる方に向けた記事となります．
- 機械学習（深層学習）の実装経験がある
- Tensorflowで言うところの，シーケンシャルAPI・関数APIは使いこなすことができる程度
- Tensorflowで言うところの，データAPIを使いこなすことができる程度
- 機械学習（深層学習）に必要な要素が理解できている
- 一般的なモデルアーキテクチャについて理解できている
- 損失関数，活性化関数，自動微分など，深層学習を支える要素について理解できている
また本記事は，PyTorchを勉強し始めて1週間目の初心者が勉強のアウトプットのためにまとめている記事になります．勉強を続けながら必要に応じて追加する予定ではありますが，拙い部分がございましたらご教示いただけますと幸いです．
本記事で利用している環境を以下に示します．
| 項目 | バージョン |
|---|---|
| Tensorflow | 2.14.0 |
| PyTorch | 2.2.0 |
TL; DR
- データのロードと加工，モデルの構築については大きな違いはない
- モデルの学習・評価の部分が大きく異なる（ココを重点的に確認すれば良い）
- PyTorchは学習過程のカスタム性が非常に高い印象
- 一方で，Tensorflowは簡潔に記述できる点が魅力的
- PyTorchは難しいという印象が先行している印象を受けた
- Tensorflowが使えるならあまり苦戦しないと思う
- GPUや分散処理が入ってくるとより難しくなるかもしれません（Tensorflowも同じですが．）
TensorflowとPyTorchの違いを確認できればそれで大丈夫という方は，第3章の学習と評価の章を確認いただければ早いと思います．
0. Tensor
PyTorchで前提になるtorch.Tensorについて確認します．既に知っているという方は第1章にスキップされてください．なお，毎回torch.Tensorとするのは冗長なので，以降はテンソルと表記します．
テンソルはnp.arrayとよく似ている一方で，以下の点で異なります．
- CPUとGPUの両方で動作可能
- 自動微分
また，テンソルはPythonの標準のデータ型やnp.array型から作成可能です．
この辺りの詳細については以下が参考になると思います．
1. データのロードと加工
今回はPyTorchチュートリアルに合わせてFashionMNISTを利用します．
まずは，Tensorflowから．
import tensorflow as tf
from tensorflow.keras.datasets import fashion_mnist
# 今回はFashionMNISTを使用
(X_train, y_train), (X_test, y_test) = fashion_mnist.load_data()
# データセットをnp.arrayからtf.data.Datasetに変換
train_dataset = tf.data.Dataset.from_tensor_slices((X_train, y_train))
test_dataset = tf.data.Dataset.from_tensor_slices((X_test, y_test))
# 今回は10種類の画像分類
NUM_TARGET = 10
# tf.dataのpipeline処理
# 画像 -> 正規化，ラベル -> one-hot encode
# データのシャッフル
# batchサイズを64に設定
def preprocess(dataset):
dataset = (
dataset
.map(lambda img, target: (img / 255, [0. if i != target else 1. for i in range(NUM_TARGET)]))
.shuffle(buffer_size=1000, seed=42)
.batch(64)
.prefetch(1)
)
return dataset
# 前処理を定義
train_dataset = preprocess(train_dataset)
test_dataset = preprocess(test_dataset)
これと（ほぼ）同じ動作をするコードをPyTorchで実装すると以下のようになります．
import torch
from torchvision.datasets import FashionMNIST
from torchvision.transforms import ToTensor, Lambda
from torch.utils.data import DataLoader
# 今回はFashionMNISTを使用
data_root_dir = './datasets'
train_data = FashionMNIST(
root=data_root_dir,
train=True,
transform=ToTensor(),
target_transform=Lambda(lambda y: torch.zeros(10, dtype=torch.float).scatter(0, torch.tensor(y), value=1))
)
test_data = FashionMNIST(
root=data_root_dir,
train=False,
transform=ToTensor()
)
# データローダの定義
train_dataloader = DataLoader(
dataset=train_data,
batch_size=64,
shuffle=True
)
test_dataloader = DataLoader(
dataset=test_data,
batch_size=64,
shuffle=True
)
ライブラリのimport
ここから各部分の詳細を見ていきます．まずは，importしたライブラリから．
import torch
from torchvision.datasets import FashionMNIST
from torchvision.transforms import ToTensor, Lambda
from torch.utils.data import DataLoader
PyTorchをimportするときはtorchと記述します．ややこしいですが，どうやらこれまでの開発経緯が関係しているようです．（本題でないのでこれ以上は触れませんが，気になる方は調べてみてください．）
また，torchvisionとありますが，このライブラリは画像認識によく利用される処理をまとめたライブラリです．PyTorchのinstall時に自動でinstallされないので，自前の環境で実行される場合は別途installが必要になります．
2024/02時点でのtorchvisionの最新バージョンは0.17.0なので，おそらく$\beta$版だと思われますが，いくつかのチュートリアルで利用されていることから，安定性やサポートのレベルにそれなりの信頼を置いても良いかもしれないと判断し，今回の実装では利用しています．
問題がある場合にはご教示いただけますと幸いです．
データセットの読み込み
data_root_dir = './datasets'
train_data = FashionMNIST(
root=data_root_dir,
train=True,
transform=ToTensor(),
target_transform=Lambda(lambda y: torch.zeros(10, dtype=torch.float).scatter(0, torch.tensor(y), value=1))
)
test_data = FashionMNIST(
root=data_root_dir,
train=False,
transform=ToTensor()
)
今回は簡単のために提供されているデータセットを用いましたが，テンソルはPythonの標準のデータ型とnp.array型から変換することが可能とのことなので，自前のデータセットでも容易に実装できます．
また，自分でカスタムしたDatasetクラスを作成することも可能です．
次に，ToTensor関数は$[0.0, 1.0]$の範囲に正規化したデータをテンソルとして渡してくれる関数です．これをFanshionMNISTクラスのtransform引数に渡すことで，画像の加工を完了した状態でデータを受け取ることができます．
最後に，torch.zeros(10, dtype=torch.float).scatter(0, torch.tensor(y), value=1)についてですが，メソッドチェーンになっており，順に，10個の0を持つテンソルを作成し，ラベルと同じインデックスの要素を1に変換しています．要するに，one-hotエンコーディングです．
データローダ
train_dataloader = DataLoader(
dataset=train_data,
batch_size=64,
shuffle=True
)
test_dataloader = DataLoader(
dataset=test_data,
batch_size=64,
shuffle=True
)
こちらはTensorflowでもお馴染みのPipelineです．コードも非常に簡潔で読みやすいので細かくは触れませんが，詳細を知りたいという方は以下を参考にされてください．
ここまでの個人的な感想ですが，TensorflowでデータAPIを使いこなすことができている人にとっては，それほど大きな違和感なく移行することができるのでは無いかと感じました．その一方で，PyTorchが初学者に優しくないと言われる部分も垣間見えるように感じました．というのも，TensorflowではデータAPIを使わなくても（効率は悪いですが）例えばlist型のデータをそのままモデルに入れることが可能ですが，PyTorchではテンソルを利用することが前提となっており，このデータ加工の部分を飛ばすことができないからです．
2. モデル構築
まずは，Tensorflowのモデル構築から．なお，モデル構造は説明のために，明示的に変数を指定しています．また，意味のある構造というわけでなく（説明のために）引数が可能な限り異なるようにしているので，違和感があるかもしれませんが，その点はご了承下さい．
from tensorflow.keras import layers
# モデル構造の定義
# 今回のモデルはシーケンスAPIでも実装できますが，個人的な好みで関数APIによる実装にしています
def build_model(
input_shape: tuple[int, int]=(28, 28, 1), # 28*28 ピクセルのモノクロ画像
output_size: int=NUM_TARGET,
lr: float=1e-3
):
input_layer = tf.keras.Input(shape=input_shape)
# 1つ目の畳み込み層
# Input: (28 * 28, 1 channel)
# Output: (26 * 26, 6 channel)
# kernelサイズ = 3, paddingなし
conv1 = layers.Conv2D(
filters=6,
kernel_size=3,
strides=(1, 1),
padding='valid',
activation=None
)
# 1つ目のプーリング層
# Input: (26 * 26, 6 channel)
# Output: (13 * 13, 6 channel)
# poolingサイズ = 2, paddingなし, MaxPooling
pool1 = layers.MaxPool2D(
pool_size=(2, 2),
padding='valid'
)
# 2つ目の畳み込み層
# Input: (13 * 13, 6 channel)
# Output: (7 * 7, 16 channel)
# kernelサイズ = 3, paddingあり, stride = 2, 活性化関数 ReLU
conv2 = layers.Conv2D(
filters=16,
kernel_size=3,
strides=(2, 2),
padding='same',
activation=tf.keras.activations.relu
)
# 2つ目のプーリング層
# Input: (7 * 7, 16 channel)
# Output: (4 * 4, 16 channel)
# poolingサイズ = 2, paddingあり, MaxPooling
pool2 = layers.MaxPool2D(
pool_size=(2, 2),
padding='same'
)
# データを1列に並べる
flatten = layers.Flatten()
# 全結合層
# Input: (4 * 4 * 16)
# Output: (1024)
# 活性化関数 ReLU
dense = layers.Dense(
units=1024,
activation=tf.keras.activations.relu
)
# 出力層（全結合）
# Input: (1024)
# Output: (10)
# 活性化関数 Softmax
output_layer = layers.Dense(
units=output_size,
activation=tf.keras.activations.softmax
)
h = conv1(input_layer)
h = pool1(h)
h = conv2(h)
h = pool2(h)
h = flatten(h)
h = dense(h)
y = output_layer(h)
model = tf.keras.Model(inputs=input_layer, outputs=y)
# 最適化手法: Adam
# Loss: Categorical Cross Entropy
# Metrics: Weighted F1-Score
model.compile(
optimizer=tf.keras.optimizers.Adam(learning_rate=lr),
loss=tf.keras.losses.CategoricalCrossentropy(),
metrics=tf.keras.metrics.F1Score(average='weighted')
)
return model
cnn = build_model()
cnn.summary()
>> 以下のような出力が得られます
Model: "model"
_________________________________________________________________
Layer (type) Output Shape Param #
=================================================================
input_1 (InputLayer) [(None, 28, 28, 1)] 0
conv2d (Conv2D) (None, 26, 26, 6) 60
max_pooling2d (MaxPooling2 (None, 13, 13, 6) 0
D)
conv2d_1 (Conv2D) (None, 7, 7, 16) 880
max_pooling2d_1 (MaxPoolin (None, 4, 4, 16) 0
g2D)
flatten (Flatten) (None, 256) 0
dense (Dense) (None, 1024) 263168
dense_1 (Dense) (None, 10) 10250
=================================================================
Total params: 274358 (1.05 MB)
Trainable params: 274358 (1.05 MB)
Non-trainable params: 0 (0.00 Byte)
_________________________________________________________________
これと（ほぼ）同じ動作をするコードをPyTorchで実装すると以下のようになります．（compileの部分については切り分けることができなかったので載せていますが，第3章で触れることにします．）
from torch import nn
import torch.nn.functional as F
NUM_TARGET = 10
# モデル構造の定義
class ConvolutionalNeuralNetwork(nn.Module):
# layerオブジェクトを作成
def __init__(self):
# nn.Moduleクラスの__init__メソッドの呼び出し
super().__init__()
# 1つ目の畳み込み層
# Input: (28 * 28, 1 channel)
# Output: (26 * 26, 6 channel)
# kernelサイズ = 3, paddingなし
self.conv1 = nn.Conv2d(
in_channels=1,
out_channels=6,
kernel_size=3,
stride=1,
padding=0
)
# 1つ目のプーリング層
# Input: (26 * 26, 6 channel)
# Output: (13 * 13, 6 channel)
# poolingサイズ = 2, paddingなし, MaxPooling
self.pool1 = nn.MaxPool2d(
kernel_size=2,
padding=0
)
# 2つ目の畳み込み層
# Input: (13 * 13, 6 channel)
# Output: (7 * 7, 16 channel)
# kernelサイズ = 3, paddingあり, stride = 2
self.conv2 = nn.Conv2d(
in_channels=6,
out_channels=16,
kernel_size=3,
stride=2,
padding=1
)
# 2つ目の畳み込み層
# Input: (13 * 13, 6 channel)
# Output: (7 * 7, 16 channel)
# kernelサイズ = 3, paddingあり, stride = 2
self.pool2 = nn.MaxPool2d(
kernel_size=2,
padding=1
)
# データを1列に並べる
self.flatten = nn.Flatten(
start_dim=1
)
# 全結合層
# Input: (4 * 4 * 16)
# Output: (1024)
self.fc = nn.Linear(
in_features=(4 ** 2) * 16,
out_features=1024
)
# 出力層（全結合）
# Input: (1024)
# Output: (10)
self.fc_out = nn.Linear(
in_features=1024,
out_features=NUM_TARGET
)
def forward(self, X):
h = self.conv1(X)
h = self.pool1(h)
# 活性化関数 ReLU の適用
h = F.relu(self.conv2(h), inplace=False)
h = self.pool2(h)
h = self.flatten(h)
# 活性化関数 ReLU の適用
h = F.relu(self.fc(h))
y = self.fc_out(h)
return F.softmax(y, dim=1) # 出力に Softmax 関数を適用
cnn = ConvolutionalNeuralNetwork()
print(cnn)
>> 以下のような出力が得られます
ConvolutionalNeuralNetwork(
(conv1): Conv2d(1, 6, kernel_size=(3, 3), stride=(1, 1))
(pool1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(conv2): Conv2d(6, 16, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
(pool2): MaxPool2d(kernel_size=2, stride=2, padding=1, dilation=1, ceil_mode=False)
(flatten): Flatten(start_dim=1, end_dim=-1)
(fc): Linear(in_features=256, out_features=1024, bias=True)
(fc_out): Linear(in_features=1024, out_features=10, bias=True)
)
以下で詳細を見ていきます．
モデルクラスの定義
PyTorchでは__init__メソッドとforwardメソッドを定義することでモデルを定義することができます．
class MyNN(nn.Module):
def __init__(self):
super().__init__()
...
def forward(self, X):
...
return y
TensorflowのクラスAPIと（ほぼ）同じなので，クラスAPIでコーディングできる方はスムーズに移行することができるのではないでしょうか？
とはいっても，Tensorflowではそれなりに複雑なモデルでも関数APIを使えば実装できるため，クラスAPIを普段使いにしている人は少ないのではないかと思います．そういったことを踏まえると，ここは躓くポイントの1つかもしれません．
layerオブジェクト
self.conv1 = nn.Conv2d(
in_channels=1,
out_channels=6,
kernel_size=3,
stride=1,
padding=0
)
Tensorflowと異なる点は何と言っても「入力形式を明示的に記述しなければならない」という点だと思います．Tensorflowは出力に注目していればライブラリが処理をしてくれていましたが，PyTorchではそうもいきません．個人的には非常に面倒に感じる点の1つでした．
forwardメソッド
def forward(self, X):
h = self.conv1(X)
h = self.pool1(h)
h = F.relu(self.conv2(h), inplace=False)
h = self.pool2(h)
h = self.flatten(h)
h = F.relu(self.fc(h))
y = self.fc_out(h)
return F.softmax(y, dim=1)
Tensorflowと（ほぼ）同じなので記法については詳細を触れません．
その代わりに，私が勉強する過程でいくつかのモデルを練習のために作成したのですが，その際に発生したエラーについて触れておきます．そのエラーとはRuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation:というもので，inplace操作（値を上書きする操作）がおこなわれたというものでした．テンソルの特性の1つにinplace操作が実行されると自動微分が使えなくなるというものがあり，逆伝播時にエラーが発生するようです．
問題はここからで，全ての関数に明示的にinplace=Falseとしても改善せず，半日を費やすことになりました．結論を述べると，
y = y + input
# エラーが発生
# y += input
というコードが原因でした．Residual Connectionを実装したときに発生したものです．
私の理解不足が原因であることは言うまでもありませんが，これまで上記の違いについて考えたことが無かったので，非常に苦戦しました．
モデル構造の出力
cnn.summary()
Model: "model"
_________________________________________________________________
Layer (type) Output Shape Param #
=================================================================
input_1 (InputLayer) [(None, 28, 28, 1)] 0
conv2d (Conv2D) (None, 26, 26, 6) 60
max_pooling2d (MaxPooling2 (None, 13, 13, 6) 0
D)
conv2d_1 (Conv2D) (None, 7, 7, 16) 880
max_pooling2d_1 (MaxPoolin (None, 4, 4, 16) 0
g2D)
flatten (Flatten) (None, 256) 0
dense (Dense) (None, 1024) 263168
dense_1 (Dense) (None, 10) 10250
=================================================================
Total params: 274358 (1.05 MB)
Trainable params: 274358 (1.05 MB)
Non-trainable params: 0 (0.00 Byte)
_________________________________________________________________
print(cnn)
ConvolutionalNeuralNetwork(
(conv1): Conv2d(1, 6, kernel_size=(3, 3), stride=(1, 1))
(pool1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(conv2): Conv2d(6, 16, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
(pool2): MaxPool2d(kernel_size=2, stride=2, padding=1, dilation=1, ceil_mode=False)
(flatten): Flatten(start_dim=1, end_dim=-1)
(fc): Linear(in_features=256, out_features=1024, bias=True)
(fc_out): Linear(in_features=1024, out_features=10, bias=True)
)
見比べて分かるようにTensorflowの方がリッチで分かりやすいです．（個人的な感想かもしれませんが．）
PyTorchもサードパーティ製のライブラリでリッチな出力をしてくれるものがあるようですが，サポートが終了しているライブラリも散見されるため，どのライブラリを用いるかについては吟味が必要になると思われます．
最終的な評価には外部ツールを用いることが多いと思いますが，簡単に確認できることのメリットも大きいと個人的に感じているので，この点については残念に感じました．
3. 学習と評価
まずは，Tensorflowのコードから．
# build_model関数内で定義した部分
# # 最適化手法: Adam
# # Loss: Categorical Cross Entropy
# # Metrics: Weighted F1-Score
# model.compile(
# optimizer=tf.keras.optimizers.Adam(learning_rate=lr),
# loss=tf.keras.losses.CategoricalCrossentropy(),
# metrics=tf.keras.metrics.F1Score(average='weighted')
# )
# 学習
# 10 epochs
cnn.fit(
train_dataset,
epochs=10
)
>> 以下のような出力が得られます
Epoch 1/10
938/938 [==============================] - 5s 5ms/step - loss: 0.5390 - f1_score: 0.8026
Epoch 2/10
938/938 [==============================] - 4s 5ms/step - loss: 0.3754 - f1_score: 0.8622
Epoch 3/10
938/938 [==============================] - 5s 5ms/step - loss: 0.3281 - f1_score: 0.8787
Epoch 4/10
938/938 [==============================] - 4s 5ms/step - loss: 0.2957 - f1_score: 0.8889
Epoch 5/10
938/938 [==============================] - 4s 5ms/step - loss: 0.2717 - f1_score: 0.8978
Epoch 6/10
938/938 [==============================] - 4s 5ms/step - loss: 0.2503 - f1_score: 0.9052
Epoch 7/10
938/938 [==============================] - 4s 5ms/step - loss: 0.2343 - f1_score: 0.9108
Epoch 8/10
938/938 [==============================] - 4s 5ms/step - loss: 0.2185 - f1_score: 0.9167
Epoch 9/10
938/938 [==============================] - 4s 5ms/step - loss: 0.2038 - f1_score: 0.9224
Epoch 10/10
938/938 [==============================] - 4s 5ms/step - loss: 0.1900 - f1_score: 0.9276
# テストデータでの評価
eval = cnn.evaluate(test_dataset, return_dict=True)
print(eval)
>> {'loss': 0.300339013338089, 'f1_score': 0.8954108953475952}
from torcheval.metrics.functional import multiclass_f1_score
# 学習
def train_loop(dataloader, model, loss_fn, optim):
# バッチごとに 予測 → lossの算出 → 勾配の初期化 → 勾配の算出 → パラメータの更新 を実行
for batch, (X, y) in enumerate(dataloader):
pred = model(X) # 予測
loss = loss_fn(pred, y) # lossの算出
optim.zero_grad() # 勾配の初期化
loss.backward() # 勾配の算出
optim.step() # パラメータの更新
# 100バッチごとにスコアを算出（無くても学習には影響しません）
if batch % 100 == 0:
pred_labels = pred.argmax(dim=1)
target_labels = y.argmax(dim=1)
score = multiclass_f1_score(
input=pred_labels,
target=target_labels,
num_classes=NUM_TARGET,
average='weighted'
)
report_str = f'loss: {loss.item():.5f}, weighted f1-score: {score:.5f}'
progress_str = f'[{batch * len(X)}/{len(dataloader.dataset)}]'
print(report_str, progress_str)
# テスト
def test_loop(dataloader, model, loss_fn):
test_loss, test_score = 0, 0
# モデルの更新をしないので勾配計算を無効化（メモリの節約）
with torch.no_grad():
for X, y in dataloader:
pred = model(X)
pred_labels = pred.argmax(dim=1)
test_loss += loss_fn(pred, y).item()
test_score += multiclass_f1_score(
input=pred_labels,
target=y,
num_classes=NUM_TARGET,
average='weighted'
)
report_str = f'Avg. loss: {test_loss / len(dataloader)}, Avg. weighted f1-score: {test_score / len(dataloader)}'
print(report_str)
learning_rate = 1e-3
epochs = 10
# loss: Cross Entropy Loss
loss_fn = nn.CrossEntropyLoss()
optim = torch.optim.Adam(cnn.parameters(), lr=learning_rate)
for t in range(epochs):
print(f'Epoch {t + 1}\n-------------------------------')
train_loop(train_dataloader, cnn, loss_fn, optim)
test_loop(test_dataloader, cnn, loss_fn)
print("Done!")
>> 以下のような出力が得られます
Epoch 1
-------------------------------
loss: 2.30322, weighted f1-score: 0.06608 [0/60000]
loss: 1.82939, weighted f1-score: 0.58820 [6400/60000]
loss: 1.71202, weighted f1-score: 0.70612 [12800/60000]
loss: 1.69892, weighted f1-score: 0.72906 [19200/60000]
loss: 1.71677, weighted f1-score: 0.68233 [25600/60000]
loss: 1.71246, weighted f1-score: 0.72033 [32000/60000]
loss: 1.66192, weighted f1-score: 0.77486 [38400/60000]
loss: 1.71226, weighted f1-score: 0.74929 [44800/60000]
loss: 1.65617, weighted f1-score: 0.79477 [51200/60000]
loss: 1.58637, weighted f1-score: 0.87646 [57600/60000]
Avg. loss: 1.6582482438178578, Avg. weighted f1-score: 0.7995069622993469
...
...
予測
PyTorchではTensorflowと異なり，学習後のモデルでないと予測できないということはありません．もっと正確に表現するなら，モデルにデータを流し込むという操作に学習と予測の区別がありません．（深層学習の理論的なことを考えるとPyTorchの方が自然かもしれません．）
pred = model(X)
ここで，model.forward()としてはいけません．PyTorchの設計原則に反するために，正常に動作しなくなるようです．（詳しくは以下を参照．）
学習
では，PyTorchで学習はどのように記述するのかというと，予測 → lossの算出 → 勾配の初期化 → 勾配の算出 → パラメータの更新，という具合に記述します．
# モデルによる予測
pred = model(X)
# lossの算出
loss = loss_fn(pred, y)
# 勾配の初期化
optim.zero_grad()
# 勾配の算出
loss.backward()
# パラメータの更新
optim.step()
上記の操作で詰まる部分として「勾配の初期化」が挙げられるのではないでしょうか？
PyTorchではこの初期化をおこなわないと勾配計算の結果が蓄積されていきます．この仕様はどうやら勾配の更新のタイミングを明示的にするためのもののようですが，パラメータの更新後に勾配の情報を残すような場面が存在するかについては分かりませんでした．
それ以外については誤差逆伝播法の手順そのままなので，特に詰まることもなくPyTorchに移行できると思います．
ちなみに，パラメータの更新が必要ない場合はwith torch.no_grad():の中で予測をおこないます．メモリの節約以外のメリットがあるかは分かりませんが，不要なものを除いた方が良いとは思うので，とりあえず付けておくという程度の認識で良いと思います．
評価
評価についてはTensorflowよりもscikit-learnに近いような気がします．
multiclass_f1_score(
input=pred_labels,
target=y,
num_classes=NUM_TARGET,
average='weighted'
)
今回はミニバッチ学習をしているので最後に平均を算出していますが，おそらくすべての予測値を保持してスコアを算出することも可能だとは思います．（メモリが許せば．）
最後に
よく巷でPyTorchは中上級者向けのライブラリであると書かれているのを目にし，TensorflowからPyTorchへの移行を躊躇していましたが，いざやってみるとそれほど難しくない印象を受けました．
一方で，TensorflowとPyTorchの使い分けが重要であるとも感じました．Tensorflowに比べて，特にモデルの学習と評価の操作が，PyTorchの方が面倒に思ったので，完全にPyTorchに移行するのではなく，簡単なモデルや前提の検証のためにTensorflowを利用し，そうでないものに対してPyTorchを利用するのが良いのではないかと感じました．
参考
以下に，私が勉強する中で特に参考になったサイトを挙げさせていただきます．本記事で説明しきれなかった部分についても以下のサイトを活用することで理解を深めることができると思います．
使用したコードの全体
PyTorch
import torch
from torchvision.datasets import FashionMNIST
from torchvision.transforms import ToTensor, Lambda
from torch.utils.data import DataLoader
from torch import nn
import torch.nn.functional as F
from torcheval.metrics.functional import multiclass_f1_score
data_root_dir = './datasets'
train_data = FashionMNIST(
root=data_root_dir,
train=True,
transform=ToTensor(),
target_transform=Lambda(lambda y: torch.zeros(10, dtype=torch.float).scatter(0, torch.tensor(y), value=1))
)
test_data = FashionMNIST(
root=data_root_dir,
train=False,
transform=ToTensor()
)
train_dataloader = DataLoader(
dataset=train_data,
batch_size=64,
shuffle=True
)
test_dataloader = DataLoader(
dataset=test_data,
batch_size=64,
shuffle=True
)
NUM_TARGET = 10
class ConvolutionalNeuralNetwork(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(
in_channels=1,
out_channels=6,
kernel_size=3,
stride=1,
padding=0
)
self.pool1 = nn.MaxPool2d(
kernel_size=2,
padding=0
)
self.conv2 = nn.Conv2d(
in_channels=6,
out_channels=16,
kernel_size=3,
stride=2,
padding=1
)
self.pool2 = nn.MaxPool2d(
kernel_size=2,
padding=1
)
self.flatten = nn.Flatten(
start_dim=1
)
self.fc = nn.Linear(
in_features=(4 ** 2) * 16,
out_features=1024
)
self.fc_out = nn.Linear(
in_features=1024,
out_features=NUM_TARGET
)
def forward(self, X):
h = self.conv1(X)
h = self.pool1(h)
h = F.relu(self.conv2(h), inplace=False)
h = self.pool2(h)
h = self.flatten(h)
h = F.relu(self.fc(h))
y = self.fc_out(h)
return F.softmax(y, dim=1)
cnn = ConvolutionalNeuralNetwork()
print(cnn)
def train_loop(dataloader, model, loss_fn, optim):
for batch, (X, y) in enumerate(dataloader):
pred = model(X)
loss = loss_fn(pred, y)
optim.zero_grad()
loss.backward()
optim.step()
if batch % 100 == 0:
pred_labels = pred.argmax(dim=1)
target_labels = y.argmax(dim=1)
score = multiclass_f1_score(
input=pred_labels,
target=target_labels,
num_classes=NUM_TARGET,
average='weighted'
)
report_str = f'loss: {loss.item():.5f}, weighted f1-score: {score:.5f}'
progress_str = f'[{batch * len(X)}/{len(dataloader.dataset)}]'
print(report_str, progress_str)
def test_loop(dataloader, model, loss_fn):
test_loss, test_score = 0, 0
with torch.no_grad():
for X, y in dataloader:
pred = model(X)
pred_labels = pred.argmax(dim=1)
test_loss += loss_fn(pred, y).item()
test_score += multiclass_f1_score(
input=pred_labels,
target=y,
num_classes=NUM_TARGET,
average='weighted'
)
report_str = f'Avg. loss: {test_loss / len(dataloader)}, Avg. weighted f1-score: {test_score / len(dataloader)}'
print(report_str)
learning_rate = 1e-3
epochs = 10
loss_fn = nn.CrossEntropyLoss()
optim = torch.optim.Adam(cnn.parameters(), lr=learning_rate)
for t in range(epochs):
print(f'Epoch {t + 1}\n-------------------------------')
train_loop(train_dataloader, cnn, loss_fn, optim)
test_loop(test_dataloader, cnn, loss_fn)
print("Done!")
Tensorflow
import tensorflow as tf
from tensorflow.keras.datasets import fashion_mnist
from tensorflow.keras import layers
(X_train, y_train), (X_test, y_test) = fashion_mnist.load_data()
train_dataset = tf.data.Dataset.from_tensor_slices((X_train, y_train))
test_dataset = tf.data.Dataset.from_tensor_slices((X_test, y_test))
NUM_TARGET = 10
def preprocess(dataset):
dataset = (
dataset
.map(lambda img, target: (img / 255, [0. if i != target else 1. for i in range(NUM_TARGET)]))
.shuffle(buffer_size=1000, seed=42)
.batch(64)
.prefetch(1)
)
return dataset
train_dataset = preprocess(train_dataset)
test_dataset = preprocess(test_dataset)
def build_model(
input_shape: tuple[int, int]=(28, 28, 1),
output_size: int=NUM_TARGET,
lr: float=1e-3
):
input_layer = tf.keras.Input(shape=input_shape)
conv1 = layers.Conv2D(
filters=6,
kernel_size=3,
strides=(1, 1),
padding='valid',
activation=None
)
pool1 = layers.MaxPool2D(
pool_size=(2, 2),
padding='valid'
)
conv2 = layers.Conv2D(
filters=16,
kernel_size=3,
strides=(2, 2),
padding='same',
activation=tf.keras.activations.relu
)
pool2 = layers.MaxPool2D(
pool_size=(2, 2),
padding='same'
)
flatten = layers.Flatten()
dense = layers.Dense(
units=1024,
activation=tf.keras.activations.relu
)
output_layer = layers.Dense(
units=output_size,
activation=tf.keras.activations.softmax
)
h = conv1(input_layer)
h = pool1(h)
h = conv2(h)
h = pool2(h)
h = flatten(h)
h = dense(h)
y = output_layer(h)
model = tf.keras.Model(inputs=input_layer, outputs=y)
model.compile(
optimizer=tf.keras.optimizers.Adam(learning_rate=lr),
loss=tf.keras.losses.CategoricalCrossentropy(),
metrics=tf.keras.metrics.F1Score(average='weighted')
)
return model
cnn = build_model()
cnn.summary()
cnn.fit(
train_dataset,
epochs=10
)
eval = cnn.evaluate(test_dataset, return_dict=True)
print(eval)