File size: 2,364 Bytes
d4cbafd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 | import argparse
from trainer import train_led_graph as led
def parse_config():
parser = argparse.ArgumentParser()
parser.add_argument("--cuda", default=True)
parser.add_argument("--learning_rate", type=int, default=0.002)
parser.add_argument("--max_epochs", type=int, default=128)
parser.add_argument('--cfg', default='led_augment')
parser.add_argument('--gpu', type=int, default=0, help='Specify which GPU to use.')
parser.add_argument('--train', type=int, default=1, help='Whether train or evaluate.')
parser.add_argument("--info", type=str, default='graph', help='Name of the experiment. '
'It will be used in file creation.')
# Graph variant knobs.
parser.add_argument('--top_n', type=int, default=5,
help='Number of sparse neighbors per agent (max 10 for NBA with A=11).')
parser.add_argument('--residual_on', type=str, default='eps', choices=['eps', 'y0'],
help='Where to apply the graph residual: directly on epsilon, or '
'on the implied y_0 estimate (then re-projected to epsilon).')
parser.add_argument('--use_sigma', action='store_true',
help='If set, pass the initializer variance_estimation to the graph '
'as per-agent uncertainty (modulates node features and edges).')
parser.add_argument('--use_v6_graph', action='store_true',
help='If set, use MoFlow V6-style RAG-scoring graph (FutureInteractionGraphV6) '
'instead of the default hand-crafted distance-based graph.')
parser.add_argument('--uncertainty_weight', type=float, default=1.0,
help='Weight for uncertainty NLL loss. Set to 0 for nosigma ablation.')
parser.add_argument('--edge_mode', type=str, default='full',
choices=['full', 'dist_only', 'relpos_only', 'heading_only', 'vel_only', 'full_relvel'],
help='Edge feature mode for RelTrajEncoder ablation.')
parser.add_argument('--resume_epoch', type=int, default=0,
help='Resume from this checkpoint epoch (0=start fresh).')
parser.add_argument('--neighbor_mode', type=str, default='rag',
choices=['rag', 'l2', 'semantic'],
help='Neighbor selection: rag (semantic+geo), l2 (closest), or semantic (learned only)')
return parser.parse_args()
def main(config):
t = led.Trainer(config)
if config.train == 1:
t.fit()
else:
t.test_single_model()
if __name__ == "__main__":
config = parse_config()
main(config)
|