Zenn
🕌

PyTorchで書籍「ゼロから作るDeep Learning」のTwoLayerNetを実装する

2024/01/10に公開
Python
機械学習
PyTorch
Apple
tech

はじめに

書籍「ゼロから作るDeep Learning」を読み終えたので、次のステップとしてPyTorchの学習を進めています。まずは簡単な題材として、書籍の第4章で登場するTwoLayerNetをPyTorchで実装することにしました。

この記事ではMacBook Proでモデルの訓練と推論を試します。小さなモデルであれば訓練が数十秒で完了することを実験します。

環境

記事の投稿にあたり、以下の環境で動作確認しました。

  • MacBook Pro（2021年に発売されたM1 PRO搭載機）
  • macOS 12.7.1 Monterey
  • Anaconda 23.7.4
  • Python 3.11.5
  • PyTorch 2.1.2

PyTorchのインストール

ローカル環境でモデルの訓練とテストを行うため、PyTorchのドキュメントに従ってPyTorchをインストールします。以下、Anacondaを利用してインストールする場合の手順を抜粋します。

次のコマンドを実行してください。

conda install pytorch torchvision -c pytorch

正常にインストールできたか確認するため、次のコマンドを実行してください。pytorchとtorchvisionのバージョンが表示されたら成功です。

conda list | grep torch
pytorch                   2.1.2                  py3.11_0    pytorch
torchvision               0.16.2                py311_cpu    pytorch

以上でインストールは完了です。

続いてMPSが有効か確認します。python3コマンドを実行してインタプリタを起動し、次のコマンドを実行してください。

>>> import torch
>>> torch.backends.mps.is_available()
True

Apple Siliconを搭載したマシンの場合はTrueが表示されます。なお、MPSが無効の場合はCPUを利用することになります。時間はかかりますが、小さなモデルであれば訓練や推論は行えます。

MPSが何者なのか、については記事の後半にて紹介します。

実装

それでは、実装を示します。作業するディレクトリはどこでも構いませんが、ここでは/tmp/torchにて作業すると仮定します。

モデルの訓練と保存

以下のコードをtrain.pyとして保存してください。

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor

class TwoLayerNet(nn.Module):
	def __init__(self, input_size, hidden_size, output_size):
		super().__init__()
		self.flatten = nn.Flatten()
		self.layers = nn.Sequential(
			nn.Linear(input_size, hidden_size),
			nn.ReLU(),
			nn.Linear(hidden_size, output_size)
		)
	def forward(self, x):
		x = self.flatten(x)
		logits = self.layers(x)
		return logits

def get_available_device():
	if torch.backends.mps.is_available():
		return 'mps'
	else:
		return 'cpu'

# どのハードウェアで訓練を行うか選択する。
device = get_available_device()

# モデルを作成する。
# 
# input_sizeには入力するデータの次元を指定する。MNISTの画像は縦横28 pxの正方形なので、28 * 28 = 784を指定する。
# hidden_sizeには隠れ層の重みの次元を指定する。この値は適当で構わないので、書籍に習って100を指定する。
# output_sizeには分類される数を指定する。MNISTのデータセットには0から9までの手書き数字の画像が含まれているため、output_sizeには10を指定する。
model = TwoLayerNet(input_size=784, hidden_size=100, output_size=10).to(device)

# 訓練データを読み込む。バッチのサイズは適当で構わないので、書籍に習って100を指定する。
batch_size = 100
training_datasets = datasets.MNIST(root="./data", train=True, download=True, transform=ToTensor())
training_dataloader = DataLoader(training_datasets, batch_size=batch_size, shuffle=True)

# 損失関数として交差エントロピー誤差、オプティマイザとして確率的勾配降下法を利用する。Learning Rateなどのハイパーパラメータは適当で構わないが、ひとまず書籍にあわせる。
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)

# エポックの数は適当なので、訓練の進み具合に応じて増減させて構わない。
for epoch in range(5):
	running_loss = 0.0
	
	for i, data in enumerate(training_dataloader, 0):
		inputs, labels = data
		inputs, labels = inputs.to(device), labels.to(device)
		optimizer.zero_grad()
		outputs = model(inputs)
		loss = criterion(outputs, labels)
		loss.backward()
		optimizer.step()
		running_loss += loss.item()
		if i % 600 == 0:
			print(f'epoch={epoch + 1}, loss={running_loss:.5f}')
			running_loss = 0.0

