al_ret/retrieval.py

import json
import logging
import os
import numpy as np
import torch

from training.distributed import is_master
from .zero_shot import zero_shot_eval
from .util import parallel_apply
from .metrics import compute_metrics, tensor_text_to_video_metrics, tensor_video_to_text_sim
from torch.nn import functional as F
try:
    import wandb
except ImportError:
    wandb = None


#
# def evaluate_al_ret(model, data, epoch, args, tb_writer=None):
#     metrics = {}
#     if not is_master(args):
#         return metrics
#     model.eval()
#
#     zero_shot_metrics = zero_shot_eval(model, data, epoch, args)
#     metrics.update(zero_shot_metrics)
#
#     if not metrics:
#         return metrics
#
#     logging.info(
#         f"Eval Epoch: {epoch} "
#         + "\t".join([f"{k}: {round(v, 4):.4f}" for k, v in metrics.items()])
#     )
#
#     if args.save_logs:
#         for name, val in metrics.items():
#             if tb_writer is not None:
#                 tb_writer.add_scalar(f"val/al_ret/{name}", val, epoch)
#         args.al_ret_output_dir = os.path.join(args.log_base_path, 'al_ret')
#         os.makedirs(args.al_ret_output_dir, exist_ok=True)
#         with open(os.path.join(args.al_ret_output_dir, "results.jsonl"), "a+") as f:
#             f.write(json.dumps(metrics))
#             f.write("\n")
#
#     if args.wandb:
#         assert wandb is not None, 'Please install wandb.'
#         for name, val in metrics.items():
#             wandb.log({f"val/{name}": val, 'epoch': epoch})
#
#     return metrics



def _run_on_single_gpu(model,
                       # batch_list_t, batch_list_v,
                       batch_sequence_output_list, batch_visual_output_list):
    sim_matrix = []
    for idx1 in range(len(batch_sequence_output_list)):
        # input_mask, segment_ids, *_tmp = b1
        sequence_output = batch_sequence_output_list[idx1]
        each_row = []
        for idx2 in range(len(batch_visual_output_list)):
            # video_mask, *_tmp = b2
            visual_output = batch_visual_output_list[idx2]
            # b1b2_logits, *_tmp = model.get_similarity_logits(sequence_output, visual_output, input_mask, video_mask,
            #                                                          loose_type=model.loose_type)
            # logging.info(f"{model.logit_scale.device}, {visual_output.device}, {sequence_output.device}")
            b1b2_logits = model.logit_scale * sequence_output @ visual_output.T
            # print(model.logit_scale.device, visual_output.device, sequence_output.device)
            # logging.info(f"{b1b2_logits.shape}, {b1b2_logits.device}")
            b1b2_logits = b1b2_logits.cpu().detach().numpy()
            each_row.append(b1b2_logits)
        each_row = np.concatenate(tuple(each_row), axis=-1)
        sim_matrix.append(each_row)
    return sim_matrix

