-
Notifications
You must be signed in to change notification settings - Fork 3
/
lenet.py
39 lines (36 loc) · 1.89 KB
/
lenet.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
33
34
35
36
37
38
39
import torch
from torch import nn
class LeNet5(nn.Module):
def __init__(self, n_channels, inputsize, outputsize):
super(LeNet5, self).__init__()
# Convolution (In LeNet-5, 32x32 images are given as input. Hence padding of 2 is done below)
self.conv1 = torch.nn.Conv2d(in_channels=n_channels, out_channels=6, kernel_size=5, stride=1, padding=int((32-inputsize)/2), bias=True)
# Max-pooling
self.max_pool_1 = torch.nn.MaxPool2d(kernel_size=2)
# Convolution
self.conv2 = torch.nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5, stride=1, padding=0, bias=True)
# Max-pooling
self.max_pool_2 = torch.nn.MaxPool2d(kernel_size=2)
# Fully connected layer
self.fc1 = torch.nn.Linear(16*5*5, 120) # convert matrix with 16*5*5 (= 400) features to a matrix of 120 features (columns)
self.fc2 = torch.nn.Linear(120, 84) # convert matrix with 120 features to a matrix of 84 features (columns)
self.fc3 = torch.nn.Linear(84, outputsize) # convert matrix with 84 features to a matrix of 10 features (columns)
def forward(self, x):
# convolve, then perform ReLU non-linearity
x = torch.nn.functional.relu(self.conv1(x))
# max-pooling with 2x2 grid
x = self.max_pool_1(x)
# convolve, then perform ReLU non-linearity
x = torch.nn.functional.relu(self.conv2(x))
# max-pooling with 2x2 grid
x = self.max_pool_2(x)
# first flatten 'max_pool_2_out' to contain 16*5*5 columns
# read through https://stackoverflow.com/a/42482819/7551231
x = x.view(-1, 16*5*5)
# FC-1, then perform ReLU non-linearity
x = torch.nn.functional.relu(self.fc1(x))
# FC-2, then perform ReLU non-linearity
x = torch.nn.functional.relu(self.fc2(x))
# FC-3
x = self.fc3(x)
return x