Update

configs.py

@@ -7,16 +7,22 @@ import torch.nn as nn
 from torchvision.models import (
     squeezenet1_0,
     SqueezeNet1_0_Weights,
+    squeezenet1_1,
+    SqueezeNet1_1_Weights,
+    shufflenet_v2_x2_0,
+    ShuffleNet_V2_X2_0_Weights,
     mobilenet_v3_small,
     MobileNet_V3_Small_Weights,
 )
-from torchvision.models import squeezenet1_0
+
+import torch.nn.functional as F
+from pytorchcv.model_provider import get_model as ptcv_get_model
 
 # Constants
 RANDOM_SEED = 123
-BATCH_SIZE = 64
-NUM_EPOCHS = 100
-LEARNING_RATE = 0.00016633192288673592
+BATCH_SIZE = 16
+NUM_EPOCHS = 40
+LEARNING_RATE = 0.00016662575248025378
 STEP_SIZE = 10
 GAMMA = 0.9
 DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
@@ -28,6 +34,8 @@ EXTERNAL_DATA_DIR = r"data/train/external/Task "
 COMBINED_DATA_DIR = r"data/train/combined/Task "
 TEMP_DATA_DIR = "data/temp/"
 NUM_CLASSES = 7
+LABEL_SMOOTHING_EPSILON = 0.1
+MIXUP_ALPHA = 0.2
 EARLY_STOPPING_PATIENCE = 20
 CLASSES = [
     "Alzheimer Disease",
@@ -42,7 +50,7 @@ MODEL_SAVE_PATH = r"output/checkpoints/model.pth"
 
 
 class SqueezeNet1_0WithDropout(nn.Module):
-    def __init__(self, num_classes=1000):
+    def __init__(self, num_classes, dropout_prob=0.5):
         super(SqueezeNet1_0WithDropout, self).__init__()
         squeezenet = squeezenet1_0(weights=SqueezeNet1_0_Weights.DEFAULT)
         self.features = squeezenet.features
@@ -52,79 +60,94 @@ class SqueezeNet1_0WithDropout(nn.Module):
             nn.ReLU(inplace=True),
             nn.AdaptiveAvgPool2d((1, 1)),
         )
+        self.dropout = nn.Dropout(
+            dropout_prob
+        )  # Add dropout layer with the specified probability
 
     def forward(self, x):
         x = self.features(x)
         x = self.classifier(x)
+        x = F.dropout(x, training=self.training)  # Apply dropout during training
         x = torch.flatten(x, 1)
         return x
 
 
-# class ShuffleNetV2WithDropout(nn.Module):
-#     def __init__(self, num_classes=1000):
-#         super(ShuffleNetV2WithDropout, self).__init__()
-#         shufflenet = shufflenet_v2_x2_0(weights=ShuffleNet_V2_X2_0_Weights)
-#         self.features = shufflenet.features
-#         self.classifier = nn.Sequential(
-#             nn.Conv2d(1024, num_classes, kernel_size=1),
-#             nn.BatchNorm2d(num_classes),  # add batch normalization
-#             nn.ReLU(inplace=True),
-#             nn.AdaptiveAvgPool2d((1, 1))
-#         )
+class SqueezeNet1_1WithDropout(nn.Module):
+    def __init__(self, num_classes, dropout_prob=0.5):
+        super(SqueezeNet1_1WithDropout, self).__init__()
+        squeezenet = squeezenet1_1(weights=SqueezeNet1_1_Weights.DEFAULT)
+        self.features = squeezenet.features
+        self.classifier = nn.Sequential(
+            nn.Conv2d(512, num_classes, kernel_size=1),
+            nn.BatchNorm2d(num_classes),  # add batch normalization
+            nn.ReLU(inplace=True),
+            nn.AdaptiveAvgPool2d((1, 1)),
+        )
+        self.dropout = nn.Dropout(
+            dropout_prob
+        )  # Add dropout layer with the specified probability
 
