File size: 2,532 Bytes
c5a22df |
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 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
# -*- coding: utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.
import argparse
import importlib
import os
import sys
from datetime import datetime
sys.dont_write_bytecode = True
from scepter.modules.solver.registry import SOLVERS
from scepter.modules.utils.config import Config
from scepter.modules.utils.distribute import we
from scepter.modules.utils.file_system import FS
from scepter.modules.utils.logger import get_logger
if os.path.exists('__init__.py'):
package_name = 'scepter_ext'
spec = importlib.util.spec_from_file_location(package_name, '__init__.py')
package = importlib.util.module_from_spec(spec)
sys.modules[package_name] = package
spec.loader.exec_module(package)
def run_task(cfg):
std_logger = get_logger(name='scepter')
solver = SOLVERS.build(cfg.SOLVER, logger=std_logger)
solver.set_up_pre()
solver.set_up()
if we.rank == 0:
FS.put_object_from_local_file(cfg.args.cfg_file, os.path.join(solver.work_dir, "train.yaml"))
if cfg.args.stage == "train":
solver.solve()
elif cfg.args.stage == "eval":
solver.run_eval()
def update_config(cfg):
if hasattr(cfg.args, 'learning_rate') and cfg.args.learning_rate:
print(
f'learning_rate change from {cfg.SOLVER.OPTIMIZER.LEARNING_RATE} to {cfg.args.learning_rate}'
)
cfg.SOLVER.OPTIMIZER.LEARNING_RATE = float(cfg.args.learning_rate)
if hasattr(cfg.args, 'max_steps') and cfg.args.max_steps:
print(
f'max_steps change from {cfg.SOLVER.MAX_STEPS} to {cfg.args.max_steps}'
)
cfg.SOLVER.MAX_STEPS = int(cfg.args.max_steps)
cfg.SOLVER.WORK_DIR = os.path.join(cfg.SOLVER.WORK_DIR, "{0:%Y%m%d%H%M%S}".format(datetime.now()))
return cfg
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Argparser for Scepter:\n')
parser.add_argument(
"--stage",
dest="stage",
help="Running stage!",
default="train",
choices=["train", "eval"]
)
parser.add_argument('--learning_rate',
dest='learning_rate',
help='The learning rate for our network!',
default=None)
parser.add_argument('--max_steps',
dest='max_steps',
help='The max steps for training!',
default=None)
cfg = Config(load=True, parser_ins=parser)
cfg = update_config(cfg)
we.init_env(cfg, logger=None, fn=run_task)
|