From 79a6e714deecb80f688871f5ad1f7fd161e1faef Mon Sep 17 00:00:00 2001 From: xrhd Date: Thu, 5 Feb 2026 17:51:39 -0300 Subject: [PATCH 1/3] feat: gat model --- bonsai/models/gat/README.md | 46 ++++++++ bonsai/models/gat/__init__.py | 0 bonsai/models/gat/modeling.py | 160 +++++++++++++++++++++++++++ bonsai/models/gat/params.py | 14 +++ bonsai/models/gat/tests/__init__.py | 0 bonsai/models/gat/tests/run_model.py | 50 +++++++++ 6 files changed, 270 insertions(+) create mode 100644 bonsai/models/gat/README.md create mode 100644 bonsai/models/gat/__init__.py create mode 100644 bonsai/models/gat/modeling.py create mode 100644 bonsai/models/gat/params.py create mode 100644 bonsai/models/gat/tests/__init__.py create mode 100644 bonsai/models/gat/tests/run_model.py 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..36a7cb74 --- /dev/null +++ b/bonsai/models/gat/modeling.py @@ -0,0 +1,160 @@ +import jax +import jax.numpy as jnp +from flax import nnx +from typing import Optional + +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) -> 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) + + # 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 = False) -> jax.Array: + h = x + # Apply dropout to input features + h = self.dropout(h) + + for i, layer in enumerate(self.layers): + h = layer(h, adj) + # Apply elu and dropout for hidden layers + if i < len(self.layers) - 1: + h = jax.nn.elu(h) + h = self.dropout(h) + + # 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..a3b04440 --- /dev/null +++ b/bonsai/models/gat/params.py @@ -0,0 +1,14 @@ +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/__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..b417a994 --- /dev/null +++ b/bonsai/models/gat/tests/run_model.py @@ -0,0 +1,50 @@ +import jax +import jax.numpy as jnp +from flax import nnx +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 + +if __name__ == "__main__": + test_gat_forward_pass() From e887705ac89e69a3478ad95fc1b7ba71c3b33622 Mon Sep 17 00:00:00 2001 From: xrhd Date: Thu, 5 Feb 2026 19:40:09 -0300 Subject: [PATCH 2/3] feat(gat): complete implementation and cora validation --- README.md | 1 + bonsai/models/gat/modeling.py | 92 +++++---- bonsai/models/gat/params.py | 3 +- .../models/gat/tests/GAT_cora_validation.py | 186 ++++++++++++++++++ bonsai/models/gat/tests/run_model.py | 56 ++++-- 5 files changed, 276 insertions(+), 62 deletions(-) create mode 100644 bonsai/models/gat/tests/GAT_cora_validation.py 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/modeling.py b/bonsai/models/gat/modeling.py index 36a7cb74..62e9a42e 100644 --- a/bonsai/models/gat/modeling.py +++ b/bonsai/models/gat/modeling.py @@ -1,7 +1,7 @@ import jax import jax.numpy as jnp from flax import nnx -from typing import Optional + class GATLayer(nnx.Module): def __init__( @@ -20,75 +20,71 @@ def __init__( 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.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) -> jax.Array: + 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) - + 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) - + + 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) - + 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) - + attention = jax.nn.softmax(attention, axis=1) # (N, N, num_heads) + # Apply dropout to attention coefficients - attention = self.dropout(attention) - + 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) - + 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 @@ -103,20 +99,20 @@ def __init__( 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 + 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( + self.layers.append( GATLayer( in_features=current_dim, out_features=hidden_features, @@ -124,10 +120,10 @@ def __init__( rngs=rngs, dropout=dropout_prob, alpha=alpha, - concat=True + concat=True, ) - ) - current_dim = hidden_features * num_heads + ) + current_dim = hidden_features * num_heads # Output Layer self.layers.append( @@ -138,23 +134,23 @@ def __init__( rngs=rngs, dropout=dropout_prob, alpha=alpha, - concat=False # Paper averages the last layer + 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 = False) -> jax.Array: + 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) - + h = self.dropout(h, deterministic=not training) + for i, layer in enumerate(self.layers): - h = layer(h, adj) - # Apply elu and dropout for hidden layers - if i < len(self.layers) - 1: - h = jax.nn.elu(h) - h = self.dropout(h) - + 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 index a3b04440..d27c55c8 100644 --- a/bonsai/models/gat/params.py +++ b/bonsai/models/gat/params.py @@ -1,9 +1,10 @@ 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)"}) diff --git a/bonsai/models/gat/tests/GAT_cora_validation.py b/bonsai/models/gat/tests/GAT_cora_validation.py new file mode 100644 index 00000000..05f6def2 --- /dev/null +++ b/bonsai/models/gat/tests/GAT_cora_validation.py @@ -0,0 +1,186 @@ +import os +import urllib.request +import jax +import jax.numpy as jnp +import numpy as np +from flax import nnx +import optax +from bonsai.models.gat.modeling import GAT +from bonsai.models.gat.params import GATConfig + +# 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" + + +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 + + +def main(): + 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%") + + +if __name__ == "__main__": + main() diff --git a/bonsai/models/gat/tests/run_model.py b/bonsai/models/gat/tests/run_model.py index b417a994..eec40c7f 100644 --- a/bonsai/models/gat/tests/run_model.py +++ b/bonsai/models/gat/tests/run_model.py @@ -1,25 +1,19 @@ import jax import jax.numpy as jnp -from flax import nnx 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 + 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 @@ -29,7 +23,7 @@ def test_gat_forward_pass(): 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}") @@ -38,13 +32,49 @@ def test_gat_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() From 64c473a271d410c95550013c3850743cabe321b1 Mon Sep 17 00:00:00 2001 From: xrhd Date: Fri, 6 Feb 2026 09:35:26 -0300 Subject: [PATCH 3/3] feat(gat): cora validation notebook --- bonsai/models/gat/.gitignore | 1 + .../gat/tests/GAT_cora_validation.ipynb | 285 ++++++++++++++++++ ...a_validation.py => GAT_cora_validation.md} | 148 +++++---- 3 files changed, 372 insertions(+), 62 deletions(-) create mode 100644 bonsai/models/gat/.gitignore create mode 100644 bonsai/models/gat/tests/GAT_cora_validation.ipynb rename bonsai/models/gat/tests/{GAT_cora_validation.py => GAT_cora_validation.md} (65%) 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/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.py b/bonsai/models/gat/tests/GAT_cora_validation.md similarity index 65% rename from bonsai/models/gat/tests/GAT_cora_validation.py rename to bonsai/models/gat/tests/GAT_cora_validation.md index 05f6def2..05023108 100644 --- a/bonsai/models/gat/tests/GAT_cora_validation.py +++ b/bonsai/models/gat/tests/GAT_cora_validation.md @@ -1,3 +1,17 @@ +--- +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 @@ -5,16 +19,25 @@ import numpy as np from flax import nnx import optax -from bonsai.models.gat.modeling import GAT -from bonsai.models.gat.params import GATConfig + +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") @@ -125,62 +148,63 @@ def eval_step(model, x, adj, labels, mask): accuracy = correct / jnp.maximum(total, 1) loss = loss_fn(model, x, adj, labels, mask, False) return loss, accuracy - - -def main(): - 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%") - - -if __name__ == "__main__": - main() +``` + +```{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 + +```