6_Temporal_Graph_Neural_Networks.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Temporal Graph Networks\n",
    "\n",
    "\n",
    "Till now, we have been focused on the *static* graphs where the graph structures and the node features are *fixed* over time. However, there are domains where the graph changes over time. \n",
    "\n",
    "Temporal graphs can be divided in two categories:\n",
    "- **Static graphs with temporal signals:** The underlying graph structure does not change over time, but features and labels evolve over time.\n",
    "        <center><img src=\"images/static_structure_dynamic_features.png\" width=\"500\"></center>\n",
    "    \n",
    "A main example is for traffic forcasting where graphs are based on traffic sensor data (e.g. the PeMS dataset) where each sensor is a node, and the edges are road connections. The geographical distribution of sensors in PeMS is shown below:\n",
    "        \n",
    "<center><img src=\"images/pems.ppm\" width=\"400\"></center>\n",
    "\n",
    "- **Dynamic graphs with temporal signals:** The topology of the graph (the presence of nodes and edges), features, and labels evolve over time.\n",
    "<center><img src=\"images/dynamic_structure_dynamic_features.png\" width=\"500\"></center>\n",
    "A main example is in a social network where new edges are added when people make new friends, existing edges are removed when people stop being friends, and node features change as people change their attributes, e.g., when they change their career assuming that career is one of the node features.\n",
    "<center><img src=\"images/Dynamic_Graphs.png\" width=\"500\"></center>\n",
    "\n",
    "> Note:\\\n",
    ">Dynamic graphs can be divided into *discrete-time* and *continuous-time* categories as well. \n",
    "\n",
    "- A discrete-time dynamic graph (DTDG) is a sequence $[G^{(1)}, G^{(2)},...,G^{(\\tau)}]$ of graph snapshots where each $G^{(t)} = \\left(V^{(t)},A^{(t)},X^{(t)}\\right)$ has vertices $V^{(t)}$, adjacency matrix $A^{(t)}$ and feature matrix $X^{(t)}$. DTDGs mainly appear in applications where data is captured at reguarly-spaced intervals.\n",
    "\n",
    "<center><img src=\"images/DTDG.png\" width=\"700\"></center>\n",
    "<center><small>Image from https://graph-neural-networks.github.io/static/file/chapter15.pdf</small></center> \n",
    "\n",
    "- A continuous-time dynamic graph (CTDG) is a pair $\\left(G^{(t_0)},O\\right)$ where ${G^{(t_0)}=\\left(V^{(t_0)},A^{(t_0)},X^{(t_0)}\\right)}$ is a static initial graph at initial state time $t_0$ and $O$ is a sequence of temporal observations/events. Each observation is a tuple of the form *(event, event type,timestamp)* where *event type* can be a node or edge addition, node or edge deletion, node feature update, etc. *event* represents the actual event that happened, and *timestamp* is the time at which the event occured:\n",
    "\n",
    "<center><img src=\"images/CTDG.png\" width=\"400\"></center>\n",
    "<center><small>Image from https://arxiv.org/pdf/2404.18211v1</small></center>\n",
    "\n",
    "We focus on DTDG in this tutorial."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Combining GNNs with sequence models\n",
    "\n",
    "DTDGs are made up of several snapshots arranged in order over time, which can be treated as sequential data. Temporal patterns in DTDGs are identified by looking at the relationships between these snapshots. Recurrent Neural Networks (RNNs) are often combined with GNNs to create dynamic models for DTDGs. These combinations are generally grouped into two types: stacked architectures and integrated architectures.\n",
    "\n",
    "- **Stacked dynamic GNNs:** The most straightforward way to model a discrete dynamic graph is to have a separate GNN handle each snapshot of the graph and feed the output of each GNN to a time series component, such as an RNN. This is illustrated in the following Figure: \n",
    "\n",
    "<center><img src=\"images/stacked_DTDG.png\" width=\"400\"></center>\n",
    "<center><small>Image from https://arxiv.org/pdf/2404.18211v1</small></center>\n",
    "\n",
    "One of most well-known approaches in this cateogry is Waterfall Dynamic-GCN. In this architectures, a GCN is stacked with an LSTM per node. More specifically, at first separate GCNs (with same parameters) handle each snapshot of the graph and next the output of each GNN is sequentially given to a LSTM. In fact, a separate LSTM is used per node (although the weights across the LSTMs are shared). The architecture is illustaretd in the following Figure:\n",
    "<center><img src=\"images/waterfall.png\" width=\"700\"></center>\n",
    "<center><small>Image from https://arxiv.org/pdf/2005.07496</small></center>\n",
    "\n",
    "The figure shows a network working on sequences of four snapshots of a graphs composed\n",
    "of five vertices. The first GCN layer acts as four copies of a regular GCN layer, each one working on a snapshot of the sequence of the graphs. The output of this first layer is processed by the LSTM layer that acts as five copies of a LSTM, each one working on a nodes of the graphs.\n",
    "The final fully-coonected (FC) layer produces the $C$-class probability vector for each nodes of every snapshot of the sequence. This layer, which produces the $C$-class probability vector for each node and for each instant of the sequence, can be seen as 5 x 4 copies of a FC layer.\n",
    "\n",
    "- **Integrated dynamic GNNs**: \n",
    "\n",
    "Integrated DGNNs are networks that combine GNNs and RNNs in one layer and thus combine modelling of the\n",
    "spatial and the temporal domain in that one layer.\n",
    "\n",
    "One major break-through approach in this category is <i>**EvolveGCN**</i>. <u>EvolveGCN applies Temporal neural networks such as RNNs to the *GCN parameters* themselves.</u> Note that GCN parameters are considered temporal and not the embeddings. In EvolveGCN, the GCN evolves over time to produce relevant temporal node embeddings. The following figure illustrates a high-level view of EvolveGCN’s architecture to produce node embeddings for a static or dynamic graph with temporal signal:\n",
    "\n",
    "<center><img src=\"images/Evolvegcn.png\" width=\"700\"></center>\n",
    "<center><small>image from Labonne, Maxime. \"Hands-On Graph Neural Networks Using Python: Practical techniques and architectures for building powerful graph and deep learning apps with PyTorch\". Packt Publishing Ltd, 2023.</small></center>\n",
    "\n",
    "but how to use RNN-based models to update the parameters of GCN according to the timesteps? EvolveGCN proposed two similar architectures, which we introduce only one of them, manily EvolveGCN-H. The main idea is shown below:\n",
    "\n",
    "<center><img src=\"images/Evolvegcn-h.png\" width=\"700\"></center>\n",
    "\n",
    "EvolveGCN-H utilizes a Gated Recurrent Unit (GRU) in place of a standard RNN. The GRU, a simplified variant of the Long Short-Term Memory (LSTM) unit, offers similar performance with fewer parameters. In this architecture, the hidden state of the GRU corresponds to the weight matrix of the GCN. \n",
    "\n",
    "Let:\n",
    "- $H_t^{(l)}$ denote the node embeddings produced at the $l$-th layer and at the timestep $t$; (Note that $H_t^{(0)}=X$) \n",
    "- $W_{(t-1)}^{(l)}$ be the weight matrix for the GCN at layer $l$ and previous timestep $t-1$. \n",
    "\n",
    "At each time step $t$, the GRU takes the node embeddings from the previous layer, $H^{(l)}_t$, as input, and uses the GCN's weight matrix, $W^{(l)}_{t-1}$, as its hidden state. It then updates the $W$ matrix for layer $l$ at time $t$ as follows:\n",
    "\n",
    "\\begin{equation*}\n",
    "W_t^{(l)} = GRU(H_t^{(l)}, W_{t-1}^{(l)})\n",
    "\\end{equation*}\n",
    "\n",
    "The updated weight matrix is used to calculate the node embeddings for the $l+1$ layer:\n",
    "\n",
    "\\begin{equation*}\n",
    "H^{(l+1)}_t = GCN(A_t, H_t^{(l)}, W_t^{(l)})\n",
    "\\end{equation*}\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example: Using Temporal Graphs for Action based recognition in Pytorch\n",
    "\n",
    "\n",
    "Skeleton-based action recognition leverages human skeletal data to identify and classify actions or gestures. This approach is particularly effective because skeleton data provides a structured and abstract representation of body movements, capturing key joint positions and their interactions. \n",
    "\n",
    "<center><img src=\"images/skeleton_as_graph.png\" width=\"300\"></center>\n",
    "\n",
    "- Skeleton Data: The input consists of sequences of joint positions, typically represented as a set of coordinates for each joint in a human skeleton. Each joint's position is recorded over time, creating a sequence of frames that shows how the skeleton moves.\n",
    "- Graph Representation: In the graph-based approach, each joint is treated as a node in a graph, and edges represent the connections between joints (e.g., bones). This graph captures the spatial relationships and constraints between different parts of the body.\n",
    "- Spatio-Temporal Modeling:\n",
    "\n",
    "    - Spatial Analysis: To analyze the spatial relationships between joints, graph convolutional layers are used. These layers capture how the joints interact with each other within a single frame.\n",
    "    - Temporal Analysis: To understand how these interactions change over time, temporal convolutional layers process sequences of frames. This helps in learning the dynamic patterns of movement across time.\n",
    "\n",
    "Spatio-Temporal Graph Convolutional Networks (ST-GCNs) are a pioneering approach designed to handle data that involves both spatial and temporal dimensions, particularly suited for tasks like action recognition from skeleton data. \n",
    "In ST-GCNs, skeleton data is represented as a graph where joints (body parts) are nodes and their connections (bones) are edges. This graph structure captures the spatial relationships between different joints in a single frame.\n",
    "ST-GCNs are composed of several layers. In each layer, ST-GCNs use GCN to apply convolution operations on the graph structure, aggregating information from neighboring nodes (joints) to learn how each joint's position relates to others. This helps in capturing the spatial patterns of body movements. Since actions involve sequences of movements over time, ST-GCNs also include one-dimensional temporal convolutional layers. These layers process sequences of frames, allowing the model to learn how the spatial configuration of joints changes over time. Temporal convolutions help capture dynamic patterns and transitions in the motion data. The overall procedure of applying several layers of STGCN is shown below:\n",
    "<center><img src=\"images/stgcn.png\" width=\"800\"></center>\n",
    "\n",
    "In this section, we will develop a basic implementation of the ST-GCN model using PyTorch Geometric, focusing on skeleton-based action recognition with the UTD-MHAD dataset."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> The skeleton joint order in UTD-MAD dataset: \\\n",
    "    head,\n",
    "    shoulder_center,\n",
    "    spine,\n",
    "    hip_center,\n",
    "    left_shoulder,\n",
    "    left_elbow,\n",
    "    left_wrist,\n",
    "    left_hand,\n",
    "    right_shoulder,\n",
    "    right_elbow,\n",
    "    right_wrist,\n",
    "    right_hand,\n",
    "    left_hip,\n",
    "    left_knee,\n",
    "    left_ankle,\n",
    "    left_foot,\n",
    "    right_hip,\n",
    "    right_knee,\n",
    "    right_ankle,\n",
    "    right_foot,\n",
    "\n",
    "> UTD-MHAD dataset consists of 27 different actions: \\\n",
    "    right arm swipe to the left,\n",
    "    right arm swipe to the right,\n",
    "    right hand wave,\n",
    "    two hand front clap,\n",
    "    right arm throw,\n",
    "    cross arms in the chest,\n",
    "    basketball shoot,\n",
    "    right hand draw x,\n",
    "    right hand draw circle (clockwise),\n",
    "    right hand draw circle (counter clockwise),\n",
    "    draw triangle,\n",
    "    bowling (right hand),\n",
    "    front boxing,\n",
    "    baseball swing from right,\n",
    "    tennis right hand forehand swing,\n",
    "    arm curl (two arms),\n",
    "    tennis serve,\n",
    "    two hand push,\n",
    "    right hand knock on door,\n",
    "    right hand catch an object,\n",
    "    right hand pick up and throw,\n",
    "    jogging in place,\n",
    "    walking in place,\n",
    "    sit to stand,\n",
    "    stand to sit,\n",
    "    forward lunge (left foot forward),\n",
    "    squat (two arms stretch out)\n",
    "\n",
    "Each skeleton data is a 20 x 3 x num_frame matrix. Each row of a skeleton frame corresponds to three spatial coordinates of a joint. \n",
    "\n",
    "The UTD-MAD dataset features a total of 8 subjects performing a variety of actions. It includes 27 distinct action categories, with each subject performing four takes (repetitions) of each action. Specifically, there are about 1,280 action sequences captured across all subjects and actions. This diversity in subjects and takes provides a comprehensive foundation for developing and evaluating action recognition algorithms.\n",
    "\n",
    "We use even subjects (S2, S4, S6, S8) to train, Odd Subjects (S1, S3, S5, S7) to test"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from scipy.io import loadmat\n",
    "import torch\n",
    "from torch_geometric.data import Data\n",
    "import torch.nn as nn\n",
    "from torch_geometric.nn import GCNConv\n",
    "from random import shuffle\n",
    "import numpy as np\n",
    "import torch.nn.functional as F\n",
    "\n",
    "\n",
    "class GCNLayer(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels):\n",
    "        super(GCNLayer, self).__init__()\n",
    "        self.conv = GCNConv(in_channels, out_channels)\n",
    "\n",
    "    def forward(self, x, edge_index):\n",
    "        x = self.conv(x, edge_index)\n",
    "        x = F.relu(x)\n",
    "        return x\n",
    "\n",
    "class TemporalConv(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1):\n",
    "        super(TemporalConv, self).__init__()\n",
    "        padding = (kernel_size - 1) // 2\n",
    "        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=(kernel_size, 1), stride=(stride, 1), padding=(padding, 0))\n",
    "\n",
    "    def forward(self, x):\n",
    "        # Input x is expected to be of shape (batch_size, in_channels, num_frames, num_nodes)\n",
    "        return self.conv(x)\n",
    "    \n",
    "    \n",
    "class STGCNLayer(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1):\n",
    "        super(STGCNLayer, self).__init__()\n",
    "        self.gcn = GCNLayer(in_channels, out_channels)\n",
    "        self.tcn1 = TemporalConv(in_channels, in_channels, kernel_size, stride)\n",
    "        self.tcn2 = TemporalConv(out_channels, out_channels, kernel_size, stride)\n",
    "        self.relu = nn.ReLU()\n",
    "        self.bn = nn.BatchNorm2d(out_channels)  # Batch normalization\n",
    "        self.dropout = nn.Dropout(0.3)  # Dropout layer\n",
    "\n",
    "    def forward(self, x, edge_index):\n",
    "        # x is expected to be of shape (batch_size, num_frames, num_nodes, in_channels)\n",
    "        batch_size, num_frames, num_nodes, in_channels = x.size()\n",
    "\n",
    "        # Apply GCN\n",
    "        x = self.gcn(x, edge_index)\n",
    "\n",
    "        # Reshape x back to (batch_size, num_frames, num_nodes, out_channels)\n",
    "        out_channels = x.size(-1)\n",
    "        x = x.view(batch_size, num_frames, num_nodes, out_channels)\n",
    "\n",
    "        # Transpose for temporal convolution: (batch_size, out_channels, num_frames, num_nodes)\n",
    "        x = x.permute(0, 3, 1, 2)\n",
    "\n",
    "        # Apply Temporal Convolution\n",
    "        x = self.tcn2(x)\n",
    "\n",
    "        x = self.bn(x)  # Batch normalization\n",
    "        x = self.dropout(x)  # Dropout\n",
    "\n",
    "\n",
    "        # Apply ReLU activation\n",
    "        x = self.relu(x)\n",
    "\n",
    "        # Transpose back: (batch_size, num_frames, num_nodes, out_channels)\n",
    "        x = x.permute(0, 2, 3, 1)\n",
    "\n",
    "        return x\n",
    "\n",
    "class STGCN(nn.Module):\n",
    "    def __init__(self, in_channels, num_classes, num_nodes, num_layers=3, hidden_dim=64, kernel_size=3, stride=1):\n",
    "        super(STGCN, self).__init__()\n",
    "        self.layers = nn.ModuleList()\n",
    "        self.layers.append(STGCNLayer(in_channels, hidden_dim, kernel_size, stride))\n",
    "        for _ in range(1, num_layers):\n",
    "            self.layers.append(STGCNLayer(hidden_dim, hidden_dim, kernel_size, stride))\n",
    "\n",
    "        self.fc = nn.Linear(hidden_dim*num_nodes, num_classes)\n",
    "        self.num_nodes = num_nodes\n",
    "\n",
    "\n",
    "    def forward(self, x, edge_index):\n",
    "        for layer in self.layers:\n",
    "            x = layer(x, edge_index)\n",
    "\n",
    "        # Global pooling (e.g., mean pooling over time and nodes)\n",
    "        x = x.view(x.size(0), -1, self.num_nodes * x.size(-1))\n",
    "        x = x.max(axis=1)[0] # Average over time\n",
    "\n",
    "        # Fully connected layer for classification\n",
    "        x = self.fc(x)\n",
    "\n",
    "        return x\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test the defined modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [],
   "source": [
    "layer = GCNLayer(3, 64)\n",
    "x = torch.randn(1, 12, 100, 3) #(batch_size, num_frames, num_nodes, in_channels)\n",
    "edge_index = torch.randint(0, 100, (2, 40))\n",
    "output = layer(x, edge_index)\n",
    "assert output.shape == (1, 12, 100, 64), output.shape\n",
    "\n",
    "\n",
    "conv = TemporalConv(3, 64)\n",
    "x = torch.randn(1, 3, 100, 20) #(batch_size, in_channels, num_frames, num_nodes)\n",
    "output = conv(x)\n",
    "assert output.shape == (1, 64, 100, 20), output.shape\n",
    "\n",
    "\n",
    "layer = STGCNLayer(3, 64)\n",
    "x = torch.randn(1, 12, 100, 3) #(batch_size, num_frames, num_nodes, in_channels)\n",
    "edge_index = torch.randint(0, 100, (2, 40))\n",
    "output = layer(x, edge_index)\n",
    "assert output.shape == (1, 12, 100, 64), output.shape\n",
    "\n",
    "model = STGCN(in_channels=3, num_classes=27, num_nodes=20)\n",
    "x = torch.randn(1, 12, 100, 3) #(batch_size, num_frames, num_nodes, in_channels)\n",
    "edge_index = torch.randint(0, 100, (2, 40))\n",
    "output = model(x, edge_index)\n",
    "assert output.shape == (1, 27)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [],
   "source": [
    "train = []\n",
    "test = []\n",
    "for file in os.listdir(\"Skeleton/\"):\n",
    "    if file.endswith(\"mat\"):\n",
    "        if \"s2\" in file or \"s4\" in file or \"s6\" in file or \"s8\" in file:\n",
    "            train.append((loadmat(\"Skeleton/\"+file)['d_skel'], file.split(\"_\")[0]))\n",
    "        else:\n",
    "            test.append((loadmat(\"Skeleton/\"+file)['d_skel'],file.split(\"_\")[0]))\n",
    "\n",
    "# Skeleton structure (joints connections) for UTD-MHAD\n",
    "skeleton_edges = [\n",
    "    (0, 1),  # head → shoulder_center\n",
    "    (1, 2),  # shoulder_center → spine\n",
    "    (2, 3),  # spine → hip_center\n",
    "\n",
    "    (1, 4),  # shoulder_center → left_shoulder\n",
    "    (4, 5),  # left_shoulder → left_elbow\n",
    "    (5, 6),  # left_elbow → left_wrist\n",
    "    (6, 7),  # left_wrist → left_hand\n",
    "\n",
    "    (1, 8),  # shoulder_center → right_shoulder\n",
    "    (8, 9),  # right_shoulder → right_elbow\n",
    "    (9, 10), # right_elbow → right_wrist\n",
    "    (10, 11),# right_wrist → right_hand\n",
    "\n",
    "    (3, 12), # hip_center → left_hip\n",
    "    (12, 13),# left_hip → left_knee\n",
    "    (13, 14),# left_knee → left_ankle\n",
    "    (14, 15),# left_ankle → left_foot\n",
    "\n",
    "    (3, 16), # hip_center → right_hip\n",
    "    (16, 17),# right_hip → right_knee\n",
    "    (17, 18),# right_knee → right_ankle\n",
    "    (18, 19) # right_ankle → right_foot\n",
    "]\n",
    "\n",
    "\n",
    "def create_graphs_from_tuple(in_tuple, use_relative_coordinates=True):\n",
    "    action = int(in_tuple[1].split('a')[1]) - 1\n",
    "    frames = in_tuple[0]  # Shape: (20, 3, num_frames)\n",
    "    \n",
    "    # Compute relative coordinates if required\n",
    "    if use_relative_coordinates:\n",
    "        central_joint_idx = 2  # Assuming the spine is the central joint\n",
    "        central_joint_coords = frames[central_joint_idx, :, :]\n",
    "        central_joint_coords = np.expand_dims(central_joint_coords, axis=0) \n",
    "        frames = frames - central_joint_coords\n",
    "\n",
    "    # Stack all frames into a single tensor of shape (num_frames, num_nodes, in_channels)\n",
    "    num_frames = frames.shape[2]\n",
    "    node_features = torch.tensor(frames.transpose(2, 0, 1), dtype=torch.float).unsqueeze(0)\n",
    "    # The transpose converts the shape to (num_frames, num_nodes, in_channels)\n",
    "    \n",
    "    # Convert edge list to tensor\n",
    "    edge_index = torch.tensor(skeleton_edges, dtype=torch.long).t().contiguous()\n",
    "\n",
    "    # Create a single PyG Data object containing the entire temporal sequence\n",
    "    temporal_graph = Data(x=node_features, edge_index=edge_index, y=action)\n",
    "    \n",
    "    return temporal_graph\n",
    "\n",
    "\n",
    "\n",
    "train_graph_snapshots = [create_graphs_from_tuple(seq) for seq in train]\n",
    "test_graph_snapshots = [create_graphs_from_tuple(seq) for seq in test]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(B, num_frames, num_nodes, in_channels)\n",
      "torch.Size([1, 72, 20, 3])\n",
      "9\n"
     ]
    }
   ],
   "source": [
    "print('(B, num_frames, num_nodes, in_channels)')\n",
    "print(train_graph_snapshots[0].x.shape)\n",
    "print(train_graph_snapshots[0].y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Showing a sample data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [],
   "source": [
    "import networkx as nx\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.animation import FuncAnimation\n",
    "\n",
    "num_frames = train_graph_snapshots[1].x.shape[1]\n",
    "node_features = train_graph_snapshots[0].x.squeeze().permute(1,2,0)\n",
    "\n",
    "G = nx.Graph()\n",
    "G.add_edges_from(skeleton_edges)\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(6, 6))\n",
    "\n",
    "def update(frame):\n",
    "    ax.clear()\n",
    "    joint_positions = node_features[:, :, frame]\n",
    "    pos = {i: [x[0].item(), x[1].item()] for i, x in enumerate(joint_positions.squeeze())}\n",
    "    nx.draw(G, pos=pos, ax=ax, node_color='red', edge_color=\"blue\", width=2, node_size=50)\n",
    "    ax.set_title(f'Frame: {frame}')\n",
    "    ax.axis('equal')\n",
    "\n",
    "ani = FuncAnimation(fig, update, frames=num_frames, interval=50, repeat=False)\n",
    "plt.close(fig)  # Prevents duplicate display in Jupyter notebooks\n",
    "ani.save('skeleton_animation.gif', writer='pillow', fps=20)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1, Loss: 3.446215401693832, Train accuracy: 0.17906976744186046\n",
      "Test accuracy: 0.20185614849187936\n",
      "Epoch 2, Loss: 2.2606151850071066, Train accuracy: 0.34186046511627904\n",
      "Test accuracy: 0.33410672853828305\n",
      "Epoch 3, Loss: 1.7693806399959464, Train accuracy: 0.413953488372093\n",
      "Test accuracy: 0.4176334106728538\n",
      "Epoch 4, Loss: 1.5146721000761487, Train accuracy: 0.4255813953488372\n",
      "Test accuracy: 0.38979118329466356\n",
      "Epoch 5, Loss: 1.2561575089884531, Train accuracy: 0.49767441860465117\n",
      "Test accuracy: 0.44779582366589327\n",
      "Epoch 6, Loss: 1.1035941395145126, Train accuracy: 0.5883720930232558\n",
      "Test accuracy: 0.5243619489559165\n",
      "Epoch 7, Loss: 0.9396904665419052, Train accuracy: 0.513953488372093\n",
      "Test accuracy: 0.42691415313225056\n",
      "Epoch 8, Loss: 0.8488747123436634, Train accuracy: 0.5581395348837209\n",
      "Test accuracy: 0.4756380510440835\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[62], line 30\u001b[0m\n\u001b[1;32m     27\u001b[0m total \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m0\u001b[39m\n\u001b[1;32m     28\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m graph_snapshots \u001b[38;5;129;01min\u001b[39;00m train_graph_snapshots:\n\u001b[1;32m     29\u001b[0m     \u001b[38;5;66;03m# Forward pass through the model\u001b[39;00m\n\u001b[0;32m---> 30\u001b[0m     out \u001b[38;5;241m=\u001b[39m \u001b[43mmodel\u001b[49m\u001b[43m(\u001b[49m\u001b[43mgraph_snapshots\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mgraph_snapshots\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43medge_index\u001b[49m\u001b[43m)\u001b[49m  \u001b[38;5;66;03m# Pass graph snapshots for one action sequence\u001b[39;00m\n\u001b[1;32m     31\u001b[0m     label \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mas_tensor([graph_snapshots\u001b[38;5;241m.\u001b[39my])\n\u001b[1;32m     32\u001b[0m     pred \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39margmax(out)\n",
      "File \u001b[0;32m~/miniforge3/envs/gnn-explainer/lib/python3.8/site-packages/torch/nn/modules/module.py:1553\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1551\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m   1552\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1553\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniforge3/envs/gnn-explainer/lib/python3.8/site-packages/torch/nn/modules/module.py:1562\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1557\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1558\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1559\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1560\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1561\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1562\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1564\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m   1565\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "Cell \u001b[0;32mIn[5], line 86\u001b[0m, in \u001b[0;36mSTGCN.forward\u001b[0;34m(self, x, edge_index)\u001b[0m\n\u001b[1;32m     84\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, x, edge_index):\n\u001b[1;32m     85\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m layer \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlayers:\n\u001b[0;32m---> 86\u001b[0m         x \u001b[38;5;241m=\u001b[39m \u001b[43mlayer\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43medge_index\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     88\u001b[0m     \u001b[38;5;66;03m# Global pooling (e.g., mean pooling over time and nodes)\u001b[39;00m\n\u001b[1;32m     89\u001b[0m     x \u001b[38;5;241m=\u001b[39m x\u001b[38;5;241m.\u001b[39mview(x\u001b[38;5;241m.\u001b[39msize(\u001b[38;5;241m0\u001b[39m), \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnum_nodes \u001b[38;5;241m*\u001b[39m x\u001b[38;5;241m.\u001b[39msize(\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m))\n",
      "File \u001b[0;32m~/miniforge3/envs/gnn-explainer/lib/python3.8/site-packages/torch/nn/modules/module.py:1553\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1551\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m   1552\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1553\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniforge3/envs/gnn-explainer/lib/python3.8/site-packages/torch/nn/modules/module.py:1562\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1557\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1558\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1559\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1560\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1561\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1562\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1564\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m   1565\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "Cell \u001b[0;32mIn[5], line 48\u001b[0m, in \u001b[0;36mSTGCNLayer.forward\u001b[0;34m(self, x, edge_index)\u001b[0m\n\u001b[1;32m     45\u001b[0m batch_size, num_frames, num_nodes, in_channels \u001b[38;5;241m=\u001b[39m x\u001b[38;5;241m.\u001b[39msize()\n\u001b[1;32m     47\u001b[0m \u001b[38;5;66;03m# Apply GCN\u001b[39;00m\n\u001b[0;32m---> 48\u001b[0m x \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mgcn\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43medge_index\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     50\u001b[0m \u001b[38;5;66;03m# Reshape x back to (batch_size, num_frames, num_nodes, out_channels)\u001b[39;00m\n\u001b[1;32m     51\u001b[0m out_channels \u001b[38;5;241m=\u001b[39m x\u001b[38;5;241m.\u001b[39msize(\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m)\n",
      "File \u001b[0;32m~/miniforge3/envs/gnn-explainer/lib/python3.8/site-packages/torch/nn/modules/module.py:1553\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1551\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m   1552\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1553\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniforge3/envs/gnn-explainer/lib/python3.8/site-packages/torch/nn/modules/module.py:1562\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1557\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1558\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1559\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1560\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1561\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1562\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1564\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m   1565\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "Cell \u001b[0;32mIn[5], line 18\u001b[0m, in \u001b[0;36mGCNLayer.forward\u001b[0;34m(self, x, edge_index)\u001b[0m\n\u001b[1;32m     17\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, x, edge_index):\n\u001b[0;32m---> 18\u001b[0m     x \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mconv\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43medge_index\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     19\u001b[0m     x \u001b[38;5;241m=\u001b[39m F\u001b[38;5;241m.\u001b[39mrelu(x)\n\u001b[1;32m     20\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m x\n",
      "File \u001b[0;32m~/miniforge3/envs/gnn-explainer/lib/python3.8/site-packages/torch/nn/modules/module.py:1553\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1551\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m   1552\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1553\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniforge3/envs/gnn-explainer/lib/python3.8/site-packages/torch/nn/modules/module.py:1562\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1557\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1558\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1559\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1560\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1561\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1562\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1564\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m   1565\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "File \u001b[0;32m~/miniforge3/envs/gnn-explainer/lib/python3.8/site-packages/torch_geometric/nn/conv/gcn_conv.py:260\u001b[0m, in \u001b[0;36mGCNConv.forward\u001b[0;34m(self, x, edge_index, edge_weight)\u001b[0m\n\u001b[1;32m    257\u001b[0m         \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    258\u001b[0m             edge_index \u001b[38;5;241m=\u001b[39m cache\n\u001b[0;32m--> 260\u001b[0m x \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlin\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    262\u001b[0m \u001b[38;5;66;03m# propagate_type: (x: Tensor, edge_weight: OptTensor)\u001b[39;00m\n\u001b[1;32m    263\u001b[0m out \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpropagate(edge_index, x\u001b[38;5;241m=\u001b[39mx, edge_weight\u001b[38;5;241m=\u001b[39medge_weight)\n",
      "File \u001b[0;32m~/miniforge3/envs/gnn-explainer/lib/python3.8/site-packages/torch/nn/modules/module.py:1553\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1551\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m   1552\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1553\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniforge3/envs/gnn-explainer/lib/python3.8/site-packages/torch/nn/modules/module.py:1562\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1557\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1558\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1559\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1560\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1561\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1562\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1564\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m   1565\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "File \u001b[0;32m~/miniforge3/envs/gnn-explainer/lib/python3.8/site-packages/torch_geometric/nn/dense/linear.py:147\u001b[0m, in \u001b[0;36mLinear.forward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m    141\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, x: Tensor) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Tensor:\n\u001b[1;32m    142\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124mr\u001b[39m\u001b[38;5;124;03m\"\"\"Forward pass.\u001b[39;00m\n\u001b[1;32m    143\u001b[0m \n\u001b[1;32m    144\u001b[0m \u001b[38;5;124;03m    Args:\u001b[39;00m\n\u001b[1;32m    145\u001b[0m \u001b[38;5;124;03m        x (torch.Tensor): The input features.\u001b[39;00m\n\u001b[1;32m    146\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m--> 147\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mF\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlinear\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mweight\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbias\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "\n",
    "model = STGCN(in_channels=3, num_classes=27, num_nodes=20,num_layers=3,kernel_size=5, hidden_dim=64)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.0001, weight_decay=5e-4)\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "\n",
    "\n",
    "train_loss = []\n",
    "for epoch in range(150):\n",
    "    model.train()\n",
    "    losses = []\n",
    "    shuffle(train_graph_snapshots)\n",
    "    for index,graph_snapshots in enumerate(train_graph_snapshots):\n",
    "        optimizer.zero_grad()\n",
    "        # Forward pass through the model\n",
    "        out = model(graph_snapshots.x, graph_snapshots.edge_index)  # Pass graph snapshots for one action sequence\n",
    "        labels = torch.as_tensor([graph_snapshots.y])\n",
    "        # Compute loss\n",
    "        loss = criterion(out, labels)\n",
    "        # Backpropagation\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        losses.append(loss.item())\n",
    "    train_loss.append(np.array(losses).mean())\n",
    "    with torch.no_grad():\n",
    "        model.eval()\n",
    "        correct = 0\n",
    "        total = 0\n",
    "        for graph_snapshots in train_graph_snapshots:\n",
    "            # Forward pass through the model\n",
    "            out = model(graph_snapshots.x, graph_snapshots.edge_index)  # Pass graph snapshots for one action sequence\n",
    "            label = torch.as_tensor([graph_snapshots.y])\n",
    "            pred = torch.argmax(out)\n",
    "            if pred.item() == label.item():\n",
    "                correct +=1\n",
    "            total +=1\n",
    "        print(f'Epoch {epoch+1}, Loss: {np.array(losses).mean()}, Train accuracy: {correct/total}')\n",
    "    with torch.no_grad():\n",
    "        model.eval()\n",
    "        correct = 0\n",
    "        total = 0\n",
    "        for graph_snapshots in test_graph_snapshots:\n",
    "            # Forward pass through the model\n",
    "            out = model(graph_snapshots.x, graph_snapshots.edge_index)  # Pass graph snapshots for one action sequence\n",
    "            label = torch.as_tensor([graph_snapshots.y])\n",
    "            # Compute loss\n",
    "            pred = torch.argmax(out)\n",
    "            if pred.item() == label.item():\n",
    "                correct +=1\n",
    "            total +=1\n",
    "        print(f'Test accuracy: {correct/total}')\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "gnn-explainer",
   "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.8.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}