39 lines
913 B
Python
39 lines
913 B
Python
import torch.nn as nn
|
|
|
|
|
|
class LeNetCustom(nn.Module):
|
|
|
|
def __init__(self,
|
|
in_channels: int=1,
|
|
out_classes: int=10):
|
|
super().__init__()
|
|
|
|
self._out_classes=out_classes
|
|
|
|
self.feature = nn.Sequential(
|
|
nn.Conv2d(in_channels, 6, kernel_size=5),
|
|
nn.Tanh(),
|
|
nn.AvgPool2d(kernel_size=2, stride=2),
|
|
|
|
nn.Conv2d(6, 16, kernel_size=5),
|
|
nn.Tanh(),
|
|
nn.AvgPool2d(kernel_size=2, stride=2),
|
|
|
|
nn.Conv2d(16, 120, kernel_size=5),
|
|
nn.Tanh(),
|
|
)
|
|
|
|
|
|
self.mlp = nn.Sequential(
|
|
nn.LazyLinear(84),
|
|
nn.Tanh(),
|
|
nn.Linear(84, out_classes)
|
|
)
|
|
|
|
def forward(self, image):
|
|
x = self.feature(image)
|
|
x = self.mlp(x.flatten(1))
|
|
|
|
return x
|
|
|