class LeNet5(nn.Module):
def __init__(self):
super(LeNet5, self).__init__()
# define a 2D convolutional layer
self.hidden_1 = torch.nn.Conv2d(1,6,kernel_size=5,stride=1,padding=2)
# define a maxpool layer
self.hidden_2 = nn.MaxPool2d(2, stride=2)
# new 2D convolutional layer
self.hidden_3 = torch.nn.Conv2d(6,16,kernel_size=5,stride=1)
# another maxpool layer
self.hidden_4 = nn.MaxPool2d(2, stride=2)
# first linear layer
self.hidden_5 = nn.Linear(16*5*5,120, bias=True)
# second linear layer
self.hidden_6 = nn.Linear(120,84, bias=True)
# final output layer
self.output = nn.Linear(84, 10, bias=False)
# activation function
self.activation = nn.ReLU()
def forward(self, x):
for i in range(1,7):
if i==5:
x=x.flatten(start_dim=1)
# activate pass through the first layer
f = getattr(self,f"hidden_{i}")
x = self.activation(f(x))
return self.output(x) # return output
x = torch.randn((1, 1, 28, 28))
model = LeNet5()
y = model(x)
print(y)
print(model)