| import torch |
| from modules.ReLoCLNet import ReLoCLNet |
| from modules.optimization import BertAdam |
| import numpy as np |
| import copy |
|
|
| def count_parameters(model, verbose=True): |
| """Count number of parameters in PyTorch model, |
| References: https://discuss.pytorch.org/t/how-do-i-check-the-number-of-parameters-of-a-model/4325/7. |
| |
| from utils.utils import count_parameters |
| count_parameters(model) |
| import sys |
| sys.exit(1) |
| """ |
| n_all = sum(p.numel() for p in model.parameters()) |
| n_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad) |
| if verbose: |
| print("Parameter Count: all {:,d}; trainable {:,d}".format(n_all, n_trainable)) |
| return n_all, n_trainable |
|
|
| def prepare_model(opt, logger): |
| model = ReLoCLNet(opt) |
| count_parameters(model) |
| if opt.device.type == "cuda": |
| logger.info("CUDA enabled.") |
| model.to(opt.device) |
| return model |
|
|
| def resume_model(logger, opt, model=None, optimizer=None, start_epoch=None): |
| checkpoint = torch.load(opt.checkpoint, map_location=opt.device) |
| if model is not None: |
| model.load_state_dict(checkpoint['model_state_dict']) |
| logger.info(f"Loading model from {opt.checkpoint} at epoch {checkpoint['epoch']}") |
|
|
| if optimizer is not None: |
| optimizer.load_state_dict(checkpoint['optimizer_state_dict']) |
| logger.info(f"Loading optimizer from {opt.checkpoint} at epoch {checkpoint['epoch']}") |
| |
| if start_epoch is not None: |
| start_epoch = checkpoint['epoch'] |
| logger.info(f"Loading start_epoch from {opt.checkpoint} at epoch {checkpoint['epoch']}") |
| |
| return model, optimizer, start_epoch, |
|
|
| def prepare_optimizer(model, opt, total_train_steps): |
| |
| param_optimizer = list(model.named_parameters()) |
| no_decay = ["bias", "LayerNorm.bias", "LayerNorm.weight"] |
| optimizer_grouped_parameters = [ |
| {"params": [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], "weight_decay": 0.01}, |
| {"params": [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], "weight_decay": 0.0}] |
|
|
| optimizer = BertAdam(optimizer_grouped_parameters, lr=opt.lr, weight_decay=opt.wd, warmup=opt.lr_warmup_proportion, |
| t_total=total_train_steps, schedule="warmup_linear") |
| return optimizer |
|
|
| def save_model(model, optimizer, epoch, path, logger): |
| data = { |
| 'epoch': epoch, |
| 'model_cfg': model.config, |
| 'model_state_dict': model.state_dict(), |
| 'optimizer_state_dict': optimizer.state_dict(), |
| } |
| torch.save(data, path) |
| logger.info(f"Save checkpoint at {path}") |
| logger.info("") |
| |
| |
| def topk_3d(tensor, k): |
| """ |
| Find the top k values and their corresponding indices in a 3D tensor. |
| |
| Args: |
| tensor (torch.Tensor): A 3D tensor of shape [v, m, n]. |
| k (int): The number of top elements to find. |
| |
| Returns: |
| topk_values (torch.Tensor): The top k values. |
| indices_3d (torch.Tensor): The indices of the top k values in the format [i, j, k]. |
| """ |
| |
| flat_tensor = tensor.view(-1) |
|
|
| |
| topk_values, topk_indices = torch.topk(flat_tensor, k) |
|
|
| |
| v, m, n = tensor.shape |
| indices_3d = torch.stack(torch.unravel_index(topk_indices, (v, m, n)), dim=1) |
|
|
| return topk_values, indices_3d |
|
|
|
|
| def generate_min_max_length_mask(array_shape, min_l, max_l): |
| """ The last two dimension denotes matrix of upper-triangle with upper-right corner masked, |
| below is the case for 4x4. |
| [[0, 1, 1, 0], |
| [0, 0, 1, 1], |
| [0, 0, 0, 1], |
| [0, 0, 0, 0]] |
| Args: |
| array_shape: np.shape??? The last two dimensions should be the same |
| min_l: int, minimum length of predicted span |
| max_l: int, maximum length of predicted span |
| Returns: |
| """ |
| single_dims = (1, ) * (len(array_shape) - 2) |
| mask_shape = single_dims + array_shape[-2:] |
| extra_length_mask_array = np.ones(mask_shape, dtype=np.float32) |
| mask_triu = np.triu(extra_length_mask_array, k=min_l) |
| mask_triu_reversed = 1 - np.triu(extra_length_mask_array, k=max_l) |
| final_prob_mask = mask_triu * mask_triu_reversed |
| return final_prob_mask |
|
|
|
|
| def extract_topk_elements(query_scores, start_probs, end_probs, video_names, k): |
|
|
| |
| topk_values, topk_indices = torch.topk(query_scores, k) |
|
|
| |
| selected_start_probs = torch.stack([start_probs[i, indices] for i, indices in enumerate(topk_indices)], dim=0) |
| selected_end_probs = torch.stack([end_probs[i, indices] for i, indices in enumerate(topk_indices)], dim=0) |
| |
| selected_video_name = [] |
| for i in range(topk_indices.shape[0]): |
| vn = copy.deepcopy(video_names) |
| tmp = [vn[idx] for idx in topk_indices[i]] |
| selected_video_name.append(tmp) |
|
|
| return topk_values, selected_start_probs, selected_end_probs, selected_video_name |
|
|
| def logger_ndcg_iou(val_ndcg_iou, logger, suffix): |
| for K, vs in val_ndcg_iou.items(): |
| for T, v in vs.items(): |
| logger.info(f"{suffix} NDCG@{K}, IoU={T}: {v:.6f}") |
| logger.info("") |
