Upload model_development.py

model_development.py

@@ -0,0 +1,647 @@
+# -*- coding: utf-8 -*-
+"""Model_Development
+
+Automatically generated by Colab.
+
+Original file is located at
+    https://colab.research.google.com/drive/1BPWPi-oYa82w42fjYxMY7SVOU2XwedyW
+"""
+
+import tensorflow as tf
+from tensorflow.keras.models import Sequential
+from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout
+from tensorflow.keras.preprocessing.image import ImageDataGenerator
+from tensorflow.keras.utils import image_dataset_from_directory
+import os
+import numpy as np
+import keras
+from keras import layers
+from tensorflow import data as tf_data
+import matplotlib.pyplot as plt
+
+!unzip "/content/drive/MyDrive/Caitlin Bodzy/Data/constellations.zip" -d "/content"
+
+
+
+# Preparing our data
+directory = "/content/content/constellations2"
+image_size = (128, 128)
+batch_size = 64
+
+train_ds, val_ds = image_dataset_from_directory(
+    directory,
+    validation_split=0.2,
+    subset="both",
+    seed=1337,
+    image_size=image_size,
+    batch_size=batch_size,
+
+)
+
+plt.figure(figsize=(10, 10))
+for images, labels in train_ds.take(1):
+    for i in range(9):
+        ax = plt.subplot(3, 3, i + 1)
+        plt.imshow(np.array(images[i]).astype("uint8"))
+        plt.title(int(labels[i]))
+        plt.axis("off")
+
+def make_model(input_shape, num_classes):
+    inputs = keras.Input(shape=input_shape)
+
+    # Entry block
+    x = layers.Rescaling(1.0 / 255)(inputs)
+    x = layers.Conv2D(128, 3, strides=2, padding="same")(x)
+    x = layers.BatchNormalization()(x)
+    x = layers.Activation("relu")(x)
+
+    previous_block_activation = x  # Set aside residual
+
+    for size in [256, 512, 728]:
+        x = layers.Activation("relu")(x)
+        x = layers.SeparableConv2D(size, 3, padding="same")(x)
+        x = layers.BatchNormalization()(x)
+
+        x = layers.Activation("relu")(x)
+        x = layers.SeparableConv2D(size, 3, padding="same")(x)
+        x = layers.BatchNormalization()(x)
+
+        x = layers.MaxPooling2D(3, strides=2, padding="same")(x)
+
+        # Project residual
+        residual = layers.Conv2D(size, 1, strides=2, padding="same")(
+            previous_block_activation
+        )
+        x = layers.add([x, residual])  # Add back residual
+        previous_block_activation = x  # Set aside next residual
+
+    x = layers.SeparableConv2D(1024, 3, padding="same")(x)
+    x = layers.BatchNormalization()(x)
+    x = layers.Activation("relu")(x)
+
+    x = layers.GlobalAveragePooling2D()(x)
+    if num_classes == 2:
+        units = 1
+    else:
+        units = num_classes
+
+    x = layers.Dropout(0.25)(x)
+    # We specify activation=None so as to return logits
+    outputs = layers.Dense(units, activation=None)(x)
+    return keras.Model(inputs, outputs)
+
+num_classes = 89
+model = make_model(input_shape=image_size + (3,), num_classes=num_classes)
+keras.utils.plot_model(model, show_shapes=True)
+
+epochs = 25
+
+callbacks = [
+    keras.callbacks.ModelCheckpoint("save_at_{epoch}.keras"),
+]
+model.compile(
+    optimizer='adam',
+    loss='sparse_categorical_crossentropy',
+    metrics=['accuracy']
+)
+model.fit(
+    train_ds,
+    epochs=epochs,
+    callbacks=callbacks,
+    validation_data=val_ds,
+
+)
+
+"""#Pytorch time"""
+
+import os
+import torch
+import pandas as pd
+from skimage import io, transform
+import numpy as np
+import matplotlib.pyplot as plt
+from torch.utils.data import Dataset, DataLoader
+from torchvision import transforms, utils
+from torchvision.datasets import ImageFolder
+from torchvision import datasets, models, transforms
+from tqdm import tqdm
+import torch.optim as optim
+import pickle
+
+transform = transforms.Compose(
+    [transforms.ToTensor(),
+     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
+
+# can use ImageFolderDataset
+path = "/content/content/constellations2"
+dataset = ImageFolder(path, transform = transform)
+
+train_dataset, test_dataset = torch.utils.data.random_split(dataset, [0.8, 0.2])
+
+batch_size = 64
+trainloader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size,
+                                          shuffle=True, num_workers=2)
+testloader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size,
+                                         shuffle=False, num_workers=2)
+
+def train_model(model, train_loader, val_loader, criterion, optimizer, num_epochs=10, device="cuda"):
+    """
+    Trains a PyTorch model and prints training + validation loss/accuracy each epoch.
+
+    Args:
+        model: nn.Module - your model
+        train_loader: DataLoader - training data
+        val_loader: DataLoader - validation data
+        criterion: loss function (e.g. nn.CrossEntropyLoss)
+        optimizer: optimizer (e.g. Adam, SGD)
+        num_epochs: int - number of training epochs
+        device: 'cuda' or 'cpu'
+    """
+
+    model.to(device)
+    # Initialize history dictionary
+    history = {
+        "train_loss": [],
+        "train_acc": [],
+        "val_loss": [],
+        "val_acc": []
+    }
+    best_accuracy = 0
+    for epoch in range(num_epochs):
+        # --- Training phase ---
+        model.train()  # set model to training mode
+        running_loss = 0.0
+        correct = 0
+        total = 0
+        print(f"\nEpoch [{epoch+1}/{num_epochs}]")
+        train_pbar = tqdm(train_loader, desc="Training", leave=False)
+        for inputs, labels in train_pbar:
+            # Move data to GPU/CPU
+            inputs, labels = inputs.to(device), labels.to(device)
+
+            # Zero gradients for this batch
+            optimizer.zero_grad()
+
+            # Forward pass
+            outputs = model(inputs)
+            loss = criterion(outputs, labels)
+
+            # Backward pass + optimization
+            loss.backward()
+            optimizer.step()
+
+            # Update training loss
+            running_loss += loss.item() * inputs.size(0)
+
+            # Get predictions (highest logit = predicted class)
+            _, predicted = torch.max(outputs, 1)
+            total += labels.size(0)
+            correct += (predicted == labels).sum().item()
+            # live update in progress bar
+            train_pbar.set_postfix(loss=loss.item())
+        epoch_train_loss = running_loss / len(train_loader.dataset)
+        epoch_train_acc = 100 * correct / total
+        history["train_loss"].append(epoch_train_loss)
+        history["train_acc"].append(epoch_train_acc)
+        # --- Validation phase ---
+        model.eval()
+        val_loss, val_correct, val_total = 0.0, 0, 0
+
+        val_pbar = tqdm(val_loader, desc="Validating", leave=False)
+        with torch.no_grad():
+            for inputs, labels in val_pbar:
+                inputs, labels = inputs.to(device), labels.to(device)
+                outputs = model(inputs)
+                loss = criterion(outputs, labels)
+
+                val_loss += loss.item() * inputs.size(0)
+                _, predicted = torch.max(outputs, 1)
+                val_total += labels.size(0)
+                val_correct += (predicted == labels).sum().item()
+
+                # live update in progress bar
+                val_pbar.set_postfix(loss=loss.item())
+
+        epoch_val_loss = val_loss / len(val_loader.dataset)
+        epoch_val_acc = 100 * val_correct / val_total
+        if epoch_val_acc > best_accuracy:
+          best_accuracy = epoch_val_acc
+          torch.save(model.state_dict(), "best_model_params.pt")
+        history["val_loss"].append(epoch_val_loss)
+        history["val_acc"].append(epoch_val_acc)
+        # --- Print results at end of epoch ---
+        print(f"Epoch [{epoch+1}/{num_epochs}] "
+              f"Train Loss: {epoch_train_loss:.4f}, Train Acc: {epoch_train_acc:.2f}% "
+              f"| Val Loss: {epoch_val_loss:.4f}, Val Acc: {epoch_val_acc:.2f}%")
+    print("\n✅ Training complete.")
+    return model, history
+
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class Net(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(3, 32, 5, stride=2, padding=2)  # downsample early
+        self.conv2 = nn.Conv2d(32, 64, 3, stride=2, padding=1)
+        self.conv3 = nn.Conv2d(64, 128, 3, stride=2, padding=1)
+        self.conv4 = nn.Conv2d(128, 256, 3, stride=2, padding=1)
+
+        # Global average pooling (output shape = batch × 256)
+        self.gap = nn.AdaptiveAvgPool2d((1, 1))
+
+        # Small fully connected head
+        self.fc1 = nn.Linear(256, 128)
+        self.fc2 = nn.Linear(128, 89)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        x = F.relu(self.conv2(x))
+        x = F.relu(self.conv3(x))
+        x = F.relu(self.conv4(x))
+        x = self.gap(x)              # -> (batch, 256, 1, 1)
+        x = torch.flatten(x, 1)      # -> (batch, 256)
+        x = F.relu(self.fc1(x))
+        x = self.fc2(x)
+        return x
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print("Using device:", device)
+model = Net().to(device)
+criterion = nn.CrossEntropyLoss()
+optimizer = optim.Adam(model.parameters(), lr=1e-3)
+num_epochs = 50
+model, history = train_model(model, trainloader, testloader, criterion, optimizer, num_epochs=num_epochs, device=device)
+
+with open('custom_cnn_history.pkl', 'wb') as f:
+    pickle.dump(history, f)
+
+"""Results for 50 epochs:
+
+Using device: cuda
+
+Epoch [1/50]
+Epoch [1/50] Train Loss: 4.3550, Train Acc: 1.81% | Val Loss: 4.0249, Val Acc: 2.54%
+
+Epoch [2/50]
+Epoch [2/50] Train Loss: 3.8108, Train Acc: 3.31% | Val Loss: 3.6244, Val Acc: 3.45%
+
+Epoch [3/50]
+Epoch [3/50] Train Loss: 3.4859, Train Acc: 6.98% | Val Loss: 3.3346, Val Acc: 5.63%
+
+Epoch [4/50]
+Epoch [4/50] Train Loss: 3.2355, Train Acc: 8.43% | Val Loss: 3.1397, Val Acc: 12.89%
+
+Epoch [5/50]
+Epoch [5/50] Train Loss: 3.0890, Train Acc: 11.51% | Val Loss: 3.0306, Val Acc: 12.16%
+
+Epoch [6/50]
+Epoch [6/50] Train Loss: 3.1845, Train Acc: 9.11% | Val Loss: 3.0034, Val Acc: 19.60%
+
+Epoch [7/50]
+Epoch [7/50] Train Loss: 2.8859, Train Acc: 15.45% | Val Loss: 2.6902, Val Acc: 15.97%
+
+Epoch [8/50]
+Epoch [8/50] Train Loss: 2.7958, Train Acc: 16.54% | Val Loss: 2.7449, Val Acc: 16.70%
+
+Epoch [9/50]
+Epoch [9/50] Train Loss: 2.5519, Train Acc: 22.34% | Val Loss: 2.7096, Val Acc: 17.79%
+
+Epoch [10/50]
+Epoch [10/50] Train Loss: 2.4776, Train Acc: 23.33% | Val Loss: 2.4054, Val Acc: 24.68%
+
+Epoch [11/50]
+Epoch [11/50] Train Loss: 2.2965, Train Acc: 29.50% | Val Loss: 2.2966, Val Acc: 23.96%
+
+Epoch [12/50]
+Epoch [12/50] Train Loss: 2.0906, Train Acc: 33.57% | Val Loss: 2.0627, Val Acc: 36.30%
+
+Epoch [13/50]
+Epoch [13/50] Train Loss: 2.0251, Train Acc: 37.52% | Val Loss: 1.9627, Val Acc: 37.57%
+
+Epoch [14/50]
+Epoch [14/50] Train Loss: 1.7805, Train Acc: 44.09% | Val Loss: 1.8681, Val Acc: 40.83%
+
+Epoch [15/50]
+Epoch [15/50] Train Loss: 1.7417, Train Acc: 43.04% | Val Loss: 1.5517, Val Acc: 51.36%
+
+Epoch [16/50]
+Epoch [16/50] Train Loss: 1.5492, Train Acc: 47.80% | Val Loss: 1.5065, Val Acc: 51.18%
+
+Epoch [17/50]
+Epoch [17/50] Train Loss: 1.5765, Train Acc: 48.39% | Val Loss: 1.4849, Val Acc: 52.99%
+
+Epoch [18/50]
+Epoch [18/50] Train Loss: 1.3129, Train Acc: 57.68% | Val Loss: 1.1928, Val Acc: 58.62%
+
+Epoch [19/50]
+Epoch [19/50] Train Loss: 1.6431, Train Acc: 46.35% | Val Loss: 1.3933, Val Acc: 54.81%
+
+Epoch [20/50]
+Epoch [20/50] Train Loss: 1.1105, Train Acc: 64.20% | Val Loss: 1.2303, Val Acc: 57.89%
+
+Epoch [21/50]
+Epoch [21/50] Train Loss: 1.1097, Train Acc: 61.53% | Val Loss: 0.9919, Val Acc: 65.88%
+
+Epoch [22/50]
+Epoch [22/50] Train Loss: 0.9873, Train Acc: 66.65% | Val Loss: 0.9837, Val Acc: 65.52%
+
+Epoch [23/50]
+Epoch [23/50] Train Loss: 0.8362, Train Acc: 70.00% | Val Loss: 1.1519, Val Acc: 62.61%
+
+Epoch [24/50]
+Epoch [24/50] Train Loss: 0.9181, Train Acc: 67.47% | Val Loss: 0.9100, Val Acc: 70.05%
+
+Epoch [25/50]
+Epoch [25/50] Train Loss: 1.0156, Train Acc: 65.52% | Val Loss: 0.9139, Val Acc: 67.70%
+
+Epoch [26/50]
+Epoch [26/50] Train Loss: 0.7847, Train Acc: 73.72% | Val Loss: 0.6174, Val Acc: 82.94%
+
+Epoch [27/50]
+Epoch [27/50] Train Loss: 0.8881, Train Acc: 70.41% | Val Loss: 1.0006, Val Acc: 60.80%
+
+Epoch [28/50]
+Epoch [28/50] Train Loss: 0.9419, Train Acc: 69.05% | Val Loss: 0.6462, Val Acc: 83.85%
+
+Epoch [29/50]
+Epoch [29/50] Train Loss: 0.6375, Train Acc: 78.84% | Val Loss: 0.6832, Val Acc: 76.23%
+
+Epoch [30/50]
+Epoch [30/50] Train Loss: 0.5708, Train Acc: 78.66% | Val Loss: 0.6045, Val Acc: 78.04%
+
+Epoch [31/50]
+Epoch [31/50] Train Loss: 0.5209, Train Acc: 81.15% | Val Loss: 0.8075, Val Acc: 76.77%
+
+Epoch [32/50]
+Epoch [32/50] Train Loss: 0.6596, Train Acc: 77.03% | Val Loss: 0.5404, Val Acc: 82.21%
+
+Epoch [33/50]
+Epoch [33/50] Train Loss: 0.4934, Train Acc: 84.05% | Val Loss: 0.3702, Val Acc: 87.66%
+
+Epoch [34/50]
+Epoch [34/50] Train Loss: 0.5477, Train Acc: 79.97% | Val Loss: 0.4291, Val Acc: 85.48%
+
+Epoch [35/50]
+Epoch [35/50] Train Loss: 0.4817, Train Acc: 83.55% | Val Loss: 0.6402, Val Acc: 76.23%
+
+Epoch [36/50]
+Epoch [36/50] Train Loss: 0.7414, Train Acc: 73.95% | Val Loss: 0.7658, Val Acc: 74.23%
+
+Epoch [37/50]
+Epoch [37/50] Train Loss: 0.4875, Train Acc: 83.64% | Val Loss: 0.5194, Val Acc: 79.13%
+
+Epoch [38/50]
+Epoch [38/50] Train Loss: 0.7001, Train Acc: 74.94% | Val Loss: 0.4669, Val Acc: 82.40%
+
+Epoch [39/50]
+Epoch [39/50] Train Loss: 0.4260, Train Acc: 86.27% | Val Loss: 0.4313, Val Acc: 84.21%
+
+Epoch [40/50]
+Epoch [40/50] Train Loss: 0.4358, Train Acc: 84.05% | Val Loss: 1.0168, Val Acc: 66.79%
+
+Epoch [41/50]
+Epoch [41/50] Train Loss: 0.4003, Train Acc: 85.50% | Val Loss: 0.4285, Val Acc: 85.12%
+
+Epoch [42/50]
+Epoch [42/50] Train Loss: 0.3837, Train Acc: 86.91% | Val Loss: 0.6751, Val Acc: 78.40%
+
+Epoch [43/50]
+Epoch [43/50] Train Loss: 0.4011, Train Acc: 85.46% | Val Loss: 1.9752, Val Acc: 55.54%
+
+Epoch [44/50]
+Epoch [44/50] Train Loss: 0.9230, Train Acc: 73.18% | Val Loss: 0.4821, Val Acc: 84.94%
+
+Epoch [45/50]
+Epoch [45/50] Train Loss: 0.3531, Train Acc: 89.13% | Val Loss: 0.3267, Val Acc: 90.38%
+
+Epoch [46/50]
+Epoch [46/50] Train Loss: 0.3086, Train Acc: 89.40% | Val Loss: 0.3978, Val Acc: 85.30%
+
+Epoch [47/50]
+Epoch [47/50] Train Loss: 0.3021, Train Acc: 89.49% | Val Loss: 0.3236, Val Acc: 88.02%
+
+Epoch [48/50]
+Epoch [48/50] Train Loss: 0.3974, Train Acc: 85.27% | Val Loss: 0.3593, Val Acc: 87.11%
+
+Epoch [49/50]
+Epoch [49/50] Train Loss: 0.2615, Train Acc: 90.39% | Val Loss: 0.2154, Val Acc: 92.38%
+
+Epoch [50/50]
+                                                                     Epoch [50/50] Train Loss: 0.2437, Train Acc: 92.03% | Val Loss: 0.3268, Val Acc: 88.38%
+
+✅ Training complete.
+"""
+
+import matplotlib.pyplot as plt
+
+# summarize history for accuracy
+plt.plot(history['train_acc'])
+plt.plot(history['val_acc'])
+plt.title('custom cnn model accuracy')
+plt.ylabel('accuracy')
+plt.xlabel('epoch')
+plt.legend(['Train', 'Validation'], loc='upper left')
+plt.show()
+# summarize history for loss
+plt.plot(history['train_loss'])
+plt.plot(history['val_loss'])
+plt.title('custom cnn model loss')
+plt.ylabel('loss')
+plt.xlabel('epoch')
+plt.legend(['Train', 'Validation'], loc='upper left')
+plt.show()
+
+"""#resnet"""
+
+#Different model
+model_ft = models.resnet18(weights='IMAGENET1K_V1')
+num_ftrs = model_ft.fc.in_features
+# Here the size of each output sample is set to 2.
+# Alternatively, it can be generalized to ``nn.Linear(num_ftrs, len(class_names))``.
+model_ft.fc = nn.Linear(num_ftrs, 89)
+
+from torchvision.datasets import ImageFolder
+import torchvision.transforms as transforms
+import torch
+
+# Standard ResNet preprocessing
+image_transforms = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406],
+                         [0.229, 0.224, 0.225])  # ImageNet stats
+])
+batch_size = 64
+# can use ImageFolderDataset
+path = "/content/content/constellations2"
+dataset = ImageFolder(path, transform = image_transforms)
+train_dataset, test_dataset = torch.utils.data.random_split(dataset, [0.8, 0.2])
+trainloader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size,
+                                          shuffle=True, num_workers=2)
+testloader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size,
+                                         shuffle=False, num_workers=2)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print("Using device:", device)
+model = Net().to(device)
+criterion = nn.CrossEntropyLoss()
+optimizer = optim.Adam(model.parameters(), lr=1e-3)
+model, history = train_model(model_ft, trainloader, testloader, criterion, optimizer, num_epochs=50, device=device)
+
+# summarize history for accuracy
+plt.plot(history['train_acc'])
+plt.plot(history['val_acc'])
+plt.title('resnet model accuracy')
+plt.ylabel('accuracy')
+plt.xlabel('epoch')
+plt.legend(['Train', 'Validation'], loc='upper left')
+plt.show()
+# summarize history for loss
+plt.plot(history['train_loss'])
+plt.plot(history['val_loss'])
+plt.title('resnet model loss')
+plt.ylabel('loss')
+plt.xlabel('epoch')
+plt.legend(['Train', 'Validation'], loc='upper left')
+plt.show()
+
+"""Chat gpt improved resnet18 model"""
+
+# ==== Imports ====
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torchvision import models, transforms
+from torchvision.datasets import ImageFolder
+from torch.utils.data import DataLoader, random_split
+from tqdm import tqdm
+import matplotlib.pyplot as plt
+
+# ==== 1. Device setup ====
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print("Using device:", device)
+
+# ==== 2. Data setup with augmentations ====
+path = "/content/content/constellations2"
+
+image_transforms = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.RandomHorizontalFlip(),
+    transforms.RandomRotation(10),
+    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406],
+                         [0.229, 0.224, 0.225])  # ImageNet normalization
+])
+
+dataset = ImageFolder(path, transform=image_transforms)
+
+# Split dataset into train/test
+train_size = int(0.8 * len(dataset))
+test_size = len(dataset) - train_size
+train_dataset, test_dataset = random_split(dataset, [train_size, test_size])
+
+# DataLoaders
+batch_size = 64
+trainloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2)
+testloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=2)
+
+# ==== 3. Model setup ====
+model_ft = models.resnet18(weights='IMAGENET1K_V1')
+num_ftrs = model_ft.fc.in_features
+model_ft.fc = nn.Linear(num_ftrs, 89)  # 89 output classes
+model_ft = model_ft.to(device)
+
+# ==== 4. Loss and optimizer ====
+criterion = nn.CrossEntropyLoss()
+optimizer = optim.Adam(model_ft.parameters(), lr=1e-4, weight_decay=1e-4)
+scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5)
+
+# ==== 5. Training function ====
+def train_model(model, train_loader, val_loader, criterion, optimizer, scheduler, num_epochs, device):
+    history = {'train_loss': [], 'val_loss': [], 'train_acc': [], 'val_acc': []}
+
+    for epoch in range(num_epochs):
+        print(f"\nEpoch [{epoch+1}/{num_epochs}]")
+        model.train()
+        train_loss, correct, total = 0.0, 0, 0
+
+        for inputs, labels in tqdm(train_loader, desc="Training", leave=False):
+            inputs, labels = inputs.to(device), labels.to(device)
+            optimizer.zero_grad()
+            outputs = model(inputs)
+            loss = criterion(outputs, labels)
+            loss.backward()
+            optimizer.step()
+
+            train_loss += loss.item() * inputs.size(0)
+            _, predicted = torch.max(outputs, 1)
+            total += labels.size(0)
+            correct += (predicted == labels).sum().item()
+
+        scheduler.step()
+
+        train_acc = 100 * correct / total
+        val_loss, val_acc = evaluate(model, val_loader, criterion, device)
+
+        history['train_loss'].append(train_loss / len(train_loader.dataset))
+        history['val_loss'].append(val_loss)
+        history['train_acc'].append(train_acc)
+        history['val_acc'].append(val_acc)
+
+        print(f"Train Loss: {train_loss/len(train_loader.dataset):.4f} | "
+              f"Train Acc: {train_acc:.2f}% | "
+              f"Val Loss: {val_loss:.4f} | "
+              f"Val Acc: {val_acc:.2f}%")
+
+    return model, history
+
+# ==== 6. Validation function ====
+def evaluate(model, loader, criterion, device):
+    model.eval()
+    loss_total, correct, total = 0.0, 0, 0
+    with torch.no_grad():
+        for inputs, labels in loader:
+            inputs, labels = inputs.to(device), labels.to(device)
+            outputs = model(inputs)
+            loss = criterion(outputs, labels)
+            loss_total += loss.item() * inputs.size(0)
+            _, predicted = torch.max(outputs, 1)
+            total += labels.size(0)
+            correct += (predicted == labels).sum().item()
+    avg_loss = loss_total / len(loader.dataset)
+    acc = 100 * correct / total
+    return avg_loss, acc
+
+# ==== 7. Train model ====
+model_ft, history = train_model(
+    model_ft,
+    trainloader,
+    testloader,
+    criterion,
+    optimizer,
+    scheduler,
+    num_epochs=30,
+    device=device
+)
+
+# ==== 8. Visualization ====
+plt.figure(figsize=(10,5))
+plt.plot(history['train_acc'], label='Train Accuracy')
+plt.plot(history['val_acc'], label='Validation Accuracy')
+plt.title('ResNet18 Model Accuracy')
+plt.xlabel('Epoch')
+plt.ylabel('Accuracy (%)')
+plt.legend()
+plt.show()
+
+plt.figure(figsize=(10,5))
+plt.plot(history['train_loss'], label='Train Loss')
+plt.plot(history['val_loss'], label='Validation Loss')
+plt.title('ResNet18 Model Loss')
+plt.xlabel('Epoch')
+plt.ylabel('Loss')
+plt.legend()
+plt.show()