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