cnn.py

import sys
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import TensorDataset, DataLoader
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt
import logging

# Configuring logging to track execution progress and results
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        # Defining feature extraction layers
        self.features = nn.Sequential(
            nn.Conv2d(1, 6, 5),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
            nn.Conv2d(6, 16, 5),
            nn.ReLU(),
            nn.MaxPool2d(2, 2)
        )
        # Defining classification layers
        self.classifier = nn.Sequential(
            nn.Linear(16 * 4 * 4, 120),
            nn.ReLU(),
            nn.Linear(120, 84),
            nn.ReLU(),
            nn.Linear(84, 10)
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.features(x)
        x = x.view(-1, 16 * 4 * 4)
        x = self.classifier(x)
        return x

def visualize_data(X: np.ndarray, T: np.ndarray) -> None:
    # Displaying sample images and class distribution
    plt.figure(figsize=(10, 5))
    for i in range(5):
        plt.subplot(1, 5, i + 1)
        plt.imshow(X[i].squeeze(), cmap='gray')
        plt.title(f"Label: {T[i]}")
    plt.show()

    unique, counts = np.unique(T, return_counts=True)
    plt.figure(figsize=(10, 5))
    plt.bar(unique, counts)
    plt.title('Class Distribution')
    plt.xlabel('Class')
    plt.ylabel('Count')
    plt.show()

def correct_data(X: torch.Tensor, T: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    # Normalizing pixel values and computing sample weights
    X = X.float() / 255.0
    return X, T

def train_and_evaluate(model: nn.Module, train_loader: DataLoader,
                       test_loader: DataLoader, criterion: nn.Module,
                       optimizer: optim.Optimizer, epochs: int) -> list[float]:
    test_accuracies = []
    
    for epoch in range(epochs):
        model.train()
        
        for inputs, labels in train_loader:
            if inputs.dim() == 5:
                inputs = inputs.squeeze(-1)
            
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
        
        # Evaluating model performance on test data
        model.eval()
        correct = 0
        total = 0
        
        with torch.no_grad():
            for inputs, labels in test_loader:
                if inputs.dim() == 5:
                    inputs = inputs.squeeze(-1)
                
                outputs = model(inputs)
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()
        
        accuracy = 100 * correct / total
        test_accuracies.append(accuracy)
        logging.info(f'Epoch {epoch + 1}, Test Accuracy: {accuracy:.2f}%')
    
    return test_accuracies

def compute_confusion_matrix(model: nn.Module,
                             test_loader: DataLoader) -> np.ndarray:
    # Computing confusion matrix for model evaluation
    model.eval()
    all_preds = []
    all_labels = []
    
    with torch.no_grad():
        for inputs, labels in test_loader:
            if inputs.dim() == 5:
                inputs = inputs.squeeze(-1)
            
            outputs = model(inputs)
            _, predicted = torch.max(outputs.data, 1)
            all_preds.extend(predicted.numpy())
            all_labels.extend(labels.numpy())
    
    cm = [[0 for _ in range(10)] for _ in range(10)]
    for t, p in zip(all_labels, all_preds):
        cm[t][p] += 1
    
    return np.array(cm)

def main(npz_file: str, test_percentage: int) -> None:
    # Loading dataset from the provided file
    data = np.load(npz_file)
    X_raw, T_raw = data['X'], data['T']
    
    logging.info(f"Dataset shape: {X_raw.shape}")
    logging.info(f"Labels shape: {T_raw.shape}")
    
    # Splitting data into training and testing sets based on the test percentage
    X_train_raw, X_test_raw, T_train_raw, T_test_raw = train_test_split(
        X_raw,
        T_raw,
        test_size=test_percentage / 100,
        random_state=42
    )
    
    # Converting data into PyTorch tensors and normalizing pixel values
    X_train_raw = torch.from_numpy(X_train_raw).float().unsqueeze(1) / 255.0
    X_test_raw = torch.from_numpy(X_test_raw).float().unsqueeze(1) / 255.0
    T_train_raw = torch.from_numpy(T_train_raw).long()
    T_test_raw = torch.from_numpy(T_test_raw).long()

    # Creating PyTorch DataLoaders for training and testing sets
    train_loader = DataLoader(TensorDataset(X_train_raw, T_train_raw), batch_size=64, shuffle=True)
    test_loader = DataLoader(TensorDataset(X_test_raw, T_test_raw), batch_size=64)

    # Initializing the CNN model and training parameters
    model = CNN()
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    # Training and evaluating the model over multiple epochs
    logging.info("Starting training...")
    train_and_evaluate(model=model,
                       train_loader=train_loader,
                       test_loader=test_loader,
                       criterion=criterion,
                       optimizer=optimizer,
                       epochs=10)

if __name__ == "__main__":
    # Checking command-line arguments to ensure proper usage
    if len(sys.argv) != 3:
        logging.error("Usage: python cnn.py <npz_file> <test_percentage>")
        sys.exit(1)

    npz_file_path = sys.argv[1]
    test_percentage_value = int(sys.argv[2])
    
    main(npz_file_path, test_percentage_value)