1import torch
2import torchvision
3import torchvision.transforms as transforms
4import torch.nn as nn
5import torch.optim as optim
6
7# データセットの読み込みと前処理
8transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
9trainset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform)
10trainloader = torch.utils.data.DataLoader(trainset, batch_size=4, shuffle=True, num_workers=2)
11
12# ニューラルネットワークの定義
13class Net(nn.Module):
14    def __init__(self):
15        super(Net, self).__init__()
16        self.conv1 = nn.Conv2d(1, 6, 5)
17        self.pool = nn.MaxPool2d(2, 2)
18        self.conv2 = nn.Conv2d(6, 16, 5)
19        self.fc1 = nn.Linear(16 * 4 * 4, 120)
20        self.fc2 = nn.Linear(120, 84)
21        self.fc3 = nn.Linear(84, 10)
22
23    def forward(self, x):
24        x = self.pool(F.relu(self.conv1(x)))
25        x = self.pool(F.relu(self.conv2(x)))
26        x = x.view(-1, 16 * 4 * 4)
27        x = F.relu(self.fc1(x))
28        x = F.relu(self.fc2(x))
29        x = self.fc3(x)
30        return x
31
32net = Net()
33
34# 損失関数と最適化アルゴリズムの定義
35criterion = nn.CrossEntropyLoss()
36optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
37
38# ニューラルネットワークの学習
39for epoch in range(2):
40    running_loss = 0.0
41    for i, data in enumerate(trainloader, 0):
42        inputs, labels = data
43        optimizer.zero_grad()
44        outputs = net(inputs)
45        loss = criterion(outputs, labels)
46        loss.backward()
47        optimizer.step()
48        running_loss += loss.item()
49        if i % 2000 == 1999:    # 2000 ミニバッチごとに表示
50            print(f'[{epoch + 1}, {i + 1}] loss: {running_loss / 2000:.3f}')
51            running_loss = 0.0
52
53print('学習が完了しました。')