forked from microsoft/SDNet
-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
33 lines (26 loc) · 935 Bytes
/
main.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
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.
import argparse
import os
import sys
import torch
from Models.SDNetTrainer import SDNetTrainer
from Utils.Arguments import Arguments
opt = None
parser = argparse.ArgumentParser(description='SDNet')
parser.add_argument('command', help='Command: train')
parser.add_argument('conf_file', help='Path to conf file.')
cmdline_args = parser.parse_args()
command = cmdline_args.command
conf_file = cmdline_args.conf_file
conf_args = Arguments(conf_file)
opt = conf_args.readArguments()
opt['cuda'] = torch.cuda.is_available()
opt['confFile'] = conf_file
opt['datadir'] = os.path.dirname(conf_file) # conf_file specifies where the data folder is
for key,val in cmdline_args.__dict__.items():
if val is not None and key not in ['command', 'conf_file']:
opt[key] = val
model = SDNetTrainer(opt)
print('Select command: ' + command)
model.train()