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

2026-09-28 15:40:05 +10:00
parent 53b9c6d854
commit 056dd87746
17 changed files with 671 additions and 733 deletions
@@ -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