LeNet-5 pytorch+torchvision+visdom

  # ====================LeNet-5_main.py===============
  # pytorch+torchvision+visdom

  1 # -*- coding: utf-8 -*-
  2 """
  3 Created on Sun May 26 22:53:52 2019
  4 
  5 @author: jiangshan
  6 """
  7 #A modified LeNet-5 [LeCun et al., 1998a] on the MNIST dataset.
  8 import torch
  9 import torch.nn as nn
 10 import torch.optim as optim
 11 from torchvision.datasets.mnist import MNIST
 12 import torchvision.transforms as transforms
 13 from torch.utils.data import DataLoader
 14 import visdom
 15 from collections import OrderedDict
 16 
 17 class LeNet5(nn.Module):
 18     """
 19     Input - 1x32x32
 20     C1 - 6@28x28 (5x5 kernel)
 21     relu
 22     S2 - 6@14x14 (2x2 kernel, stride 2) Subsampling
 23     C3 - 16@10x10 (5x5 kernel, complicated shit)
 24     relu
 25     S4 - 16@5x5 (2x2 kernel, stride 2) Subsampling
 26     C5 - 120@1x1 (5x5 kernel)
 27     F6 - 84
 28     relu
 29     F7 - 10 (Output)
 30     """
 31     def __init__(self):
 32         super(LeNet5, self).__init__()
 33 
 34         self.convnet = nn.Sequential(OrderedDict([
 35             ('c1', nn.Conv2d(1, 6, kernel_size=(5, 5))),
 36             ('relu1', nn.ReLU()),
 37             ('s2', nn.MaxPool2d(kernel_size=(2, 2), stride=2)),
 38             ('c3', nn.Conv2d(6, 16, kernel_size=(5, 5))),
 39             ('relu3', nn.ReLU()),
 40             ('s4', nn.MaxPool2d(kernel_size=(2, 2), stride=2)),
 41             ('c5', nn.Conv2d(16, 120, kernel_size=(5, 5))),
 42             ('relu5', nn.ReLU())
 43         ]))
 44 
 45         self.fc = nn.Sequential(OrderedDict([
 46             ('f6', nn.Linear(120, 84)),
 47             ('relu6', nn.ReLU()),
 48             ('f7', nn.Linear(84, 10)),
 49             ('sig7', nn.LogSoftmax(dim=-1))
 50         ]))
 51 
 52     def forward(self, img):
 53         output = self.convnet(img)
 54         output = output.view(img.size(0), -1)
 55         output = self.fc(output)
 56         return output
 57 
 58 
 59 viz = visdom.Visdom()
 60 data_train = MNIST('./data/mnist',
 61                    download=True,
 62                    transform=transforms.Compose([
 63                        transforms.Resize((32, 32)),
 64                        transforms.ToTensor()]))
 65 data_test = MNIST('./data/mnist',
 66                   train=False,
 67                   download=True,
 68                   transform=transforms.Compose([
 69                       transforms.Resize((32, 32)),
 70                       transforms.ToTensor()]))
 71 data_train_loader = DataLoader(data_train, batch_size=256, shuffle=True, num_workers=8)
 72 data_test_loader = DataLoader(data_test, batch_size=1024, num_workers=8)
 73 
 74 net = LeNet5()
 75 criterion = nn.CrossEntropyLoss()
 76 optimizer = optim.Adam(net.parameters(), lr=2e-3)
 77 
 78 cur_batch_win = None
 79 cur_batch_win_opts = {
 80     'title': 'Epoch Loss Trace',
 81     'xlabel': 'Batch Number',
 82     'ylabel': 'Loss',
 83     'width': 1200,
 84     'height': 600,
 85 }
 86 
 87 
 88 def train(epoch):
 89     global cur_batch_win
 90     net.train()
 91     loss_list, batch_list = [], []
 92     for i, (images, labels) in enumerate(data_train_loader):
 93         optimizer.zero_grad()
 94 
 95         output = net(images)
 96 
 97         loss = criterion(output, labels)
 98 
 99         loss_list.append(loss.detach().cpu().item())
100         batch_list.append(i+1)
101 
102         if i % 10 == 0:
103             print('Train - Epoch %d, Batch: %d, Loss: %f' % (epoch, i, loss.detach().cpu().item()))
104 
105         # Update Visualization
106         if viz.check_connection():
107             cur_batch_win = viz.line(torch.Tensor(loss_list), torch.Tensor(batch_list),
108                                      win=cur_batch_win, name='current_batch_loss',
109                                      update=(None if cur_batch_win is None else 'replace'),
110                                      opts=cur_batch_win_opts)
111         loss.backward()
112         optimizer.step()
113 
114 
115 def test():
116     net.eval()
117     total_correct = 0
118     avg_loss = 0.0
119     for i, (images, labels) in enumerate(data_test_loader):
120         output = net(images)
121         avg_loss += criterion(output, labels).sum()
122         pred = output.detach().max(1)[1]
123         total_correct += pred.eq(labels.view_as(pred)).sum()
124 
125     avg_loss /= len(data_test)
126     print('Test Avg. Loss: %f, Accuracy: %f' % (avg_loss.detach().cpu().item(), float(total_correct) / len(data_test)))
127 
128 
129 def train_and_test(epoch):
130     train(epoch)
131     test()
132 
133 
134 def main():
135     for e in range(1, 16):
136         train_and_test(e)
137 
138 
139 if __name__ == '__main__':
140     main()