Basic_GCN.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8fff9a1c5c727638",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Get a simple GCN Working"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca653fc46f8da2e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-31T13:01:01.906294Z",
     "start_time": "2024-05-31T13:00:57.445915Z"
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import pickle\n",
    "import numpy as np\n",
    "import networkx as nx\n",
    "import torch.nn.functional as F\n",
    "from torch_geometric.nn import GINConv, global_mean_pool\n",
    "from torch_geometric.loader import DataLoader\n",
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.metrics import confusion_matrix, classification_report"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac8795e392cef1da",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-31T13:01:04.107907Z",
     "start_time": "2024-05-31T13:01:03.808267Z"
    }
   },
   "outputs": [],
   "source": [
    "# Load data splits\n",
    "# dataset = 'fer2013'\n",
    "dataset = 'ck'\n",
    "train_data_path = dataset + '_data/train_data_70_20_10.pkl'\n",
    "val_data_path = dataset + '_data/val_data_70_20_10.pkl'\n",
    "test_data_path = dataset + '_data/test_data_70_20_10.pkl'\n",
    "\n",
    "with open(train_data_path, 'rb') as f:\n",
    "    train_data = pickle.load(f)\n",
    "with open(val_data_path, 'rb') as f:\n",
    "    val_data = pickle.load(f)\n",
    "with open(test_data_path, 'rb') as f:\n",
    "    test_data = pickle.load(f)\n",
    "\n",
    "adjacency_matrix = np.loadtxt('standard_mesh_adj_matrix.csv', delimiter=',')\n",
    "G = nx.from_numpy_array(adjacency_matrix)\n",
    "\n",
    "# Add batch attribute to each data object\n",
    "for data in train_data:\n",
    "    data.batch = torch.zeros(data.x.size(0), dtype=torch.long)\n",
    "for data in val_data:\n",
    "    data.batch = torch.zeros(data.x.size(0), dtype=torch.long)\n",
    "for data in test_data:\n",
    "    data.batch = torch.zeros(data.x.size(0), dtype=torch.long)\n",
    "\n",
    "class EarlyStopping:\n",
    "    def __init__(self, patience=10, delta=0):\n",
    "        self.patience = patience\n",
    "        self.delta = delta\n",
    "        self.best_score = None\n",
    "        self.early_stop = False\n",
    "        self.counter = 0\n",
    "\n",
    "    def __call__(self, val_loss, model):\n",
    "        if self.best_score is None:\n",
    "            self.best_score = val_loss\n",
    "            self.save_checkpoint(val_loss, model)\n",
    "        elif val_loss > self.best_score + self.delta:\n",
    "            self.counter += 1\n",
    "            if self.counter >= self.patience:\n",
    "                self.early_stop = True\n",
    "        else:\n",
    "            self.best_score = val_loss\n",
    "            self.save_checkpoint(val_loss, model)\n",
    "            self.counter = 0\n",
    "\n",
    "    def save_checkpoint(self, val_loss, model):\n",
    "        torch.save(model.state_dict(), 'checkpoint.pt')\n",
    "        self.val_loss_min = val_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22a5770c567363a6",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-31T13:03:48.299736Z",
     "start_time": "2024-05-31T13:01:21.419573Z"
    }
   },
   "outputs": [],
   "source": [
    "class SimpleGCN(torch.nn.Module):\n",
    "    def __init__(self, input_dim, hidden_dim, output_dim):\n",
    "        super(SimpleGCN, self).__init__()\n",
    "        nn1 = torch.nn.Sequential(torch.nn.Linear(input_dim, hidden_dim), torch.nn.ReLU(), torch.nn.Linear(hidden_dim, hidden_dim))\n",
    "        self.conv1 = GINConv(nn1)\n",
    "        nn2 = torch.nn.Sequential(torch.nn.Linear(hidden_dim, hidden_dim), torch.nn.ReLU(), torch.nn.Linear(hidden_dim, hidden_dim))\n",
    "        self.conv2 = GINConv(nn2)\n",
    "        nn2 = torch.nn.Sequential(torch.nn.Linear(hidden_dim, hidden_dim), torch.nn.ReLU(), torch.nn.Linear(hidden_dim, hidden_dim))\n",
    "        self.conv2 = GINConv(nn2)\n",
    "        self.lin = torch.nn.Linear(hidden_dim, output_dim)\n",
    "\n",
    "    def forward(self, data):\n",
    "        x, edge_index, batch = data.x.to(torch.float), data.edge_index.to(torch.int64), data.batch\n",
    "        x = self.conv1(x, edge_index)\n",
    "        x = F.relu(x)\n",
    "        x = self.conv2(x, edge_index)\n",
    "        x = F.relu(x)\n",
    "        x = global_mean_pool(x, batch)\n",
    "        x = self.lin(x)\n",
    "        return F.log_softmax(x, dim=1)\n",
    "\n",
    "# Convert data lists to DataLoader\n",
    "batch_size = 32\n",
    "\n",
    "train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True)\n",
    "val_loader = DataLoader(val_data, batch_size=batch_size, shuffle=False)\n",
    "test_loader = DataLoader(test_data, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "# Define the training and evaluation functions\n",
    "def train():\n",
    "    model.train()\n",
    "    total_loss = 0\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    for data in train_loader:\n",
    "        data = data.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        out = model(data)\n",
    "        loss = criterion(out, data.y)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        total_loss += loss.item()\n",
    "        pred = out.argmax(dim=1)\n",
    "        correct += pred.eq(data.y).sum().item()\n",
    "        total += data.y.size(0)\n",
    "    return total_loss / len(train_loader), correct / total\n",
    "\n",
    "def evaluate(loader):\n",
    "    model.eval()\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    val_loss = 0\n",
    "    all_preds = []\n",
    "    all_labels = []\n",
    "    with torch.no_grad():\n",
    "        for data in loader:\n",
    "            data = data.to(device)\n",
    "            out = model(data)\n",
    "            pred = out.argmax(dim=1)\n",
    "            correct += pred.eq(data.y).sum().item()\n",
    "            total += data.y.size(0)\n",
    "            val_loss += criterion(out, data.y).item()\n",
    "            all_preds.extend(pred.cpu().numpy())\n",
    "            all_labels.extend(data.y.cpu().numpy())\n",
    "    return correct / total, val_loss / len(loader), all_labels, all_preds\n",
    "\n",
    "# Get number of classes\n",
    "output_dim = len(np.unique([data.y.item() for data in train_data]))\n",
    "\n",
    "# Initialize model, optimizer, and criterion\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "model = SimpleGCN(input_dim=3, hidden_dim=64, output_dim=output_dim).to(device) \n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n",
    "\n",
    "# Calculate class weights\n",
    "label_counts = np.bincount([data.y.item() for data in train_data])\n",
    "class_weights = 1.0 / label_counts\n",
    "class_weights = class_weights / class_weights.sum()\n",
    "class_weights = torch.tensor(class_weights, dtype=torch.float).to(device)\n",
    "criterion = torch.nn.CrossEntropyLoss(weight=class_weights)\n",
    "\n",
    "early_stopping = EarlyStopping(patience=20, delta=0.001)\n",
    "\n",
    "train_losses, val_losses = [], []\n",
    "train_accuracies, val_accuracies = [], []\n",
    "\n",
    "# Training loop\n",
    "for epoch in range(1, 501):\n",
    "    train_loss, train_acc = train()\n",
    "    val_acc, val_loss, _, _ = evaluate(val_loader)\n",
    "    train_losses.append(train_loss)\n",
    "    val_losses.append(val_loss)\n",
    "    train_accuracies.append(train_acc)\n",
    "    val_accuracies.append(val_acc)\n",
    "    print(f'Epoch: {epoch:03d}, Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}, '\n",
    "          f'Train Acc: {int(100 * train_acc):02d}%, Val Acc: {int(100 * val_acc):02d}%')\n",
    "    \n",
    "    early_stopping(val_loss, model)\n",
    "    if early_stopping.early_stop:\n",
    "        print(\"Early stopping\")\n",
    "        break\n",
    "\n",
    "# Load the last checkpoint with the best model\n",
    "model.load_state_dict(torch.load('checkpoint.pt'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cce6309c00f4b40c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-31T13:03:55.638485Z",
     "start_time": "2024-05-31T13:03:54.130111Z"
    }
   },
   "outputs": [],
   "source": [
    "# Evaluation on test set\n",
    "test_acc, test_loss, test_labels, test_preds = evaluate(test_loader)\n",
    "print(f'Test Accuracy: {100 * test_acc:.2f}%')\n",
    "\n",
    "plt.figure(figsize=(20, 5))\n",
    "\n",
    "# Plotting training and validation loss\n",
    "plt.subplot(1, 3, 1)\n",
    "plt.plot(train_losses, label='Training Loss')\n",
    "plt.plot(val_losses, label='Validation Loss')\n",
    "plt.xlabel('Epochs')\n",
    "plt.ylabel('Loss')\n",
    "plt.title('Training and Validation Loss Over Epochs')\n",
    "plt.legend()\n",
    "plt.grid(True)\n",
    "\n",
    "# Plotting training and validation accuracy\n",
    "plt.subplot(1, 3, 2)\n",
    "plt.plot(train_accuracies, label='Training Accuracy')\n",
    "plt.plot(val_accuracies, label='Validation Accuracy')\n",
    "plt.xlabel('Epochs')\n",
    "plt.ylabel('Accuracy')\n",
    "plt.title('Training and Validation Accuracy Over Epochs')\n",
    "plt.legend()\n",
    "plt.grid(True)\n",
    "\n",
    "if dataset == 'ck':\n",
    "    label_mapping = {'neutral': 0, 'happiness': 1, 'sadness': 2, 'surprise': 3, 'fear': 4, 'disgust': 5, 'anger': 6, 'contempt': 7}\n",
    "elif dataset == 'fer2013':\n",
    "    label_mapping = {'neutral': 0, 'happy': 1, 'sad': 2, 'surprise': 3, 'fear': 4, 'disgust': 5, 'angry': 6}\n",
    "else:\n",
    "    raise ValueError('Invalid dataset')\n",
    "class_names = [name for name in label_mapping.keys()]\n",
    "\n",
    "# Plot confusion matrix\n",
    "conf_matrix = confusion_matrix(test_labels, test_preds, normalize='true')\n",
    "plt.subplot(1, 3, 3)\n",
    "sns.heatmap(conf_matrix, annot=True, fmt='.2f', cmap='Blues', xticklabels=class_names, yticklabels=class_names)\n",
    "plt.xlabel('Predicted')\n",
    "plt.ylabel('True')\n",
    "plt.title('Confusion Matrix')\n",
    "plt.show()\n",
    "\n",
    "# Print classification report\n",
    "class_report = classification_report(test_labels, test_preds, target_names=class_names, zero_division=0)\n",
    "print(class_report)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d162a52541f8e4ea",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}