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.
29 lines
451 B
YAML
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
|