lenet实现

Author: yuangu

lenet/main.py

@@ -0,0 +1,103 @@
+#-*- coding:utf-8-*-
+
+import torch
+from torch.utils.data import Dataset, DataLoader
+
+from moudle import LeNet
+from torchvision.datasets.mnist import MNIST
+import torchvision.transforms as transforms
+import time
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+batch_size=1 
+epoch_num=5 #2-0.8323, 5-0.9545
+LR = 0.001
+
+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()]))
+
+
+
+train_loader = DataLoader( 
+  dataset= data_train,  #CustomDataset(), 
+  batch_size=batch_size, # 批大小 
+ 
+  shuffle=True, # 是否随机打乱顺序 
+  num_workers=8, # 多线程读取数据的线程数 
+  ) 
+
+test_loader = DataLoader( 
+  dataset= data_test,  #CustomDataset(), 
+  batch_size=batch_size, # 批大小 
+ 
+  shuffle=True, # 是否随机打乱顺序 
+  num_workers=8, # 多线程读取数据的线程数 
+  ) 
+
+
+net = LeNet().to(device)
+
+opt = torch.optim.SGD(net.parameters(), lr=LR)
+
+loss_function = torch.nn.CrossEntropyLoss()
+
+def train():
+  net.train()
+  for epoch in range(epoch_num):
+      total_loss = 0
+      epoch_step = 0
+      tic = time.time()
+
+      for batch_image, batch_label in train_loader:
+          batch_image = batch_image.to(device)
+          batch_label  = batch_label.to(device)
+
+          opt.zero_grad() 
+          output = net(batch_image)
+          loss = loss_function(output, batch_label) 
+          
+          total_loss += loss
+          epoch_step += 1
+          
+          loss.backward() 
+          opt.step()
+        
+      toc = time.time()    
+      print("one epoch does take approximately " + str((toc - tic)) + " seconds),average loss: " + str(total_loss/epoch_step))
+
+#torch.save(net.state_dict(), "./moudle/moudle")      
+
+def test():
+  net.eval()
+  total_correct = 0
+  for batch_image, batch_label in test_loader:
+    batch_image = batch_image.to(device)
+    batch_label  = batch_label.to(device)
+
+    output = net(batch_image)
+    
+    pred = output.detach().max(1)[1]
+    total_correct += pred.eq(batch_label.view_as(pred)).sum()
+
+  print("total_correct:", float(total_correct) / len(data_test))
+
+
+
+if __name__ == "__main__":
+  train()
+  test()
+     

lenet/moudle.py

@@ -0,0 +1,44 @@
+#-*-coding:utf-8-*-
+
+'''
+《深度学习-卷积神经网络从入门到精通》中的lenet5实现（p43 - p44），
+但是原书公式的效果极差，所以相对原书做了如下修改：
+1. 将sigmoid改换成Relu
+2. 增加一个84 -> 10的全连接层
+3. 将平均池化换成最大池化
+'''
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.autograd import Variable
+
+class LeNet(nn.Module):
+    def __init__(self):
+        super(LeNet, self).__init__()
+
+        self.conv1 = nn.Conv2d(1, 6, 5, padding=2)
+        self.conv2 = nn.Conv2d(6, 16, 5)
+        self.conv3 = nn.Conv2d(16, 120, 5)
+        
+        self.fc1 =  nn.Linear(120, 84)
+        self.fc2 =  nn.Linear(84, 10)
+
+    def forward(self, input):
+        # 28 * 28 - > 28 * 28 -> 14 * 14
+        x = F.max_pool2d(F.relu(self.conv1(input)), 2, stride=2)
+        # 14 * 14 -> 10 * 10 - > 5 * 5
+        x = F.max_pool2d(F.relu(self.conv2(x)), 2, stride=2)
+        
+        x = F.relu(self.conv3(x))
+
+        x = x.view(x.size(0),-1)
+        
+        # 84 * 84 -> 10*10
+        x = self.fc2(F.relu(self.fc1(x)))
+        
+        x = F.softmax(x,dim=1)
+        #print(x.size)
+        return x
+        
+