def evaluate_al_ret(model, data, epoch, args, tb_writer=None):
    if is_master(args) and (args.val_frequency and ((epoch % args.val_frequency) == 0 or epoch == args.epochs)):
        # print(data)
        val_al_ret_data = list(data.keys())
        # print(val_vl_ret_data)
        assert len(val_al_ret_data) == 1
        val_al_ret_data = val_al_ret_data[0]
        test_dataloader = data[val_al_ret_data]
        # print(len(test_dataloader))
        # print(len(test_dataloader))
        # print(len(test_dataloader))
        # print(len(test_dataloader))
        device = model.device
        n_gpu = torch.cuda.device_count()
        logging.info(f"\nEval Epoch: {epoch}, eval Audio-Text Retrieval under {val_al_ret_data.upper()} test data")
        if hasattr(model, 'module'):
            model = model.module.to(device)
        else:
            model = model.to(device)
        # #################################################################
        ## below variables are used to multi-sentences retrieval
        # multi_sentence_: important tag for eval
        # cut_off_points: used to tag the label when calculate the metric
        # sentence_num: used to cut the sentence representation
        # video_num: used to cut the video representation
        # #################################################################
        multi_sentence_ = False
        cut_off_points_, sentence_num_, video_num_ = [], -1, -1
        if hasattr(test_dataloader.dataset, 'multi_sentence_per_audio') and test_dataloader.dataset.multi_sentence_per_audio:
        # if False:
            multi_sentence_ = True
            cut_off_points_ = test_dataloader.dataset.cut_off_points
            sentence_num_ = test_dataloader.dataset.sentence_num
            video_num_ = test_dataloader.dataset.audio_num
            cut_off_points_ = [itm - 1 for itm in cut_off_points_]

        if multi_sentence_:
            print("Eval under the multi-sentence per audio clip setting.")
            print("sentence num: {}, video num: {}".format(sentence_num_, video_num_))
            logging.info("Eval under the multi-sentence per audio clip setting.")
            logging.info("sentence num: {}, video num: {}".format(sentence_num_, video_num_))

        model.eval()
        with torch.no_grad():
            # batch_list_t = []
            # batch_list_v = []
            batch_sequence_output_list, batch_visual_output_list = [], []
            total_video_num = 0

            # ----------------------------
            # 1. cache the features
            # ----------------------------
            for bid, batch in enumerate(test_dataloader):
                # batch = tuple(t.to(device) for t in batch)
                video, input_ids, attention_mask = batch
                # print(input_ids.shape, video.shape, video.dtype)
                input_ids = input_ids.squeeze().to(device)
                attention_mask = attention_mask.squeeze().to(device)
                # video = video.squeeze().permute(0, 2, 1, 3, 4).float().to(device)
                video = video.float().to(device)



                # print(input_ids.shape, video.shape, video.dtype)
                # print(input_ids.shape, video.shape)
                if multi_sentence_:
                    # multi-sentences retrieval means: one clip has two or more descriptions.
                    b, *_t = video.shape
                    sequence_output = model.encode_text(input_ids, attention_mask)
                    # logging.info(f'multi: {sequence_output.shape}')
                    # sequence_output = model.get_sequence_output(input_ids, segment_ids, input_mask)
                    batch_sequence_output_list.append(sequence_output)
                    # batch_list_t.append((input_mask, segment_ids,))
                    # 0 16
                    s_, e_ = total_video_num, total_video_num + b
                    filter_inds = [itm - s_ for itm in cut_off_points_ if itm >= s_ and itm < e_] # cut_off_points_ [0 4 9 14]

                    if len(filter_inds) > 0:
                        # video, video_mask = video[filter_inds, ...], video_mask[filter_inds, ...]
                        # print('before', video.shape)
                        video = video[filter_inds, ...]
                        # print('after', video.shape)
                        # visual_output = model.get_visual_output(video, video_mask)
                        visual_output = model.encode_image(video)
                        batch_visual_output_list.append(visual_output)
                        # batch_list_v.append((video_mask,))
                    total_video_num += b
                else:
                    sequence_output = model.encode_text(input_ids, attention_mask)
                    visual_output = model.encode_image(video)
                    # sequence_output, visual_output = model.get_sequence_visual_output(input_ids, segment_ids, input_mask, video, video_mask)

                    batch_sequence_output_list.append(sequence_output)
                    # batch_list_t.append((input_mask, segment_ids,))

                    batch_visual_output_list.append(visual_output)
                    # batch_list_v.append((video_mask,))

                print(f"Process {val_al_ret_data.upper()}: {bid}/{len(test_dataloader)}\r", end='')
            # ----------------------------------
            # 2. calculate the similarity
            # ----------------------------------
            n_gpu = torch.cuda.device_count()
            if n_gpu > 1:
                # print('n_gpu > 1')
                device_ids = list(range(n_gpu))
                # print('device_ids', device_ids)
                batch_t_output_splits = []
                batch_v_output_splits = []
                bacth_len = len(batch_sequence_output_list)
                # print(bacth_len)
                split_len = (bacth_len + n_gpu - 1) // n_gpu
                for dev_id in device_ids:
                    s_, e_ = dev_id * split_len, (dev_id + 1) * split_len
                    if dev_id == 0:

                        batch_t_output_splits.append(batch_sequence_output_list[s_:e_])
                        batch_v_output_splits.append(batch_visual_output_list)
                        # print(len(batch_sequence_output_list[s_:e_]), len(batch_visual_output_list))
                    else:
                        devc = torch.device('cuda:{}'.format(str(dev_id)))

                        devc_batch_list = [b.to(devc) for b in batch_sequence_output_list[s_:e_]]
                        batch_t_output_splits.append(devc_batch_list)
                        devc_batch_list = [b.to(devc) for b in batch_visual_output_list]
                        batch_v_output_splits.append(devc_batch_list)
                        # print(len(devc_batch_list), len(devc_batch_list))
                parameters_tuple_list = [(
                                          batch_t_output_splits[dev_id], batch_v_output_splits[dev_id]) for dev_id in device_ids]
                parallel_outputs = parallel_apply(_run_on_single_gpu, model, parameters_tuple_list, device_ids)
                sim_matrix = []
                for idx in range(len(parallel_outputs)):
                    sim_matrix += parallel_outputs[idx]
                sim_matrix = np.concatenate(tuple(sim_matrix), axis=0)
            else:
                sim_matrix = _run_on_single_gpu(model,
                                                # batch_list_t, batch_list_v,
                                                batch_sequence_output_list, batch_visual_output_list)
                sim_matrix = np.concatenate(tuple(sim_matrix), axis=0)
        #####################################################################
        if multi_sentence_:

            logging.info(f"{val_al_ret_data.upper()} before reshape, sim matrix size: {sim_matrix.shape}")
            cut_off_points2len_ = [itm + 1 for itm in cut_off_points_]
            max_length = max([e_-s_ for s_, e_ in zip([0]+cut_off_points2len_[:-1], cut_off_points2len_)])
            sim_matrix_new = []
            for s_, e_ in zip([0] + cut_off_points2len_[:-1], cut_off_points2len_):
                sim_matrix_new.append(np.concatenate((sim_matrix[s_:e_],
                                                      np.full((max_length-e_+s_, sim_matrix.shape[1]), -np.inf)), axis=0))
            sim_matrix = np.stack(tuple(sim_matrix_new), axis=0)
            logging.info(f"{val_al_ret_data.upper()} after reshape, sim matrix size: {sim_matrix.shape}")

            tv_metrics = tensor_text_to_video_metrics(sim_matrix)
            # vt_metrics = compute_metrics(tensor_video_to_text_sim(sim_matrix))
        else:
            logging.info(f"{val_al_ret_data.upper()} sim matrix size: {sim_matrix.shape[0]}, {sim_matrix.shape[1]}")
            t2v_sim_matrix = torch.from_numpy(sim_matrix).cuda()
            # t2v_sim_matrix = t2v_sim_matrix * F.softmax(t2v_sim_matrix*10, dim=0) * len(t2v_sim_matrix)
            tv_metrics = compute_metrics(t2v_sim_matrix.cpu().numpy())


            # vt_metrics = compute_metrics(t2v_sim_matrix.T.cpu().numpy())

            logging.info('\t Length-T: {}, Length-V:{}'.format(len(sim_matrix), len(sim_matrix[0])))

        logging.info(f"{val_al_ret_data.upper()} Text-to-Audio:")
        logging.info('\t>>>  R@1: {:.1f} - R@5: {:.1f} - R@10: {:.1f} - Median R: {:.1f} - Mean R: {:.1f}'.
                    format(tv_metrics['R1'], tv_metrics['R5'], tv_metrics['R10'], tv_metrics['MR'], tv_metrics['MeanR']))
        # logging.info(f"{val_al_ret_data.upper()} Text-to-Audio:")
        # logging.info('\t>>>  V2T$R@1: {:.1f} - V2T$R@5: {:.1f} - V2T$R@10: {:.1f} - V2T$Median R: {:.1f} - V2T$Mean R: {:.1f}'.
        #             format(vt_metrics['R1'], vt_metrics['R5'], vt_metrics['R10'], vt_metrics['MR'], vt_metrics['MeanR']))


        if args.save_logs:
            for name, val in tv_metrics.items():
                if tb_writer is not None:
                    tb_writer.add_scalar(f"val/al_ret/{val_al_ret_data}/t2a/{name}", val, epoch)
            # for name, val in vt_metrics.items():
            #     if tb_writer is not None:
            #         tb_writer.add_scalar(f"val/al_ret/{val_al_ret_data}/v2t/{name}", val, epoch)

            args.al_ret_output_dir = os.path.join(args.log_base_path, f'al_ret/{val_al_ret_data}')
            os.makedirs(args.al_ret_output_dir, exist_ok=True)
            with open(os.path.join(args.al_ret_output_dir, "results.jsonl"), "a+") as f:
                f.write(json.dumps({'t2a': tv_metrics}))
                f.write("\n")
                # f.write(json.dumps({'v2t': vt_metrics}))
                # f.write("\n")