Обнавлена структура проекта
This commit is contained in:
@@ -0,0 +1,38 @@
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user