26.pdf

    1. import torch
    2. import torch.nn as nn
    3. import torch.nn.functional as F
    4. import torch.optim as optim
    5. from torchvision import datasets, transforms
    6. batch_size=200
    7. learning_rate=0.01
    8. epochs=10
    9. train_loader = torch.utils.data.DataLoader(
    10. datasets.MNIST('../data', train=True, download=True,
    11. transform=transforms.Compose([
    12. transforms.ToTensor(),
    13. transforms.Normalize((0.1307,), (0.3081,))
    14. ])),
    15. batch_size=batch_size, shuffle=True)
    16. test_loader = torch.utils.data.DataLoader(
    17. datasets.MNIST('../data', train=False, transform=transforms.Compose([
    18. transforms.ToTensor(),
    19. transforms.Normalize((0.1307,), (0.3081,))
    20. ])),
    21. batch_size=batch_size, shuffle=True)
    22. w1, b1 = torch.randn(200, 784, requires_grad=True),\
    23. torch.zeros(200, requires_grad=True)
    24. w2, b2 = torch.randn(200, 200, requires_grad=True),\
    25. torch.zeros(200, requires_grad=True)
    26. w3, b3 = torch.randn(10, 200, requires_grad=True),\
    27. torch.zeros(10, requires_grad=True)
    28. torch.nn.init.kaiming_normal_(w1)
    29. torch.nn.init.kaiming_normal_(w2)
    30. torch.nn.init.kaiming_normal_(w3)
    31. def forward(x):
    32. x = x@w1.t() + b1
    33. x = F.relu(x)
    34. x = x@w2.t() + b2
    35. x = F.relu(x)
    36. x = x@w3.t() + b3
    37. x = F.relu(x)
    38. return x
    39. optimizer = optim.SGD([w1, b1, w2, b2, w3, b3], lr=learning_rate)
    40. criteon = nn.CrossEntropyLoss()
    41. for epoch in range(epochs):
    42. for batch_idx, (data, target) in enumerate(train_loader):
    43. data = data.view(-1, 28*28)
    44. logits = forward(data)
    45. loss = criteon(logits, target)
    46. optimizer.zero_grad()
    47. loss.backward()
    48. # print(w1.grad.norm(), w2.grad.norm())
    49. optimizer.step()
    50. if batch_idx % 100 == 0:
    51. print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
    52. epoch, batch_idx * len(data), len(train_loader.dataset),
    53. 100. * batch_idx / len(train_loader), loss.item()))
    54. test_loss = 0
    55. correct = 0
    56. for data, target in test_loader:
    57. data = data.view(-1, 28 * 28)
    58. logits = forward(data)
    59. test_loss += criteon(logits, target).item()
    60. pred = logits.data.max(1)[1]
    61. correct += pred.eq(target.data).sum()
    62. test_loss /= len(test_loader.dataset)
    63. print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
    64. test_loss, correct, len(test_loader.dataset),
    65. 100. * correct / len(test_loader.dataset)))