rokati/focal_gnn_one_graph

Heterogeneous GNN xG Prediction Model (Focal Loss, One Graph Architecture)

This is a Heterogeneous Graph Neural Network (GNN) model trained to predict Expected Goals (xG) in football/soccer using a single persistent graph with shot-indexed edges.

Model Description

• —Architecture: Heterogeneous GNN with GAT attention layers
• —Graph Structure: One persistent graph with all nodes, shot-specific edge masking
• —Nodes: shooter (player embeddings, 496 total), goal (learnable embedding), goalkeeper (learnable embedding)
• —Edges: 4 edge types (bidirectional shooter↔goal and shooter↔goalkeeper)
• —Global Features: 18 shot-level contextual features
• —Hidden Dimensions: 64
• —Number of Layers: 3 GAT layers
• —Attention Heads: 4
• —Dropout Rate: 0.3
• —Framework: PyTorch Geometric
• —Loss Function: Focal Loss (alpha=0.8773, gamma=2.0)

Performance Metrics

• —Test Loss: 0.08266454190015793
• —Test AUC: 0.609370231628418
• —Test Accuracy: 0.14332698285579681
• —Test F1: 0.2105599045753479

Global Features (18 features)

The model uses the following contextual features for each shot:

• —ballcloserthan_gk
• —bodypartname_Left Foot
• —bodypartname_Other
• —bodypartname_Right Foot
• —goaldistto_gk
• —minute
• —nearestopponentdist
• —nearestteammatedist
• —opponentswithin5m
• —playpatternname_From Counter
• —playpatternname_From Free Kick
• —playpatternname_From Goal Kick
• —playpatternname_From Keeper
• —playpatternname_From Kick Off
• —playpatternname_From Throw In
• —playpatternname_Other
• —playpatternname_Regular Play
• —teammateswithin5m

Graph Structure

Persistent Nodes

1. 1.Shooter Nodes: 496 player embeddings (dimension: 64)
2. 2.Goal Node: 1 learnable goal representation (dimension: 64)
3. 3.Goalkeeper Node: 1 learnable goalkeeper representation (dimension: 64)

Edge Types (with attributes, shot-indexed)

1. 1.shooter → goal: Distance (meters) + Angle to goal (radians)
2. 2.goal → shooter: Reverse edges with same attributes
3. 3.shooter → goalkeeper: Distance to goalkeeper (meters)
4. 4.goalkeeper → shooter: Reverse edges with same attributes

All edges are indexed by shot_idx for efficient masking during prediction.

Usage

import torch
from torch_geometric.data import HeteroData
from huggingface_hub import hf_hub_download
import importlib.util
import json

# Download files
model_path = hf_hub_download(repo_id="rokati/focal_gnn_one_graph", filename="best_gnn_model.pth")
architecture_path = hf_hub_download(repo_id="rokati/focal_gnn_one_graph", filename="model_architecture.py")
config_path = hf_hub_download(repo_id="rokati/focal_gnn_one_graph", filename="config.json")

# Load configuration
with open(config_path, 'r') as f:
    config = json.load(f)

# Load architecture
spec = importlib.util.spec_from_file_location("model_architecture", architecture_path)
model_module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(model_module)

# Create model instance
model = model_module.XGNet(
    num_players=config['num_players'],
    hid=config['hidden_dim'],
    p=config['dropout_rate'],
    heads=config['num_heads'],
    num_layers=config['num_layers'],
    use_norm=config['use_norm'],
    num_global_features=config['num_global_features']
)

# Load weights
model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
model.eval()

# Prepare persistent graph (build once with all players/shots)
# Then predict by passing the graph and shot_idx:
with torch.no_grad():
    xg_prediction = torch.sigmoid(model(graph, shot_idx)).item()

Training Details

The model was trained with:

• —Loss Function: Focal Loss (addresses class imbalance)
• —Optimizer: Adam (lr=1e-3, weight_decay=1e-4)
• —Scheduler: ReduceLROnPlateau
• —Batch Size: 32
• —Max Epochs: 100
• —Early Stopping: Patience=15 on validation F1
• —Framework: PyTorch Lightning

Model Architecture

The heterogeneous GNN uses:

1. 1.Node Embeddings: Learnable embeddings for shooters, goal, and goalkeeper
2. 2.Edge-Conditioned Attention: GAT layers with edge attributes
3. 3.Bidirectional Message Passing: 3 layers with forward and reverse edges
4. 4.Shot-Indexed Masking: Efficient prediction via edge masking
5. 5.Global Context: Shot-level features encoded and combined
6. 6.Readout: Concatenate shooter + goal + goalkeeper representations → MLP classifier

Key Innovation

Unlike traditional approaches that create separate graphs for each shot, this model uses:

• —One persistent graph with all nodes
• —Shot-indexed edges for efficient masking
• —Bidirectional message passing for richer representations
• —This architecture is more memory-efficient and allows for better batch processing

License

MIT