【PyTorch】LeNet5を実装してみる

# -*- coding: utf-8 -*-
import torch
from torch import nn

class LeNet5(nn.Module):
    def __init__(self) -> None:
        super(LeNet5, self).__init__()
        self.conv1 = nn.Conv2d(
            in_channels=1, out_channels=6, kernel_size=5,
            stride=1, padding=0)
        self.avgpool1 = nn.AvgPool2d(2, 2)
        self.conv2 = nn.Conv2d(
            in_channels=6, out_channels=16, kernel_size=5,
            stride=1, padding=0
        )
        self.avgpool2 = nn.AvgPool2d(2, 2)
        self.conv3 = nn.Conv2d(
            in_channels=16, out_channels=120, kernel_size=5,
            stride=1, padding=0
        )
        self.fc1 = nn.Linear(120, 84)
        self.fc2 = nn.Linear(84, 10)

    def forward(self, x):
        x = torch.tanh(self.conv1(x))
        x = torch.tanh(self.avgpool1(x))
        x = torch.tanh(self.conv2(x))
        x = torch.tanh(self.avgpool2(x))
        x = torch.tanh(self.conv3(x))
        x = torch.tanh(self.fc1(x.view(-1, 120*1*1)))
        x = self.fc2(x)
        
        return x

import torch
import torchvision
import torchvision.transforms as transforms
from torch import nn, optim
from LeNet5 import LeNet5

preprocess = transforms.Compose([
    transforms.Resize(32),
    transforms.ToTensor(),
])

trainset = torchvision.datasets.MNIST(
    root='../data',
    train=True,
    download=True,
    transform=preprocess
)

testset = torchvision.datasets.MNIST(
    root='../data',
    train=False,
    download=True,
    transform=preprocess
) 

batch_size = 100

trainloader = torch.utils.data.DataLoader(
    trainset,
    batch_size=batch_size,
    shuffle=True,
)

testloader = torch.utils.data.DataLoader(
    testset,
    batch_size=batch_size,
    shuffle=False,
)

# 使用するデバイスの設定
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)

# 損失関数とオプティマイザの定義
net = LeNet5()
net = net.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.01)

# ネットワークの訓練
epochs = 100

net.train()
for epoch in range(epochs):
    loss = None
    count = 0
    train_acc = 0
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data[0].to(device), data[1].to(device)
        
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        
        # 訓練データの精度を計算
        pred = outputs.argmax(dim=1, keepdim=True)
        train_acc += pred.eq(labels.view_as(pred)).sum().item()
        count += inputs.shape[0]
        
    train_acc = (train_acc/count) * 100
    print("epoch{:d}/{:d} train_loss:{:.3f} train_acc:{:.2f}".format(epoch+1, epochs, loss, train_acc))

# モデルを保存
torch.save(net.state_dict(), './LeNet_mnist.pth')

# testデータで評価
test_acc = 0
count = 0
net.eval()
with torch.no_grad():
    for data in testloader:
        inputs, labels = data[0].to(device), data[1].to(device)
        outputs = net(inputs)
        pred = outputs.argmax(dim=1, keepdim=True)
        test_acc += pred.eq(labels.view_as(pred)).sum().item()
        count += inputs.shape[0]

acc = (test_acc / count) * 100
print(acc)