| | import os, json |
| | import torch |
| | 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.setup import set_seed, get_args |
| | from utils.run_utils import prepare_optimizer, prepare_model, logger_ndcg_iou, save_model, resume_model |
| |
|
| | def main(): |
| | opt = get_args() |
| | logger = get_logger(opt.results_path, opt.exp_id) |
| | set_seed(opt.seed) |
| | logger.info("Arguments:\n%s", json.dumps(vars(opt), indent=4)) |
| | opt.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
| | logger.info(f"device: {opt.device}") |
| | |
| | |
| | |
| | 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 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 >= 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)): |
| | 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() |
| | |
| | optimizer.step() |
| |
|
| | if step % opt.log_step == 0: |
| | 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) |
| |
|
| | logger_ndcg_iou(val_ndcg_iou, logger, "VAL") |
| | logger_ndcg_iou(test_ndcg_iou, logger, "TEST") |
| |
|
| | if val_ndcg_iou[20][0.5] > best_val_ndcg: |
| | best_val_ndcg = val_ndcg_iou[20][0.5] |
| | logger_ndcg_iou(val_ndcg_iou, logger, "BEST VAL") |
| | logger_ndcg_iou(test_ndcg_iou, logger, "BEST TEST") |
| | |
| | bestmodel_path = os.path.join(opt.results_path, "best_model.pt") |
| | save_model(model, optimizer, epoch, bestmodel_path, logger) |
| |
|
| | if __name__ == '__main__': |
| | main() |
| |
|
