File size: 1,297 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 | 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)
|