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