Deep-Learning-Architectures/LeNet/LeNet.py at main · GSNCodes/Deep-Learning-Architectures · GitHub

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
import torch
import torch.nn as nn

class LeNet(nn.Module):

	def __init__(self):

		super(LeNet, self).__init__()
		self.relu = nn.ReLU()
		self.pool = nn.AvgPool2d(kernel_size=2, stride=2)
		self.conv1 = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5, stride=1, padding=0)
		self.conv2 = nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5, stride=1, padding=0)
		self.conv3 = nn.Conv2d(in_channels=16, out_channels=120, kernel_size=5, stride=1, padding=0)
		self.linear1 = nn.Linear(120, 84)
		self.linear2 = nn.Linear(84, 10)

	def forward(self, x):

		x = self.relu(self.conv1(x))
		x = self.pool(x)
		x = self.relu(self.conv2(x))
		x = self.pool(x)
		x = self.relu(self.conv3(x))
		x = x.reshape(x.shape[0], -1)
		x = self.relu(self.linear1(x))
		x = self.linear2(x)

		return x

if __name__ == '__main__':

	test_tensor = torch.randn(100, 1, 32, 32)
	model = LeNet()
	print(model(test_tensor).shape) # Expected Output:- torch.Size([100, 10])