LED/main_sport_led.py

import argparse
from trainer import train_sport_led as sl


def parse_config():
    parser = argparse.ArgumentParser()
    parser.add_argument('--cuda', default=True)
    parser.add_argument('--learning_rate', type=float, default=2e-3)
    parser.add_argument('--cfg', default='soccer', help='Sport cfg name (soccer / football)')
    parser.add_argument('--gpu', type=int, default=0)
    parser.add_argument('--train', type=int, default=1)
    parser.add_argument('--info', type=str, default='baseline')

    parser.add_argument('--use_graph', action='store_true',
                        help='If set, use the FutureInteractionGraph residual '
                             '(default: baseline LED).')
    parser.add_argument('--residual_on', type=str, default='y0',
                        choices=['eps', 'y0'],
                        help='Where the graph residual is applied (only used with --use_graph).')
    parser.add_argument('--use_v6_graph', action='store_true',
                        help='If set, use MoFlow V6-style RAG-scoring graph instead of '
                             'the hand-crafted distance-based one.')
    return parser.parse_args()


def main(config):
    t = sl.Trainer(config)
    t.fit()


if __name__ == '__main__':
    config = parse_config()
    main(config)