diff --git a/README.md b/README.md index 1cfd167f..85c8b4ed 100644 --- a/README.md +++ b/README.md @@ -27,6 +27,7 @@ These are listed based on status and then alphabetically. | [ConvNeXT](bonsai/models/convnext/) | Image classification | ✅ | | | [Densenet](bonsai/models/densenet121/) | Image classification | ✅ | | | [EfficientNet](bonsai/models/efficientnet/) | Image classification | ✅ | | +| [GAT](bonsai/models/gat/) | Graph Attention | ✅ | | | [Qwen 3](bonsai/models/qwen3) | LLM | ✅ | | | [ResNet50](bonsai/models/resnet) | Image classification | ✅ | | | [VGG](bonsai/models/vgg19) | Image classification | ✅ | | diff --git a/bonsai/models/gat/.gitignore b/bonsai/models/gat/.gitignore new file mode 100644 index 00000000..4173ae8d --- /dev/null +++ b/bonsai/models/gat/.gitignore @@ -0,0 +1 @@ +tests/data diff --git a/bonsai/models/gat/README.md b/bonsai/models/gat/README.md new file mode 100644 index 00000000..611b038a --- /dev/null +++ b/bonsai/models/gat/README.md @@ -0,0 +1,46 @@ +# Graph Attention Network (GAT) + +A JAX/Flax implementation of the Graph Attention Network (GAT) based on [Graph Attention Networks](https://arxiv.org/abs/1710.10903) by Veličković et al. + +This implementation uses `flax.nnx` for explicit state management. + +## Usage + +```python +import jax +import jax.numpy as jnp +from flax import nnx +from bonsai.models.gat.modeling import GAT + +# Configuration +key = jax.random.key(0) +N, F, C = 10, 5, 2 + +# Instantiate Model +model = GAT( + in_features=F, + hidden_features=8, + out_features=C, + num_heads=2, + dropout_rng=key, + dropout_prob=0.6, + alpha=0.2 +) + +# Dummy Data +key, k1, k2 = jax.random.split(key, 3) +x = jax.random.normal(k1, (N, F)) +adj = jax.random.bernoulli(k2, 0.3, (N, N)).astype(jnp.float32) + jnp.eye(N) +adj = jnp.clip(adj, 0.0, 1.0) + +# Forward Pass +logits = model(x, adj, training=False) +``` + +## Validation + +To reproduce the results on the Cora dataset (~83% accuracy): + +```bash +python bonsai/models/gat/tests/GAT_cora_validation.py +``` diff --git a/bonsai/models/gat/__init__.py b/bonsai/models/gat/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/bonsai/models/gat/modeling.py b/bonsai/models/gat/modeling.py new file mode 100644 index 00000000..62e9a42e --- /dev/null +++ b/bonsai/models/gat/modeling.py @@ -0,0 +1,156 @@ +import jax +import jax.numpy as jnp +from flax import nnx + + +class GATLayer(nnx.Module): + def __init__( + self, + in_features: int, + out_features: int, + num_heads: int, + rngs: nnx.Rngs, + dropout: float = 0.6, + alpha: float = 0.2, + concat: bool = True, + ): + self.in_features = in_features + self.out_features = out_features + self.num_heads = num_heads + self.dropout_prob = dropout + self.alpha = alpha + self.concat = concat + + self.W = nnx.Linear(in_features, num_heads * out_features, use_bias=False, rngs=rngs) + # Attention parameter 'a' + # shape: (num_heads, 2 * out_features) + self.a = nnx.Param(nnx.initializers.glorot_uniform()(rngs.params(), (num_heads, 2 * out_features), jnp.float32)) + + self.bias = nnx.Param(jnp.zeros((num_heads * out_features if concat else out_features,))) + + self.leaky_relu = lambda x: jax.nn.leaky_relu(x, negative_slope=alpha) + self.dropout = nnx.Dropout(dropout, rngs=rngs) + + def __call__(self, h: jax.Array, adj: jax.Array, training: bool = True) -> jax.Array: + # h: (N, in_features) + # adj: (N, N) + N = h.shape[0] + + # Linear transformation + Wh = self.W(h) # (N, num_heads * out_features) + Wh = Wh.reshape(N, self.num_heads, self.out_features) # (N, num_heads, out_features) + + # Prepare attention mechanism + # Concatenate Wh_i and Wh_j -> (N, N, num_heads, 2 * out_features) + # However, to avoid O(N^2) memory for features, we can do: + # e_ij = LeakyReLU(a^T [Wh_i || Wh_j]) + # = LeakyReLU((a_1)^T Wh_i + (a_2)^T Wh_j) + + a1 = self.a.value[:, : self.out_features] # (num_heads, out_features) + a2 = self.a.value[:, self.out_features :] # (num_heads, out_features) + + # Calculate (a_1)^T Wh_i -> (N, num_heads) + # Wh: (N, num_heads, out_features) + # a1: (num_heads, out_features) + # We want inner product over out_features for each head + attn_1 = jnp.einsum("nho,ho->nh", Wh, a1) # (N, num_heads) + attn_2 = jnp.einsum("nho,ho->nh", Wh, a2) # (N, num_heads) + + # Broadcast add -> (N, N, num_heads) + e = attn_1[:, None, :] + attn_2[None, :, :] + e = self.leaky_relu(e) + + # Masked attention + # adj is assumed to be 0 for no edge, 1 for edge (including self-loop) + # We want to mask where adj is 0 + zero_vec = -9e15 * jnp.ones_like(e) + attention = jnp.where(adj[..., None] > 0, e, zero_vec) + + # Softmax over neighbors (dim 1) + attention = jax.nn.softmax(attention, axis=1) # (N, N, num_heads) + + # Apply dropout to attention coefficients + attention = self.dropout(attention, deterministic=not training) + + # Aggregation: h'_i = sum_j alpha_ij W h_j + # attention: (N, N, num_heads) + # Wh: (N, num_heads, out_features) + # Result: (N, num_heads, out_features) + h_prime = jnp.einsum("ijh,jho->iho", attention, Wh) + + if self.concat: + # Concatenate heads -> (N, num_heads * out_features) + output = h_prime.reshape(N, self.num_heads * self.out_features) + else: + # Average heads -> (N, out_features) + output = jnp.mean(h_prime, axis=1) + + return output + self.bias.value + + +class GAT(nnx.Module): + def __init__( + self, + in_features: int, + hidden_features: int, + out_features: int, + num_heads: int, + dropout_rng: jax.Array, + dropout_prob: float = 0.6, + alpha: float = 0.2, + concat_hidden: bool = True, + num_layers: int = 2, # Defaults to 2 layers as per paper/spec + num_out_heads: int = 1, + ): + self.dropout_prob = dropout_prob + self.layers = nnx.List([]) + rngs = nnx.Rngs(dropout_rng) + + # Input/Hidden Layers + # Usually GAT has layers 1 to N-1 + current_dim = in_features + + # First layer (and subsequent hidden layers if num_layers > 2) + for _ in range(num_layers - 1): + self.layers.append( + GATLayer( + in_features=current_dim, + out_features=hidden_features, + num_heads=num_heads, + rngs=rngs, + dropout=dropout_prob, + alpha=alpha, + concat=True, + ) + ) + current_dim = hidden_features * num_heads + + # Output Layer + self.layers.append( + GATLayer( + in_features=current_dim, + out_features=out_features, + num_heads=num_out_heads, + rngs=rngs, + dropout=dropout_prob, + alpha=alpha, + concat=False, # Paper averages the last layer + ) + ) + + self.dropout = nnx.Dropout(dropout_prob, rngs=rngs) + + def __call__(self, x: jax.Array, adj: jax.Array, training: bool = True) -> jax.Array: + h = x + # Apply dropout to input features + h = self.dropout(h, deterministic=not training) + + for i, layer in enumerate(self.layers): + h = layer(h, adj, training=training) + # Apply elu and dropout for hidden layers + if i < len(self.layers) - 1: + h = jax.nn.elu(h) + h = self.dropout(h, deterministic=not training) + + # Final layer usually is softmax for classification, but we return logits + return h diff --git a/bonsai/models/gat/params.py b/bonsai/models/gat/params.py new file mode 100644 index 00000000..d27c55c8 --- /dev/null +++ b/bonsai/models/gat/params.py @@ -0,0 +1,15 @@ +from dataclasses import dataclass, field + + +@dataclass +class GATConfig: + """Configuration class for Graph Attention Network (GAT).""" + + in_features: int = field(metadata={"help": "Dimension of input node features"}) + hidden_features: int = field(default=8, metadata={"help": "Dimension of hidden features PER HEAD"}) + out_features: int = field(default=7, metadata={"help": "Dimension of output features (classes)"}) + num_heads: int = field(default=8, metadata={"help": "Number of attention heads for hidden layers"}) + num_out_heads: int = field(default=1, metadata={"help": "Number of attention heads for output layer"}) + num_layers: int = field(default=2, metadata={"help": "Number of GAT layers"}) + dropout_prob: float = field(default=0.6, metadata={"help": "Dropout probability"}) + alpha: float = field(default=0.2, metadata={"help": "LeakyReLU negative slope"}) diff --git a/bonsai/models/gat/tests/GAT_cora_validation.ipynb b/bonsai/models/gat/tests/GAT_cora_validation.ipynb new file mode 100644 index 00000000..71ece6eb --- /dev/null +++ b/bonsai/models/gat/tests/GAT_cora_validation.ipynb @@ -0,0 +1,285 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "ab873b8a", + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "import urllib.request\n", + "import jax\n", + "import jax.numpy as jnp\n", + "import numpy as np\n", + "from flax import nnx\n", + "import optax\n", + "\n", + "try:\n", + " from bonsai.models.gat.modeling import GAT\n", + " from bonsai.models.gat.params import GATConfig\n", + "except ImportError:\n", + " try:\n", + " !pip insetall -e .\n", + " except:\n", + " !pip install -q git+https://github.com/jax-ml/bonsai@main\n", + "\n", + "\n", + "# Configuration\n", + "# Use a temporary directory for data to avoid polluting the repo\n", + "DATA_DIR = os.path.join(os.getcwd(), \"data\", \"cora\")\n", + "CORA_CONTENT_URL = \"https://raw.githubusercontent.com/tkipf/pygcn/master/data/cora/cora.content\"\n", + "CORA_CITES_URL = \"https://raw.githubusercontent.com/tkipf/pygcn/master/data/cora/cora.cites\"" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "1e0aee81", + "metadata": {}, + "outputs": [], + "source": [ + "def download_cora():\n", + " os.makedirs(DATA_DIR, exist_ok=True)\n", + " content_path = os.path.join(DATA_DIR, \"cora.content\")\n", + " cites_path = os.path.join(DATA_DIR, \"cora.cites\")\n", + "\n", + " if not os.path.exists(content_path):\n", + " print(f\"Downloading {CORA_CONTENT_URL}...\")\n", + " urllib.request.urlretrieve(CORA_CONTENT_URL, content_path)\n", + "\n", + " if not os.path.exists(cites_path):\n", + " print(f\"Downloading {CORA_CITES_URL}...\")\n", + " urllib.request.urlretrieve(CORA_CITES_URL, cites_path)\n", + "\n", + " return content_path, cites_path\n", + "\n", + "\n", + "def load_data():\n", + " content_path, cites_path = download_cora()\n", + "\n", + " # Load content\n", + " # Format: paper_id, word_attributes..., label\n", + " content = np.genfromtxt(content_path, dtype=np.dtype(str))\n", + " idx = content[:, 0].astype(np.int32)\n", + " features = content[:, 1:-1].astype(np.float32)\n", + " labels_str = content[:, -1]\n", + "\n", + " # Map labels to integers\n", + " unique_labels = sorted(set(labels_str))\n", + " label_map = {label: i for i, label in enumerate(unique_labels)}\n", + " labels = np.array([label_map[l] for l in labels_str], dtype=np.int32)\n", + "\n", + " # Map paper IDs to 0-N\n", + " idx_map = {id: i for i, id in enumerate(idx)}\n", + "\n", + " # Load cites\n", + " # Format: cited_paper_id, citing_paper_id\n", + " edges_unordered = np.genfromtxt(cites_path, dtype=np.int32)\n", + "\n", + " edges = []\n", + " for edge in edges_unordered:\n", + " if edge[0] in idx_map and edge[1] in idx_map:\n", + " edges.append([idx_map[edge[0]], idx_map[edge[1]]])\n", + " edges = np.array(edges)\n", + "\n", + " N = features.shape[0]\n", + " adj = np.zeros((N, N), dtype=np.float32)\n", + " adj[edges[:, 0], edges[:, 1]] = 1.0\n", + " # Symmetric adjacency matrix\n", + " adj = adj + adj.T\n", + " adj = np.clip(adj, 0, 1)\n", + "\n", + " # Add self-loops\n", + " adj = adj + np.eye(N)\n", + "\n", + " # Row-normalize features\n", + " features_sum = features.sum(axis=1, keepdims=True)\n", + " features_sum = np.where(features_sum == 0, 1, features_sum)\n", + " features = features / features_sum\n", + "\n", + " # Standard split indices (based on common Cora implementations)\n", + " # We'll use 20 nodes per class for training, 500 for val, 1000 for test\n", + " train_mask = np.zeros(N, dtype=bool)\n", + " val_mask = np.zeros(N, dtype=bool)\n", + " test_mask = np.zeros(N, dtype=bool)\n", + "\n", + " # For reproducibility, we use a class-balanced split\n", + " np.random.seed(42)\n", + " for i in range(len(unique_labels)):\n", + " indices = np.where(labels == i)[0]\n", + " np.random.shuffle(indices)\n", + " train_mask[indices[:20]] = True\n", + " val_mask[indices[20 : 20 + 71]] = True # 71 * 7 = 497\n", + " test_mask[indices[20 + 71 : 20 + 71 + 143]] = True # 143 * 7 = 1001\n", + "\n", + " return (\n", + " jnp.array(features),\n", + " jnp.array(adj),\n", + " jnp.array(labels),\n", + " jnp.array(train_mask),\n", + " jnp.array(val_mask),\n", + " jnp.array(test_mask),\n", + " )\n", + "\n", + "\n", + "def loss_fn(model, x, adj, labels, mask, training):\n", + " logits = model(x, adj, training=training)\n", + " log_probs = jax.nn.log_softmax(logits)\n", + " one_hot = jax.nn.one_hot(labels, num_classes=logits.shape[-1])\n", + " # Cross entropy only on masked nodes\n", + " loss = -jnp.sum(one_hot * log_probs, axis=-1)\n", + " return jnp.sum(loss * mask) / jnp.maximum(jnp.sum(mask), 1)\n", + "\n", + "\n", + "@nnx.jit\n", + "def train_step(model, optimizer, x, adj, labels, mask):\n", + " grad_fn = nnx.value_and_grad(loss_fn)\n", + " loss, grads = grad_fn(model, x, adj, labels, mask, True)\n", + " optimizer.update(model, grads)\n", + " return loss\n", + "\n", + "\n", + "@nnx.jit\n", + "def eval_step(model, x, adj, labels, mask):\n", + " logits = model(x, adj, training=False)\n", + " preds = jnp.argmax(logits, axis=-1)\n", + " correct = jnp.sum((preds == labels) * mask)\n", + " total = jnp.sum(mask)\n", + " accuracy = correct / jnp.maximum(total, 1)\n", + " loss = loss_fn(model, x, adj, labels, mask, False)\n", + " return loss, accuracy" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "3db7dd83", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Loading Cora data...\n", + "Downloading https://raw.githubusercontent.com/tkipf/pygcn/master/data/cora/cora.content...\n", + "Downloading https://raw.githubusercontent.com/tkipf/pygcn/master/data/cora/cora.cites...\n", + "Data loaded. Features: (2708, 1433), Nodes: 2708, Edges: 13264\n", + "Starting training...\n", + "Epoch 10: Loss = 1.8906, Val Loss = 1.9067, Val Acc = 0.7485\n", + "Epoch 20: Loss = 1.8212, Val Loss = 1.8453, Val Acc = 0.7827\n", + "Epoch 30: Loss = 1.7223, Val Loss = 1.7531, Val Acc = 0.7827\n", + "Epoch 40: Loss = 1.4675, Val Loss = 1.6313, Val Acc = 0.7746\n", + "Epoch 50: Loss = 1.3681, Val Loss = 1.4850, Val Acc = 0.8169\n", + "Epoch 60: Loss = 1.1953, Val Loss = 1.3345, Val Acc = 0.8028\n", + "Epoch 70: Loss = 1.1052, Val Loss = 1.1954, Val Acc = 0.8089\n", + "Epoch 80: Loss = 1.0032, Val Loss = 1.0757, Val Acc = 0.8109\n", + "Epoch 90: Loss = 0.9063, Val Loss = 0.9835, Val Acc = 0.8089\n", + "Epoch 100: Loss = 0.8257, Val Loss = 0.9084, Val Acc = 0.8149\n", + "Epoch 110: Loss = 0.7812, Val Loss = 0.8464, Val Acc = 0.8169\n", + "Epoch 120: Loss = 0.6745, Val Loss = 0.7981, Val Acc = 0.8149\n", + "Epoch 130: Loss = 0.6951, Val Loss = 0.7553, Val Acc = 0.8089\n", + "Epoch 140: Loss = 0.6721, Val Loss = 0.7244, Val Acc = 0.8089\n", + "Epoch 150: Loss = 0.7320, Val Loss = 0.6954, Val Acc = 0.8229\n", + "Epoch 160: Loss = 0.5734, Val Loss = 0.6827, Val Acc = 0.8109\n", + "Epoch 170: Loss = 0.5050, Val Loss = 0.6635, Val Acc = 0.7988\n", + "Epoch 180: Loss = 0.5997, Val Loss = 0.6547, Val Acc = 0.8089\n", + "Epoch 190: Loss = 0.7087, Val Loss = 0.6414, Val Acc = 0.8008\n", + "Epoch 200: Loss = 0.4626, Val Loss = 0.6308, Val Acc = 0.7948\n", + "\n", + "Final Results:\n", + "Test Loss: 0.6163\n", + "Test Accuracy: 0.8032\n", + "SUCCESS: Accuracy is above 80%\n" + ] + } + ], + "source": [ + "print(\"Loading Cora data...\")\n", + "x, adj, labels, train_mask, val_mask, test_mask = load_data()\n", + "print(f\"Data loaded. Features: {x.shape}, Nodes: {x.shape[0]}, Edges: {jnp.sum(adj > 0)}\")\n", + "\n", + "key = jax.random.key(42)\n", + "model_key, _ = jax.random.split(key)\n", + "\n", + "config = GATConfig(\n", + " in_features=x.shape[1],\n", + " hidden_features=8,\n", + " out_features=int(jnp.max(labels) + 1),\n", + " num_heads=8,\n", + " num_out_heads=1,\n", + " dropout_prob=0.6,\n", + " alpha=0.2,\n", + ")\n", + "\n", + "model = GAT(\n", + " in_features=config.in_features,\n", + " hidden_features=config.hidden_features,\n", + " out_features=config.out_features,\n", + " num_heads=config.num_heads,\n", + " dropout_rng=model_key,\n", + " dropout_prob=config.dropout_prob,\n", + " alpha=config.alpha,\n", + " num_out_heads=config.num_out_heads,\n", + ")\n", + "\n", + "# Standard Adam optimizer for GAT\n", + "# Paper uses lr=0.005 and weight_decay=5e-4\n", + "optimizer = nnx.Optimizer(model, optax.adam(learning_rate=0.005), wrt=nnx.Param)\n", + "\n", + "print(\"Starting training...\")\n", + "best_val_acc = 0\n", + "for epoch in range(1, 201):\n", + " loss = train_step(model, optimizer, x, adj, labels, train_mask)\n", + "\n", + " if epoch % 10 == 0:\n", + " val_loss, val_acc = eval_step(model, x, adj, labels, val_mask)\n", + " print(f\"Epoch {epoch:3d}: Loss = {loss:.4f}, Val Loss = {val_loss:.4f}, Val Acc = {val_acc:.4f}\")\n", + " if val_acc > best_val_acc:\n", + " best_val_acc = val_acc\n", + "\n", + "test_loss, test_acc = eval_step(model, x, adj, labels, test_mask)\n", + "print(\"\\nFinal Results:\")\n", + "print(f\"Test Loss: {test_loss:.4f}\")\n", + "print(f\"Test Accuracy: {test_acc:.4f}\")\n", + "\n", + "if test_acc >= 0.80:\n", + " print(\"SUCCESS: Accuracy is above 80%\")\n", + "else:\n", + " print(\"FAILURE: Accuracy is below 80%\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9e91c18a", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "jupytext": { + "default_lexer": "ipython3" + }, + "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.12.6" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/bonsai/models/gat/tests/GAT_cora_validation.md b/bonsai/models/gat/tests/GAT_cora_validation.md new file mode 100644 index 00000000..05023108 --- /dev/null +++ b/bonsai/models/gat/tests/GAT_cora_validation.md @@ -0,0 +1,210 @@ +--- +jupytext: + text_representation: + extension: .md + format_name: myst + format_version: 0.13 + jupytext_version: 1.18.1 +kernelspec: + display_name: .venv + language: python + name: python3 +--- + +```{code-cell} ipython3 +import os +import urllib.request +import jax +import jax.numpy as jnp +import numpy as np +from flax import nnx +import optax + +try: + from bonsai.models.gat.modeling import GAT + from bonsai.models.gat.params import GATConfig +except ImportError: + try: + !pip insetall -e . + except: + !pip install -q git+https://github.com/jax-ml/bonsai@main + + +# Configuration +# Use a temporary directory for data to avoid polluting the repo +DATA_DIR = os.path.join(os.getcwd(), "data", "cora") +CORA_CONTENT_URL = "https://raw.githubusercontent.com/tkipf/pygcn/master/data/cora/cora.content" +CORA_CITES_URL = "https://raw.githubusercontent.com/tkipf/pygcn/master/data/cora/cora.cites" +``` + +```{code-cell} ipython3 +def download_cora(): + os.makedirs(DATA_DIR, exist_ok=True) + content_path = os.path.join(DATA_DIR, "cora.content") + cites_path = os.path.join(DATA_DIR, "cora.cites") + + if not os.path.exists(content_path): + print(f"Downloading {CORA_CONTENT_URL}...") + urllib.request.urlretrieve(CORA_CONTENT_URL, content_path) + + if not os.path.exists(cites_path): + print(f"Downloading {CORA_CITES_URL}...") + urllib.request.urlretrieve(CORA_CITES_URL, cites_path) + + return content_path, cites_path + + +def load_data(): + content_path, cites_path = download_cora() + + # Load content + # Format: paper_id, word_attributes..., label + content = np.genfromtxt(content_path, dtype=np.dtype(str)) + idx = content[:, 0].astype(np.int32) + features = content[:, 1:-1].astype(np.float32) + labels_str = content[:, -1] + + # Map labels to integers + unique_labels = sorted(set(labels_str)) + label_map = {label: i for i, label in enumerate(unique_labels)} + labels = np.array([label_map[l] for l in labels_str], dtype=np.int32) + + # Map paper IDs to 0-N + idx_map = {id: i for i, id in enumerate(idx)} + + # Load cites + # Format: cited_paper_id, citing_paper_id + edges_unordered = np.genfromtxt(cites_path, dtype=np.int32) + + edges = [] + for edge in edges_unordered: + if edge[0] in idx_map and edge[1] in idx_map: + edges.append([idx_map[edge[0]], idx_map[edge[1]]]) + edges = np.array(edges) + + N = features.shape[0] + adj = np.zeros((N, N), dtype=np.float32) + adj[edges[:, 0], edges[:, 1]] = 1.0 + # Symmetric adjacency matrix + adj = adj + adj.T + adj = np.clip(adj, 0, 1) + + # Add self-loops + adj = adj + np.eye(N) + + # Row-normalize features + features_sum = features.sum(axis=1, keepdims=True) + features_sum = np.where(features_sum == 0, 1, features_sum) + features = features / features_sum + + # Standard split indices (based on common Cora implementations) + # We'll use 20 nodes per class for training, 500 for val, 1000 for test + train_mask = np.zeros(N, dtype=bool) + val_mask = np.zeros(N, dtype=bool) + test_mask = np.zeros(N, dtype=bool) + + # For reproducibility, we use a class-balanced split + np.random.seed(42) + for i in range(len(unique_labels)): + indices = np.where(labels == i)[0] + np.random.shuffle(indices) + train_mask[indices[:20]] = True + val_mask[indices[20 : 20 + 71]] = True # 71 * 7 = 497 + test_mask[indices[20 + 71 : 20 + 71 + 143]] = True # 143 * 7 = 1001 + + return ( + jnp.array(features), + jnp.array(adj), + jnp.array(labels), + jnp.array(train_mask), + jnp.array(val_mask), + jnp.array(test_mask), + ) + + +def loss_fn(model, x, adj, labels, mask, training): + logits = model(x, adj, training=training) + log_probs = jax.nn.log_softmax(logits) + one_hot = jax.nn.one_hot(labels, num_classes=logits.shape[-1]) + # Cross entropy only on masked nodes + loss = -jnp.sum(one_hot * log_probs, axis=-1) + return jnp.sum(loss * mask) / jnp.maximum(jnp.sum(mask), 1) + + +@nnx.jit +def train_step(model, optimizer, x, adj, labels, mask): + grad_fn = nnx.value_and_grad(loss_fn) + loss, grads = grad_fn(model, x, adj, labels, mask, True) + optimizer.update(model, grads) + return loss + + +@nnx.jit +def eval_step(model, x, adj, labels, mask): + logits = model(x, adj, training=False) + preds = jnp.argmax(logits, axis=-1) + correct = jnp.sum((preds == labels) * mask) + total = jnp.sum(mask) + accuracy = correct / jnp.maximum(total, 1) + loss = loss_fn(model, x, adj, labels, mask, False) + return loss, accuracy +``` + +```{code-cell} ipython3 +print("Loading Cora data...") +x, adj, labels, train_mask, val_mask, test_mask = load_data() +print(f"Data loaded. Features: {x.shape}, Nodes: {x.shape[0]}, Edges: {jnp.sum(adj > 0)}") + +key = jax.random.key(42) +model_key, _ = jax.random.split(key) + +config = GATConfig( + in_features=x.shape[1], + hidden_features=8, + out_features=int(jnp.max(labels) + 1), + num_heads=8, + num_out_heads=1, + dropout_prob=0.6, + alpha=0.2, +) + +model = GAT( + in_features=config.in_features, + hidden_features=config.hidden_features, + out_features=config.out_features, + num_heads=config.num_heads, + dropout_rng=model_key, + dropout_prob=config.dropout_prob, + alpha=config.alpha, + num_out_heads=config.num_out_heads, +) + +# Standard Adam optimizer for GAT +# Paper uses lr=0.005 and weight_decay=5e-4 +optimizer = nnx.Optimizer(model, optax.adam(learning_rate=0.005), wrt=nnx.Param) + +print("Starting training...") +best_val_acc = 0 +for epoch in range(1, 201): + loss = train_step(model, optimizer, x, adj, labels, train_mask) + + if epoch % 10 == 0: + val_loss, val_acc = eval_step(model, x, adj, labels, val_mask) + print(f"Epoch {epoch:3d}: Loss = {loss:.4f}, Val Loss = {val_loss:.4f}, Val Acc = {val_acc:.4f}") + if val_acc > best_val_acc: + best_val_acc = val_acc + +test_loss, test_acc = eval_step(model, x, adj, labels, test_mask) +print("\nFinal Results:") +print(f"Test Loss: {test_loss:.4f}") +print(f"Test Accuracy: {test_acc:.4f}") + +if test_acc >= 0.80: + print("SUCCESS: Accuracy is above 80%") +else: + print("FAILURE: Accuracy is below 80%") +``` + +```{code-cell} ipython3 + +``` diff --git a/bonsai/models/gat/tests/__init__.py b/bonsai/models/gat/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/bonsai/models/gat/tests/run_model.py b/bonsai/models/gat/tests/run_model.py new file mode 100644 index 00000000..eec40c7f --- /dev/null +++ b/bonsai/models/gat/tests/run_model.py @@ -0,0 +1,80 @@ +import jax +import jax.numpy as jnp +from bonsai.models.gat.modeling import GAT + + +def test_gat_forward_pass(): + print("Initializing GAT model...") + # 1. Configuration + key = jax.random.key(0) + N, F, C = 10, 5, 2 # 10 nodes, 5 features, 2 classes + + # 2. Instantiate Model + model = GAT( + in_features=F, hidden_features=8, out_features=C, num_heads=2, dropout_rng=key, dropout_prob=0.6, alpha=0.2 + ) + + print("Model initialized successfully.") + + # 3. Create Dummy Data + key, k1, k2 = jax.random.split(key, 3) + x = jax.random.normal(k1, (N, F)) + # Adjacency matrix (binary, symmetric + self-loops) + adj = jax.random.bernoulli(k2, 0.3, (N, N)).astype(jnp.float32) + adj = adj + jnp.eye(N) + adj = jnp.clip(adj, 0.0, 1.0) + + print(f"Input features shape: {x.shape}") + print(f"Adjacency matrix shape: {adj.shape}") + + # 4. Forward Pass + print("Running forward pass...") + try: + logits = model(x, adj, training=False) + print(f"Logits shape: {logits.shape}") + + assert logits.shape == (N, C), f"Expected logits shape {(N, C)}, but got {logits.shape}" + print("Forward pass successful!") + + except Exception as e: + print(f"Forward pass failed: {e}") + raise e + + +def test_gat_edge_cases(): + print("\nTesting edge cases...") + key = jax.random.key(1) + N, F, C = 5, 4, 3 + + model = GAT(in_features=F, hidden_features=4, out_features=C, num_heads=2, dropout_rng=key) + + # 1. Zero-edge graph (only self-loops) + print("Case 1: Zero-edge graph (only self-loops)...") + x = jax.random.normal(key, (N, F)) + adj_zero = jnp.eye(N) + logits = model(x, adj_zero, training=False) + assert logits.shape == (N, C) + print("Success: Zero-edge graph handled.") + + # 2. Dimension mismatch (features) + print("Case 2: Feature dimension mismatch...") + x_wrong = jax.random.normal(key, (N, F + 1)) + try: + model(x_wrong, adj_zero, training=False) + print("Failure: Model should have raised a dimension mismatch error.") + except Exception as e: + print(f"Success: Correctly caught error: {e}") + + # 3. Dimension mismatch (adjacency) + print("Case 3: Adjacency dimension mismatch...") + adj_wrong = jnp.eye(N + 1) + try: + model(x, adj_wrong, training=False) + print("Failure: Model should have raised an adjacency dimension mismatch error.") + except Exception as e: + print(f"Success: Correctly caught error: {e}") + + +if __name__ == "__main__": + test_gat_forward_pass() + test_gat_edge_cases()