【Pytorch】mnistを用いたクラス分類モデルをCNNで学習するコード

conv_net = nn.Sequential(
    Conv((20, 1, 5, 5), F.relu),     # 28x28x 1 -> 24x24x20
    Pooling((2, 2)),                 # 24x24x20 -> 12x12x20
    Conv((50, 20, 5, 5), F.relu),    # 12x12x20 ->  8x 8x50
    Pooling((2, 2)),                 #  8x 8x50 ->  4x 4x50
    Flatten(),
    Dense(4*4*50, 10)
)

batch_size = 100
n_epochs = 10
lr = 0.02
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


conv_net.to(device)
optimizer = optim.SGD(conv_net.parameters(), lr=lr)

# torchvisionのdatasetsを使ってMNISTのデータを取得
# ミニバッチ化や前処理などの処理を行ってくれるDataLoaderを定義

dataloader_train = torch.utils.data.DataLoader(
    datasets.MNIST('./data/mnist', train=True, download=True, transform=transforms.ToTensor()),
    batch_size=batch_size,
    shuffle=True
)

dataloader_valid = torch.utils.data.DataLoader(
    datasets.MNIST('./data/mnist', train=False, download=True, transform=transforms.ToTensor()),
    batch_size=batch_size,
    shuffle=False
)

for epoch in range(n_epochs):
    losses_train = []
    losses_valid = []

    conv_net.train()
    n_train = 0
    acc_train = 0
    for x, t in dataloader_train:
        n_train += t.size()[0]

        conv_net.zero_grad()  # 勾配の初期化

        x = x.to(device)  # テンソルをGPUに移動

        t_hot = torch.eye(10)[t]  # 正解ラベルをone-hot vector化

        t_hot = t_hot.to(device)  # 正解ラベルとone-hot vectorをそれぞれGPUに移動

        y = conv_net.forward(x)  # 順伝播

        loss = -(t_hot*torch.log_softmax(y, dim=-1)).sum(axis=1).mean()  # 誤差(クロスエントロピー誤差関数)の計算

        loss.backward()  # 誤差の逆伝播

        optimizer.step()  # パラメータの更新

        pred = y.argmax(1)  # 最大値を取るラベルを予測ラベルとする

        acc_train += (pred.to("cpu") == t).float().sum().item()
        losses_train.append(loss.tolist())

    conv_net.eval()
    n_val = 0
    acc_val = 0
    for x, t in dataloader_valid:
        n_val += t.size()[0]

        x = x.to(device)  # テンソルをGPUに移動

        t_hot = torch.eye(10)[t]  # 正解ラベルをone-hot vector化

        t_hot = t_hot.to(device)  # 正解ラベルとone-hot vectorをそれぞれGPUに移動

        y = conv_net.forward(x)  # 順伝播

        loss = -(t_hot*torch.log_softmax(y, dim=-1)).sum(axis=1).mean()  # 誤差(クロスエントロピー誤差関数)の計算

        pred = y.argmax(1)  # 最大値を取るラベルを予測ラベルとする

        acc_val += (pred.to("cpu") == t).float().sum().item()
        losses_valid.append(loss.tolist())

    print('EPOCH: {}, Train [Loss: {:.3f}, Accuracy: {:.3f}], Valid [Loss: {:.3f}, Accuracy: {:.3f}]'.format(
        epoch,
        np.mean(losses_train),
        acc_train/n_train,
        np.mean(losses_valid),
        acc_val/n_val
    ))

# 出力
# EPOCH: 0, Train [Loss: 0.471, Accuracy: 0.866], Valid [Loss: 0.244, Accuracy: 0.928]
# EPOCH: 1, Train [Loss: 0.225, Accuracy: 0.935], Valid [Loss: 0.168, Accuracy: 0.952]
# EPOCH: 2, Train [Loss: 0.167, Accuracy: 0.953], Valid [Loss: 0.130, Accuracy: 0.963]
# EPOCH: 3, Train [Loss: 0.135, Accuracy: 0.961], Valid [Loss: 0.112, Accuracy: 0.968]
# EPOCH: 4, Train [Loss: 0.116, Accuracy: 0.967], Valid [Loss: 0.102, Accuracy: 0.968]
# EPOCH: 5, Train [Loss: 0.101, Accuracy: 0.971], Valid [Loss: 0.080, Accuracy: 0.976]
# EPOCH: 6, Train [Loss: 0.092, Accuracy: 0.974], Valid [Loss: 0.076, Accuracy: 0.976]
# EPOCH: 7, Train [Loss: 0.084, Accuracy: 0.975], Valid [Loss: 0.072, Accuracy: 0.978]
# EPOCH: 8, Train [Loss: 0.078, Accuracy: 0.978], Valid [Loss: 0.061, Accuracy: 0.982]
# EPOCH: 9, Train [Loss: 0.074, Accuracy: 0.979], Valid [Loss: 0.060, Accuracy: 0.981]