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)
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)
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)