-#     def forward(self, x):
-#         x = self.features(x)
-#         x = self.classifier(x)
-#         x = torch.flatten(x, 1)
-#         return x
+    def forward(self, x):
+        x = self.features(x)
+        x = self.classifier(x)
+        x = F.dropout(x, training=self.training)  # Apply dropout during training
+        x = torch.flatten(x, 1)
+        return x
 
 
-class MobileNetV3SmallWithDropout(nn.Module):
-    def __init__(self, num_classes=1000):
-        super(MobileNetV3SmallWithDropout, self).__init__()
-        mobilenet = mobilenet_v3_small(weights=MobileNet_V3_Small_Weights)
-        self.features = mobilenet.features
+class ShuffleNetV2WithDropout(nn.Module):
+    def __init__(self, num_classes, dropout_prob=0.5):
+        super(ShuffleNetV2WithDropout, self).__init__()
+        shufflenet = shufflenet_v2_x2_0(weights=ShuffleNet_V2_X2_0_Weights.DEFAULT)
+        self.features = shufflenet.features
         self.classifier = nn.Sequential(
-            nn.Conv2d(576, num_classes, kernel_size=1),
+            nn.Conv2d(1024, num_classes, kernel_size=1),
             nn.BatchNorm2d(num_classes),  # add batch normalization
             nn.ReLU(inplace=True),
             nn.AdaptiveAvgPool2d((1, 1)),
         )
+        self.dropout = nn.Dropout(
+            dropout_prob
+        )  # Add dropout layer with the specified probability
 
     def forward(self, x):
         x = self.features(x)
         x = self.classifier(x)
+        x = F.dropout(x, training=self.training)  # Apply dropout during training
         x = torch.flatten(x, 1)
         return x
 
 
-class ResNet18WithNorm(nn.Module):
-    def __init__(self, num_classes=1000):
-        super(ResNet18WithNorm, self).__init__()
-        resnet = resnet18(pretrained=False)
-        self.features = nn.Sequential(
-            *list(resnet.children())[:-2]
-        )  # Remove last 2 layers (avgpool and fc)
+class MobileNetV3SmallWithDropout(nn.Module):
+    def __init__(self, num_classes, dropout_prob=0.5):
+        super(MobileNetV3SmallWithDropout, self).__init__()
+        mobilenet = mobilenet_v3_small(weights=MobileNet_V3_Small_Weights.DEFAULT)
+        self.features = mobilenet.features
         self.classifier = nn.Sequential(
+            nn.Conv2d(576, num_classes, kernel_size=1),
+            nn.BatchNorm2d(num_classes),  # add batch normalization
+            nn.ReLU(inplace=True),
             nn.AdaptiveAvgPool2d((1, 1)),
-            nn.Flatten(),
-            nn.Linear(512, num_classes),
-            nn.BatchNorm1d(num_classes),  # Add batch normalization
         )
+        self.dropout = nn.Dropout(
+            dropout_prob
+        )  # Add dropout layer with the specified probability
 
     def forward(self, x):
         x = self.features(x)
         x = self.classifier(x)
+        x = F.dropout(x, training=self.training)  # Apply dropout during training
         x = torch.flatten(x, 1)
         return x
 
 
 MODEL = SqueezeNet1_0WithDropout(num_classes=7)
