shemalfoy/eegnet-gnn-features

EEGNet-GNN feature extractor (Option C)

Standard EEGNet with one change: Layer 2's depthwise spatial conv is replaced by a graph convolution over the electrode montage, so electrodes mix according to how close they are on the scalp rather than as a flat, order-agnostic channel list.

Temporal conv -> GRAPH conv -> Separable conv -> Avg-pool + flatten -> flat vector
  (Layer 1)      (Layer 2)      (Layer 3)          (Layer 4)           OUTPUT
  IDENTICAL      THE SWAP       IDENTICAL          IDENTICAL

This model stops at Layer 4 and returns a flat feature vector — there is no classifier head. Attach your own head (a linear layer, an MLP, or a variational quantum circuit) downstream. The output size is exposed as model.flat_dim (496 with the defaults).

The swap

Standard EEGNet's Layer 2 learns, per output map, a single weighted sum over all electrodes at once — topology is ignored. Here that becomes a graph convolution H = Â X W, where Â is the symmetric-normalised adjacency of the electrode montage. Each electrode aggregates only from its physical neighbours (e.g. C3 from FC3, FC1, C5, C1, CP3, CP1), then a learned per-node readout collapses the electrodes to the same (F2, 1, T) shape the depthwise conv produced — so Layers 1/3/4 are unchanged. A spatial="depthwise" flag restores the original EEGNet for a controlled comparison.

The adjacency is built from 2-D positions for the 22 channels of BCI Competition IV-2a and travels with the checkpoint (a saved buffer), so from_pretrained restores the exact graph you trained on.

Usage

The model is a custom PyTorchModelHubMixin module, so you need its class definition (eegnet_gnn.py, included in this repo) alongside the weights.

from eegnet_gnn import EEGNetGNN
import torch

model = EEGNetGNN.from_pretrained("shemalfoy/eegnet-gnn-features").eval()

x = torch.randn(1, 1, 22, 1000)     # (batch, 1, channels, time)
features = model(x)                 # (1, model.flat_dim) == (1, 496)

# attach your own classifier
head = torch.nn.Linear(model.flat_dim, 4)
logits = head(features)

Pull the class straight from the repo if you don't have the file locally:

import importlib.util
from huggingface_hub import hf_hub_download
path = hf_hub_download("shemalfoy/eegnet-gnn-features", "eegnet_gnn.py")
spec = importlib.util.spec_from_file_location("eegnet_gnn", path)
mod = importlib.util.module_from_spec(spec); spec.loader.exec_module(mod)
model = mod.EEGNetGNN.from_pretrained("shemalfoy/eegnet-gnn-features")

Inputs / outputs

• Input: (batch, 1, 22, T) float tensor — 22 EEG channels in BCI IV-2a order, T samples.
• Output: (batch, flat_dim) feature vector (flat_dim = 496 with defaults).

Dependencies

torch, huggingface_hub, safetensors. No quantum / PennyLane dependency.

Notes and limitations

• A different electrode set requires rebuilding the adjacency (build_adjacency(coords=...)) and retraining.
• Built on EEGNet (Lawhern et al., 2018) and graph convolution (Kipf & Welling, 2017).