【PyTorch】元祖CNNのLeNetを実装してみた

# -*- coding: utf-8 -*-
from torch import nn
import torch.nn.functional as F

class LeNet(nn.Module):
    def __init__(self) -> None:
        super(LeNet, self).__init__()
        self.conv1 = nn.Conv2d(
            in_channels=1, out_channels=2, kernel_size=3,
            stride=2, padding=1)
        self.conv2 = nn.Conv2d(
            in_channels=2, out_channels=4, kernel_size=5,
            stride=1, padding=0
        )
        self.fc = nn.Linear(4*4*4, 10)

    def forward(self, x):
        x = F.tanh(self.conv1(x))
        x = F.tanh(self.conv2(x))
        x = F.tanh(self.fc(x.view(-1, 4*4*4)))
        
        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(16),
    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)