notebooks/02_model_training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Malware Image Classification - Model Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Google Colab Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!git clone https://github.com/Silvano315/Malware-classification-model.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# change working directory\n",
    "\n",
    "import os \n",
    "\n",
    "os.chdir(\"Malware-classification-model\")\n",
    "os.getcwd()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from google.colab import drive\n",
    "drive.mount('/content/drive')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Move your Kaggle API to /root/.config/kaggle and /root/.kaggle/kaggle.json\n",
    "\n",
    "os.makedirs('/root/.kaggle', exist_ok = True)\n",
    "\n",
    "!cp /content/drive/MyDrive/Kaggle_api/kaggle.json /root/.config/kaggle.json\n",
    "!cp /content/drive/MyDrive/Kaggle_api/kaggle.json /root/.kaggle/kaggle.json"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Setup & Import"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'/Users/silvanoquarto/Desktop/PROJECTS/Malware-classification-model'"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# If you run this notebook locally --> change working directory\n",
    "\n",
    "import os \n",
    "\n",
    "os.chdir(\"..\")\n",
    "os.getcwd()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import yaml\n",
    "import logging\n",
    "import torch\n",
    "import pandas as pd\n",
    "from pathlib import Path\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import transforms, datasets\n",
    "\n",
    "from src.utils.kaggle_downloader import setup_dataset\n",
    "from src.utils.logger_setup import get_logger\n",
    "from src.models.model_factory import create_model, validate_model_config, get_available_models\n",
    "from src.training.trainer import ModelTrainer\n",
    "from src.training.experiment import Experiment\n",
    "from src.training.callbacks import EarlyStopping, ModelCheckpoint, ReduceLROnPlateau\n",
    "from src.visualization.plot_results import scatter_plot_metrics, plot_confusion_matrix, plot_misclassified_images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setup logging\n",
    "\n",
    "logger = get_logger(ch_log_level=logging.INFO, fh_log_level=logging.DEBUG)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load config\n",
    "\n",
    "with open('config/config.yaml', 'r') as f:\n",
    "    config = yaml.safe_load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using device: cpu\n"
     ]
    }
   ],
   "source": [
    "# Setup CUDA\n",
    "\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "print(f\"Using device: {device}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setup paths\n",
    "\n",
    "DATA_PATH = Path(config['paths']['data']['raw_path'])\n",
    "MODELS_PATH = Path(config['paths']['data']['models_path'])\n",
    "MODELS_PATH.mkdir(parents=True, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training parameters\n",
    "\n",
    "BATCH_SIZE = config['training']['batch_size']\n",
    "NUM_EPOCHS = config['training']['num_epochs']\n",
    "LEARNING_RATE = config['training']['learning_rate']\n",
    "NUM_CLASSES = config['dataset']['num_classes']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Dataset e DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download Malimg dataset from Kaggle\n",
    "\n",
    "dataset_path = setup_dataset()\n",
    "print(f\"Dataset found at: {dataset_path}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define transforms\n",
    "\n",
    "mean_tf = config['preprocessing']['image']['mean'] * 3\n",
    "std_tf = config['preprocessing']['image']['std'] * 3\n",
    "\n",
    "train_transform = transforms.Compose([\n",
    "    transforms.Resize(config['preprocessing']['image']['size']),\n",
    "    transforms.Grayscale(num_output_channels=3),\n",
    "    #transforms.RandomRotation(20),\n",
    "    #transforms.RandomHorizontalFlip(),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(\n",
    "        #mean=config['preprocessing']['image']['mean'],\n",
    "        #std=config['preprocessing']['image']['std']\n",
    "        mean=mean_tf,\n",
    "        std=std_tf\n",
    "    )\n",
    "])\n",
    "\n",
    "val_transform = transforms.Compose([\n",
    "    transforms.Resize(config['preprocessing']['image']['size']),\n",
    "    transforms.Grayscale(num_output_channels=3),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(\n",
    "        #mean=config['preprocessing']['image']['mean'],\n",
    "        #std=config['preprocessing']['image']['std']\n",
    "        mean=mean_tf,\n",
    "        std=std_tf\n",
    "    )\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create datasets\n",
    "\n",
    "train_dataset = datasets.ImageFolder(root=f\"{config['paths']['data']['raw_path']}/train\", transform=train_transform)\n",
    "\n",
    "val_dataset = datasets.ImageFolder(root=f\"{config['paths']['data']['raw_path']}/val\", transform=val_transform)\n",
    "\n",
    "test_dataset = datasets.ImageFolder(root=f\"{config['paths']['data']['raw_path']}/test\", transform=val_transform)\n",
    "\n",
    "CLASS_NAMES = test_dataset.classes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create datasets\n",
    "\n",
    "train_loader = DataLoader(\n",
    "    train_dataset,\n",
    "    batch_size = BATCH_SIZE,\n",
    "    shuffle=True,\n",
    "    #num_workers=4,\n",
    "    #pin_memory=True\n",
    ")\n",
    "\n",
    "val_loader = DataLoader(\n",
    "    val_dataset,\n",
    "    batch_size = BATCH_SIZE,\n",
    "    shuffle=False,\n",
    "    #num_workers=4,\n",
    "    #pin_memory=True\n",
    ")\n",
    "\n",
    "test_loader = DataLoader(\n",
    "    test_dataset,\n",
    "    batch_size = BATCH_SIZE,\n",
    "    shuffle=False,\n",
    "    #num_workers=4,\n",
    "    #pin_memory=True\n",
    ")\n",
    "\n",
    "print(f\"Train samples: {len(train_dataset)}\")\n",
    "print(f\"Validation samples: {len(val_dataset)}\")\n",
    "print(f\"Test samples: {len(test_dataset)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Model Definition & Setup Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model configuration --> Baseline model\n",
    "\n",
    "model_config = {\n",
    "    'type': 'baseline',\n",
    "    'num_classes': 25,\n",
    "    'input_channels': 1\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model configuration --> Transfer learning model\n",
    "\n",
    "model_config = {\n",
    "    'type': 'transfer',\n",
    "    'model_name': 'resnet50', \n",
    "    'num_classes': NUM_CLASSES,\n",
    "    'pretrained': True,\n",
    "    'use_custom_classifier': True\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Give a look at every avialable model\n",
    "\n",
    "get_available_models()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Validate configuration\n",
    "\n",
    "validate_model_config(model_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create model\n",
    "\n",
    "model = create_model(model_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optimizer e Loss\n",
    "\n",
    "optimizer = torch.optim.Adam(\n",
    "    model.parameters(),\n",
    "    lr=LEARNING_RATE\n",
    ")\n",
    "\n",
    "criterion = torch.nn.CrossEntropyLoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-12-14 11:57:40,888 - training - INFO - Creating new experiment\n"
     ]
    }
   ],
   "source": [
    "# Setup experiment\n",
    "\n",
    "experiment = Experiment(\n",
    "    name=\"malware_classification_v1\",       # maybe use the name of the model \n",
    "    root=\"/content/drive/MyDrive/Malimg_dataset_and_data/experiments\",\n",
    "    logger=logger\n",
    ")\n",
    "experiment.init()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setup callbacks\n",
    "\n",
    "callbacks = [\n",
    "    EarlyStopping(\n",
    "        monitor='val_loss',\n",
    "        patience=config['training']['early_stopping']['patience'],\n",
    "        min_delta=config['training']['early_stopping']['min_delta'],\n",
    "        verbose=True\n",
    "    ),\n",
    "    ModelCheckpoint(\n",
    "        filepath='best_baseline_model.pth',\n",
    "        monitor='val_loss',\n",
    "        save_best_only=True,\n",
    "        verbose=1\n",
    "    ),\n",
    "    ReduceLROnPlateau(\n",
    "        optimizer=optimizer,\n",
    "        mode='min',\n",
    "        patience=5,\n",
    "        factor=0.1,\n",
    "        verbose=True\n",
    "    )\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize trainer\n",
    "\n",
    "trainer = ModelTrainer(\n",
    "    model=model,\n",
    "    optimizer=optimizer,\n",
    "    criterion=criterion,\n",
    "    experiment=experiment,\n",
    "    device=device,\n",
    "    logger=logger\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Model Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train the model\n",
    "\n",
    "trained_model = trainer.train(\n",
    "    train_loader=train_loader,\n",
    "    val_loader=val_loader,\n",
    "    num_epochs=NUM_EPOCHS,\n",
    "    callbacks=callbacks\n",
    ")\n",
    "\n",
    "torch.save(trained_model.state_dict(), experiment.root / 'final_model.pth')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Validation on Test Set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evaluation on test set\n",
    "\n",
    "test_logs = trainer.validate(test_loader)\n",
    "experiment.save_history('test', **test_logs)\n",
    "logger.info(f\"Test Results: {test_logs}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get predictions on test set\n",
    "\n",
    "test_targets, test_predictions = trainer.predict(test_loader)\n",
    "\n",
    "plot_confusion_matrix(test_targets, test_predictions, classes = CLASS_NAMES)\n",
    "logger.info(\"Confusion matrix saved as 'confusion_matrix.png'\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save test results\n",
    "test_results = {\n",
    "    'targets': test_targets.tolist(),\n",
    "    'predictions': test_predictions.tolist()\n",
    "}\n",
    "\n",
    "with open(f\"{experiment.results_dir}/test_results.json\", 'w') as f:\n",
    "    json.dump(test_results, f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Analysis of Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate and Save plots training history\n",
    "\n",
    "experiment.plot_history()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evaluation train and validation results\n",
    "\n",
    "scatter_plot_metrics(f'{experiment.root}/history/train.csv', \n",
    "                     f'{experiment.root}/history/val.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Replace missing values with 0 in column lr from val.csv and test.csv (TO BE REFACTORED)\n",
    "\n",
    "val = pd.read_csv(f\"{experiment.root}/history/val.csv\")\n",
    "val['lr'] = val['lr'].fillna(0).to_numpy()\n",
    "val.to_csv(f\"{experiment.root}/history/val.csv\", index=False)\n",
    "\n",
    "test = pd.read_csv(f\"{experiment.root}/history/test.csv\")\n",
    "test['lr'] = test['lr'].fillna(0).to_numpy()\n",
    "test.to_csv(f\"{experiment.root}/history/test.csv\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Calculate average metrics for last n epochs\n",
    "\n",
    "experiment = Experiment(\"malware_classification_baseline_v1\", \"/content/drive/MyDrive/Malimg_dataset_and_data/experiments\")\n",
    "experiment.load_history_from_file(\"val\")\n",
    "experiment.load_history_from_file(\"train\")\n",
    "experiment.load_history_from_file(\"test\")\n",
    "\n",
    "avg_metrics = experiment.calculate_average_metrics('val', last_n_epochs=5)\n",
    "print(\"Average validation metrics:\", avg_metrics)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Export results in JSON\n",
    "\n",
    "experiment.export_results_to_json(\"/content/drive/MyDrive/Malimg_dataset_and_data/experiments/malware_classification_baseline_v1/results/results.json\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Find best epoch according to validation accuracy \n",
    "\n",
    "metric = 'accuracy'\n",
    "\n",
    "best_epoch = experiment.get_best_epoch(metric, mode='max')\n",
    "print(f\"Best validation accuracy was achieved at epoch {best_epoch} with\n",
    "       {100*experiment.history['val'][metric][best_epoch-1]:.1f}%\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot learning rate\n",
    "\n",
    "experiment.plot_learning_rate(experiment.history['train']['lr'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot misclassified images with ground truth and prediction\n",
    "\n",
    "plot_misclassified_images(\n",
    "    model=trained_model,\n",
    "    dataloader=test_loader,\n",
    "    device=device,\n",
    "    num_images=16,\n",
    "    class_names=CLASS_NAMES,\n",
    "    mean=(0.5, 0.5, 0.5),\n",
    "    std=(0.5, 0.5, 0.5)\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}