lesson5.pdf

    1. import  torch
    2. from    torch import nn
    3. from    torch.nn import functional as F
    4. from    torch import optim
    6. import  torchvision
    7. from    matplotlib import pyplot as plt
    9. from    utils import plot_image, plot_curve, one_hot
    13. batch_size = 512
    15. # step1. load dataset
    16. train_loader = torch.utils.data.DataLoader(
    17.     torchvision.datasets.MNIST('mnist_data', train=True, download=True,
    18.                                transform=torchvision.transforms.Compose([
    19.                                    torchvision.transforms.ToTensor(),
    20.                                    torchvision.transforms.Normalize(
    21.                                        (0.1307,), (0.3081,))
    22.                                ])),
    23.     batch_size=batch_size, shuffle=True)
    25. test_loader = torch.utils.data.DataLoader(
    26.     torchvision.datasets.MNIST('mnist_data/', train=False, download=True,
    27.                                transform=torchvision.transforms.Compose([
    28.                                    torchvision.transforms.ToTensor(),
    29.                                    torchvision.transforms.Normalize(
    30.                                        (0.1307,), (0.3081,))
    31.                                ])),
    32.     batch_size=batch_size, shuffle=False)
    34. x, y = next(iter(train_loader))
    35. print(x.shape, y.shape, x.min(), x.max())
    36. plot_image(x, y, 'image sample')
    40. class Net(nn.Module):
    42.     def __init__(self):
    43.         super(Net, self).__init__()
    45.         # xw+b
    46.         self.fc1 = nn.Linear(28*28, 256)
    47.         self.fc2 = nn.Linear(256, 64)
    48.         self.fc3 = nn.Linear(64, 10)
    50.     def forward(self, x):
    51.         # x: [b, 1, 28, 28]
    52.         # h1 = relu(xw1+b1)
    53.         x = F.relu(self.fc1(x))
    54.         # h2 = relu(h1w2+b2)
    55.         x = F.relu(self.fc2(x))
    56.         # h3 = h2w3+b3
    57.         x = self.fc3(x)
    59.         return x
    63. net = Net()
    64. # [w1, b1, w2, b2, w3, b3]
    65. optimizer = optim.SGD(net.parameters(), lr=0.01, momentum=0.9)
    68. train_loss = []
    70. for epoch in range(3):
    72.     for batch_idx, (x, y) in enumerate(train_loader):
    74.         # x: [b, 1, 28, 28], y: [512]
    75.         # [b, 1, 28, 28] => [b, 784]
    76.         x = x.view(x.size(0), 28*28)
    77.         # => [b, 10]
    78.         out = net(x)
    79.         # [b, 10]
    80.         y_onehot = one_hot(y)
    81.         # loss = mse(out, y_onehot)
    82.         loss = F.mse_loss(out, y_onehot)
    84.         optimizer.zero_grad()
    85.         loss.backward()
    86.         # w' = w - lr*grad
    87.         optimizer.step()
    89.         train_loss.append(loss.item())
    91.         if batch_idx % 10==0:
    92.             print(epoch, batch_idx, loss.item())
    94. plot_curve(train_loss)
    95. # we get optimal [w1, b1, w2, b2, w3, b3]
    98. total_correct = 0
    99. for x,y in test_loader:
    100.     x  = x.view(x.size(0), 28*28)
    101.     out = net(x)
    102.     # out: [b, 10] => pred: [b]
    103.     pred = out.argmax(dim=1)
    104.     correct = pred.eq(y).sum().float().item()
    105.     total_correct += correct
    107. total_num = len(test_loader.dataset)
    108. acc = total_correct / total_num
    109. print('test acc:', acc)
    111. x, y = next(iter(test_loader))
    112. out = net(x.view(x.size(0), 28*28))
    113. pred = out.argmax(dim=1)
    114. plot_image(x, pred, 'test')
    1. import  torch
    2. from    matplotlib import pyplot as plt
    5. def plot_curve(data):
    6.     fig = plt.figure()
    7.     plt.plot(range(len(data)), data, color='blue')
    8.     plt.legend(['value'], loc='upper right')
    9.     plt.xlabel('step')
    10.     plt.ylabel('value')
    11.     plt.show()
    15. def plot_image(img, label, name):
    17.     fig = plt.figure()
    18.     for i in range(6):
    19.         plt.subplot(2, 3, i + 1)
    20.         plt.tight_layout()
    21.         plt.imshow(img[i][0]*0.3081+0.1307, cmap='gray', interpolation='none')
    22.         plt.title("{}: {}".format(name, label[i].item()))
    23.         plt.xticks([])
    24.         plt.yticks([])
    25.     plt.show()
    28. def one_hot(label, depth=10):
    29.     out = torch.zeros(label.size(0), depth)
    30.     idx = torch.LongTensor(label).view(-1, 1)
    31.     out.scatter_(dim=1, index=idx, value=1)
    32.     return out