def handle_training_batches(self):
"""
Training loop for one mini-epoch.
:return: training loss for the current mini-epoch
"""
# setting the model to training mode
self.model.train()
# initializing a list object to hold losses from each iteration
epoch_loss = []
# for the first epoch, only the linear layer is trained.
# starting from the second epoch, all the parameters of the model are trained.
if self.current_epoch == 1:
for param in self.model.parameters():
param.requires_grad = True
# training loop
for batch_idx, batch in enumerate(self.data_loader):
# if overfit_batch == 1, only the same batch is trained.
# this helps to see whether there are any issues with optimization.
# a fast over-fitting behaviour is expected.
if self.cfg['overfit_batch'] == 1:
if batch_idx == 0:
overfit_batch = batch
else:
batch = overfit_batch
# making sure the data and labels are in the correct device and in float32 type
items, labels = batch
items = handle_device(items, self.device)
labels = handle_device(labels, self.device)
# forward pass of the model
# obtaining the embeddings of each item in the batch
embs = self.model(items)
# calculating the loss value for the iteration
loss = LOSS_DICT[self.cfg['loss']](
data=embs,
labels=labels,
emb_size=self.model.fin_emb_size,
proxies=self.proxies,
margin=self.cfg['margin'],
mining_strategy=self.cfg['mining_strategy'])
# setting gradients of the optimizer to zero
self.optimizer.zero_grad()
# calculating gradients with backpropagation
loss.backward()
# updating the weights
self.optimizer.step()
# logging the loss value of the current batch
epoch_loss.append(loss.detach().item())
# logging the loss value of the current mini-epoch
return np.mean(epoch_loss)
def handle_training_batches(self):
"""
Training loop for one mini-epoch.
:return: training loss for the current mini-epoch
"""
# setting the model to training mode
self.model.train()
# initializing a list object to hold losses from each iteration
epoch_loss = []
# training loop
for batch_idx, batch in enumerate(self.data_loader):
# if overfit_batch == 1, only the same batch is trained.
# this helps to see whether there are any issues with optimization.
# a fast over-fitting behaviour is expected.
if self.cfg['overfit_batch'] == 1:
if batch_idx == 0:
overfit_batch = batch
else:
batch = overfit_batch
# making sure the data and labels are in the correct device and in float32 type
items, labels = batch
items = handle_device(items, self.device)
labels = handle_device(labels, self.device)
# forward pass of the model
# obtaining the embeddings of each item in the batch
embs = self.model(items)
# calculating the loss value for the iteration
loss = triplet_loss(data=embs,
labels=labels,
emb_size=self.cfg['emb_size'],
margin=self.cfg['margin'],
mining_strategy=self.cfg['mining_strategy'])
# setting gradients of the optimizer to zero
self.optimizer.zero_grad()
# calculating gradients with backpropagation
loss.backward()
# updating the weights
self.optimizer.step()
# applying the zero-mask to the selected indices
if self.prune_iteration > 0:
self.apply_mask()
# logging the loss value of the current batch
epoch_loss.append(loss.detach().item())
# logging the loss value of the current mini-epoch
return np.mean(epoch_loss)
def select_indices_to_prune(self):
"""
Selecting which indices to prune based on the trained model.
:return:
"""
self.indices_to_prune = torch.topk(
torch.abs(self.model.lin1.weight).mean(dim=1),
k=NUM_OF_ROWS_TO_PRUNE[self.prune_iteration],
largest=False).indices
# creating a mask of ones and zeros
mask = torch.ones(self.model.lin1.weight.shape)
zero_row = torch.zeros(1, self.model.lin1.weight.shape[1])
# sending the tensors to the proper device
mask = handle_device(mask, self.device)
zero_row = handle_device(zero_row, self.device)
# finalizing the mask based on the selected indices
mask[self.indices_to_prune] = zero_row
self.mask = mask
def handle_validation_batches(self):
"""
Validation loop.
"""
# setting the model to evaluation mode
self.model.eval()
# disabling gradient tracking
with torch.no_grad():
# initializing an empty tensor for storing the embeddings
embed_all = torch.tensor([], device=self.device)
# iterating through the data loader
for batch_idx, items in enumerate(self.data_loader):
# sending the items to the proper device
items = handle_device(items, self.device)
# forward pass of the model
# obtaining the embeddings of each item in the batch
embs = self.model(items)
embed_all = torch.cat((embed_all, embs))
# if Triplet or ProxyNCA loss is used, the distance function is Euclidean distance
if self.cfg['loss'] in [0, 1]:
dist_all = pairwise_euclidean_distance(embed_all)
dist_all /= self.model.fin_emb_size
# if NormalizedSoftmax loss is used, the distance function is cosine distance
elif self.cfg['loss'] == 2:
dist_all = -1 * pairwise_cosine_similarity(embed_all)
# if Group loss is used, the distance function is Pearson correlation coefficient
else:
dist_all = -1 * pairwise_pearson_coef(embed_all)
# computing evaluation metrics from the obtained distances
val_map_score = average_precision(
-1 * dist_all.cpu().float().clone() +
torch.diag(torch.ones(len(self.data)) * float('-inf')),
dataset=0,
path=self.cfg['main_path'],
print_metrics=False)
# logging the mean average precision value of the current validation run
self.val_loss_log.append(val_map_score.item())
def handle_training_batches(self):
"""
Training loop for one mini-epoch.
:return: training loss for the current mini-epoch
"""
# setting the model to training mode
self.model.train()
# initializing a list object to hold losses from each iteration
epoch_loss = []
# training loop
for batch_idx, batch in enumerate(self.data_loader):
# if overfit_batch == 1, only the same batch is trained.
# this helps to see whether there are any issues with optimization.
# a fast over-fitting behaviour is expected.
if self.cfg['overfit_batch'] == 1:
if batch_idx == 0:
overfit_batch = batch
else:
batch = overfit_batch
# making sure the data and labels are in the correct device and in float32 type
items, labels = batch
items = handle_device(items, self.device)
labels = handle_device(labels, self.device)
# forward pass of the student model
# obtaining the embeddings of each item in the batch
embs_s = self.model(items)
# if the distance-based KD loss is chosen,
# we obtain the embeddings of each item from the teacher model
with torch.no_grad():
embs_t = self.teacher(
items) if self.cfg['kd_loss'] == 'distance' else None
# calculating the KD loss for the iteration
kd_loss = KD_LOSS_DICT[self.cfg['kd_loss']](
embs_s=embs_s,
embs_t=embs_t,
emb_size=self.cfg['emb_size'],
lp_layer=self.lp_layer,
labels=labels,
centroids=self.centroids)
# calculating the triplet loss for the iteration
main_loss = triplet_loss(
data=embs_s,
labels=labels,
emb_size=self.cfg['emb_size'],
margin=self.cfg['margin'],
mining_strategy=self.cfg['mining_strategy'])
# summing KD and triplet loss values
loss = kd_loss + main_loss
# setting gradients of the optimizer to zero
self.optimizer.zero_grad()
# calculating gradients with backpropagation
loss.backward()
# updating the weights
self.optimizer.step()
# logging the loss value of the current batch
epoch_loss.append(loss.detach().item())
# logging the loss value of the current mini-epoch
return np.mean(epoch_loss)
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
model.to(device)
remove_items = []
with torch.no_grad(): # disabling gradient tracking
model.eval() # setting the model to evaluation mode
# initializing an empty tensor for storing the embeddings
embed_all = torch.tensor([], device=device)
# iterating through the data loader
for batch_idx, item in tqdm(enumerate(test_loader)):
try:
# sending the items to the proper device
item = handle_device(item, device)
# forward pass of the model
# obtaining the embeddings of each item in the batch
emb = model(item)
# appending the current embedding to the collection of embeddings
embed_all = torch.cat((embed_all, emb))
except Exception as e:
print("Error: {}, input shape: {}, index".format(
e, item.shape, batch_idx))
remove_items.append(audio_files[batch_idx])
continue
for re_item in remove_items:
audio_files.remove(re_item)
print("Number of audios: {}".format(len(audio_files)))
def evaluate(exp_name,
exp_type,
main_path,
emb_size,
loss,
data_dir):
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
print('Evaluating model {}.'.format(exp_name))
file_list = enumerate_h5_files(data_dir)
file_list.sort(key=lambda x: os.path.splitext(os.path.basename(x))[0])
print("Number feature files: {}".format(len(file_list)))
data = []
name = list(map(lambda x: os.path.splitext(os.path.relpath(x, data_dir))[0], file_list))
print("name: {}".format(name))
#image_with_index_list = dict(zip(name, range(len(name))))
#print("image_with_index_list: {}".format(image_with_index_list))
for file in tqdm(file_list):
temp_crema = dd.io.load(file)["crema"]
#print("crema shape: {}".format(temp_crema.shape))
idxs = np.arange(0, temp_crema.shape[0], 8)
temp_tensor = torch.from_numpy(temp_crema[idxs].T)
data.append(torch.cat((temp_tensor, temp_tensor))[:23].unsqueeze(0))
#name.append(os.path.splitext(os.path.basename(file))[0])
test_set = FullSizeInstanceDataset(data=data)
test_loader = DataLoader(test_set, batch_size=1, shuffle=False)
print("Initializing model")
# initializing the model
model = MOVEModel(emb_size=emb_size)
# loading a pre-trained model
model_name = os.path.join(main_path, 'saved_models', '{}_models'.format(exp_type), 'model_{}.pt'.format(exp_name))
model.load_state_dict(torch.load(model_name, map_location='cpu'))
# sending the model to gpu, if available
model.to(device)
remove_items = []
with torch.no_grad(): # disabling gradient tracking
model.eval() # setting the model to evaluation mode
# initializing an empty tensor for storing the embeddings
embed_all = torch.tensor([], device=device)
# iterating through the data loader
for batch_idx, item in tqdm(enumerate(test_loader)):
try:
# sending the items to the proper device
item = handle_device(item, device)
# forward pass of the model
# obtaining the embeddings of each item in the batch
emb = model(item)
# appending the current embedding to the collection of embeddings
embed_all = torch.cat((embed_all, emb))
except Exception as e:
print("Error: {}, input shape: {}, index".format(e, item.shape, batch_idx))
remove_items.append(name[batch_idx])
continue
for re_item in remove_items:
name.remove(re_item)
print("name length: {}".format(len(name)))
image_with_index_list = dict(zip(name, range(len(name))))
embed_all = F.normalize(embed_all, p=2, dim=1)
return embed_all.cpu(), image_with_index_list
def evaluate(exp_name, exp_type, main_path, emb_size, loss):
"""
Main evaluation function of MOVE. For a detailed explanation of parameters,
please check 'python move_main.py -- help'
:param main_path: main working directory
:param exp_name: name to save model and experiment summary
:param exp_type: type of experiment
:param emb_size: the size of the final embeddings produced by the model
:param loss: the loss used for training the model
"""
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
eval_dataset = os.path.join(main_path, 'data/benchmark_crema.pt')
print('Evaluating model {} on dataset {}.'.format(exp_name, eval_dataset))
# initializing the model
model = MOVEModel(emb_size=emb_size)
# loading a pre-trained model
model_name = os.path.join(main_path, 'saved_models',
'{}_models'.format(exp_type),
'model_{}.pt'.format(exp_name))
model.load_state_dict(torch.load(model_name, map_location='cpu'))
# sending the model to gpu, if available
model.to(device)
# loading test data, initializing the dataset object and the data loader
test_data, test_labels = import_dataset_from_pt(filename=eval_dataset,
suffix=False)
test_set = FullSizeInstanceDataset(data=test_data)
test_loader = DataLoader(test_set, batch_size=1, shuffle=False)
start_time = time.monotonic()
with torch.no_grad(): # disabling gradient tracking
model.eval() # setting the model to evaluation mode
# initializing an empty tensor for storing the embeddings
embed_all = torch.tensor([], device=device)
# iterating through the data loader
for batch_idx, item in enumerate(test_loader):
# sending the items to the proper device
item = handle_device(item, device)
# forward pass of the model
# obtaining the embeddings of each item in the batch
emb = model(item)
# appending the current embedding to the collection of embeddings
embed_all = torch.cat((embed_all, emb))
# if Triplet or ProxyNCA loss is used, the distance function is Euclidean distance
if loss in [0, 1]:
dist_all = pairwise_euclidean_distance(embed_all)
dist_all /= model.fin_emb_size
# if NormalizedSoftmax loss is used, the distance function is cosine distance
elif loss == 2:
dist_all = -1 * pairwise_cosine_similarity(embed_all)
# if Group loss is used, the distance function is Pearson correlation coefficient
else:
dist_all = -1 * pairwise_pearson_coef(embed_all)
# computing evaluation metrics from the obtained distances
average_precision(-1 * dist_all.cpu().float().clone() +
torch.diag(torch.ones(len(test_data)) * float('-inf')),
dataset=1)
test_time = time.monotonic() - start_time
print('Total time: {:.0f}m{:.0f}s.'.format(test_time // 60,
test_time % 60))