-
Notifications
You must be signed in to change notification settings - Fork 1
/
lenet5.py
32 lines (26 loc) · 914 Bytes
/
lenet5.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
import torch.nn as nn
class LeNet5(nn.Module):
def __init__(self):
super(LeNet5, self).__init__()
self.convnet = nn.Sequential(
nn.Conv2d(1, 6, kernel_size=5, stride=1, padding=0),
nn.BatchNorm2d(6),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),
nn.Conv2d(6, 16, kernel_size=5, stride=1, padding=0),
nn.BatchNorm2d(16),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),
nn.Conv2d(16, 120, kernel_size=5, stride=1, padding=0),
nn.BatchNorm2d(120),
nn.ReLU())
self.fc = nn.Sequential(
nn.Linear(120, 84),
nn.ReLU(),
nn.Linear(84, 10),
nn.LogSoftmax(dim=-1))
def forward(self, x):
out = self.convnet(x)
out = out.view(x.size(0), -1)
out = self.fc(out)
return out