PyTorchに入門してみる　part 5（最終回）　MNISTを単純なLSTMで解く

PyTorch

はじめに

PyTorch入門シリーズ5回目になります。これで最終回にしようと思います。

今まで4回は公開されているモデルを使用させて頂きました。

今回は自分でスクリプトを書きます。

題材は入門にうってつけの「MNIST」です。

テーマは『できるだけ簡潔に』としました。


LSTMを使ってMNISTにチャレンジします。画像分類にLSTM？と疑問に思う方もいるかもしれません。
詳細はこちらを参照して下さい。
MNISTを単純なLSTMで解く（MXNet） - パソコン関連もろもろ

Pythonスクリプト

できる限り簡潔に書いたつもりです。ただPyTorchでこのようなことをするのは初めてなので改善点があるかもしれません。

import torch
import torch.nn as nn
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

device = 'cuda' if torch.cuda.is_available() else 'cpu' 

#データの取得
transform = transforms.Compose([transforms.ToTensor()])
trainset = datasets.MNIST(root='.', download = True, transform = transform, train = True)
testset = datasets.MNIST(root='.', download = True, transform = transform, train = False)

trainloader = DataLoader(trainset, batch_size = 100, shuffle = True)
testloader = DataLoader(testset, batch_size = 100, shuffle = False)

#モデルの定義
class MyModel(nn.Module):
    def __init__(self, input_size = 28, hidden_layer = 128):
        super().__init__()
        self.lstm = nn.LSTM(input_size, hidden_layer, batch_first = True)
        self.flat = nn.Flatten()
        self.linear = nn.Linear(input_size * hidden_layer, 10)
        torch.nn.init.xavier_uniform_(self.linear.weight)
        self.linear.bias.data.fill_(0.01)
        self.softmax = nn.Softmax(dim =1)
        
    def forward(self, x):
        x = self.lstm(x)
        x = self.flat(x[0])
        x = self.linear(x)
        x = self.softmax(x)
        return x

model = MyModel().to(device)

loss_func = nn.CrossEntropyLoss()
opt = torch.optim.Adam(model.parameters())

#評価のための関数定義
def evaluate_accuracy(data, net):
    total_count = 0
    correct_count = 0
    net.eval()
    for batch_data in data:
        x = torch.squeeze(batch_data[0], dim = 1).to(device)
        y = batch_data[1]
        out = net(x)
        out = out.argmax(1).to('cpu')
        total_count += y.shape[0]
        correct_count += (out == y).sum().item()
    acc = correct_count / total_count
    return acc

#学習ループ
for epoch in range(10):
    model.train()
    for batch_data in trainloader:
        opt.zero_grad()
        x = torch.squeeze(batch_data[0], dim = 1).to(device)
        y = batch_data[1].to(device)

        output = model(x)

        loss = loss_func(output, y)
        loss.backward()
        opt.step()
    
    test_acc = evaluate_accuracy(testloader, model)
    print('%d epoch test_acc = %f' %(epoch + 1, test_acc))

設定したハイパーパラメーターは以下の4つだけです。

LSTMのlayer数：128
全結合層のバイアス初期化：0.01
batchサイズ：100
epoch数：10

結果

1 epoch test_acc = 0.940400
2 epoch test_acc = 0.970900
3 epoch test_acc = 0.975200
4 epoch test_acc = 0.979000
5 epoch test_acc = 0.984600
6 epoch test_acc = 0.986100
7 epoch test_acc = 0.986200
8 epoch test_acc = 0.983900
9 epoch test_acc = 0.985100
10 epoch test_acc = 0.988800

それなりの結果が得られたので書いたスクリプトに大きな間違いはなさそうです。

参考にさせて頂いたサイト

tzmi.hatenablog.com

さいごに

間違いや改善点があればコメント頂けましたら幸いです。

PyTorch入門シリーズ（前回までの記事）

touch-sp.hatenablog.com
touch-sp.hatenablog.com
touch-sp.hatenablog.com
touch-sp.hatenablog.com