# 訓練が終わったら、モデルをmnist.pthとしてファイルに保存する。
torch.save(model.state_dict(), './mnist.pth')

print("\nFinished!")

モデルの読み込みと評価

以下のコードをtest.pyとして保存してください。

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor

class TwoLayerNet(nn.Module):
	def __init__(self, input_size, hidden_size, output_size):
		super().__init__()
		self.flatten = nn.Flatten()
		self.layers = nn.Sequential(
			nn.Linear(input_size, hidden_size),
			nn.ReLU(),
			nn.Linear(hidden_size, output_size)
			# nn.Softmax(dim=1)
		)
	def forward(self, x):
		x = self.flatten(x)
		logits = self.layers(x)
		return logits

def get_available_device():
	if torch.backends.mps.is_available():
		return 'mps'
	else:
		return 'cpu'

device = get_available_device()

# 訓練済みのモデルを読み込む。
model = TwoLayerNet(input_size=784, hidden_size=100, output_size=10).to(device)
model_path = './mnist.pth'
model.load_state_dict(torch.load(model_path))

# テスト用のデータを読み込む。
batch_size = 100
test_datasets = datasets.MNIST(root="./data", train=False, download=True, transform=ToTensor())
test_dataloader = DataLoader(test_datasets, batch_size=batch_size, shuffle=True)

# モデルを評価する。
correct = 0
total = 0

with torch.no_grad():
	for data in test_dataloader:
		inputs, labels = data
		inputs = inputs.to(device)
		labels = labels.to(device)
		outputs = model(inputs)
		_, predicted = torch.max(outputs.data, 1)
		total += labels.size(0)
		correct += (predicted == labels).sum().item()

print(f'Accuracy of the network on the 10000 test images: {100 * correct / total:.2f} %')

動作確認

それでは動作確認を行いましょう。まずはモデルの訓練からです。

次のコマンドを実行してください。

python3 train.py

初回は画像データをダウンロードするため完了まで数分かかります。２回目以降の実行であれば、10秒ほどで完了します。

以下のように損失（loss）が徐々に小さくなっていることが確認できたら成功です。

epoch=1, loss=2.28737
epoch=2, loss=0.11533
epoch=3, loss=0.09539
epoch=4, loss=0.09099
epoch=5, loss=0.06506

Finished!

、訓練したモデルを評価しましょう。次のコマンドを実行してください。

python3 test.py

結果は以下のようになります。この結果と全く同じになるとは限りませんが、おおむね95 %以上になるはずです。

Accuracy of the network on the 10000 test images: 97.65 %

以上で動作確認は完了です！

（おまけ）MPSは何者だ？

Apple製デバイスにもGPUは搭載されています。MPS（Metal Performance Shaders）はGPUを利用するためのAPIです。行列の演算を多用する処理の高速化が期待できます。

余談になりますが、AppleのMPS紹介ページにはPyTorchのインストールコマンドが掲載されています。内容が古いのでその部分は無視してください。

（おまけ）PyTorchのチュートリアル、どの順番で読み進めるか？

PyTorchのドキュメントには初学者向けのチュートリアルがいくつか用意されています。どこから手をつけるか迷うところですが、素直にLearn the Basicsを1から7まで順番に沿って読み進めるのがおすすめです。

  1. Tensors
  2. Datasets & DataLoaders
  3. Transforms
  4. Build the Neural Network
  5. Automatic Differentiation with torch.autograd
  6. Optimizing Model Parameters
  7. Save and Load the Model

Learn the Basicsを読み終えたら、Deep Learning with PyTorch: A 60 Minute Blitzに進むのがおすすめです。

なぜかと言えば、TensorsはLearn the Basicsと同じ内容なので読み飛ばすことができるからです。さらに、A Gentle Introduction to torch.autogradは書籍「ゼロから作るDeep Learning」の第5章と内容が重複しているからです。

Discussion