|
import shutil |
|
import os |
|
import multiprocessing as mp |
|
if mp.get_start_method(allow_none=True) != 'spawn': |
|
mp.set_start_method('spawn') |
|
import argparse |
|
from core.distributed_utils import dist_init |
|
from core.config import Config_Hulk |
|
from core.testers import tester_entry |
|
import torch |
|
import yaml |
|
import re |
|
import torch.distributed as dist |
|
|
|
parser = argparse.ArgumentParser(description='Multi-Task Training Framework') |
|
parser.add_argument('--spec_ginfo_index', type=int, required=True) |
|
parser.add_argument('--load-path', default='', type=str) |
|
parser.add_argument('--ignore', nargs='+', default=[], type=str) |
|
parser.add_argument('--recover', action='store_true') |
|
parser.add_argument('--load-single', action='store_true') |
|
parser.add_argument('--port', default='23456', type=str) |
|
parser.add_argument('--config', default='', type=str) |
|
parser.add_argument('--test_config', default='', type=str) |
|
parser.add_argument('--expname', type=str, default=None, help='experiment name, output folder') |
|
parser.add_argument('--auto-resume', type=str, default=None, help='jobs auto resume from the pattern_path or the folder') |
|
|
|
loader = yaml.SafeLoader |
|
loader.add_implicit_resolver( |
|
u'tag:yaml.org,2002:float', |
|
re.compile(u'''^(?: |
|
[-+]?(?:[0-9][0-9_]*)\\.[0-9_]*(?:[eE][-+]?[0-9]+)? |
|
|[-+]?(?:[0-9][0-9_]*)(?:[eE][-+]?[0-9]+) |
|
|\\.[0-9_]+(?:[eE][-+][0-9]+)? |
|
|[-+]?[0-9][0-9_]*(?::[0-5]?[0-9])+\\.[0-9_]* |
|
|[-+]?\\.(?:inf|Inf|INF) |
|
|\\.(?:nan|NaN|NAN))$''', re.X), |
|
list(u'-+0123456789.')) |
|
|
|
def main(): |
|
args = parser.parse_args() |
|
dist_init() |
|
|
|
|
|
if args.auto_resume is not None: |
|
if os.path.isdir(args.auto_resume): |
|
max_iter = 0 |
|
filename = os.listdir(args.auto_resume) |
|
for file in filename: |
|
if file.startswith('ckpt_task0') and file.endswith('.pth.tar'): |
|
cur_iter = int(file.split('_')[-1].split('.')[0]) |
|
max_iter = max(max_iter, cur_iter) |
|
if max_iter > 0: |
|
args.load_path = os.path.join(args.auto_resume, |
|
'ckpt_task_iter_{}.pth.tar'.format(str(max_iter))) |
|
args.recover = True |
|
args.ignore = [] |
|
print('auto-resume from: {}'.format(args.load_path)) |
|
elif args.auto_resume.endswith('.pth.tar'): |
|
tmpl = args.auto_resume.replace('ckpt_task_', 'ckpt_task*_') |
|
import glob |
|
ckpt = glob.glob(tmpl) |
|
if len(ckpt) > 0: |
|
args.load_path = args.auto_resume |
|
args.recover = True |
|
args.ignore = [] |
|
print('auto-resume from: {}'.format(args.load_path)) |
|
else: |
|
print('auto-resume not work:{}'.format(args.auto_resume)) |
|
|
|
C_train = Config_Hulk(args.config, spec_ginfo_index=args.spec_ginfo_index) |
|
|
|
with open(args.test_config) as f: |
|
test_config = yaml.load(f, Loader=loader) |
|
num_test_tasks = len(test_config['tasks']) |
|
|
|
for test_spec_ginfo_index in range(num_test_tasks): |
|
C_test = Config_Hulk(args.test_config, spec_ginfo_index=test_spec_ginfo_index) |
|
if args.expname is not None: |
|
C_train.config['expname'] = args.expname |
|
|
|
S = tester_entry(C_train, C_test) |
|
config_save_to = os.path.join(S.ckpt_path, 'config.yaml') |
|
test_config_save_to = os.path.join(S.ckpt_path, 'test_config_task{}.yaml'.format(test_spec_ginfo_index)) |
|
if not os.path.exists(config_save_to): |
|
shutil.copy(args.config, config_save_to) |
|
shutil.copy(args.test_config, test_config_save_to) |
|
|
|
S.initialize(args) |
|
|
|
S.run() |
|
|
|
|
|
if __name__ == '__main__': |
|
main() |
|
|
