Files
solver/configs/model/gat.yaml
forbes fe41fa3b00 feat(models): implement GNN model layer for assembly constraint analysis
Add complete model layer with GIN and GAT encoders, 5 task-specific
prediction heads, uncertainty-weighted multi-task loss (Kendall et al.
2018), and config-driven factory functions.

New modules:
- graph_conv.py: datagen dict -> PyG Data conversion (22D node/edge features)
- encoder.py: GINEncoder (3-layer, 128D) and GATEncoder (4-layer, 256D, 8-head)
- heads.py: edge classification, graph classification, joint type,
  DOF regression, per-body DOF tracking heads
- assembly_gnn.py: AssemblyGNN wiring encoder + configurable heads
- losses.py: MultiTaskLoss with learnable log-variance per task
- factory.py: build_model() and build_loss() from YAML configs

Supporting changes:
- generator.py: serialize anchor_a, anchor_b, pitch in joint dicts
- configs: fix joint_type num_classes 12 -> 11 (matches JointType enum)

92 tests covering shapes, gradients, edge cases, and end-to-end
datagen-to-model pipeline.
2026-02-07 10:14:19 -06:00

29 lines
451 B
YAML

# Advanced GAT model config
name: gat_solver
architecture: gat
encoder:
num_layers: 4
hidden_dim: 256
num_heads: 8
dropout: 0.1
residual: true
node_features_dim: 22
edge_features_dim: 22
heads:
edge_classification:
enabled: true
hidden_dim: 128
graph_classification:
enabled: true
num_classes: 4
joint_type:
enabled: true
num_classes: 11
dof_regression:
enabled: true
dof_tracking:
enabled: true