LeNet-5 pytorch+torchvision+visdom

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

 # -*- coding: utf-8 -*-
 """
 Created on Sun May 26 22:53:52 2019

 @author: jiangshan
 """
 #A modified LeNet-5 [LeCun et al., 1998a] on the MNIST dataset.
 import torch
 import torch.nn as nn
 import torch.optim as optim
 from torchvision.datasets.mnist import MNIST
 import torchvision.transforms as transforms
 from torch.utils.data import DataLoader
 import visdom
 from collections import OrderedDict

 class LeNet5(nn.Module):
     """
     Input - 1x32x32
     C1 - 6@28x28 (5x5 kernel)
     relu
     S2 - 6@14x14 (2x2 kernel, stride 2) Subsampling
     C3 - 16@10x10 (5x5 kernel, complicated shit)
     relu
     S4 - 16@5x5 (2x2 kernel, stride 2) Subsampling
     C5 - 120@1x1 (5x5 kernel)
     F6 - 84
     relu
     F7 - 10 (Output)
     """
     def __init__(self):
         super(LeNet5, self).__init__()

         self.convnet = nn.Sequential(OrderedDict([
             ('c1', nn.Conv2d(1, 6, kernel_size=(5, 5))),
             ('relu1', nn.ReLU()),
             ('s2', nn.MaxPool2d(kernel_size=(2, 2), stride=2)),
             ('c3', nn.Conv2d(6, 16, kernel_size=(5, 5))),
             ('relu3', nn.ReLU()),
             ('s4', nn.MaxPool2d(kernel_size=(2, 2), stride=2)),
             ('c5', nn.Conv2d(16, 120, kernel_size=(5, 5))),
             ('relu5', nn.ReLU())
         ]))

         self.fc = nn.Sequential(OrderedDict([
             ('f6', nn.Linear(120, 84)),
             ('relu6', nn.ReLU()),
             ('f7', nn.Linear(84, 10)),
             ('sig7', nn.LogSoftmax(dim=-1))
         ]))

     def forward(self, img):
         output = self.convnet(img)
         output = output.view(img.size(0), -1)
         output = self.fc(output)
         return output

 viz = visdom.Visdom()
 data_train = MNIST('./data/mnist',
                    download=True,
                    transform=transforms.Compose([
                        transforms.Resize((32, 32)),
                        transforms.ToTensor()]))
 data_test = MNIST('./data/mnist',
                   train=False,
                   download=True,
                   transform=transforms.Compose([
                       transforms.Resize((32, 32)),
                       transforms.ToTensor()]))
 data_train_loader = DataLoader(data_train, batch_size=256, shuffle=True, num_workers=8)
 data_test_loader = DataLoader(data_test, batch_size=1024, num_workers=8)

 net = LeNet5()
 criterion = nn.CrossEntropyLoss()
 optimizer = optim.Adam(net.parameters(), lr=2e-3)

 cur_batch_win = None
 cur_batch_win_opts = {
     'title': 'Epoch Loss Trace',
     'xlabel': 'Batch Number',
     'ylabel': 'Loss',
     'width': 1200,
     'height': 600,
 }

 def train(epoch):
     global cur_batch_win
     net.train()
     loss_list, batch_list = [], []
     for i, (images, labels) in enumerate(data_train_loader):
         optimizer.zero_grad()

         output = net(images)

         loss = criterion(output, labels)

         loss_list.append(loss.detach().cpu().item())
         batch_list.append(i+1)

         if i % 10 == 0:
             print('Train - Epoch %d, Batch: %d, Loss: %f' % (epoch, i, loss.detach().cpu().item()))

         # Update Visualization
         if viz.check_connection():
             cur_batch_win = viz.line(torch.Tensor(loss_list), torch.Tensor(batch_list),
                                      win=cur_batch_win, name='current_batch_loss',
                                      update=(None if cur_batch_win is None else 'replace'),
                                      opts=cur_batch_win_opts)
         loss.backward()
         optimizer.step()

 def test():
     net.eval()
     total_correct = 0
     avg_loss = 0.0
     for i, (images, labels) in enumerate(data_test_loader):
         output = net(images)
         avg_loss += criterion(output, labels).sum()
         pred = output.detach().max(1)[1]
         total_correct += pred.eq(labels.view_as(pred)).sum()

     avg_loss /= len(data_test)
     print('Test Avg. Loss: %f, Accuracy: %f' % (avg_loss.detach().cpu().item(), float(total_correct) / len(data_test)))

 def train_and_test(epoch):
     train(epoch)
     test()

 def main():
     for e in range(1, 16):
         train_and_test(e)

 if __name__ == '__main__':
     main()

先开启visdom 进行可视化

python -m visdom.server

运行程序

python LeNet-5_main.py

打开浏览器查看live graph

http://localhost:8097