-
Notifications
You must be signed in to change notification settings - Fork 2
/
launch.py
24 lines (21 loc) · 1.04 KB
/
launch.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
import argparse
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Launcher")
parser.add_argument("--task", default="train", type=str,
choices=["train", "val", "export", "visualize"], help="task to run")
parser.add_argument("--config", required=True, help="config to load")
parser.add_argument("--epoch", type=int, help="epoch to val or test")
parser.add_argument("--image", type=str, default="", help="image info")
args = parser.parse_args()
if args.task == "train":
from poi.apis.train import train_net
train_net(config=args.config)
elif args.task == "val":
from poi.apis.test import test_net
test_net(config=args.config, task="val", epoch=args.epoch)
elif args.task == "export":
from poi.apis.export import export_net
export_net(config=args.config, epoch=args.epoch)
elif args.task == "visualize":
from poi.apis.visualize import visualize_net
visualize_net(config=args.config, epoch=args.epoch, image=args.image)