+# MODEL = ptcv_get_model("sqnxt23v5_w2", pretrained=False, num_classes=7)
 print(CLASSES)
 
 preprocess = transforms.Compose(
     [
-        transforms.Resize((112, 112)), # Resize to 112x112
+        transforms.Resize((274, 274)),  # Resize to 112x112
         transforms.ToTensor(),  # Convert to tensor
         transforms.Grayscale(num_output_channels=3),  # Convert to 3 channels
         # Normalize 3 channels

eval.py

@@ -1,30 +1,31 @@
 import os
 import torch
-from torchvision.transforms import transforms
+import numpy as np
 import pathlib
 from PIL import Image
-from torchmetrics import ConfusionMatrix, Accuracy, F1Score
 import matplotlib.pyplot as plt
 from sklearn.metrics import (
     classification_report,
     precision_recall_curve,
+    accuracy_score,
+    f1_score,
+    confusion_matrix,
+    ConfusionMatrixDisplay,
     roc_curve,
     auc,
-    confusion_matrix,
 )
+from sklearn.preprocessing import label_binarize
 from configs import *
 from data_loader import load_data  # Import the load_data function
-import numpy as np
 
 # Constants
 DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
 # Load the model
 MODEL = MODEL.to(DEVICE)
-MODEL.load_state_dict(torch.load(r"C:\Users\User\Documents\Flutter-test\handetect\assets\model.pt", map_location=DEVICE))
+MODEL.load_state_dict(torch.load(MODEL_SAVE_PATH, map_location=DEVICE))
 MODEL.eval()
 
-
 def predict_image(image_path, model, transform):
     model.eval()
     correct_predictions = 0
@@ -38,9 +39,6 @@ def predict_image(image_path, model, transform):
     predicted_labels = []
     predicted_scores = []  # To store predicted class probabilities
 
-    accuracy_metric = Accuracy(num_classes=NUM_CLASSES, task="multiclass")
-    f1_metric = F1Score(num_classes=NUM_CLASSES, task="multiclass")
-
     with torch.no_grad():
         for image_file in images:
             print("---------------------------")
@@ -67,22 +65,20 @@ def predict_image(image_path, model, transform):
                 correct_predictions += 1
 
     # Calculate accuracy and f1 score
-    accuracy = correct_predictions / total_predictions
+    accuracy = accuracy_score(true_classes, predicted_labels)
     print("Accuracy:", accuracy)
-    f1 = f1_metric(torch.tensor(predicted_labels), torch.tensor(true_classes)).item()
+    f1 = f1_score(true_classes, predicted_labels, average="weighted")
     print("Weighted F1 Score:", f1)
 
     # Convert the lists to tensors
     predicted_labels_tensor = torch.tensor(predicted_labels)
     true_classes_tensor = torch.tensor(true_classes)
 
-    # Create a confusion matrix
-    conf_matrix = ConfusionMatrix(num_classes=NUM_CLASSES, task="multiclass")
-    conf_matrix(predicted_labels_tensor, true_classes_tensor)
+    # Calculate the confusion matrix
+    conf_matrix = confusion_matrix(true_classes, predicted_labels)
 
     # Plot the confusion matrix
-    conf_matrix.compute()
-    conf_matrix.plot()
+    ConfusionMatrixDisplay(confusion_matrix=conf_matrix, display_labels=CLASSES).plot(cmap=plt.cm.Blues)
     plt.title("Confusion Matrix")
     plt.show()
 
@@ -92,8 +88,20 @@ def predict_image(image_path, model, transform):
         true_classes, predicted_labels, target_names=class_names
     )
     print("Classification Report:\n", report)
-    
-    
+
+    # Calculate precision and recall for each class
+    true_classes_binary = label_binarize(true_classes, classes=range(NUM_CLASSES))
+    precision, recall, _ = precision_recall_curve(
+        true_classes_binary.ravel(), np.array(predicted_scores).ravel()
+    )
+
+    # Plot precision-recall curve
+    plt.figure(figsize=(10, 6))
+    plt.plot(recall, precision)
+    plt.title("Precision-Recall Curve")
+    plt.xlabel("Recall")
+    plt.ylabel("Precision")
+    plt.show()
 
 
 # Call predict_image function with your image path

train.py

@@ -7,12 +7,44 @@ from models import *
 from torch.utils.tensorboard import SummaryWriter
 from configs import *
 import data_loader
+import torch.nn.functional as F
+import numpy as np
+
+
+class LabelSmoothingLoss(nn.Module):
+    def __init__(self, epsilon=0.1, num_classes=2):
+        super(LabelSmoothingLoss, self).__init__()
+        self.epsilon = epsilon
+        self.num_classes = num_classes
+
+    def forward(self, input, target):
+        target_smooth = (1 - self.epsilon) * target + self.epsilon / self.num_classes
+        return nn.CrossEntropyLoss()(input, target_smooth)
 
 
 def setup_tensorboard():
     return SummaryWriter(log_dir="output/tensorboard/training")
 
 
+def mixup_data(x, y, alpha=1.0):
+    """Returns mixed inputs, pairs of targets, and lambda"""
+    if alpha > 0:
+        lam = np.random.beta(alpha, alpha)
+    else:
+        lam = 1
+
+    batch_size = x.size()[0]
+    index = torch.randperm(batch_size)
+
+    mixed_x = lam * x + (1 - lam) * x[index, :]
+    y_a, y_b = y, y[index]
+    return mixed_x, y_a, y_b, lam
+
+
+def mixup_criterion(criterion, pred, y_a, y_b, lam):
+    return lam * criterion(pred, y_a) + (1 - lam) * criterion(pred, y_b)
+
+
 def load_and_preprocess_data():
     return data_loader.load_data(
         RAW_DATA_DIR + str(TASK),
@@ -26,7 +58,7 @@ def initialize_model_optimizer_scheduler():
     model = MODEL.to(DEVICE)
     criterion = nn.CrossEntropyLoss()
     optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
-    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=STEP_SIZE, gamma=GAMMA)
+    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=NUM_EPOCHS)
     return model, criterion, optimizer, scheduler
 
 
