Initial commit

architectures/DenseNet121.py

@@ -0,0 +1,15 @@
+
+import torch.nn as nn
+from torchvision import models
+from architectures.Model import Model
+
+class DenseNet121(Model):
+    def get(self):
+
+        # load pretrained
+        m = models.densenet121(weights=models.DenseNet121_Weights.DEFAULT)
+        # will modify only the final layers
+        num_ftrs = m.classifier.in_features
+        m.classifier = nn.Linear(num_ftrs, self.size)
+        
+        return m.to(self.device)

architectures/Inception.py

@@ -0,0 +1,47 @@
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torchvision import models
+import time
+
+# Base model
+from architectures.Model import Model
+
+class Inception(Model):
+    def get(self):
+        model = models.inception_v3(weights=models.Inception_V3_Weights.DEFAULT)
+        #for param in model.parameters():
+        #    param.requires_grad = False
+        model.fc = nn.Linear(model.fc.in_features, self.size)
+        model.AuxLogits.fc = nn.Linear(model.AuxLogits.fc.in_features, self.size)
+        return model.to(self.device)
+
+    def train(self, epochs, loader, rate):
+        # Override because Inception returns a tuple (main, aux)
+        criterion = nn.CrossEntropyLoss()
+        optimizer = optim.Adam(filter(lambda p: p.requires_grad, self.model.parameters()), lr=rate)
+
+        print(f"Starting training on {self.device}...")
+        start_time = time.time()
+
+        self.model.train()
+
+        for epoch in range(epochs):
+            total_loss = 0.0
+            for inputs, labels in loader:
+                
+                inputs, labels = inputs.to(self.device), labels.to(self.device)
+                optimizer.zero_grad()
+                
+                outputs, aux_outputs = self.model(inputs)
+                loss = criterion(outputs, labels) + 0.3 * criterion(aux_outputs, labels)
+                
+                loss.backward()
+                optimizer.step()
+                total_loss += loss.item()
+
+            print(f"Epoch {epoch+1}/{epochs} | Loss: {total_loss/len(loader):.4f}")
+        
+        if self.device.type == 'cuda': torch.cuda.synchronize()
+        print(f"Training completed in: {time.time() - start_time:.2f}s")

architectures/Model.py

@@ -0,0 +1,99 @@
+from abc import ABC, abstractmethod
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import time
+import numpy as np
+from sklearn.metrics import classification_report
+
+class Model(ABC):
+    def __init__(self, device, size):
+        self.device = device
+        self.size = size
+        self.model = self.get()
+
+    @abstractmethod
+    def get(self):
+        # return the model
+        return self.model
+
+    def train(self, epochs, loader, rate):
+        criterion = nn.CrossEntropyLoss()
+        optimizer = optim.Adam(filter(lambda p: p.requires_grad, self.model.parameters()), lr=rate)
+        
+        print(f"Starting training on {self.device}...")
+        start_time = time.time()
+        self.model.train()
+
+        for epoch in range(epochs):
+            total_loss = 0.0
+            for inputs, labels in loader:
+                inputs, labels = inputs.to(self.device), labels.to(self.device)
+                optimizer.zero_grad()
+                outputs = self.model(inputs)
+                loss = criterion(outputs, labels)
+                loss.backward()
+                optimizer.step()
+                total_loss += loss.item()
+            
+            print(f"Epoch {epoch+1}/{epochs} | Loss: {total_loss / len(loader):.4f}")
+
+        if self.device.type == 'cuda': torch.cuda.synchronize()
+        print(f"Training completed in: {time.time() - start_time:.2f}s")
+
+    def evaluate(self, loader):
+        self.model.eval()
+        all_preds, all_labels = [], []
+        print("\nEvaluating...")
+        
+        with torch.no_grad():
+            for inputs, labels in loader:
+                inputs, labels = inputs.to(self.device), labels.to(self.device)
+                outputs = self.model(inputs)
+                _, predicted = torch.max(outputs, 1)
+                all_preds.extend(predicted.cpu().numpy())
+                all_labels.extend(labels.cpu().numpy())
+
+        accuracy = 100 * (np.array(all_preds) == np.array(all_labels)).sum() / len(all_labels)
+        print(f"Test Accuracy: {accuracy:.2f}%")
+        print(classification_report(all_labels, all_preds, zero_division=0))
+
+
+    # Using the factory patern here
+    @staticmethod
+    def create(arch, device, size):
+        print(f'>> MODEL ARCHITECTURE  >> {arch.name}.')
+
+        match arch:
+
+            # ResNet18
+            case Architecture.RESNET18: 
+                from architectures.ResNet18 import ResNet18
+                return ResNet18(device, size)
+            
+            # ResNet50
+            case Architecture.RESNET50:
+                from architectures.ResNet18 import ResNet18
+                return ResNet18(device, size)
+            
+            # INCEPTION
+            case Architecture.INCEPTION:
+                from architectures.Inception import Inception
+                return Inception(device, size)
+            
+            # DENSENET121
+            case Architecture.DENSENET121:
+                from architectures.DenseNet121 import DenseNet121
+                return DenseNet121(device, size)
+            case _:
+                raise ValueError(f"Unknown model: {arch}")
+
+ 
+# model architectures
+from enum import Enum, auto
+
+class Architecture(Enum):
+    RESNET18 = auto()
+    RESNET50 = auto()
+    INCEPTION = auto()
+    DENSENET121 = auto()

architectures/ResNet18.py

@@ -0,0 +1,22 @@
+
+import torch.nn as nn
+from torchvision import models
+
+# Base model
+from architectures.Model import Model
+
+class ResNet18(Model):
+    
+    def get(self):
+        m = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
+
+        # freez all layers
+        for param in m.parameters():
+            param.requires_grad = False
+
+        # unfreez the last two
+        for param in m.layer3.parameters(): param.requires_grad = True
+        for param in m.layer4.parameters(): param.requires_grad = True
+        
+        m.fc = nn.Linear(m.fc.in_features, self.size)
+        return m.to(self.device)

architectures/ResNet50.py

@@ -0,0 +1,22 @@
+
+import torch.nn as nn
+from torchvision import models
+
+# Base model
+from architectures.Model import Model
+
+class ResNet50(Model):
+    
+    def get(self):
+        m = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
+        
+        # freez all layers
+        for param in m.parameters():
+            param.requires_grad = False
+
+        # unfreez the last two
+        for param in m.layer3.parameters(): param.requires_grad = True
+        for param in m.layer4.parameters(): param.requires_grad = True
+
+        m.fc = nn.Linear(m.fc.in_features, self.size)
+        return m.to(self.device)