Add letnet.py

letnet.py

@@ -0,0 +1,61 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torchvision import datasets, transforms
+from torch.autograd import Variable
+
+class LeNet5(nn.Module):
+    def __init__(self):
+        super(LeNet5,self).__init__()
+        self.layer1 = nn.Sequential(
+            nn.Conv2d(3,16,5,1),
+            nn.ReLU(),
+            nn.MaxPool2d(2,2,0,1),
+            nn.Conv2d(16,32,5,1),
+            nn.ReLU(),
+            nn.MaxPool2d(2,2,0,1)
+        )
+        self.fc = nn.Sequential(
+            nn.Linear(32*5*5,120),
+            nn.Linear(120,84),
+            nn.Linear(84, 10)
+        )
+        
+    def forward(self, x):
+        x = self.layer1(x)
+        x = x.view(-1, 800)
+        x = self.fc(x)
+        return x
+
+
+lr = 0.01
+momentum = 0.5
+
+log_interval = 10
+epochs = 10
+batch_size = 64
+test_batch_size = 1000
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+model = LeNet5()
+optimizer = torch.optim.SGD(model.parameters(),lr=lr,momentum=momentum)
+
+train_loader = torch.utils.data.DataLoader(
+    datasets.MNIST('../data', train=True, download=True,
+                   transform=transforms.Compose([
+                       transforms.ToTensor(),
+                       transforms.Normalize((0.1307,), (0.3081,))
+                   ])),
+    batch_size=batch_size, shuffle=True)
+
+def train(epoch):  # 定义每个epoch的训练细节
+    model.train()  # 设置为trainning模式
+    for batch_idx, (data, target) in enumerate(train_loader):
+        data = data.to(device)
+        target = target.to(device)
+        data, target = Variable(data), Variable(target)
+        optimizer.zero_grad()  # 优化器梯度初始化为零
+        output = model(data)  # 把数据输入网络并得到输出，即进行前向传播
+        loss = F.cross_entropy(output, target)  # 交叉熵损失函数
+        loss.backward()  # 反向传播梯度
+        optimizer.step()  # 结束一次前传+反传之后，更新参数