def prepare(self, pin_memory: bool = False): """ Prepares the task: - dataloaders - model - copy model to correct device - meters - loss - optimizer - LR schedulers - AMP state - resume from a checkpoint if available """ self.dataloaders = self.build_dataloaders(pin_memory=pin_memory) self.phases = self._build_phases() train_phases = [phase for phase in self.phases if phase["train"]] num_train_phases = len(train_phases) self.base_model = self._build_model() self._set_ddp_options() self.base_loss = self._build_loss() self.meters = self._build_meters() self.optimizer = self._build_optimizer() self.optimizer_schedulers = self._build_optimizer_schedulers() self.num_train_phases = num_train_phases self.base_loss = self.base_loss.to(self.device) if self.device.type == "cuda": self.base_model = copy_model_to_gpu(self.base_model) # initialize the pytorch optimizer now since the model has been moved to # the appropriate device. self.prepare_optimizer() # Enable mixed precision grad scalers if self.amp_type == AmpType.APEX: # Allow Apex Amp to perform casts as specified by the amp_args. # This updates the model and the PyTorch optimizer (which is wrapped # by the ClassyOptimizer in self.optimizer). # NOTE: this must happen before loading the checkpoint. See # https://nvidia.github.io/apex/amp.html#checkpointing for more details. self.base_model, self.optimizer.optimizer = apex.amp.initialize( self.base_model, self.optimizer.optimizer, **self.amp_args) # Restore an hypothetical checkpoint vissl_state_dict = None if self.checkpoint_path is not None: self.checkpoint = CheckpointLoader.load_and_broadcast_checkpoint( checkpoint_folder=self.checkpoint_folder, checkpoint_path=self.checkpoint_path, device=torch.device("cpu"), ) self.iteration = self.checkpoint["iteration"] self.local_iteration_num = self.checkpoint["iteration_num"] vissl_state_dict = self.checkpoint.get("classy_state_dict") if "loss" in self.checkpoint: self.base_loss.load_state_dict(self.checkpoint["loss"]) logging.info("======Loaded loss state from checkpoint======") return self._update_classy_state(vissl_state_dict)
def prepare(self): """Prepares task for training, populates all derived attributes """ pin_memory = self.use_gpu and torch.cuda.device_count() > 1 self.phases = self._build_phases() self.train = False if self.test_only else self.train self.dataloaders = self.build_dataloaders( pin_memory=pin_memory, multiprocessing_context=mp.get_context(self.dataloader_mp_context), ) if self.batch_norm_sync_mode == BatchNormSyncMode.PYTORCH: self.base_model = nn.SyncBatchNorm.convert_sync_batchnorm( self.base_model) elif self.batch_norm_sync_mode == BatchNormSyncMode.APEX: self.base_model = apex.parallel.convert_syncbn_model( self.base_model) # move the model and loss to the right device if self.use_gpu: self.base_model, self.loss = copy_model_to_gpu( self.base_model, self.loss) else: self.loss.cpu() self.base_model.cpu() if self.optimizer is not None: # initialize the pytorch optimizer now since the model has been moved to # the appropriate device self.optimizer.init_pytorch_optimizer(self.base_model, loss=self.loss) if self.amp_args is not None: # Initialize apex.amp. This updates the model and the PyTorch optimizer ( # if training, which is wrapped by the ClassyOptimizer in self.optimizer). # Please note this must happen before loading the checkpoint, cause # there's amp state to be restored. if self.optimizer is None: self.base_model = apex.amp.initialize(self.base_model, optimizers=None, **self.amp_args) else: self.base_model, self.optimizer.optimizer = apex.amp.initialize( self.base_model, self.optimizer.optimizer, **self.amp_args) if self.checkpoint_path: self.checkpoint_dict = load_and_broadcast_checkpoint( self.checkpoint_path) classy_state_dict = (None if self.checkpoint_dict is None else self.checkpoint_dict["classy_state_dict"]) if classy_state_dict is not None: state_load_success = update_classy_state(self, classy_state_dict) assert (state_load_success ), "Update classy state from checkpoint was unsuccessful." self.init_distributed_data_parallel_model()
def prepare(self, num_dataloader_workers=0, dataloader_mp_context=None): """Prepares task for training, populates all derived attributes Args: num_dataloader_workers: Number of dataloading processes. If 0, dataloading is done on main process dataloader_mp_context: Determines how processes are spawned. Value must be one of None, "spawn", "fork", "forkserver". If None, then context is inherited from parent process """ pin_memory = self.use_gpu and torch.cuda.device_count() > 1 self.phases = self._build_phases() self.dataloaders = self.build_dataloaders( num_workers=num_dataloader_workers, pin_memory=pin_memory, multiprocessing_context=dataloader_mp_context, ) if self.batch_norm_sync_mode == BatchNormSyncMode.PYTORCH: self.base_model = nn.SyncBatchNorm.convert_sync_batchnorm( self.base_model) elif self.batch_norm_sync_mode == BatchNormSyncMode.APEX: self.base_model = apex.parallel.convert_syncbn_model( self.base_model) # move the model and loss to the right device if self.use_gpu: self.base_model, self.loss = copy_model_to_gpu( self.base_model, self.loss) else: self.loss.cpu() self.base_model.cpu() # initialize the pytorch optimizer now since the model has been moved to # the appropriate device self.optimizer.init_pytorch_optimizer(self.base_model, loss=self.loss) if self.amp_args is not None: # Initialize apex.amp. This updates the model and the PyTorch optimizer ( # which is wrapped by the ClassyOptimizer in self.optimizer). # Please note this must happen before loading the checkpoint, cause # there's amp state to be restored. self.base_model, self.optimizer.optimizer = apex.amp.initialize( self.base_model, self.optimizer.optimizer, **self.amp_args) classy_state_dict = (None if self.checkpoint is None else self.checkpoint.get("classy_state_dict")) if classy_state_dict is not None: state_load_success = update_classy_state(self, classy_state_dict) assert (state_load_success ), "Update classy state from checkpoint was unsuccessful." self.init_distributed_data_parallel_model()
def prepare_extraction(self, pin_memory: bool = False): """ Prepares a light-weight task for feature extraction on multi-gpu. The model runs in eval mode only. """ self.dataloaders = self.build_dataloaders(pin_memory=pin_memory) self.base_model = self._build_model() if self.device.type == "cuda": self.base_model = copy_model_to_gpu(self.base_model) return self
def prepare(self): """Prepares task for training, populates all derived attributes """ self.phases = self._build_phases() self.train = False if self.test_only else self.train if self.batch_norm_sync_mode == BatchNormSyncMode.PYTORCH: self.base_model = nn.SyncBatchNorm.convert_sync_batchnorm( self.base_model) elif self.batch_norm_sync_mode == BatchNormSyncMode.APEX: sync_bn_process_group = apex.parallel.create_syncbn_process_group( self.batch_norm_sync_group_size) self.base_model = apex.parallel.convert_syncbn_model( self.base_model, process_group=sync_bn_process_group) # move the model and loss to the right device if self.use_gpu: self.base_model, self.base_loss = copy_model_to_gpu( self.base_model, self.base_loss) else: self.base_loss.cpu() self.base_model.cpu() if self.optimizer is not None: self.prepare_optimizer(optimizer=self.optimizer, model=self.base_model, loss=self.base_loss) if self.amp_args is not None: # Initialize apex.amp. This updates the model and the PyTorch optimizer ( # if training, which is wrapped by the ClassyOptimizer in self.optimizer). # Please note this must happen before loading the checkpoint, cause # there's amp state to be restored. if self.optimizer is None: self.base_model = apex.amp.initialize(self.base_model, optimizers=None, **self.amp_args) else: self.base_model, self.optimizer.optimizer = apex.amp.initialize( self.base_model, self.optimizer.optimizer, **self.amp_args) if self.checkpoint_path: self.checkpoint_dict = load_and_broadcast_checkpoint( self.checkpoint_path) classy_state_dict = (None if self.checkpoint_dict is None else self.checkpoint_dict["classy_state_dict"]) if classy_state_dict is not None: state_load_success = update_classy_state(self, classy_state_dict) assert (state_load_success ), "Update classy state from checkpoint was unsuccessful." self.init_distributed_data_parallel_model()
def prepare( self, num_dataloader_workers=0, pin_memory=False, use_gpu=False, dataloader_mp_context=None, ): """Prepares task for training, populates all derived attributes Args: num_dataloader_workers: Number of dataloading processes. If 0, dataloading is done on main process pin_memory: if true pin memory on GPU use_gpu: if true, load model, optimizer, loss, etc on GPU dataloader_mp_context: Determines how processes are spawned. Value must be one of None, "spawn", "fork", "forkserver". If None, then context is inherited from parent process """ self.phases = self._build_phases() self.dataloaders = self.build_dataloaders( num_workers=num_dataloader_workers, pin_memory=pin_memory, multiprocessing_context=dataloader_mp_context, ) # move the model and loss to the right device if use_gpu: self.base_model, self.loss = copy_model_to_gpu(self.base_model, self.loss) else: self.loss.cpu() self.base_model.cpu() # initialize the pytorch optimizer now since the model has been moved to # the appropriate device self.optimizer.init_pytorch_optimizer(self.base_model, loss=self.loss) classy_state_dict = ( None if self.checkpoint is None else self.checkpoint.get("classy_state_dict") ) if classy_state_dict is not None: state_load_success = update_classy_state(self, classy_state_dict) assert ( state_load_success ), "Update classy state from checkpoint was unsuccessful." if self.amp_opt_level is not None: # Initialize apex.amp. This updates the model and the PyTorch optimizer ( # which is wrapped by the ClassyOptimizer in self.optimizer) self.base_model, self.optimizer.optimizer = apex.amp.initialize( self.base_model, self.optimizer.optimizer, opt_level=self.amp_opt_level ) self.init_distributed_data_parallel_model()
def extract(self): """ Extract workflow supports multi-gpu feature extraction. Since we are only extracting features, only the model is built (and initialized from some model weights file if specified by user). The model is set to the eval mode fully. The features are extracted for whatever data splits (train, val, test) etc that user wants. """ # support feature extraction on gpu only. assert self.task.device.type == "cuda", "Set MACHINE.DEVICE = gpu" self.task.prepare_extraction(pin_memory=self.cfg.DATA.PIN_MEMORY) # in case of feature evaluation mode, if we are freezing both trunk and # head, DDP won't work as there are no parameters in the model. Adding # the dummy head will lead to features being not right. So we rather # add the dummy layer to the model and use DDP. We copy the model to # gpu (if using gpus) after the new dummy layer addition. fully_frozen_model = self.task.base_model.is_fully_frozen_model() if fully_frozen_model: self.task.base_model.dummy_layer = torch.nn.Linear(4, 4) if self.task.device.type == "cuda": self.task.base_model = copy_model_to_gpu(self.task.base_model) self.task.init_distributed_data_parallel_model() if is_primary(): logging.info("Model is:\n {}".format(self.task.model)) # Get the names of the features that we are extracting. If user doesn't # specify the features to evaluate, we get the full model output and freeze # head/trunk both as caution. feat_names = get_trunk_output_feature_names(self.cfg.MODEL) if len(feat_names) == 0: feat_names = ["heads"] features = {} for split in self.task.available_splits: logging.info(f"Extracting features for partition: {split.lower()}") self.task.data_iterator = iter( self.task.dataloaders[split.lower()]) features[split.lower()] = self._get_split_features( feat_names, self.cfg, self.task) logging.info( f"Done getting features for partition: {split.lower()}") if hasattr(self.task, "data_iterator"): del self.task.data_iterator gc.collect() if hasattr(self.task, "dataloaders"): del self.task.dataloaders gc.collect() return features
def add_dummy_layer(self): """ In case of feature evaluation mode, if we are freezing both trunk and head, DDP won't work as there are no parameters in the model. Adding the dummy head will lead to features being not right. So we rather add the dummy layer to the model and use DDP. We copy the model to gpu (if using gpus) after the new dummy layer addition. """ fully_frozen_model = self.base_model.is_fully_frozen_model() if fully_frozen_model: self.base_model.dummy_layer = torch.nn.Linear(4, 4) if self.device.type == "cuda": self.base_model = copy_model_to_gpu(self.base_model)
def prepare_extraction(self, pin_memory: bool = False): """ Prepares a light-weight task for feature extraction on multi-gpu. The model runs in eval mode only. """ self.datasets, self.data_and_label_keys = self.build_datasets() self.dataloaders = self.build_dataloaders(pin_memory=pin_memory) # build the meters in case the extraction is for predictions. self.meters = self._build_meters() self.base_model = self._build_model(strict_load=True) if self.device.type == "cuda": self.base_model = copy_model_to_gpu(self.base_model) return self
def prepare(self, pin_memory: bool = False): """ Prepares the task: - dataloaders - model - copy model to correct device - meters - loss - optimizer - LR schedulers - AMP state - resume from a checkpoint if available """ self.phases = self._build_phases() self.num_phases = len(self.phases) self.base_model = self._build_model() self._set_ddp_options() self.meters = self._build_meters() self.optimizer = self._build_optimizer() self.optimizer_schedulers = self._build_optimizer_schedulers() if self.device.type == "cuda": self.base_model = copy_model_to_gpu(self.base_model) # initialize the pytorch optimizer now since the model has been moved to # the appropriate device. self.prepare_optimizer() # Enable mixed precision grad scalers if self.amp_type == AmpType.APEX: # Allow Apex Amp to perform casts as specified by the amp_args. # This updates the model and the PyTorch optimizer (which is wrapped # by the ClassyOptimizer in self.optimizer). # NOTE: this must happen before loading the checkpoint. See # https://nvidia.github.io/apex/amp.html#checkpointing for more details. self.base_model, self.optimizer.optimizer = apex.amp.initialize( self.base_model, self.optimizer.optimizer, **self.amp_args ) # Create EMA average of the model if hook is specified. ema_config = self.config["HOOKS"]["EMA_MODEL"] if ema_config["ENABLE_EMA_METERS"] or ema_config["SAVE_EMA_MODEL"]: self._create_ema_model() # Restore an hypothetical checkpoint vissl_state_dict = None if self.checkpoint_path is not None: self.checkpoint = CheckpointLoader.load_and_broadcast_checkpoint( checkpoint_folder=self.checkpoint_folder, checkpoint_path=self.checkpoint_path, device=torch.device("cpu"), ) if self.checkpoint is not None: self.iteration = self.checkpoint["iteration"] self.local_iteration_num = self.checkpoint["iteration_num"] vissl_state_dict = self.checkpoint.get("classy_state_dict") else: raise ValueError(f"Could not load checkpoint: {self.checkpoint_path}") current_train_phase_idx = ( vissl_state_dict["train_phase_idx"] + 1 if vissl_state_dict else 0 ) self.datasets, self.data_and_label_keys = self.build_datasets( current_train_phase_idx ) # set dataset state before building dataloader, in order to capture checkpoint info. if vissl_state_dict and "train" in self.datasets: self.datasets["train"].set_classy_state( vissl_state_dict.get("train_dataset_iterator") ) self.dataloaders = self.build_dataloaders( pin_memory=pin_memory, current_train_phase_idx=current_train_phase_idx ) # Build base loss, move to device, and load from checkpoint if applicable self.base_loss = self._build_loss() self.base_loss = self.base_loss.to(self.device) if self.checkpoint and "loss" in self.checkpoint: self.base_loss.load_state_dict(self.checkpoint["loss"]) logging.info("======Loaded loss state from checkpoint======") return self._update_classy_state(vissl_state_dict)
def instance_retrieval_test(args, cfg): # We require 1-gpu for feature extraction. Hence check CUDA is available. assert torch.cuda.is_available(), "CUDA not available, Exit!" train_dataset_name = cfg.IMG_RETRIEVAL.TRAIN_DATASET_NAME eval_dataset_name = cfg.IMG_RETRIEVAL.EVAL_DATASET_NAME spatial_levels = cfg.IMG_RETRIEVAL.SPATIAL_LEVELS resize_img = cfg.IMG_RETRIEVAL.RESIZE_IMG eval_binary_path = cfg.IMG_RETRIEVAL.EVAL_BINARY_PATH root_dataset_path = cfg.IMG_RETRIEVAL.DATASET_PATH temp_dir = f"{cfg.IMG_RETRIEVAL.TEMP_DIR}/{str(uuid.uuid4())}" logging.info(f"Temp directory: {temp_dir}") ############################################################################ # Step 1: Prepare the train/eval datasets, create model and load weights # We only create the train dataset if we need PCA/whitening otherwise # train_dataset is None train_dataset = get_train_dataset(cfg, root_dataset_path, train_dataset_name, eval_binary_path) # create the eval dataset. INSTRE data evaluation requires whitening. eval_dataset = get_eval_dataset(cfg, root_dataset_path, eval_dataset_name, eval_binary_path) # Setup the data transforms (basic) that we apply on the train/eval dataset. transforms = get_transforms(cfg, eval_dataset_name) # Create the image helper image_helper = InstanceRetrievalImageLoader(S=resize_img, transforms=transforms) # Build the model on gpu and set in the eval mode model = build_retrieval_model(cfg) model = copy_model_to_gpu(model) logging.info("Freezing the model.....") model.eval() model.freeze_head_and_trunk() ############################################################################ # Step 2: Extract the features for the train dataset, calculate PCA or # whitening and save if cfg.IMG_RETRIEVAL.SHOULD_TRAIN_PCA_OR_WHITENING: logging.info("Extracting training features...") # the features are already processed based on type: rmac | GeM | l2 norm train_features = get_train_features( cfg, temp_dir, train_dataset_name, resize_img, spatial_levels, image_helper, train_dataset, model, ) ######################################################################## # Train PCA on the train features pca_out_fname = f"{temp_dir}/{train_dataset_name}_S{resize_img}_PCA.pickle" if PathManager.exists(pca_out_fname): logging.info("Loading PCA...") pca = load_pca(pca_out_fname) else: logging.info("Training and saving PCA...") pca = train_and_save_pca(train_features, cfg.IMG_RETRIEVAL.N_PCA, pca_out_fname) else: pca = None ############################################################################ # Step 4: Extract db_features and q_features for the eval dataset logging.info("Extracting Queries features...") features_queries = get_queries_features( cfg, temp_dir, eval_dataset_name, resize_img, spatial_levels, image_helper, eval_dataset, model, pca, ) logging.info("Extracting Dataset features...") features_dataset = get_dataset_features( cfg, temp_dir, eval_dataset_name, resize_img, spatial_levels, image_helper, eval_dataset, model, pca, ) ############################################################################ # Step 5: Compute similarity and score logging.info("Calculating similarity and score...") sim = features_queries.dot(features_dataset.T) logging.info(f"Similarity tensor: {sim.shape}") eval_dataset.score(sim, temp_dir) ############################################################################ # Step 6: cleanup the temp directory logging.info(f"Cleaning up temp directory: {temp_dir}") cleanup_dir(temp_dir) logging.info("All done!!")
def instance_retrieval_test(args, cfg): if (cfg.IMG_RETRIEVAL.USE_FEATURE_EXTRACTION_ENGINE and not cfg.IMG_RETRIEVAL.FEATURE_EXTRACTION_DIR): # We require 1-gpu for feature extraction. Hence check CUDA is available. # If we provide FEATURE_EXTRACTION_DIR, we have already extracted the features # and do not require GPU. assert torch.cuda.is_available(), "CUDA not available, Exit!" train_dataset_name = cfg.IMG_RETRIEVAL.TRAIN_DATASET_NAME eval_dataset_name = cfg.IMG_RETRIEVAL.EVAL_DATASET_NAME spatial_levels = cfg.IMG_RETRIEVAL.SPATIAL_LEVELS resize_img = cfg.IMG_RETRIEVAL.RESIZE_IMG eval_binary_path = cfg.IMG_RETRIEVAL.EVAL_BINARY_PATH root_dataset_path = cfg.IMG_RETRIEVAL.DATASET_PATH save_features = cfg.IMG_RETRIEVAL.SAVE_FEATURES use_feature_extractor = cfg.IMG_RETRIEVAL.USE_FEATURE_EXTRACTION_ENGINE temp_dir = None if save_features: temp_dir = os.path.join(get_checkpoint_folder(cfg), "features") logging.info(f"Temp directory: {temp_dir}") ############################################################################ # Step 1: Prepare the train/eval datasets, create model and load weights # We only create the train dataset if we need PCA/whitening otherwise # train_dataset is None train_dataset = get_train_dataset(cfg, root_dataset_path, train_dataset_name, eval_binary_path) # create the eval dataset. INSTRE data evaluation requires whitening. eval_dataset = get_eval_dataset(cfg, root_dataset_path, eval_dataset_name, eval_binary_path) # Setup the data transforms (basic) that we apply on the train/eval dataset. transforms = get_transforms(cfg, eval_dataset_name) # Create the image helper image_helper = InstanceRetrievalImageLoader( S=resize_img, transforms=transforms, center_crop=cfg.IMG_RETRIEVAL.CENTER_CROP) model = None if not use_feature_extractor: # Build the model on gpu and set in the eval mode model = build_retrieval_model(cfg) model = copy_model_to_gpu(model) logging.info("Freezing the model.....") model.eval() model.freeze_head_and_trunk() ############################################################################ # Step 2: Extract the features for the train dataset, calculate PCA or # whitening and save if cfg.IMG_RETRIEVAL.TRAIN_PCA_WHITENING: logging.info("Extracting training features...") # the features are already processed based on type: rmac | GeM | l2 norm with PerfTimer("get_train_features", PERF_STATS): # TODO: encapsulate the approach "WithFeatureExtractor" from the other one. if use_feature_extractor: input_dir = (cfg.IMG_RETRIEVAL.FEATURE_EXTRACTION_DIR or get_checkpoint_folder(cfg)) input_dir = os.path.join(input_dir, "train_database") train_features = load_and_process_features( cfg, input_dir, "train") else: train_features = extract_train_features( cfg, temp_dir, train_dataset_name, resize_img, spatial_levels, image_helper, train_dataset, model, ) train_features = np.vstack( [x.reshape(-1, x.shape[-1]) for x in train_features]) ######################################################################## # Train PCA on the train features pca_out_fname = None if temp_dir: pca_out_fname = f"{temp_dir}/{train_dataset_name}_S{resize_img}_PCA.pickle" if pca_out_fname and g_pathmgr.exists(pca_out_fname): logging.info("Loading PCA...") pca = load_pca(pca_out_fname) else: logging.info("Training and saving PCA...") pca = train_and_save_pca(train_features, cfg.IMG_RETRIEVAL.N_PCA, pca_out_fname) else: pca = None ############################################################################ # Step 4: Extract db_features and q_features for the eval dataset with PerfTimer("get_query_features", PERF_STATS): logging.info("Extracting Queries features...") # TODO: encapsulate the approach "WithFeatureExtractor" from the other one. if use_feature_extractor: input_dir = (cfg.IMG_RETRIEVAL.FEATURE_EXTRACTION_DIR or get_checkpoint_folder(cfg)) input_dir = os.path.join(input_dir, "query") features_queries = load_and_process_features( cfg, input_dir, "test", pca) else: features_queries = get_queries_features( cfg, temp_dir, eval_dataset_name, resize_img, spatial_levels, image_helper, eval_dataset, model, pca, ) features_queries = np.vstack(features_queries) with PerfTimer("get_dataset_features", PERF_STATS): logging.info("Extracting Dataset features...") # TODO: encapsulate the approach "WithFeatureExtractor" from the other one. if use_feature_extractor: input_dir = (cfg.IMG_RETRIEVAL.FEATURE_EXTRACTION_DIR or get_checkpoint_folder(cfg)) input_dir = os.path.join(input_dir, "train_database") features_dataset = load_and_process_features( cfg, input_dir, "test", pca) else: features_dataset = get_dataset_features( cfg, temp_dir, eval_dataset_name, resize_img, spatial_levels, image_helper, eval_dataset, model, pca, ) features_dataset = np.vstack(features_dataset) ############################################################################ # Step 5: Compute similarity, score, and save results with PerfTimer("scoring_results", PERF_STATS): logging.info("Calculating similarity and score...") if cfg.IMG_RETRIEVAL.SIMILARITY_MEASURE == "cosine_similarity": sim = features_queries.dot(features_dataset.T) elif cfg.IMG_RETRIEVAL.SIMILARITY_MEASURE == "l2": sim = -compute_l2_distance_matrix(features_queries, features_dataset) else: raise ValueError( f"{ cfg.IMG_RETRIEVAL.SIMILARITY_MEASURE } not supported.") logging.info(f"Similarity tensor: {sim.shape}") results = eval_dataset.score(sim, temp_dir) ############################################################################ # Step 6: save results and cleanup the temp directory if cfg.IMG_RETRIEVAL.SAVE_RETRIEVAL_RANKINGS_SCORES: checkpoint_folder = get_checkpoint_folder(cfg) # Save the rankings sim = sim.T ranks = np.argsort(-sim, axis=0) save_file(ranks.T.tolist(), os.path.join(checkpoint_folder, "rankings.json")) # Save the similarity scores save_file(sim.tolist(), os.path.join(checkpoint_folder, "similarity_scores.json")) # Save the result metrics save_file( results, os.path.join(checkpoint_folder, "metrics.json"), append_to_json=False, ) logging.info("All done!!")
def instance_retrieval_test(args, cfg): # We require 1-gpu for feature extraction. Hence check CUDA is available. assert torch.cuda.is_available(), "CUDA not available, Exit!" train_dataset_name = cfg.IMG_RETRIEVAL.TRAIN_DATASET_NAME eval_dataset_name = cfg.IMG_RETRIEVAL.EVAL_DATASET_NAME spatial_levels = cfg.IMG_RETRIEVAL.SPATIAL_LEVELS resize_img = cfg.IMG_RETRIEVAL.RESIZE_IMG eval_binary_path = cfg.IMG_RETRIEVAL.EVAL_BINARY_PATH root_dataset_path = cfg.IMG_RETRIEVAL.DATASET_PATH save_features = cfg.IMG_RETRIEVAL.SAVE_FEATURES temp_dir = None if save_features: temp_dir = os.path.join(get_checkpoint_folder(cfg), "features") logging.info(f"Temp directory: {temp_dir}") ############################################################################ # Step 1: Prepare the train/eval datasets, create model and load weights # We only create the train dataset if we need PCA/whitening otherwise # train_dataset is None train_dataset = get_train_dataset(cfg, root_dataset_path, train_dataset_name, eval_binary_path) # create the eval dataset. INSTRE data evaluation requires whitening. eval_dataset = get_eval_dataset(cfg, root_dataset_path, eval_dataset_name, eval_binary_path) # Setup the data transforms (basic) that we apply on the train/eval dataset. transforms = get_transforms(cfg, eval_dataset_name) # Create the image helper image_helper = InstanceRetrievalImageLoader(S=resize_img, transforms=transforms) # Build the model on gpu and set in the eval mode model = build_retrieval_model(cfg) model = copy_model_to_gpu(model) logging.info("Freezing the model.....") model.eval() model.freeze_head_and_trunk() ############################################################################ # Step 2: Extract the features for the train dataset, calculate PCA or # whitening and save if cfg.IMG_RETRIEVAL.TRAIN_PCA_WHITENING: logging.info("Extracting training features...") # the features are already processed based on type: rmac | GeM | l2 norm train_features = get_train_features( cfg, temp_dir, train_dataset_name, resize_img, spatial_levels, image_helper, train_dataset, model, ) ######################################################################## # Train PCA on the train features pca_out_fname = None if temp_dir: pca_out_fname = f"{temp_dir}/{train_dataset_name}_S{resize_img}_PCA.pickle" if pca_out_fname and PathManager.exists(pca_out_fname): logging.info("Loading PCA...") pca = load_pca(pca_out_fname) else: logging.info("Training and saving PCA...") pca = train_and_save_pca(train_features, cfg.IMG_RETRIEVAL.N_PCA, pca_out_fname) else: pca = None ############################################################################ # Step 4: Extract db_features and q_features for the eval dataset logging.info("Extracting Queries features...") features_queries = get_queries_features( cfg, temp_dir, eval_dataset_name, resize_img, spatial_levels, image_helper, eval_dataset, model, pca, ) logging.info("Extracting Dataset features...") features_dataset = get_dataset_features( cfg, temp_dir, eval_dataset_name, resize_img, spatial_levels, image_helper, eval_dataset, model, pca, ) ############################################################################ # Step 5: Compute similarity, score, and save results logging.info("Calculating similarity and score...") sim = features_queries.dot(features_dataset.T) logging.info(f"Similarity tensor: {sim.shape}") results = eval_dataset.score(sim, temp_dir) ############################################################################ # Step 6: save results and cleanup the temp directory if cfg.IMG_RETRIEVAL.SAVE_RETRIEVAL_RANKINGS_SCORES: checkpoint_folder = get_checkpoint_folder(cfg) # Save the rankings sim = sim.T ranks = np.argsort(-sim, axis=0) save_file(ranks.T.tolist(), os.path.join(checkpoint_folder, "rankings.json")) # Save the similarity scores save_file( sim.tolist(), os.path.join(checkpoint_folder, "similarity_scores.json"), ) # Save the result metrics save_file(results, os.path.join(checkpoint_folder, "metrics.json")) logging.info("All done!!")