shemalfoy/eegnet-gnn-features-f2is12

EEGNet-GNN feature extractor (F1=6, F2=12)

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

The model stops at Layer 4 and returns a flat feature vector; EEGNetGNNClassifier attaches a linear head. Output size is exposed as model.flat_dim (372 with the defaults: F1=6, F2=12, pool1=4, pool2=8, n_times=1000).

The swap

Standard EEGNet's Layer 2 learns, per output map, a single signed weighted sum over all electrodes — topology is ignored, but the signs let it compute contrasts. Here that becomes H = Â X W + X W_res, where Â is the symmetric-normalised adjacency of the montage. Each electrode aggregates from its physical neighbours (e.g. C3 from FC3, FC1, C5, C1, CP3, CP1); a residual path preserves un-smoothed focal detail; then a learned signed, per-feature-map readout (F2, N) collapses electrodes to the same (F2, 1, T) shape the depthwise conv produced, so Layers 1/3/4 are unchanged.

The readout must stay signed. An earlier version used softmax(node_weight), giving a non-negative convex combination shared across all feature maps. That cannot represent a difference between electrodes, and the C3-vs-C4 ERD asymmetry is the dominant motor imagery feature — so the model sat at chance (~25% on 4-class) no matter how long it trained.

spatial="depthwise" restores the original EEGNet for a controlled A/B.

Options

Usage

from eegnet_gnn import EEGNetGNN, EEGNetGNNClassifier
import torch

backbone = EEGNetGNN(spatial="graph", F1=6, F2=12)
model = EEGNetGNNClassifier(backbone, n_classes=4)

x = torch.randn(1, 1, 22, 1000)     # (batch, 1, channels, time)
logits = model(x)                   # (1, 4)

Train the whole thing end-to-end. Do not freeze the backbone unless you are loading genuinely pretrained weights — freezing a random init reduces the model to logistic regression on random projections.