PyTorchの基本的な使い方(十)——神経ネットワークを巻き込む

import torch
import torchvision
import torch.nn as nn
import torch.utils.data as Data
import matplotlib.pyplot as plt
from torch.autograd import Variable

#  
EPOCH = 1
LR = 0.01
BATCH_SIZE = 64

#  MNIST 
train_data = torchvision.datasets.MNIST(
    root = './mnist/',
    #  
    train = True,
    #  (0, 255) -> (0, 1)
    transform = torchvision.transforms.ToTensor(),
    #  MNIST 
    download = True
)

test_data = torchvision.datasets.MNIST(
    root = './mnist/',
    #  
    train = False,
    #  (0, 255) -> (0, 1)
    transform = torchvision.transforms.ToTensor(),
    #  MNIST 
    download = True
)

print train_data.train_data.size()
print train_data.train_labels.size()
print test_data.test_data.size()
print test_data.test_labels.size()

torch.Size([60000, 28, 28])
torch.Size([60000])
torch.Size([10000, 28, 28])
torch.Size([10000])

#  
plt.imshow(train_data.train_data[0].numpy(), cmap = 'gray')
plt.title('%i' % train_data.train_labels[0])
plt.show()

plt.imshow(test_data.test_data[0].numpy(), cmap = 'gray')
plt.title('%i' % test_data.test_labels[0])
plt.show()

#  
train_loader = Data.DataLoader(dataset = train_data, batch_size = BATCH_SIZE, shuffle = True, num_workers = 2)
test_loader = Data.DataLoader(dataset = test_data, batch_size = BATCH_SIZE, shuffle = False, num_workers = 1)

#  
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(
                in_channels = 1,
                out_channels = 16,
                kernel_size = 5,
                stride = 1,
                padding = 2
            ),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size = 2)
        )
        # conv1 (16, 14, 14)
        self.conv2 = nn.Sequential(
            nn.Conv2d(16, 32, 5, 1, 2),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        # conv2 (32, 7, 7)
        self.output = nn.Linear(32 * 7 * 7, 10)
        
    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size(0), -1)
        prediction = self.output(x)
        return prediction

cnn = CNN()
print cnn

CNN (
  (conv1): Sequential (
    (0): Conv2d(1, 16, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
    (1): ReLU ()
    (2): MaxPool2d (size=(2, 2), stride=(2, 2), dilation=(1, 1))
  )
  (conv2): Sequential (
    (0): Conv2d(16, 32, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
    (1): ReLU ()
    (2): MaxPool2d (size=(2, 2), stride=(2, 2), dilation=(1, 1))
  )
  (output): Linear (1568 -> 10)
)

#  
optimizer = torch.optim.Adam(cnn.parameters(), lr = LR, betas= (0.9, 0.999))

#  
loss_func = nn.CrossEntropyLoss()

#  
for epoch in xrange(EPOCH):
    for step, (x, y) in enumerate(train_loader):
        x_var = Variable(x)
        y_var = Variable(y)
        prediction = cnn(x_var)
        loss = loss_func(prediction, y_var)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        if step % 100 == 0:
            correct = 0.0
            for step_test, (test_x, test_y) in enumerate(test_loader):
                test_x = Variable(test_x)
                test_output = cnn(test_x)
                pred_y = torch.max(test_output, 1)[1].data.squeeze()
                correct += sum(pred_y == test_y)
            accuracy = correct / test_data.test_data.size(0)
            print 'Epoch: ', epoch, '| train loss: %.4f' % loss.data[0], '| accuracy: ', accuracy

Epoch:  0 | train loss: 2.2787 | accuracy:  0.0982
Epoch:  0 | train loss: 0.0788 | accuracy:  0.9592
Epoch:  0 | train loss: 0.0587 | accuracy:  0.9626
Epoch:  0 | train loss: 0.0188 | accuracy:  0.9745
Epoch:  0 | train loss: 0.0707 | accuracy:  0.9759
Epoch:  0 | train loss: 0.0564 | accuracy:  0.9775
Epoch:  0 | train loss: 0.0489 | accuracy:  0.9779
Epoch:  0 | train loss: 0.0925 | accuracy:  0.9791
Epoch:  0 | train loss: 0.0566 | accuracy:  0.9834

参考資料