@@ -35,7 +67,7 @@ def plot_and_log_metrics(metrics_dict, step, writer, prefix="Train"):
         writer.add_scalar(f"{prefix}/{metric_name}", metric_value, step)
 
 
-def train_one_epoch(model, criterion, optimizer, train_loader, epoch):
+def train_one_epoch(model, criterion, optimizer, train_loader, epoch, alpha):
     model.train()
     running_loss = 0.0
     total_train = 0
@@ -44,11 +76,15 @@ def train_one_epoch(model, criterion, optimizer, train_loader, epoch):
     for i, (inputs, labels) in enumerate(train_loader, 0):
         inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)
         optimizer.zero_grad()
-        if model.__class__.__name__ == "GoogLeNet":
-            outputs = model(inputs).logits
-        else:
-            outputs = model(inputs)
-        loss = criterion(outputs, labels)
+
+        # Apply mixup
+        inputs, targets_a, targets_b, lam = mixup_data(inputs, labels, alpha)
+
+        outputs = model(inputs)
+
+        # Calculate mixup loss
+        loss = mixup_criterion(criterion, outputs, targets_a, targets_b, lam)
+
         loss.backward()
         optimizer.step()
         running_loss += loss.item()
@@ -107,7 +143,7 @@ def main_training_loop():
         print("Learning rate:", scheduler.get_last_lr()[0])
 
         avg_train_loss, train_accuracy = train_one_epoch(
-            model, criterion, optimizer, train_loader, epoch
+            model, criterion, optimizer, train_loader, epoch, MIXUP_ALPHA
         )
         AVG_TRAIN_LOSS_HIST.append(avg_train_loss)
         TRAIN_ACC_HIST.append(train_accuracy)
@@ -154,7 +190,7 @@ def main_training_loop():
                 )
             )
             break
-
+    MODEL_SAVE_PATH = "output/checkpoints/model.pth"
     # Ensure the parent directory exists
     os.makedirs(os.path.dirname(MODEL_SAVE_PATH), exist_ok=True)
     torch.save(model.state_dict(), MODEL_SAVE_PATH)

tuning.py

@@ -12,7 +12,7 @@ from torch.utils.tensorboard import SummaryWriter
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 EPOCHS = 10
 N_TRIALS = 1000
-TIMEOUT = 1800
+TIMEOUT = 14400  
 
 # Create a TensorBoard writer
 writer = SummaryWriter(log_dir="output/tensorboard/tuning")