-
Notifications
You must be signed in to change notification settings - Fork 5
/
main_mix.py
56 lines (45 loc) · 1.36 KB
/
main_mix.py
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
55
56
from motion_vqvae_mix import MoQ
import argparse
import os
import yaml
from pprint import pprint
from easydict import EasyDict
def parse_args():
parser = argparse.ArgumentParser(
description='Pytorch implementation of Music2Dance')
parser.add_argument('--config', default='')
# exclusive arguments
group = parser.add_mutually_exclusive_group(required=True)
group.add_argument('--train', action='store_true')
group.add_argument('--eval', action='store_true')
group.add_argument('--visgt', action='store_true')
group.add_argument('--anl', action='store_true')
group.add_argument('--sample', action='store_true')
return parser.parse_args()
def main():
# parse arguments and load config
args = parse_args()
with open(args.config) as f:
config = yaml.load(f)
for k, v in vars(args).items():
config[k] = v
pprint(config)
config = EasyDict(config)
agent = MoQ(config)
print(config)
if args.train:
agent.train()
elif args.eval:
agent.eval()
elif args.visgt:
# print('Wula!')
agent.visgt()
elif args.anl:
# print('Wula!')
agent.analyze_code()
elif args.sample:
config.update({'need_not_train_data':True})
config.update({'need_not_test_data':True})
agent.sample()
if __name__ == '__main__':
main()