Обнавлена структура проекта

custom_arhitecture/__init__.py

@@ -0,0 +1,9 @@
+from .lenet import LeNetCustom
+from .alexnet import AlexNetCustom
+from .vgg import VGGCustom
+
+__all__ = [
+    'LeNetCustom',
+    'AlexNetCustom',
+    'VGGCustom'
+]

custom_arhitecture/alexnet.py

@@ -0,0 +1,53 @@
+import torch
+import torch.nn as nn
+
+
+class AlexNetCustom(nn.Module):
+    def __init__(self,
+                 in_channels: int = 3,
+                 out_classes: int=1_000):
+        super().__init__()
+        
+        self._out_classes = out_classes
+        
+        self.feature = nn.Sequential(
+            nn.Conv2d(in_channels, 96, kernel_size=11, stride=4, padding=2),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=3, stride=2),
+            
+            nn.Conv2d(96, 256, kernel_size=5, padding=2),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=3, stride=2),
+            
+            nn.Conv2d(256, 384, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            
+            nn.Conv2d(384, 384, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            
+            nn.Conv2d(384, 256, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=3, stride=2)
+        )
+        
+        self.avgpool = nn.AdaptiveAvgPool2d((6, 6)) 
+        
+        self.mlp = nn.Sequential(
+            nn.Dropout(p=0.5),
+            nn.Linear(256*6*6, 4096),
+            nn.ReLU(inplace=True),
+            
+            nn.Dropout(p=0.5),
+            nn.Linear(4096, 4096),
+            nn.ReLU(inplace=True),
+            
+            nn.Linear(4096, self._out_classes)
+        )
+        
+    def forward(self, image):
+        x = self.feature(image)
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+        x = self.mlp(x)
+        
+        return x

custom_arhitecture/lenet.py

@@ -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
+    

custom_arhitecture/vgg.py

@@ -0,0 +1,84 @@
+import torch
+import torch.nn as nn
+
+
+VGG_CONFIG = {
+    "VGG16": [64, 64, "M",
+              128, 128, "M",
+              256, 256, 256, "M",
+              512, 512, 512, "M",
+              512, 512, 512, "M"],
+    "VGG19": [64, 64, "M",
+              128, 128, "M",
+              256, 256, 256, 256, "M",
+              512, 512, 512, 512, "M",
+              512, 512, 512, 512, "M"]
+}
+
+
+class VGGCustom(nn.Module):
+    def __init__(self,
+                 architecture:str="VGG16",
+                 num_channels:int = 3,
+                 num_classes:int=10,
+                 batch_normalization=False,
+                 dropout:float=0.5) -> None:
+
+        super().__init__()
+
+        if architecture not in VGG_CONFIG:
+            raise ValueError(f"Architecture {architecture} is invalid. ",
+                             f"Available architectures are :{list(VGG_CONFIG.keys())}")
+
+
+        self.batch_normalization = batch_normalization
+
+        self.features = self._make_layers(
+            VGG_CONFIG[architecture],
+            in_channels=num_channels
+        )
+
+        self.avgpool = nn.AdaptiveAvgPool2d((7, 7))
+
+        self.classifier = nn.Sequential(
+            nn.Linear(512*7*7, 4096),
+            nn.ReLU(True),
+            nn.Dropout(p=dropout),
+
+            nn.Linear(4096, 4096),
+            nn.ReLU(True),
+            nn.Dropout(p=dropout),
+
+            nn.Linear(4096, num_classes)
+        )
+
+    def _make_layers(self, cfg, in_channels):
+        layers = []
+
+        for layer in cfg:
+            if layer == 'M':
+                layers.append(nn.MaxPool2d(kernel_size=2, stride=2))
+            else:
+
+                layers.append(nn.Conv2d(in_channels=in_channels,
+                                        out_channels=layer,
+                                        kernel_size=3,
+                                        stride=1,
+                                        padding=1))
+
+                if self.batch_normalization:
+                    layers.append(nn.BatchNorm2d(layer))
+
+                layers.append(nn.ReLU(True))
+
+                in_channels = layer
+
+        return nn.Sequential(*layers)
+
+    def forward(self, image):
+        x = self.features(image)
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+        x = self.classifier(x)
+
+        return x