PyTorch入門！実践的な機械学習のはじめ方

from sklearn.datasets import fetch_openml # scikit-learnのデータ取得ライブラリ

import numpy as np
from sklearn.model_selection import train_test_split # Testデータ生成用ライブラリ 

import torch # PyTorch
import torch.nn as nn # ニューラルネットワーク生成用ライブラリ
import torch.nn.functional as F # ニューラルネットワーク用の関数
import torch.optim as optim # 最適化関数のライブラリ

# Xに画像データ、tに対応する正解ラベル(0~9)が取得されます。
X, t = fetch_openml('mnist_784', version=1, return_X_y=True)

# 取得したデータを確認
print(X.shape, t.shape) # (70000, 784) (70000,)
print(type(X), type(t)) # <class 'numpy.ndarray'> <class 'numpy.ndarray'>

class Model(nn.Module):
    # 使用するニューラルネットの層を定義していきます。
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(num_features=20)
        self.conv3 = nn.Conv2d(20, 20, kernel_size=3, padding=1)
        self.conv4 = nn.Conv2d(20, 40, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(num_features=40)
        self.conv5 = nn.Conv2d(40, 40, kernel_size=3, padding=0)
        self.conv6 = nn.Conv2d(40, 40, kernel_size=3, padding=0)
        self.bn3 = nn.BatchNorm2d(num_features=40)

        self.fc1 = nn.Linear(40 * 3 * 3, 120)  # 6*6 from image dimension
        self.fc2 = nn.Linear(120, 10)

    # __init__で定義した層を用いて、順伝播の処理を記載していきます。
    def forward(self, x):
        # import pdb; pdb.set_trace()
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(F.relu(self.bn1(self.conv2(x))), (2,2))
        x = F.relu(self.conv3(x))
        x = F.max_pool2d(F.relu(self.bn2(self.conv4(x))), (2,2))
        x = F.relu(self.conv5(x))
        x = F.relu(self.bn2(self.conv6(x)))

        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

# 損失関数の定義
criterion = nn.CrossEntropyLoss()

# 評価関数の定義
# 最適化関数(optimizer)の定義
# 第一引数にoptによって最適化するパラメータ群を指定する。モデルのパラメータ群になるはず。
lr = 1e-3
opt = optim.Adam(model.parameters(), lr=lr)

class MyDataset(torch.utils.data.Dataset):
    def __init__(self, dataset):
        
        super(MyDataset, self).__init__()
        self.dataset = dataset

    # データセットの数
    def __len__(self):
        return len(self.dataset[0])
    
    # DataLoaderに呼ばれた際に取得するデータセット
    def __getitem__(self, idx):
        
        # idx単位で取得するようにするのがポイントな気がする。__getitem__として動作させるには。
        x = self.dataset[0][idx]
        label = self.dataset[1][idx]
        
        x = x[np.newaxis,:,:].astype(np.float32) / 255

        return x, label

# 学習用のデータセット
train_dataset = MyDataset([train_X, train_t])
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=4,shuffle=True, num_workers=1)

# 評価用のデータセット
test_dataset = MyDataset([test_X, test_t])
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=4, shuffle=False, num_workers=1)

best_loss = 1e5 # 誤差の初期値
n_epochs = 10 # エポック数

for epoch in range(1, n_epochs):
    print(f'Epoch = {epoch: 03d}')

    model.train()
    train_loss = 0

    # Training
    for _xs, _ts in train_loader:
        _xs, _ts = _xs.to(device), _ts.to(device)

        # stack overflow
        # https://stackoverflow.com/questions/48001598/why-do-we-need-to-call-zero-grad-in-pytorch
        opt.zero_grad() # 勾配を初期化
        
        _ys = model(_xs)

        loss = criterion(_ys, _ts)
        train_loss += loss.item() 

        # back propagation
        loss.backward() # モデルの勾配を計算
        
        opt.step() # モデルのパラメータを更新

        _xs = _xs.to('cpu').detach().numpy()
        _ts = _ts.to('cpu').detach().numpy()

    train_loss /= len(train_loader)

    # Evaluation
    model.eval()

    valid_loss = 0
    
    # 評価の場合は勾配を保存しないように設定
    with torch.no_grad():
        ts, ys = [], []
        for _xs, _ts in test_loader:
            _xs, _ts = _xs.to(device), _ts.to(device)
            _ys = model(_xs)
            valid_loss += criterion(_ys,_ts)
            ys.append(_ys.to('cpu').detach().numpy())
            ts.append(_ts.to('cpu').detach().numpy())
        ys = np.concatenate(ys, axis=0)

    valid_loss /= len(test_loader)

    # 評価時の誤差関数の結果が一番小さいモデルを保存する
    if valid_loss <= best_loss:
        torch.save(model.state_dict(), './model_00.pt')
        best_loss = valid_loss

    # 学習状況の確認。学習時、評価時、一番良い誤差を出力。
    print(f'Loss Train={train_loss:.4f}, Test={valid_loss:.4f}, Best={best_loss:.4f}')

# モデルのLOAD
# 公式サイト：https://pytorch.org/tutorials/beginner/saving_loading_models.html
model.load_state_dict(torch.load('./model_00.pt'))
model.eval()

# 推論する画像の表示(確認用に表示しています。)
input_id = 0
_x = test_X[input_id,:,:].astype(np.float32) /255

fig, ax = plt.subplots()
ax.axes.xaxis.set_visible(False) # X軸を非表示にする
ax.axes.yaxis.set_visible(False) # Y軸を非表示にする
ax.imshow(_x, cmap='Greys')

# 推論する画像をモデルに入力
_x = torch.from_numpy(_x[np.newaxis, np.newaxis,:,:]).clone() # 推論するためにndarray型のデータをTensor型に変換。
_x = _x.to(device) # GPUに載せる
_y = model(_x) # 推論
_y = _y.to('cpu').detach().numpy() # CPUに載せて、Tensor型からndarray型へ

# 推論結果を元にクラス分類
pred = np.argmax(_y, axis=1)

# 推論結果の表示
print(f'推論結果：{pred}')