예제 #1
0
def train(args):
    with open(os.path.join('config', 'procgen.yaml')) as f:
        procgen_config = yaml.load(f, Loader=yaml.SafeLoader)

    with open(os.path.join('config', args.agent+'.yaml')) as f:
        agent_config = yaml.load(f, Loader=yaml.SafeLoader)

    device = torch.device(
        'cuda' if (torch.cuda.is_available() and args.cuda) else 'cpu')

    time = datetime.now().strftime("%Y%m%d-%H%M")
    log_dir = os.path.join(
        args.log_dir,
        args.agent,
        args.env_name.split('NoFrameskip')[0],
        args.agent+f'-{args.seed}-{time}')

    # fix seed
    fix_seed(args.seed)

    # make envs
    envs = make_env(args.env_name, args.num_processes, device, **procgen_config)
    eval_envs = make_env(args.env_name, 1, device, **procgen_config)

    # make agent
    agent, param_file = AGENT_CLASS[args.agent]

    # run
    agent(envs, eval_envs, device, log_dir, args.num_processes,
          **agent_config).run()
def pretrain(cfg):
    print(cfg.pretty())
    pretrain_config_validator(cfg)
    fix_seed(cfg.seed)

    controller = load_pretrained_weights(
        NAO(**cfg.controller).to(0), cfg.pretrained_model_path)
    models = {'trunk': controller}
    dataset = get_dataset(seed=cfg.seed, **cfg.dataset)
    optimizers = {
        'trunk_optimizer':
        get_optimizer(parameters=models['trunk'].parameters(), **cfg.optimizer)
    }
    lr_schedulers = {
        'trunk_scheduler_by_iteration':
        get_scheduler(optimizer=optimizers['trunk_optimizer'], **cfg.scheduler)
    }
    loss_funcs = {
        'reconstruction_loss': torch.nn.NLLLoss(),
        'metric_loss': get_loss(**cfg.loss)
    }
    mining_funcs = {"tuple_miner": get_miner(**cfg.miner)}
    visualizers = [umap.UMAP(**params) for params in cfg.visualizers]
    end_of_iteration_hook = TensorboardHook(visualizers).end_of_iteration_hook
    end_of_epoch_hook = ModelSaverHook().end_of_epoch_hook
    get_trainer(
        models=models,
        optimizers=optimizers,
        lr_schedulers=lr_schedulers,
        loss_funcs=loss_funcs,
        mining_funcs=mining_funcs,
        dataset=dataset,
        end_of_iteration_hook=end_of_iteration_hook,
        end_of_epoch_hook=end_of_epoch_hook,
        **cfg.trainer,
    ).train()
예제 #3
0
def train(cfg):
    print(cfg.pretty())
    train_config_validator(cfg)
    fix_seed(cfg.seed)

    writer = SummaryWriter(log_dir='logs')
    controller = load_pretrained_weights(
        NAO(**cfg.controller).to(0), cfg.pretrained_model_path)
    dataset = get_dataset(writer=writer, seed=cfg.seed, **cfg.dataset)
    optimizer = get_optimizer(parameters=_get_target_parameters(
        controller, cfg.freeze_encoder_decoder),
                              **cfg.optimizer)
    lr_scheduler = get_scheduler(optimizer=optimizer, **cfg.scheduler)
    end_of_epoch_hook = ModelSaverHook().end_of_epoch_hook

    get_trainer(
        controller=controller,
        dataset=dataset,
        optimizer=optimizer,
        lr_scheduler=lr_scheduler,
        writer=writer,
        end_of_epoch_hook=end_of_epoch_hook,
        **cfg.trainer,
    ).train()
예제 #4
0
        '~/kaggle/landmark_recognition_2020/logs/Landmarks/4a293h13/checkpoints'
    )
    SUBMISSION_PATH = 'submission.csv'
    DEVICE = 'cuda:0'
    BATCH_SIZE = 512
    NUM_WORKERS = 20
else:
    raise ValueError("Unknown environment exception")

CHECKPOINT_NAME = 'epoch_1.ckpt'
NORMALIZE_VECTORS = True
LOAD_VECTORS_FROM_CHECKPOINT = False
TOPK = 5
SEED = 17
DEVICE = torch.device(DEVICE)
fix_seed(SEED)


def main():
    start_time = datetime.datetime.now()
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        level=logging.DEBUG,
    )
    logger = logging.getLogger(__name__)

    # load config file
    model_args, training_args = load_or_parse_args(
        (ModelArgs, TrainingArgs),
        verbose=True,