Exemplo n.º 1
0
def main(cfg):
    log = logging.getLogger(__name__)

    # Get device
    device = torch.device("cuda" if (
        torch.cuda.is_available() and cfg.eval.cuda) else "cpu")
    print("DEVICE : {}".format(device))

    # Get task and model_name
    tested_task = cfg.data.task
    tested_model_name = cfg.model_name

    # Find and create associated model
    model_config = getattr(cfg.models, tested_model_name, None)

    cfg_eval = set_format(model_config, cfg.eval)

    # Enable CUDNN BACKEND
    torch.backends.cudnn.enabled = cfg_eval.enable_cudnn

    # Find and create associated dataset
    dataset_config = cfg.data
    tested_dataset_name = dataset_config.name
    dataset_config.dataroot = hydra.utils.to_absolute_path(
        dataset_config.dataroot)
    dataset = find_dataset_using_name(tested_dataset_name,
                                      tested_task)(dataset_config, cfg_eval)

    # Find and create associated model
    resolve_model(model_config, dataset, tested_task)
    model_config = merges_in_sub(model_config, [cfg_eval, dataset_config])
    model = find_model_using_name(model_config.architecture, tested_task,
                                  model_config, dataset)

    log.info(model)

    model.eval()
    if cfg_eval.enable_dropout:
        model.enable_dropout_in_eval()

    # Set sampling / search strategies
    dataset.set_strategies(
        model, precompute_multi_scale=cfg_eval.precompute_multi_scale)

    model = model.to(device)
    model_parameters = filter(lambda p: p.requires_grad, model.parameters())
    params = sum([np.prod(p.size()) for p in model_parameters])
    log.info("Model size = %i", params)

    tracker: BaseTracker = dataset.get_tracker(model, tested_task, dataset,
                                               cfg.wandb, cfg.tensorboard)

    checkpoint = get_model_checkpoint(model, cfg_eval.checkpoint_dir,
                                      tested_model_name, True,
                                      cfg_eval.weight_name, "test")

    # Run training / evaluation
    run(cfg, model, dataset, device, tracker, checkpoint, log)
Exemplo n.º 2
0
def main(cfg):
    log = logging.getLogger(__name__)

    # Get device
    device = torch.device("cuda" if (
        torch.cuda.is_available() and cfg.training.cuda) else "cpu")
    print("DEVICE : {}".format(device))

    # Get task and model_name
    tested_task = cfg.data.task
    tested_model_name = cfg.model_name

    # Find and create associated model
    model_config = getattr(cfg.models, tested_model_name, None)

    # Find which dataloader to use
    cfg_training = set_format(model_config, cfg.training)

    # Enable CUDNN BACKEND
    torch.backends.cudnn.enabled = cfg_training.enable_cudnn

    # Find and create associated dataset
    dataset_config = cfg.data
    tested_dataset_name = dataset_config.name
    dataset_config.dataroot = hydra.utils.to_absolute_path(
        dataset_config.dataroot)
    dataset = find_dataset_using_name(tested_dataset_name,
                                      tested_task)(dataset_config,
                                                   cfg_training)

    # Find and create associated model
    resolve_model(model_config, dataset, tested_task)
    model_config = merges_in_sub(model_config, [cfg_training, dataset_config])
    model = find_model_using_name(model_config.architecture, tested_task,
                                  model_config, dataset)

    log.info(model)

    # Optimizer
    lr_params = cfg_training.learning_rate
    model.set_optimizer(getattr(torch.optim, cfg_training.optimizer, None),
                        lr_params=lr_params)

    # Set sampling / search strategies
    dataset.set_strategies(
        model, precompute_multi_scale=cfg_training.precompute_multi_scale)

    model = model.to(device)
    model_parameters = filter(lambda p: p.requires_grad, model.parameters())
    params = sum([np.prod(p.size()) for p in model_parameters])
    log.info("Model size = %i", params)

    # metric tracker
    if cfg.wandb.log:
        import wandb

        wandb.init(project=cfg.wandb.project)
        # wandb.watch(model)

    tracker: BaseTracker = dataset.get_tracker(model, tested_task, dataset,
                                               cfg.wandb, cfg.tensorboard)

    checkpoint = get_model_checkpoint(
        model,
        cfg_training.checkpoint_dir,
        tested_model_name,
        cfg_training.resume,
        cfg_training.weight_name,
        "val" if dataset.has_val_loader else "test",
    )

    # Run training / evaluation
    run(cfg, model, dataset, device, tracker, checkpoint, log)
Exemplo n.º 3
0
def _find_model_using_name(model_class, task, model_config, dataset):
    resolve_model(model_config, dataset, task)
    return find_model_using_name(model_class, task, model_config, dataset)