correct ndcg-iou

.gitignore

@@ -7,8 +7,6 @@ __pycache__/
 *.so
 
 unused
-
-results
 # Distribution / packaging
 .Python
 build/
@@ -59,7 +57,6 @@ coverage.xml
 *.pot
 
 # Django stuff:
-*.log
 local_settings.py
 db.sqlite3
 db.sqlite3-journal

README.md

@@ -1,14 +1,9 @@
----
-license: cc
-language:
-- en
----
 # Video Moment Retrieval in Practical Setting: A Dataset of Ranked Moments for Imprecise  Queries
 
 The benchmark and dataset for the paper "Video Moment Retrieval in Practical Settings: A Dataset of Ranked Moments for Imprecise Queries" is coming soon.
 
-We recommend cloning the code, data, and feature files from the Hugging Face repository at [TVR-Ranking](https://huggingface.co/axgroup/TVR-Ranking).  
-  
+We recommend cloning the code, data, and feature files from the Hugging Face repository at [TVR-Ranking](https://huggingface.co/axgroup/TVR-Ranking).  This repository only includes the code for ReLoCLNet. You can download the other baseline models from [XML](https://huggingface.co/LiangRenjie/XML_RVMR) and [CONQUER](https://huggingface.co/LiangRenjie/CONQUER_RVMR).
+
 ![TVR_Ranking_overview](./figures/taskComparisonV.png)  
 
 
@@ -57,30 +52,55 @@ tar -xf tvr_feature_release.tar.gz -C data/TVR_Ranking/feature
 # modify the data path first 
 sh run_top20.sh
 ```
+### 5. Inferring
+The checkpoint can all be accessed from Hugging Face [TVR-Ranking](https://huggingface.co/axgroup/TVR-Ranking).
+```shell
+sh infer_top20.sh
+```
 
-## Baseline
-(ToDo: running the new version...) \
-The baseline performance of  $NDGC@20$ was shown as follows.
+## Experiment Results
+### Baseline
+The baseline performance of  $NDGC@40$ was shown as follows.
 Top $N$ moments were comprised of a pseudo training set by the query-caption similarity.
-| Model          | $N$ | IoU = 0.3, val | IoU = 0.3, test | IoU = 0.5, val | IoU = 0.5, test | IoU = 0.7, val | IoU = 0.7, test |
-|----------------|-----|----------------|-----------------|----------------|-----------------|----------------|-----------------|
-| **XML**        | 1   | 0.1050         | 0.1047          | 0.0767         | 0.0751          | 0.0287         | 0.0314          |
-|                | 20  | 0.1948         | 0.1964          | 0.1417         | 0.1434          | 0.0519         | 0.0583          |
-|                | 40  | 0.2101         | 0.2110          | 0.1525         | 0.1533          | 0.0613         | 0.0617          |
-| **CONQUER**    | 1   | 0.0979         | 0.0830          | 0.0817         | 0.0686          | 0.0547         | 0.0479          |
-|                | 20  | 0.2007         | 0.1935          | 0.1844         | 0.1803          | 0.1391         | 0.1341          |
-|                | 40  | 0.2094         | 0.1943          | 0.1930         | 0.1825          | 0.1481         | 0.1334          |
-| **ReLoCLNet**  | 1   | 0.1306         | 0.1299          | 0.1169         | 0.1154          | 0.0738         | 0.0789          |
-|                | 20  | 0.3264         | 0.3214          | 0.3007         | 0.2956          | 0.2074         | 0.2084          |
-|                | 40  | 0.3479         | 0.3473          | 0.3221         | 0.3217          | 0.2218         | 0.2275          |
+
+| **Model**      | **Train Set Top N** | **IoU=0.3**  | |**IoU=0.5**  | |**IoU=0.7**  | |
+|----------------|---------------------|--------------|--------------|--------------|--------------|--------------|--------------|
+|                |                     | **Val** | **Test** | **Val** | **Test** | **Val** | **Test** |
+| **XML**        | 1                   | 0.1077 | 0.1016 | 0.0775 | 0.0727 | 0.0273 | 0.0294 |
+|                | 20                  | 0.2580 | 0.2512 | 0.1874 | 0.1853 | 0.0705 | 0.0753 |
+|                | 40                  | 0.2408 | 0.2432 | 0.1740 | 0.1791 | 0.0666 | 0.0720 |
+| **ReLoCLNet**  | 1                   | 0.1533 | 0.1489 | 0.1321 | 0.1304 | 0.0878 | 0.0869 |
+|                | 20                  | 0.4039 | 0.4031 | 0.3656 | 0.3648 | 0.2542 | 0.2567 |
+|                | 40                  | 0.4725 | 0.4735 | 0.4337 | 0.4337 | 0.3015 | 0.3079 |
+
+
+###  ReLoCLNet Performance
+
+| **Model**  | **Train Set Top N** | **IoU=0.3**  | |**IoU=0.5**  | |**IoU=0.7**  | |
+|------------|---------------------|--------------|--------------|--------------|--------------|--------------|--------------|
+|            |                     | **Val** | **Test** | **Val** | **Test** | **Val** | **Test** |
+| **NDCG@10** |                     |              |              |              |              |              |              |
+| ReLoCLNet  | 1                   | 0.1575 | 0.1525 | 0.1358 | 0.1349 | 0.0908 | 0.0916 |
+| ReLoCLNet  | 20                  | 0.3751 | 0.3751 | 0.3407 | 0.3397 | 0.2316 | 0.2338 |
+| ReLoCLNet  | 40                  | 0.4339 | 0.4353 | 0.3984 | 0.3986 | 0.2693 | 0.2807 |
+| **NDCG@20** |                     |              |              |              |              |              |              |
+| ReLoCLNet  | 1                   | 0.1504 | 0.1439 | 0.1303 | 0.1269 | 0.0866 | 0.0849 |
+| ReLoCLNet  | 20                  | 0.3815 | 0.3792 | 0.3462 | 0.3427 | 0.2381 | 0.2386 |
+| ReLoCLNet  | 40                  | 0.4418 | 0.4439 | 0.4060 | 0.4059 | 0.2787 | 0.2877 |
+| **NDCG@40** |                     |              |              |              |              |              |              |
+| ReLoCLNet  | 1                   | 0.1533 | 0.1489 | 0.1321 | 0.1304 | 0.0878 | 0.0869 |
+| ReLoCLNet  | 20                  | 0.4039 | 0.4031 | 0.3656 | 0.3648 | 0.2542 | 0.2567 |
+| ReLoCLNet  | 40                  | 0.4725 | 0.4735 | 0.4337 | 0.4337 | 0.3015 | 0.3079 |
+
+
+
+
 
 
-### 4. Inferring
-[ToDo] The checkpoint can all be accessed from Hugging Face [TVR-Ranking](https://huggingface.co/axgroup/TVR-Ranking).
 
 
 ## Citation
 If you feel this project helpful to your research, please cite our work.
 ```
 
-```
+```

infer.py

@@ -5,9 +5,9 @@ from tqdm import tqdm
 from modules.dataset_init import prepare_dataset
 from modules.infer_lib import grab_corpus_feature, eval_epoch 
 
-from utils.basic_utils import AverageMeter, get_logger
+from utils.basic_utils import get_logger
 from utils.setup import set_seed, get_args
-from utils.run_utils import prepare_optimizer, prepare_model, logger_ndcg_iou
+from utils.run_utils import prepare_optimizer, prepare_model, logger_ndcg_iou, resume_model
 
 def main():
     opt = get_args()
@@ -21,7 +21,11 @@ def main():
 
     model = prepare_model(opt, logger)
     # optimizer = prepare_optimizer(model, opt, len(train_loader) * opt.n_epoch)
+    # start_epoch = 0
+    # model, optimizer, start_epoch = resume_model(logger, opt, model, optimizer, start_epoch)
+    model, _, _ = resume_model(logger, opt, model)
 
+    model.eval()
     corpus_feature = grab_corpus_feature(model, corpus_loader, opt.device)
     val_ndcg_iou = eval_epoch(model, corpus_feature, val_loader, val_gt, opt, corpus_video_list)
     test_ndcg_iou = eval_epoch(model, corpus_feature, test_loader, test_gt, opt, corpus_video_list)

infer_top20.sh

@@ -1,17 +1,14 @@
 python infer.py \
     --results_path      results/tvr_ranking \
-    --checkpoint        results/tvr_ranking/best_model.pt \
     --train_path        data/TVR_Ranking/train_top20.json \
     --val_path          data/TVR_Ranking/val.json \
     --test_path         data/TVR_Ranking/test.json \
     --corpus_path       data/TVR_Ranking/video_corpus.json \
-    --desc_bert_path    /home/renjie.liang/datasets/TVR_Ranking/features/query_bert.h5 \
-    --video_feat_path   /home/share/czzhang/Dataset/TVR/TVR_feature/video_feature/tvr_i3d_rgb600_avg_cl-1.5.h5 \
-    --sub_bert_path     /home/share/czzhang/Dataset/TVR/TVR_feature/bert_feature/sub_query/tvr_sub_pretrained_w_sub_query_max_cl-1.5.h5 \
-    --exp_id infer
+    --desc_bert_path    data/features/query_bert.h5 \
+    --video_feat_path   data/features/tvr_i3d_rgb600_avg_cl-1.5.h5 \
+    --sub_bert_path     data/features/tvr_sub_pretrained_w_sub_query_max_cl-1.5.h5 \
+    --checkpoint        results/tvr_ranking/top20/best_model.pt \
+    --exp_id            top20_infer
 
 # qsub -I -l select=1:ngpus=1 -P gs_slab -q slab_gpu8
-# cd /home/renjie.liang/11_TVR-Ranking/ReLoCLNet; conda activate py11; sh infer_top20.sh 
-    # --hard_negative_start_epoch 0 \
-    # --no_norm_vfeat \
-    # --use_hard_negative
+# cd /home/renjie.liang/11_TVR-Ranking/ReLoCLNet; conda activate py11; sh infer_top20.sh

modules/ReLoCLNet.py

@@ -201,7 +201,8 @@ class ReLoCLNet(nn.Module):
         feat = input_proj_layer(feat)
         feat = pos_embed_layer(feat)
         mask = mask.unsqueeze(1)  # (N, 1, L), torch.FloatTensor
-        return encoder_layer(feat, mask)  # (N, L, D_hidden)
+        feat = encoder_layer(feat, mask)  # (N, L, D_hidden)
+        return feat
 
     def get_modularized_queries(self, encoded_query, query_mask, return_modular_att=False):
         """

modules/dataset_init.py

@@ -41,6 +41,9 @@ def collate_fn(batch, task):
         batch_data["sub_feat"] = sub_feat_mask[0]
         batch_data["sub_mask"] = sub_feat_mask[1]
         
+        # batch_data["video_name"] = [e["video_name"] for e in batch]
+        
+        
     if task == "eval":
         query_feat_mask = pad_sequences_1d([e["query_feat"] for e in batch], dtype=torch.float32, fixed_length=None)
         batch_data["query_feat"] = query_feat_mask[0]

modules/dataset_tvrr.py

@@ -23,7 +23,7 @@ class TrainDataset(Dataset):
         # prepare desc data
         self.use_video = "video" in ctx_mode
         self.use_sub = "sub" in ctx_mode
-
+        
         self.desc_bert_h5 = h5py.File(desc_bert_path, "r")
         if self.use_video:
             self.vid_feat_h5 = h5py.File(video_feat_path, "r")
@@ -56,6 +56,7 @@ class TrainDataset(Dataset):
         query_id=raw_data["query_id"]
         video_name=raw_data["video_name"]
         timestamp = raw_data["timestamp"]
+        duration = raw_data["duration"]
         
         model_inputs = dict()
         model_inputs["simi"] = raw_data["similarity"]
@@ -80,12 +81,8 @@ class TrainDataset(Dataset):
         else:
             model_inputs["sub_feat"] = torch.zeros((2, 2))
 
-        # print(ctx_l)
-        # print(timestamp)
         model_inputs["st_ed_indices"] = self.get_st_ed_label(timestamp, max_idx=ctx_l - 1)
-        # print(model_inputs["st_ed_indices"])
         return model_inputs
-        # return dict(meta=meta, model_inputs=model_inputs)
 
     def get_st_ed_label(self, ts, max_idx):
         """
@@ -175,6 +172,7 @@ class CorpusEvalDataset(Dataset):
 
         self.use_video = "video" in ctx_mode
         self.use_sub = "sub" in ctx_mode
+
         if self.use_video:
             self.vid_feat_h5 = h5py.File(video_feat_path, "r")
         if self.use_sub:
@@ -187,6 +185,8 @@ class CorpusEvalDataset(Dataset):
         """No need to batch, since it has already been batched here"""
         raw_data = self.video_data[index]
         # initialize with basic data
+        duration = raw_data["duration"]
+        video_name = raw_data["vid_name"]
         meta = dict(vid_name=raw_data["vid_name"], duration=raw_data["duration"])
         model_inputs = dict()
 

modules/infer_lib.py

@@ -10,16 +10,18 @@ def grab_corpus_feature(model, corpus_loader, device):
     model.eval()
     all_video_feat, all_video_mask = [], []
     all_sub_feat, all_sub_mask = [], []
-        
-    for batch_input in tqdm(corpus_loader, desc="Compute Corpus Feature: ", total=len(corpus_loader)):
-        batch_input = {k: v.to(device) for k, v in batch_input.items()}
-        _video_feat, _sub_feat = model.encode_context(batch_input["video_feat"], batch_input["video_mask"],
+    
+    # all_video_name = []
+    with torch.no_grad():
+        for batch_input in tqdm(corpus_loader, desc="Compute Corpus Feature: ", total=len(corpus_loader)):
+            batch_input = {k: v.to(device) for k, v in batch_input.items()}
+            _video_feat, _sub_feat = model.encode_context(batch_input["video_feat"], batch_input["video_mask"],
                                                       batch_input["sub_feat"], batch_input["sub_mask"])
-        
-        all_video_feat.append(_video_feat.detach().cpu())
-        all_video_mask.append(batch_input["video_mask"].detach().cpu())
-        all_sub_feat.append(_sub_feat.detach().cpu())
-        all_sub_mask.append(batch_input["sub_mask"].detach().cpu())
+            
+            all_video_feat.append(_video_feat.detach().cpu())
+            all_video_mask.append(batch_input["video_mask"].detach().cpu())
+            all_sub_feat.append(_sub_feat.detach().cpu())
+            all_sub_mask.append(batch_input["sub_mask"].detach().cpu())
         
     all_video_feat = torch.cat(all_video_feat, dim=0)
     all_video_mask = torch.cat(all_video_mask, dim=0)
@@ -41,7 +43,7 @@ def eval_epoch(model, corpus_feature, eval_loader, eval_gt, opt, corpus_video_li
     all_video_mask = corpus_feature["all_video_mask"].to(device)
     all_sub_feat = corpus_feature["all_sub_feat"].to(device)
     all_sub_mask = corpus_feature["all_sub_mask"].to(device)
-    all_query_score, all_end_prob, all_start_prob = [], [], []
+    all_query_score, all_end_prob, all_start_prob, all_top_video_name = [], [], [], []
     for batch_input in tqdm(eval_loader, desc="Compute Query Scores: ", total=len(eval_loader)):
         batch_input = {k: v.to(device) for k, v in batch_input.items()}
         query_scores, start_probs, end_probs = model.get_pred_from_raw_query(
@@ -56,34 +58,36 @@ def eval_epoch(model, corpus_feature, eval_loader, eval_gt, opt, corpus_video_li
         start_probs = F.softmax(start_probs, dim=-1) 
         end_probs = F.softmax(end_probs, dim=-1)
         
-        query_scores, start_probs, end_probs = extract_topk_elements(query_scores, start_probs, end_probs, topn_video)
+        query_scores, start_probs,  end_probs, video_name_top = extract_topk_elements(query_scores, start_probs, end_probs, corpus_video_list, topn_video)
         
         all_query_id.append(batch_input["query_id"].detach().cpu())
         all_query_score.append(query_scores.detach().cpu())
         all_start_prob.append(start_probs.detach().cpu())
         all_end_prob.append(end_probs.detach().cpu())
-        
+        all_top_video_name.extend(video_name_top)
+
     all_query_id = torch.cat(all_query_id, dim=0)
     all_query_id = all_query_id.tolist()
     
     all_query_score = torch.cat(all_query_score, dim=0)
     all_start_prob = torch.cat(all_start_prob, dim=0)
     all_end_prob = torch.cat(all_end_prob, dim=0)
-    average_ndcg = calculate_average_ndcg(all_query_id, all_start_prob, all_query_score, all_end_prob, corpus_video_list, eval_gt, opt)
+    average_ndcg = calculate_average_ndcg(all_query_id, all_start_prob, all_query_score, all_end_prob, all_top_video_name, eval_gt, opt)
     return average_ndcg
 
-def calculate_average_ndcg(all_query_id, all_start_prob, all_query_score, all_end_prob, corpus_video_list, eval_gt, opt):
+def calculate_average_ndcg(all_query_id, all_start_prob, all_query_score, all_end_prob, all_top_video_name, eval_gt, opt):
     topn_moment = max(opt.ndcg_topk)
     
     all_2D_map = torch.einsum("qvm,qv,qvn->qvmn", all_start_prob, all_query_score, all_end_prob)
     map_mask = generate_min_max_length_mask(all_2D_map.shape, min_l=opt.min_pred_l, max_l=opt.max_pred_l)
     all_2D_map = all_2D_map * map_mask
     all_pred = {}
-    for i in trange(len(all_2D_map), desc="Collect Predictions: "):
-        query_id = all_query_id[i]
-        score_map = all_2D_map[i]
+    for idx in trange(len(all_2D_map), desc="Collect Predictions: "):
+        query_id = all_query_id[idx]
+        score_map = all_2D_map[idx]
         top_score, top_idx = topk_3d(score_map, topn_moment)
-        pred_videos = [corpus_video_list[i[0]] for i in top_idx]
+        top_video_name  = all_top_video_name[idx]
+        pred_videos = [top_video_name[i[0]] for i in top_idx]
         pre_start_time = [i[1].item() * opt.clip_length for i in top_idx]
         pre_end_time   = [i[2].item() * opt.clip_length for i in top_idx]
         
@@ -94,7 +98,7 @@ def calculate_average_ndcg(all_query_id, all_start_prob, all_query_score, all_en
                 "timestamp": [s, e],
                 "model_scores": score
             })
-        print(pred_result)
+        # print(pred_result)
         all_pred[query_id] = pred_result
         
     average_ndcg = calculate_ndcg_iou(eval_gt, all_pred, opt.iou_threshold, opt.ndcg_topk)

modules/ndcg_iou.py

@@ -25,7 +25,7 @@ def calculate_ndcg(pred_scores, true_scores):
 def calculate_ndcg_iou(all_gt, all_pred, TS, KS):
     performance = defaultdict(lambda: defaultdict(list))
     performance_avg = defaultdict(lambda: defaultdict(float))
-    for k in all_pred.keys():
+    for k in tqdm(all_pred.keys(), desc="Calculate NDCG"):
         one_pred = all_pred[k]
         one_gt = all_gt[k]  
 

results/ReLoCLNet/top01/best_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:10d055d1d72eac6aef2937d38422ec3ff5760aba541b1f9de1b3a6127925550b
+size 83802857

results/ReLoCLNet/top20/best_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9c40bf96c586b463cf324a42101c92a2e7b5e5692a0caa25d234b977403dc8bd
+size 83802857

results/ReLoCLNet/top40/best_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:05a50f5bd49514b72669755b5d08e39906f9bc2bdedb4b1b53a519ec66b5e980
+size 83802857

run_top01.sh

@@ -0,0 +1,17 @@
+python train.py \
+    --results_path      results/tvr_ranking \
+    --train_path        data/TVR_Ranking/train_top01.json \
+    --val_path          data/TVR_Ranking/val.json \
+    --test_path         data/TVR_Ranking/test.json \
+    --corpus_path       data/TVR_Ranking/video_corpus.json \
+    --desc_bert_path    data/features/query_bert.h5 \
+    --video_feat_path   data/features/tvr_i3d_rgb600_avg_cl-1.5.h5 \
+    --sub_bert_path     data/features/tvr_sub_pretrained_w_sub_query_max_cl-1.5.h5 \
+    --n_epoch               4000 \
+    --eval_num_per_epoch    0.05 \
+    --seed                  2024 \
+    --bsz                   512 \
+    --exp_id                top01
+
+# qsub -I -l select=1:ngpus=1 -P gs_slab -q gpu8
+# cd /home/renjie.liang/11_TVR-Ranking/ReLoCLNet; conda activate py11; sh run_top01.sh 

run_top20.sh

@@ -4,11 +4,14 @@ python train.py \
     --val_path          data/TVR_Ranking/val.json \
     --test_path         data/TVR_Ranking/test.json \
     --corpus_path       data/TVR_Ranking/video_corpus.json \
-    --desc_bert_path    data/TVR_Ranking/features/query_bert.h5 \
-    --video_feat_path   data/TVR_Ranking/features/tvr_i3d_rgb600_avg_cl-1.5.h5 \
-    --sub_bert_path     data/TVR_Ranking/features/tvr_sub_pretrained_w_sub_query_max_cl-1.5.h5 \
-    --n_epoch 100 \
-    --eval_num_per_epoch 1 \
-    --seed 2024 \
-    --exp_id new_version
+    --desc_bert_path    data/features/query_bert.h5 \
+    --video_feat_path   data/features/tvr_i3d_rgb600_avg_cl-1.5.h5 \
+    --sub_bert_path     data/features/tvr_sub_pretrained_w_sub_query_max_cl-1.5.h5 \
+    --n_epoch               200 \
+    --eval_num_per_epoch    1 \
+    --seed                  2024 \
+    --bsz                   512 \
+    --exp_id                top20
 
+# qsub -I -l select=1:ngpus=1 -P gs_slab -q gpu8
+# cd /home/renjie.liang/11_TVR-Ranking/ReLoCLNet; conda activate py11; sh run_top20.sh 

run_top40.sh

@@ -0,0 +1,17 @@
+python train.py \
+    --results_path      results/tvr_ranking \
+    --train_path        data/TVR_Ranking/train_top40.json \
+    --val_path          data/TVR_Ranking/val.json \
+    --test_path         data/TVR_Ranking/test.json \
+    --corpus_path       data/TVR_Ranking/video_corpus.json \
+    --desc_bert_path    data/features/query_bert.h5 \
+    --video_feat_path   data/features/tvr_i3d_rgb600_avg_cl-1.5.h5 \
+    --sub_bert_path     data/features/tvr_sub_pretrained_w_sub_query_max_cl-1.5.h5 \
+    --n_epoch               100 \
+    --eval_num_per_epoch    2 \
+    --seed                  2024 \
+    --bsz                   512 \
+    --exp_id                top40
+
+# qsub -I -l select=1:ngpus=1 -P gs_slab -q gpu8
+# cd /home/renjie.liang/11_TVR-Ranking/ReLoCLNet; conda activate py11; sh run_top40.sh 

train.py

@@ -7,7 +7,7 @@ from modules.infer_lib import grab_corpus_feature, eval_epoch
 
 from utils.basic_utils import AverageMeter, get_logger
 from utils.setup import set_seed, get_args
-from utils.run_utils import prepare_optimizer, prepare_model, logger_ndcg_iou
+from utils.run_utils import prepare_optimizer, prepare_model, logger_ndcg_iou, save_model, resume_model
 
 def main():
     opt = get_args()
@@ -20,32 +20,39 @@ def main():
     
     
     train_loader, corpus_loader, corpus_video_list, val_loader, test_loader, val_gt, test_gt = prepare_dataset(opt)
-
+    
     model = prepare_model(opt, logger)
     optimizer = prepare_optimizer(model, opt, len(train_loader) * opt.n_epoch)
     
+    start_epoch = 0
+    if opt.checkpoint is not None:
+        model, optimizer, start_epoch = resume_model(logger, opt, model, optimizer, start_epoch)
+    
     eval_step = len(train_loader) // opt.eval_num_per_epoch
     best_val_ndcg = 0
-    for epoch_i in range(0, opt.n_epoch):
-        logger.info(f"TRAIN EPOCH: {epoch_i}|{opt.n_epoch}")
+    for epoch in range(start_epoch, opt.n_epoch):
+        logger.info(f"TRAIN EPOCH: {epoch}|{opt.n_epoch}")
         model.train()
-        if opt.hard_negative_start_epoch != -1 and epoch_i >= opt.hard_negative_start_epoch:
+        if opt.hard_negative_start_epoch != -1 and epoch >= opt.hard_negative_start_epoch:
             model.set_hard_negative(True, opt.hard_pool_size)
-
         model.train()
+
         for step, batch_input in tqdm(enumerate(train_loader), desc="Training", total=len(train_loader)):
-            step += 1
+            global_step = epoch * len(train_loader) + step + 1
             batch_input = {k: v.to(opt.device) for k, v in batch_input.items()}
             loss = model(**batch_input)
             optimizer.zero_grad()
             loss.backward()
             # nn.utils.clip_grad_norm_(model.parameters())
             optimizer.step()
-            
+
             if step % opt.log_step == 0:
-                logger.info(f"EPOCH {epoch_i}/{opt.n_epoch} | STEP: {step}|{len(train_loader)} | Loss: {loss.item():.6f}")
-                
-            if step % eval_step == 0 or step == len(train_loader):
+                logger.info(f"EPOCH {epoch}/{opt.n_epoch} | STEP: {step}|{len(train_loader)} | Loss: {loss.item():.6f}")
+                for i in range(torch.cuda.device_count()):
+                    print(f"Memory Allocated on GPU {i}: {torch.cuda.memory_allocated(i) / 1024**3:.2f} GB")
+                    print(f"Memory Cached on GPU {i}: {torch.cuda.memory_reserved(i) / 1024**3:.2f} GB")
+                print("-------------------------")
+            if global_step % eval_step == 0 or step == len(train_loader):
                 corpus_feature = grab_corpus_feature(model, corpus_loader, opt.device)
                 val_ndcg_iou = eval_epoch(model, corpus_feature, val_loader, val_gt, opt, corpus_video_list)
                 test_ndcg_iou = eval_epoch(model, corpus_feature, test_loader, test_gt, opt, corpus_video_list)
@@ -58,12 +65,8 @@ def main():
                     logger_ndcg_iou(val_ndcg_iou, logger, "BEST VAL")
                     logger_ndcg_iou(test_ndcg_iou, logger, "BEST TEST")
                     
-                    checkpoint = {"model": model.state_dict(), "model_cfg": model.config, "epoch": epoch_i}
-
                     bestmodel_path = os.path.join(opt.results_path, "best_model.pt")
-                    torch.save(checkpoint, bestmodel_path)
-                    logger.info(f"Save checkpoint at {bestmodel_path}")
-                    logger.info("")
+                    save_model(model, optimizer, epoch, bestmodel_path, logger)
 
 if __name__ == '__main__':
     main()

utils/run_utils.py

@@ -2,6 +2,7 @@ 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,
@@ -21,18 +22,27 @@ def count_parameters(model, verbose=True):
 def prepare_model(opt, logger):
     model = ReLoCLNet(opt)
     count_parameters(model)
-
-    if opt.checkpoint is not None:
-        checkpoint = torch.load(opt.checkpoint, map_location=opt.device)  
-        model.load_state_dict(checkpoint['model'])
-        logger.info(f"Loading checkpoint from {opt.checkpoint}")
-
-    # Prepare optimizer (unchanged)
     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())
@@ -43,10 +53,20 @@ def prepare_optimizer(model, opt, total_train_steps):
 
     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.
@@ -94,7 +114,7 @@ def generate_min_max_length_mask(array_shape, min_l, max_l):
     return final_prob_mask  # with valid bit to be 1
 
 
-def extract_topk_elements(query_scores, start_probs, end_probs, k):
+def extract_topk_elements(query_scores, start_probs, end_probs, video_names, k):
 
     # Step 1: Find the top k values and their indices in query_scores
     topk_values, topk_indices = torch.topk(query_scores, k)
@@ -102,8 +122,14 @@ def extract_topk_elements(query_scores, start_probs, end_probs, k):
     # Step 2: Use these indices to select the corresponding elements from start_probs and end_probs
     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
+    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():

utils/setup.py

@@ -84,12 +84,13 @@ def get_args():
     parser.add_argument("--ndcg_topk", type=int, nargs='+', default=[10, 20, 40], help="List of NDCG top k values")
     args = parser.parse_args()
     
-    
+    args.results_path = os.path.join(args.results_path, args.exp_id)
     os.makedirs(args.results_path, exist_ok=True)
+    
     if args.hard_negative_start_epoch != -1:
         if args.hard_pool_size > args.bsz:
             print("[WARNING] hard_pool_size is larger than bsz")
-    
+
     return args