Code release for "EgoVLPv2: Egocentric Video-Language Pre-training with Fusion in the Backbone" [ICCV, 2023]
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import json
import torch
import os
import pdb

######################################################################################################
#                     Load data
######################################################################################################
annotation_path = ""  # Change to your own path containing canonical annotation files
feat_path = "saved_features/EgoVLPv2"  # Change to your own path containing features of canonical videos
info_path = 'jsons/ego4d.json'
annot_path_train = 'jsons/moments_train.json'
annot_path_val = 'jsons/moments_val.json'
annot_path_test = 'jsons/moments_test_unannotated.json'

with open(annot_path_train, 'r') as f:
    v_annot_train = json.load(f)

with open(annot_path_val, 'r') as f:
    v_annot_val = json.load(f)

with open(annot_path_test, 'r') as f:
    v_annot_test = json.load(f)

with open(info_path, 'r') as f:
    feat_info=json.load(f)

v_all_duration = {}
for video in feat_info['videos']:
    v_id = video['video_uid']
    v_dur = video['duration_sec']
    v_all_duration[v_id] = v_dur

v_annot = {}
v_annot['videos'] = v_annot_train['videos'] + v_annot_val['videos'] + v_annot_test['videos']

######################################################################################################
#                     Convert video annotations to clip annotations: clip_annot_1
######################################################################################################
clip_annot_1 = {}
for video in v_annot['videos']:
    vid = video['video_uid']
    clips = video['clips']
    v_duration = v_all_duration[vid] #feat_info[feat_info.video_uid == vid].canonical_video_duration_sec.values[0]
    try:
        feats = torch.load(os.path.join(feat_path, clips[0]['clip_uid'] + '.pt'))
    except:
        print(f'{vid} features do not exist!')
        continue
    fps = feats.shape[0] / v_duration
    for clip in clips:
        clip_id = clip['clip_uid']

        if clip_id not in clip_annot_1.keys():
            clip_annot_1[clip_id] = {}
            clip_annot_1[clip_id]['video_id'] = vid
            clip_annot_1[clip_id]['clip_id'] = clip_id
            clip_annot_1[clip_id]['parent_start_sec'] = clip['video_start_sec']
            clip_annot_1[clip_id]['parent_end_sec'] = clip['video_end_sec']
            clip_annot_1[clip_id]['v_duration'] = v_duration
            clip_annot_1[clip_id]['fps'] = fps
            clip_annot_1[clip_id]['annotations'] = []
            clip_annot_1[clip_id]['subset'] = video['split']

        if video['split'] != 'test':
            annotations = clip['annotations']
            for cnt, annot in enumerate(annotations):
                for label in annot['labels']:
                    if label['primary']:
                        clip_annot_1[clip_id]['annotations'].append(label)

#######################################################################
## If there are no remaining annotations for a clip ###
## Remove the clip ###

remove_list = []
for k, v in clip_annot_1.items():
    if v['subset']!='test' and len(v['annotations']) == 0:
        print(f'NO annotations: video {k}')
        remove_list.append(k)

for item in remove_list:
    del clip_annot_1[item]

cnt_train = 0
cnt_val = 0
for k, v in clip_annot_1.items():
    if v['subset'] == 'train':
        cnt_train += 1
    elif v['subset'] == 'val':
        cnt_val += 1

print(f"Number of clips in training: {cnt_train}")
print(f"Number of clips in validation: {cnt_val}")

with open("Evaluation/ego4d/annot/clip_annotations.json", "w") as fp:
    json.dump(clip_annot_1, fp)


EgoMQ/Convert_annotations.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


from Evaluation.ego4d.generate_detection import gen_detection_multicore as gen_det_ego4d
from Evaluation.ego4d.get_detect_performance import evaluation_detection as eval_det_ego4d
from Evaluation.ego4d.generate_retrieval import gen_retrieval_multicore as gen_retrieval
from Evaluation.ego4d.get_retrieval_performance import evaluation_retrieval as eval_retrieval

import Utils.opts as opts
import os
if __name__ == '__main__':

    opt = opts.parse_opt()
    opt = vars(opt)

    print(opt)

    if not os.path.exists(opt["output_path"]):
        print('No predictions! Please run inference first!')

    if opt['infer_datasplit'] == 'test':
        if opt['eval_stage'] == 'eval_detection' or opt['eval_stage'] == 'all':
            print("---------------------------------------------------------------------------------------------")
            print("2. Detection evaluation starts!")
            print("---------------------------------------------------------------------------------------------")

            print("a. Generate detections!")
            gen_det_ego4d(opt)   # Not knowing video categories
            print("Detection evaluation finishes! \n")

        if opt['eval_stage'] == 'eval_retrieval' or opt['eval_stage'] == 'all':

            print("---------------------------------------------------------------------------------------------")
            print("3. Retrieval evaluation starts!")
            print("---------------------------------------------------------------------------------------------")

            print("a. Generate retrieval!")
            gen_retrieval(opt)
            print("Detection evaluation finishes! \n")

    else:
        if opt['eval_stage'] == 'eval_detection' or opt['eval_stage'] == 'all':
            print("---------------------------------------------------------------------------------------------")
            print("2. Detection evaluation starts!")
            print("---------------------------------------------------------------------------------------------")

            print("a. Generate detections!")
            gen_det_ego4d(opt)   # Not knowing video categories

            print("b. Evaluate the detection results!")
            eval_det_ego4d(opt)

            print("Detection evaluation finishes! \n")

        if opt['eval_stage'] == 'eval_retrieval' or opt['eval_stage'] == 'all':

            print("---------------------------------------------------------------------------------------------")
            print("3. Retrieval evaluation starts!")
            print("---------------------------------------------------------------------------------------------")

            print("a. Generate retrieval!")
            gen_retrieval(opt)

            print("b. Evaluate the retrieval results!")
            eval_retrieval(opt)
            print("Detection evaluation finishes! \n")

EgoMQ/Eval.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import logging
import numpy as np
import pandas as pd

from utils import interpolated_prec_rec
from utils import segment_iou

from joblib import Parallel, delayed

class EvalActionDetection(object):

    def __init__(self, ground_truth_filename=None, 
                 tiou_thresholds=np.linspace(0.5, 0.95, 10), 
                 verbose=True):
        if not ground_truth_filename:
            raise IOError('Please input a valid ground truth file.')
        self.tiou_thresholds = tiou_thresholds
        self.verbose = verbose
        self.ap = None

        self.ground_truth, self.activity_labels = self._import_ground_truth(
            ground_truth_filename)

        if self.verbose:
            logging.info('[INIT] Loaded annotations')
            nr_gt = len(np.unique(self.ground_truth['gt-id']))
            logging.info('\tNumber of ground truth instances: {}'.format(nr_gt))
            logging.info('\tFixed threshold for tiou score: {}'.format(self.tiou_thresholds))

    def _import_ground_truth(self, ground_truth_filename):
        """Reads ground truth file, checks if it is well formatted, and returns
           the ground truth instances and the activity classes.

        Parameters
        ----------
        ground_truth_filename : str
            Full path to the ground truth CSV file.

        Outputs
        -------
        ground_truth : df
            Data frame containing the ground truth instances.
        activity_labels : list
            List of the activity classes
        """
        ground_truth = pd.read_csv(ground_truth_filename)
        # remove instances of length <=0 
        ground_truth = ground_truth.loc[ground_truth['t-start'].values < ground_truth['t-end'].values].reset_index(drop=True)

        activity_labels = np.unique(ground_truth['label'])
        ground_truth['gt-id'] = range(len(ground_truth))

        return ground_truth, activity_labels

    def wrapper_compute_average_precision(self, prediction):
        """Computes average precision for each class in the subset.
        """
        ap = np.zeros((len(self.tiou_thresholds), len(self.activity_labels)))
        ground_truth_by_label = self.ground_truth.groupby(by='label')

        results = Parallel(n_jobs=len(self.activity_labels))(
                    delayed(compute_average_precision_detection)(
                        ground_truth=ground_truth_by_label.get_group(l).reset_index(drop=True),
                        prediction=prediction[prediction['label']==l].reset_index(drop=True),
                        tiou_thresholds=self.tiou_thresholds,
                    ) for l in self.activity_labels)
        
        for i in range(len(self.activity_labels)):
            ap[:,i]= results[i]

        return ap

    def evaluate(self, prediction):
        """Evaluates a prediction df. For the detection task we measure the
        interpolated mean average precision to measure the performance of a
        method.
        """
        logging.info('\tNumber of predictions: {}. Number of unique videos: {}'.format(len(prediction),len(np.unique(prediction['video-name']))))

        self.ap = self.wrapper_compute_average_precision(prediction)
        self.mAP = self.ap.mean(axis=1)
        self.average_mAP = self.mAP.mean()

        if self.verbose:
            logging.info('[RESULTS] Performance on ActivityNet detection task.')
            logging.info('\tmAP: {}'.format(self.mAP))
            logging.info('\tAverage-mAP: {}'.format(self.average_mAP))
        return self.average_mAP

def compute_average_precision_detection(ground_truth, prediction, tiou_thresholds=np.linspace(0.5, 0.95, 10)):
    """Compute average precision (detection task) between ground truth and
    predictions data frames. If multiple predictions occurs for the same
    predicted segment, only the one with highest score is matches as
    true positive. This code is greatly inspired by Pascal VOC devkit.

    Parameters
    ----------
    ground_truth : df
        Data frame containing the ground truth instances.
        Required fields: ['video-name', 't-start', 't-end']
    prediction : df
        Data frame containing the prediction instances.
        Required fields: ['video-name, 't-start', 't-end', 'score']
    tiou_thresholds : 1darray, optional
        Temporal intersection over union threshold.

    Outputs
    -------
    ap : float
        Average precision score.
    """
    ap = np.zeros(len(tiou_thresholds))
    if prediction.empty:
        return ap

    gt_id_lst = np.unique(ground_truth['gt-id'].values)
    gt_id_to_index = dict(zip(gt_id_lst, range(len(gt_id_lst))))
    lock_gt = np.ones((len(tiou_thresholds),len(gt_id_to_index))) * -1
    
    npos = float(len(gt_id_lst))

    # Sort predictions by decreasing score order.
    sort_idx = prediction['score'].values.argsort()[::-1]
    prediction = prediction.loc[sort_idx].reset_index(drop=True)

    # Initialize true positive and false positive vectors.
    tp = np.zeros((len(tiou_thresholds), len(prediction)))
    fp = np.zeros((len(tiou_thresholds), len(prediction)))


    # Adaptation to query faster
    ground_truth_gbvn = ground_truth.groupby(by='video-name')

    # Assigning true positive to truly grount truth instances.
    for idx, this_pred in prediction.iterrows():

        try:
            # Check if there is at least one ground truth in the video associated.
            ground_truth_videoid = ground_truth_gbvn.get_group(this_pred['video-name'])
        except Exception as e:
            fp[:, idx] = 1
            continue

        this_gt = ground_truth_videoid.reset_index()
        tiou_arr = segment_iou(this_pred[['t-start', 't-end']].values,
                               this_gt[['t-start', 't-end']].values)
        # We would like to retrieve the predictions with highest tiou score.
        tiou_sorted_idx = tiou_arr.argsort()[::-1]
        for tidx, tiou_thr in enumerate(tiou_thresholds):
            for jdx in tiou_sorted_idx:
                if tiou_arr[jdx] < tiou_thr:
                    fp[tidx, idx] = 1
                    break
                if lock_gt[tidx, gt_id_to_index[this_gt.loc[jdx]['gt-id']]] >= 0:
                    continue
                # Assign as true positive after the filters above.
                tp[tidx, idx] = 1
                lock_gt[tidx, gt_id_to_index[this_gt.loc[jdx]['gt-id']]] = idx
                break

            if fp[tidx, idx] == 0 and tp[tidx, idx] == 0:
                fp[tidx, idx] = 1

    tp_cumsum = np.cumsum(tp, axis=1).astype(np.float)
    fp_cumsum = np.cumsum(fp, axis=1).astype(np.float)
    recall_cumsum = tp_cumsum / npos

    precision_cumsum = recall_cumsum * npos / (recall_cumsum * npos + fp_cumsum)
    
    for tidx in range(len(tiou_thresholds)):
        ap[tidx] = interpolated_prec_rec(precision_cumsum[tidx,:], recall_cumsum[tidx,:])

    return ap

EgoMQ/Evaluation/ego4d/eval_action_detection.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import json
import numpy as np
import pandas as pd
from joblib import Parallel, delayed

from ..utils import get_blocked_videos
from ..utils import interpolated_prec_rec
from ..utils import segment_iou

import warnings

warnings.filterwarnings("ignore", message="numpy.dtype size changed")
warnings.filterwarnings("ignore", message="numpy.ufunc size changed")


class ANETdetection(object):
    GROUND_TRUTH_FIELDS = ['database']
    # GROUND_TRUTH_FIELDS = ['database', 'taxonomy', 'version']
    PREDICTION_FIELDS = ['results', 'version', 'external_data']

    def __init__(self, ground_truth_filename=None, prediction_filename=None,
                 ground_truth_fields=GROUND_TRUTH_FIELDS,
                 prediction_fields=PREDICTION_FIELDS,
                 tiou_thresholds=np.linspace(0.5, 0.95, 10),
                 subset='test', verbose=True,
                 check_status=False):
        if not ground_truth_filename:
            raise IOError('Please input a valid ground truth file.')
        if not prediction_filename:
            raise IOError('Please input a valid prediction file.')
        self.subset = subset
        self.tiou_thresholds = tiou_thresholds
        self.verbose = verbose
        self.gt_fields = ground_truth_fields
        self.pred_fields = prediction_fields
        self.ap = None
        self.check_status = check_status
        # Retrieve blocked videos from server.

        if self.check_status:
            self.blocked_videos = get_blocked_videos()
        else:
            self.blocked_videos = list()

        # Import ground truth and predictions.
        self.ground_truth, self.activity_index = self._import_ground_truth(
            ground_truth_filename)
        self.prediction = self._import_prediction(prediction_filename)

        if self.verbose:
            print('[INIT] Loaded annotations from {} subset.'.format(subset))
            nr_gt = len(self.ground_truth)
            print('\tNumber of ground truth instances: {}'.format(nr_gt))
            nr_pred = len(self.prediction)
            print('\tNumber of predictions: {}'.format(nr_pred))
            print('\tFixed threshold for tiou score: {}'.format(self.tiou_thresholds))

    def _import_ground_truth(self, ground_truth_filename):
        """Reads ground truth file, checks if it is well formatted, and returns
           the ground truth instances and the activity classes.

        Parameters
        ----------
        ground_truth_filename : str
            Full path to the ground truth json file.

        Outputs
        -------
        ground_truth : df
            Data frame containing the ground truth instances.
        activity_index : dict
            Dictionary containing class index.
        """
        print(ground_truth_filename)
        with open(ground_truth_filename, 'r') as fobj:
            data = json.load(fobj)
        # Checking format
        # if not all([field in data.keys() for field in self.gt_fields]):
        #     raise IOError('Please input a valid ground truth file.')

        # Read ground truth data.
        activity_index, cidx = {}, 0
        video_lst, t_start_lst, t_end_lst, label_lst = [], [], [], []
        for videoid, v in data.items():
            # print(v)
            if not v['subset'] in self.subset:
                continue
            if videoid in self.blocked_videos:
                continue
            for ann in v['annotations']:
                if ann['label'] not in activity_index:
                    activity_index[ann['label']] = cidx
                    cidx += 1
                video_lst.append(videoid)
                t_start_lst.append(float(ann['start_time']))
                t_end_lst.append(float(ann['end_time']))
                label_lst.append(activity_index[ann['label']])

        ground_truth = pd.DataFrame({'video-id': video_lst,
                                     't-start': t_start_lst,
                                     't-end': t_end_lst,
                                     'label': label_lst})
        if self.verbose:
            print(activity_index)

        with open('activity_idx.json', 'w') as f:
            json.dump(activity_index, f)

        return ground_truth, activity_index

    def _import_prediction(self, prediction_filename):
        """Reads prediction file, checks if it is well formatted, and returns
           the prediction instances.

        Parameters
        ----------
        prediction_filename : str
            Full path to the prediction json file.

        Outputs
        -------
        prediction : df
            Data frame containing the prediction instances.
        """
        with open(prediction_filename, 'r') as fobj:
            data = json.load(fobj)
        # Checking format...
        if not all([field in data.keys() for field in self.pred_fields]):
            raise IOError('Please input a valid prediction file.')

        # Read predictions.
        video_lst, t_start_lst, t_end_lst = [], [], []
        label_lst, score_lst = [], []
        for videoid, v in data['results'].items():
            if videoid in self.blocked_videos:
                continue
            for result in v:
                try:
                    label = self.activity_index[result['label']]
                except:
                    label_words = result['label']
                    print(f'Category {label_words} is not in groundtruth set!')
                    continue
                video_lst.append(videoid)
                t_start_lst.append(float(result['segment'][0]))
                t_end_lst.append(float(result['segment'][1]))
                label_lst.append(label)
                score_lst.append(result['score'])
        prediction = pd.DataFrame({'video-id': video_lst,
                                   't-start': t_start_lst,
                                   't-end': t_end_lst,
                                   'label': label_lst,
                                   'score': score_lst})
        return prediction

    def _get_predictions_with_label(self, prediction_by_label, label_name, cidx):
        """Get all predicitons of the given label. Return empty DataFrame if there
        is no predcitions with the given label.
        """
        try:
            return prediction_by_label.get_group(cidx).reset_index(drop=True)
        except:
            if self.verbose:
                print('Warning: No predictions of label \'%s\' were provdied.' % label_name)
            return pd.DataFrame()

    def wrapper_compute_average_precision(self):
        """Computes average precision for each class in the subset.
        """
        ap = np.zeros((len(self.tiou_thresholds), len(self.activity_index)))

        # Adaptation to query faster
        ground_truth_by_label = self.ground_truth.groupby('label')
        prediction_by_label = self.prediction.groupby('label')

        results = Parallel(n_jobs=len(self.activity_index))(
            delayed(compute_average_precision_detection)(
                ground_truth=ground_truth_by_label.get_group(cidx).reset_index(drop=True),
                prediction=self._get_predictions_with_label(prediction_by_label, label_name, cidx),
                tiou_thresholds=self.tiou_thresholds,
            ) for label_name, cidx in self.activity_index.items())

        # for label_name, cidx in self.activity_index.items():
        #     compute_average_precision_detection(
        #                 ground_truth=ground_truth_by_label.get_group(cidx).reset_index(drop=True),
        #                 prediction=self._get_predictions_with_label(prediction_by_label, label_name, cidx),
        #                 tiou_thresholds=self.tiou_thresholds
        #             )

        for i, cidx in enumerate(self.activity_index.values()):
            ap[:, cidx] = results[i]

        np.save('ap.npy', ap)

        return ap

    def evaluate(self):
        """Evaluates a prediction file. For the detection task we measure the
        interpolated mean average precision to measure the performance of a
        method.
        """
        self.ap = self.wrapper_compute_average_precision()

        self.mAP = self.ap.mean(axis=1)
        self.average_mAP = self.mAP.mean()

        if self.verbose:
            print('[RESULTS] Performance on ActivityNet detection task.')
            print('Average-mAP: {}'.format(self.average_mAP))
            print('mAPs are {}'.format(self.mAP))

        return self.mAP, self.average_mAP


def compute_average_precision_detection(ground_truth, prediction, tiou_thresholds=np.linspace(0.5, 0.95, 10)):
    """Compute average precision (detection task) between ground truth and
    predictions data frames. If multiple predictions occurs for the same
    predicted segment, only the one with highest score is matches as
    true positive. This code is greatly inspired by Pascal VOC devkit.

    Parameters
    ----------
    ground_truth : df
        Data frame containing the ground truth instances.
        Required fields: ['video-id', 't-start', 't-end']
    prediction : df
        Data frame containing the prediction instances.
        Required fields: ['video-id, 't-start', 't-end', 'score']
    tiou_thresholds : 1darray, optional
        Temporal intersection over union threshold.

    Outputs
    -------
    ap : float
        Average precision score.
    """
    ap = np.zeros(len(tiou_thresholds))
    if prediction.empty:
        return ap

    npos = float(len(ground_truth))
    lock_gt = np.ones((len(tiou_thresholds), len(ground_truth))) * -1
    # Sort predictions by decreasing score order.
    sort_idx = prediction['score'].values.argsort()[::-1]
    prediction = prediction.loc[sort_idx].reset_index(drop=True)

    # Initialize true positive and false positive vectors.
    tp = np.zeros((len(tiou_thresholds), len(prediction)))
    fp = np.zeros((len(tiou_thresholds), len(prediction)))

    # Adaptation to query faster
    ground_truth_gbvn = ground_truth.groupby('video-id')

    # Assigning true positive to truly grount truth instances.
    for idx, this_pred in prediction.iterrows():

        if this_pred['video-id'] == 'c424ecf4-c654-4770-8af6-de97adb8241f':
            a = 1
        try:
            # Check if there is at least one ground truth in the video associated.
            ground_truth_videoid = ground_truth_gbvn.get_group(this_pred['video-id'])
        except Exception as e:
            fp[:, idx] = 1
            continue

        this_gt = ground_truth_videoid.reset_index()
        tiou_arr = segment_iou(this_pred[['t-start', 't-end']].values,
                               this_gt[['t-start', 't-end']].values)
        # We would like to retrieve the predictions with highest tiou score.
        tiou_sorted_idx = tiou_arr.argsort()[::-1]
        for tidx, tiou_thr in enumerate(tiou_thresholds):
            for jdx in tiou_sorted_idx:
                if tiou_arr[jdx] < tiou_thr:
                    fp[tidx, idx] = 1
                    break
                if lock_gt[tidx, this_gt.loc[jdx]['index']] >= 0:
                    continue
                # Assign as true positive after the filters above.
                tp[tidx, idx] = 1
                lock_gt[tidx, this_gt.loc[jdx]['index']] = idx
                break

            if fp[tidx, idx] == 0 and tp[tidx, idx] == 0:
                fp[tidx, idx] = 1

    tp_cumsum = np.cumsum(tp, axis=1).astype(np.float)
    fp_cumsum = np.cumsum(fp, axis=1).astype(np.float)
    recall_cumsum = tp_cumsum / npos

    precision_cumsum = tp_cumsum / (tp_cumsum + fp_cumsum)

    for tidx in range(len(tiou_thresholds)):
        ap[tidx] = interpolated_prec_rec(precision_cumsum[tidx, :], recall_cumsum[tidx, :])

    return ap

EgoMQ/Evaluation/ego4d/eval_detection.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import numpy as np
import pandas as pd
import json
import os
from joblib import Parallel, delayed


def IOU(s1, e1, s2, e2):
    if (s2 > e1) or (s1 > e2):
        return 0
    Aor = max(e1, e2) - min(s1, s2)
    Aand = min(e1, e2) - max(s1, s2)
    return float(Aand) / Aor

def Soft_NMS(df, nms_threshold=1e-5, num_prop=200):

    df = df.sort_values(by="score", ascending=False)

    tstart = list(df.xmin.values[:])
    tend = list(df.xmax.values[:])
    tscore = list(df.score.values[:])
    tlabel = list(df.label.values[:])

    rstart = []
    rend = []
    rscore = []
    rlabel = []

    while len(tscore) > 1 and len(rscore) < num_prop and max(tscore)>0:
        max_index = tscore.index(max(tscore))
        if tlabel[max_index] == 0:
            tstart.pop(max_index)
            tend.pop(max_index)
            tscore.pop(max_index)
            tlabel.pop(max_index)
            continue

        for idx in range(0, len(tscore)):
            if idx != max_index:
                tmp_iou = IOU(tstart[max_index], tend[max_index], tstart[idx], tend[idx])
                if tmp_iou > 0:
                    tscore[idx] = tscore[idx] * np.exp(-np.square(tmp_iou) / nms_threshold)

        rstart.append(tstart[max_index])
        rend.append(tend[max_index])
        rscore.append(tscore[max_index])
        rlabel.append(tlabel[max_index])

        tstart.pop(max_index)
        tend.pop(max_index)
        tscore.pop(max_index)
        tlabel.pop(max_index)

    newDf = pd.DataFrame()
    newDf['score'] = rscore
    newDf['xmin'] = rstart
    newDf['xmax'] = rend
    newDf['label'] = rlabel
    return newDf

def _gen_detection_video(video_name, result, idx_classes, opt, num_prop=200, topk = 2):
    path = os.path.join(opt["output_path"], opt["prop_path"]) + "/"
    files = [path + f for f in os.listdir(path) if  video_name in f]
    if len(files) == 0:
        print('Missing result for video {}'.format(video_name))
        return {video_name:[]}

    dfs = []  # merge pieces of video
    for snippet_file in files:
        snippet_df = pd.read_csv(snippet_file)
        dfs.append(snippet_df)
    df = pd.concat(dfs)

    df = Soft_NMS(df, nms_threshold=opt['nms_alpha_detect'], num_prop=num_prop)
    df = df.sort_values(by="score", ascending=False)

    duration = result['duration']
    proposal_list = []
    for j in range(min(num_prop, len(df))):
        tmp_proposal = {}
        tmp_proposal["label"] = idx_classes[int(df.label.values[j])]
        tmp_proposal["score"] = float(round(df.score.values[j], 6))
        tmp_proposal["segment"] = [float(round(max(0, df.xmin.values[j]), 1)),
                                   float(round(min(duration, df.xmax.values[j]), 1))]
        proposal_list.append(tmp_proposal)

    return {video_name:proposal_list}

def gen_detection_multicore(opt):

    with open(opt['clip_anno'], 'r') as f:
        ego4d_gt = json.load(f)

    video_list = []
    for k, v in ego4d_gt.items():
        if v['subset'] in opt['infer_datasplit']:
            video_list.append(k)


    # load all categories
    if os.path.exists(opt["moment_classes"]):
        with open(opt["moment_classes"], 'r') as f:
            classes = json.load(f)

    idx_classes = {}
    for key, value in classes.items():
        idx_classes[value] = key

    result = {
        video:
            {
                'duration': ego4d_gt[video]['parent_end_sec'] - ego4d_gt[video]['parent_start_sec']
            }
        for video in video_list
    }

    parallel = Parallel(n_jobs=16, prefer="processes")
    detection = parallel(delayed(_gen_detection_video)(video_name, result[video_name], idx_classes, opt)
                        for video_name in video_list)
    detection_dict = {}
    [detection_dict.update(d) for d in detection]
    output_dict = {"version": "EGO4D", "results": detection_dict, "external_data": {}}

    with open(os.path.join(opt["output_path"], opt["detect_result_file"]), "w") as out:
        json.dump(output_dict, out)



EgoMQ/Evaluation/ego4d/generate_detection.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import numpy as np
import pandas as pd
import json
import os
import sys
from joblib import Parallel, delayed

def IOU(s1, e1, s2, e2):
    if (s2 > e1) or (s1 > e2):
        return 0
    Aor = max(e1, e2) - min(s1, s2)
    Aand = min(e1, e2) - max(s1, s2)
    return float(Aand) / Aor

def Soft_NMS(df, nms_threshold=1e-5, num_prop=200):

    df = df.sort_values(by="score", ascending=False)

    tstart = list(df.xmin.values[:])
    tend = list(df.xmax.values[:])
    tscore = list(df.score.values[:])
    tlabel = list(df.label.values[:])

    rstart = []
    rend = []
    rscore = []
    rlabel = []


    while len(tscore) > 1 and len(rscore) < num_prop and max(tscore)>0:
        max_index = tscore.index(max(tscore))
        if tlabel[max_index] == 0:
            tstart.pop(max_index)
            tend.pop(max_index)
            tscore.pop(max_index)
            tlabel.pop(max_index)
            continue

        for idx in range(0, len(tscore)):
            if idx != max_index:
                tmp_iou = IOU(tstart[max_index], tend[max_index], tstart[idx], tend[idx])
                if tmp_iou > 0:
                    tscore[idx] = tscore[idx] * np.exp(-np.square(tmp_iou) / nms_threshold)

        rstart.append(tstart[max_index])
        rend.append(tend[max_index])
        rscore.append(tscore[max_index])
        rlabel.append(tlabel[max_index])

        tstart.pop(max_index)
        tend.pop(max_index)
        tscore.pop(max_index)
        tlabel.pop(max_index)

    newDf = pd.DataFrame()
    newDf['score'] = rscore
    newDf['xmin'] = rstart
    newDf['xmax'] = rend
    newDf['label'] = rlabel
    return newDf


def rm_other_category(df, annots, classes):
    df_v = []
    list_label = []
    for ann in annots:
        label = classes[ann['label']]

        if not label in list_label:
            df_v.append(df[df['label'] == label])
            list_label.append(label)

    df_v = pd.concat(df_v)
    return df_v

def _gen_retrieval_video(video_name, result, classes, idx_classes, opt, test_anno, num_prop=200,  topk = 2):
    path = os.path.join(opt["output_path"], opt["prop_path"]) + "/"
    files = [path + f for f in os.listdir(path) if  video_name in f]
    if len(files) == 0:
        print('Missing result for video {}'.format(video_name))
        return {video_name:[]}

    dfs = []  # merge pieces of video
    for snippet_file in files:
        snippet_df = pd.read_csv(snippet_file)
        dfs.append(snippet_df)
    df = pd.concat(dfs)

    if opt['infer_datasplit'] != 'test':
        df = rm_other_category(df, test_anno['annotations'], classes)

    df = df.sort_values(by="score", ascending=False)

    duration = result['duration']
    proposal_list = []
    for j in range(min(num_prop, len(df))):
        tmp_proposal = {}
        tmp_proposal["label"] = idx_classes[int(df.label.values[j])]
        tmp_proposal["score"] = float(round(df.score.values[j], 6))
        tmp_proposal["segment"] = [float(round(max(0, df.xmin.values[j]), 1)),
                                   float(round(min(duration, df.xmax.values[j]), 1))]
        proposal_list.append(tmp_proposal)

    return {video_name:proposal_list}

def gen_retrieval_multicore(opt):

    with open(opt['clip_anno'], 'r') as f:
        ego4d_gt = json.load(f)

    video_list = []
    for k, v in ego4d_gt.items():
        if v['subset'] in opt['infer_datasplit']:
            video_list.append(k)

    # load all categories
    if os.path.exists(opt["moment_classes"]):
        with open(opt["moment_classes"], 'r') as f:
            classes = json.load(f)

    idx_classes = {}
    for key, value in classes.items():
        idx_classes[value] = key


    result = {
        video:
            {
                'duration': ego4d_gt[video]['parent_end_sec'] - ego4d_gt[video]['parent_start_sec']
            }
        for video in video_list
    }

    parallel = Parallel(n_jobs=1, prefer="processes")
    detection = parallel(delayed(_gen_retrieval_video)(video_name, result[video_name], classes, idx_classes, opt, ego4d_gt[video_name])
                        for video_name in video_list)

    detection_dict = {}
    [detection_dict.update(d) for d in detection]
    output_dict = {"version": "EGO4D", "results": detection_dict, "external_data": {}}

    with open(os.path.join(opt["output_path"], opt["retrieval_result_file"]), "w") as out:
        json.dump(output_dict, out)



EgoMQ/Evaluation/ego4d/generate_retrieval.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import os
from .eval_detection import ANETdetection

def run_evaluation(ground_truth_filename, prediction_filename,
         subset='test', tiou_thresholds=[0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7],
         verbose=True):

    anet_detection = ANETdetection(ground_truth_filename, prediction_filename,
                                   subset=subset, tiou_thresholds=tiou_thresholds,
                                   verbose=verbose, check_status=False)

    mAPs, average_mAP = anet_detection.evaluate()

    for (tiou, mAP) in zip(tiou_thresholds, mAPs):
        print("mAP at tIoU {} is {}".format(tiou, mAP))



def evaluation_detection(opt):

    run_evaluation(ground_truth_filename = opt["clip_anno"],
                   prediction_filename = os.path.join(opt["output_path"], opt["detect_result_file"]),
                   subset=opt['infer_datasplit'], tiou_thresholds=opt['tIoU_thr'])

EgoMQ/Evaluation/ego4d/get_detect_performance.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import numpy as np
import os
import json
class Moment_Retrieval(object):
    GROUND_TRUTH_FIELDS = ['database']
    PREDICTION_FIELDS = ['results', 'version', 'external_data']

    def __init__(self, ground_truth_filename=None, prediction_filename=None,
                 ground_truth_fields=GROUND_TRUTH_FIELDS,
                 prediction_fields=PREDICTION_FIELDS,
                 tiou_thresholds=np.linspace(0.5, 0.95, 10),
                 subset='test', verbose=False,
                 check_status=False):
        if not ground_truth_filename:
            raise IOError('Please input a valid ground truth file.')
        if not prediction_filename:
            raise IOError('Please input a valid prediction file.')
        self.subset = subset
        self.tiou_thresholds = tiou_thresholds
        self.verbose = verbose
        self.gt_fields = ground_truth_fields
        self.pred_fields = prediction_fields
        self.ap = None
        self.check_status = check_status
        # Retrieve blocked videos from server.

        # Import ground truth and predictions.
        self.ground_truth = self._import_ground_truth(
            ground_truth_filename)
        self.prediction = self._import_prediction(prediction_filename)

        if self.verbose:
            print('[INIT] Loaded annotations from {} subset.'.format(subset))
            nr_gt = len(self.ground_truth)
            print('\tNumber of ground truth instances: {}'.format(nr_gt))
            nr_pred = len(self.prediction)
            print('\tNumber of predictions: {}'.format(nr_pred))
            print('\tFixed threshold for tiou score: {}'.format(self.tiou_thresholds))

    def _import_ground_truth(self, ground_truth_filename):

        with open(ground_truth_filename, 'r') as fobj:
            data = json.load(fobj)


        ground_truth = {}
        for videoid, v in data.items():

            if not v['subset'] in self.subset:
                continue

            annotations = {}
            for ann in v['annotations']:

                if ann['label'] not in annotations.keys():
                    annotations[ann['label']] = []
                annotations[ann['label']].append([ann['start_time'], ann['end_time']])

            ground_truth[v['clip_id']] = annotations

        return ground_truth

    def _import_prediction(self, prediction_filename):

        with open(prediction_filename, 'r') as fobj:
            data = json.load(fobj)
        # Checking format...
        if not all([field in data.keys() for field in self.pred_fields]):
            raise IOError('Please input a valid prediction file.')


        prediction = {}
        for videoid, v in data['results'].items():

            pred_props = {}
            for prop in v:
                if prop['label'] not in pred_props.keys():
                    pred_props[prop['label']] = []
                pred_props[prop['label']].append([prop['segment'][0], prop['segment'][1], prop['score']])

            prediction[videoid] = pred_props

        return prediction

    def evaluate(self):
        tious = [0.3, 0.5, 0.7]
        recalls = [1, 2, 3, 4, 5]

        # eval_result = [[[ [] for _ in range(len(self.ground_truth))] for _ in recalls] for _ in tious]
        eval_result = [[ [] for _ in recalls] for _ in tious]

        # v_cnt = 0
        for key_v, value_v in self.ground_truth.items():
            gt_v = value_v
            pred_v = self.prediction[key_v]

            for key_label, value_c in gt_v.items():
                gt_v_c = value_c
                num_gt_v_c = len(gt_v_c)
                if key_label in pred_v.keys():
                    pred_v_c = pred_v[key_label]
                    overlap = iou(pred_v_c, gt_v_c)

                    for i, t in enumerate(tious):
                        for j, r in enumerate(recalls):

                            is_retrieved = [(overlap > t)[:r*num_gt_v_c][:,i].any() for i in range(num_gt_v_c)]
                            eval_result[i][j].extend(is_retrieved)
                else:
                    for i, t in enumerate(tious):
                        for j, r in enumerate(recalls):
                            eval_result[i][j].extend([False] * len(gt_v_c))


        eval_result = np.array(eval_result).mean(axis=-1)

        for i, t in enumerate(tious):
            for j, r in enumerate(recalls):
                recall = eval_result[i, j]
                print(f'Rank {r}x @ tIoU {t} is {recall}')

        return eval_result


def iou(pred, gt): # require pred and gt is numpy
    assert isinstance(pred, list) and isinstance(gt,list)
    pred_is_list = isinstance(pred[0],list)
    gt_is_list = isinstance(gt[0],list)
    if not pred_is_list: pred = [pred]
    if not gt_is_list: gt = [gt]
    pred, gt = np.array(pred), np.array(gt)
    inter_left = np.maximum(pred[:,0,None], gt[None,:,0])
    inter_right = np.minimum(pred[:,1,None], gt[None,:,1])
    inter = np.maximum(0.0, inter_right - inter_left)
    union_left = np.minimum(pred[:,0,None], gt[None,:,0])
    union_right = np.maximum(pred[:,1,None], gt[None,:,1])
    union = np.maximum(0.0, union_right - union_left)
    overlap = 1.0 * inter / union
    if not gt_is_list:
        overlap = overlap[:,0]
    if not pred_is_list:
        overlap = overlap[0]
    return overlap

def evaluation_retrieval(opt):

    ego4d_MR = Moment_Retrieval(ground_truth_filename = opt["clip_anno"],
                   prediction_filename = os.path.join(opt["output_path"], opt["retrieval_result_file"]),
                   subset=opt['infer_datasplit'], tiou_thresholds=opt['tIoU_thr'],
                   verbose=True, check_status=False)

    eval_result = ego4d_MR.evaluate()

    print(eval_result)

EgoMQ/Evaluation/ego4d/get_retrieval_performance.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import json
import urllib.request as request

import numpy as np

API = 'http://ec2-52-11-11-89.us-west-2.compute.amazonaws.com/challenge16/api.py'

def get_blocked_videos(api=API):
    api_url = '{}?action=get_blocked'.format(api)
    # req = request(api_url)
    response = request.urlopen(api_url)
    return json.loads(response.read())

def interpolated_prec_rec(prec, rec):
    """Interpolated AP - VOCdevkit from VOC 2011.
    """
    mprec = np.hstack([[0], prec, [0]])
    mrec = np.hstack([[0], rec, [1]])
    for i in range(len(mprec) - 1)[::-1]:
        mprec[i] = max(mprec[i], mprec[i + 1])
    idx = np.where(mrec[1::] != mrec[0:-1])[0] + 1
    ap = np.sum((mrec[idx] - mrec[idx - 1]) * mprec[idx])
    return ap

def segment_iou(target_segment, candidate_segments):
    """Compute the temporal intersection over union between a
    target segment and all the test segments.

    Parameters
    ----------
    target_segment : 1d array
        Temporal target segment containing [starting, ending] times.
    candidate_segments : 2d array
        Temporal candidate segments containing N x [starting, ending] times.

    Outputs
    -------
    tiou : 1d array
        Temporal intersection over union score of the N's candidate segments.
    """
    tt1 = np.maximum(target_segment[0], candidate_segments[:, 0])
    tt2 = np.minimum(target_segment[1], candidate_segments[:, 1])
    # Intersection including Non-negative overlap score.
    segments_intersection = (tt2 - tt1).clip(0)
    # Segment union.
    segments_union = (candidate_segments[:, 1] - candidate_segments[:, 0]) \
      + (target_segment[1] - target_segment[0]) - segments_intersection
    # Compute overlap as the ratio of the intersection
    # over union of two segments.
    tIoU = segments_intersection.astype(float) / segments_union
    return tIoU

def wrapper_segment_iou(target_segments, candidate_segments):
    """Compute intersection over union btw segments
    Parameters
    ----------
    target_segments : ndarray
        2-dim array in format [m x 2:=[init, end]]
    candidate_segments : ndarray
        2-dim array in format [n x 2:=[init, end]]
    Outputs
    -------
    tiou : ndarray
        2-dim array [n x m] with IOU ratio.
    Note: It assumes that candidate-segments are more scarce that target-segments
    """
    if candidate_segments.ndim != 2 or target_segments.ndim != 2:
        raise ValueError('Dimension of arguments is incorrect')

    n, m = candidate_segments.shape[0], target_segments.shape[0]
    tiou = np.empty((n, m))
    for i in xrange(m):
        tiou[:, i] = segment_iou(target_segments[i,:], candidate_segments)

    return tiou

    # TEST

EgoMQ/Evaluation/utils.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import sys
sys.dont_write_bytecode = True
import torch
import torch.nn.parallel
from Utils.dataset import VideoDataSet
from Models.VSGN import VSGN
import pandas as pd
from joblib import Parallel, delayed
import sys
sys.dont_write_bytecode = True
import torch.nn.parallel
import os
import Utils.opts as opts

import datetime
import numpy as np

torch.manual_seed(21)


# Infer all data
def Infer_SegTAD(opt):
    model = VSGN(opt)
    model = torch.nn.DataParallel(model).cuda()
    if not os.path.exists(opt["checkpoint_path"] + "/best.pth.tar"):
        print("There is no checkpoint. Please train first!!!")
    else:
        checkpoint = torch.load(opt["checkpoint_path"] + "/best.pth.tar")
        model.load_state_dict(checkpoint['state_dict'])
    model.eval()
    proposal_path = os.path.join(opt["output_path"], opt["prop_path"])

    if not os.path.exists(proposal_path):
        os.mkdir(proposal_path)

    test_loader = torch.utils.data.DataLoader(VideoDataSet(opt, subset=opt['infer_datasplit'], mode="inference"),
                                              batch_size=opt["batch_size"], shuffle=False,
                                              num_workers=0, pin_memory=True, drop_last=False)

    with torch.no_grad():
        for i, (index_list, input_data, num_frms) in enumerate(test_loader):

            infer_batch_selectprop(model, index_list, input_data, test_loader, proposal_path, num_frms)


# Infer one batch of data, for the model with proposal selection
def infer_batch_selectprop(model,
                           index_list,
                           input_data,
                           test_loader,
                           proposal_path,
                           num_frms):

    loc_dec, score_dec, loc_adjusted, pred_action, pred_start, pred_end = model(input_data.cuda(), num_frms)

    # Move variables to output to CPU
    loc_adjusted_batch = loc_adjusted.detach().cpu().numpy()
    score_dec_batch = score_dec.detach().cpu().numpy()
    pred_action_batch = pred_action.detach().cpu().numpy()
    pred_start_batch = pred_start.detach().cpu().numpy()
    pred_end_batch = pred_end.detach().cpu().numpy()
    num_frms_batch = num_frms.detach().cpu().numpy()


    Parallel(n_jobs=len(index_list))(
        delayed(infer_v_asis)(
            opt,
            video=test_loader.dataset.clip_dict[test_loader.dataset.clip_list[full_idx]],
            score_dec_v = score_dec_batch[batch_idx],
            loc_dec_v = loc_adjusted_batch[batch_idx],
            pred_action_v = pred_action_batch[batch_idx],
            pred_start_v = pred_start_batch[batch_idx],
            pred_end_v = pred_end_batch[batch_idx],
            proposal_path = proposal_path,
            num_frms_v = num_frms_batch[batch_idx]

        ) for batch_idx, full_idx in enumerate(index_list))



def infer_v_asis(*args, **kwargs):
    tscale = args[0]["temporal_scale"]
    loc_pred_v = kwargs['loc_dec_v']
    score_dec_v = kwargs['score_dec_v']
    pred_start_v = kwargs['pred_start_v']
    pred_end_v = kwargs['pred_end_v']
    proposal_path = kwargs['proposal_path']
    num_frms_v = kwargs['num_frms_v']
    thresh = 0.000000005

    clip_name = kwargs['video']['clip_id']
    fps = kwargs['video']['fps']

    loc_pred_v[:, 0] = loc_pred_v[:, 0].clip(min=0, max=tscale - 1)
    loc_pred_v[:, 1] = loc_pred_v[:, 1].clip(min=0, max=tscale - 1)

    start_score = (pred_start_v[np.ceil(loc_pred_v[:, 0]).astype('int32')] + pred_start_v[
        np.floor(loc_pred_v[:, 0]).astype('int32')]) / 2
    end_score = (pred_end_v[np.ceil(loc_pred_v[:, 1]).astype('int32')] + pred_end_v[
        np.floor(loc_pred_v[:, 1]).astype('int32')]) / 2

    score_stage2 = start_score * end_score

    loc_pred_v[:, 0] = loc_pred_v[:, 0].clip(min=0, max=num_frms_v - 1)
    loc_pred_v[:, 1] = loc_pred_v[:, 1].clip(min=0, max=num_frms_v - 1)

    new_props = []
    for j in range(1, opt['decoder_num_classes']):
        inds = (score_dec_v[:, j] > thresh)
        scores = (score_dec_v[:, j] * score_stage2)[inds]
        locations = loc_pred_v[inds, :]
        labels = np.array([j] * locations.shape[0])
        cls_dets = np.concatenate((locations, scores[:, None], labels[:, None]), axis=1)
        keep = nms(cls_dets, opt['nms_thr'])
        if (len(keep) > 0):
            cls_dets = cls_dets[keep]

        new_props.append(cls_dets)

    new_props = np.concatenate(new_props, axis=0)
    new_props[:, :2] = (new_props[:, :2]) / fps

    col_name = ["xmin", "xmax", "score", "label"]
    new_df = pd.DataFrame(new_props, columns=col_name)
    path = proposal_path + "/" + clip_name + ".csv"
    new_df.to_csv(path, index=False)


def nms(dets, thresh=0.4):
    """Pure Python NMS baseline."""
    if len(dets) == 0: return []
    x1 = dets[:, 0]
    x2 = dets[:, 1]
    scores = dets[:, 2]
    lengths = x2 - x1
    order = scores.argsort()[::-1]
    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)
        xx1 = np.maximum(x1[i], x1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        inter = np.maximum(0.0, xx2 - xx1)
        ovr = inter / (lengths[i] + lengths[order[1:]] - inter)
        inds = np.where(ovr <= thresh)[0]
        order = order[inds + 1]
    return keep


if __name__ == '__main__':

    opt = opts.parse_opt()
    opt = vars(opt)

    print(opt)

    if not os.path.exists(opt["output_path"]):
        os.makedirs(opt["output_path"])

    print(datetime.datetime.now())
    print("---------------------------------------------------------------------------------------------")
    print("1. Inference starts!")
    print("---------------------------------------------------------------------------------------------")

    Infer_SegTAD(opt)

    print("Inference finishes! \n")


EgoMQ/Infer.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import os
import json
import argparse

parser = argparse.ArgumentParser()
parser.add_argument(
    '--output_path',
    type=str)

args = parser.parse_args()

result_det, result_rev = None, None

if os.path.exists(f'{args.output_path}/detections_postNMS.json'):
    result_det = f'{args.output_path}/detections_postNMS.json'
else:
    raise ValueError('Detection results missing.')

if os.path.exists(f'{args.output_path}/retreival_postNMS.json'):
    result_rev = f'{args.output_path}/retreival_postNMS.json'
else:
    raise ValueError('Retrieval results missing.')


with open(result_det, 'r') as fobj:
    data_det = json.load(fobj)

with open(result_rev, 'r') as fobj:
    data_rev = json.load(fobj)

data_submission = {"version": "1.0","challenge": "ego4d_moment_queries"}
data_submission['detect_results'] = data_det['results']
data_submission['retrieve_results'] = data_rev['results']

with open(f"{args.output_path}/submission.json", "w") as fp:
    json.dump(data_submission,fp)

print('Submission succesfully packed. Good luck!')

EgoMQ/Merge_detection_retrieval.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import torch
from .BoxCoder import BoxCoder
import torch.nn.functional as F

class Pred_loc_cls(object):
    def __init__(self, opt):
        super(Pred_loc_cls, self).__init__()

        self.pre_nms_thresh = 0.00
        self.pre_nms_top_n = 10000
        self.num_classes = opt['decoder_num_classes']

        self.box_coder = BoxCoder(opt)

    def __call__(self, cls_pred_dec, reg_pred_dec, anchors):
        bs = cls_pred_dec[0].shape[0]
        anchors = [anchor.unsqueeze(0).repeat(bs, 1, 1).to(device=cls_pred_dec[0].device) for anchor in anchors]
        all_anchors = torch.cat(anchors, dim=1)          # bs, levels*positions*scales, left-right

        cls_pred_dec = [cls_pred_dec[i] for i in range(len(cls_pred_dec)-1, -1, -1)]
        reg_pred_dec = [reg_pred_dec[i] for i in range(len(reg_pred_dec)-1, -1, -1)]

        loc_dec, _ = self._call_one_stage(cls_pred_dec, reg_pred_dec, all_anchors)
        _, score_dec = self._call_one_stage(cls_pred_dec, reg_pred_dec, torch.stack(loc_dec, dim=0))


        return torch.stack(score_dec, dim=0), torch.stack(loc_dec, dim=0)

    def _call_one_stage(self, cls_pred, reg_pred, all_anchors):

        N = cls_pred[0].shape[0]

        cls_pred = F.softmax(torch.cat(cls_pred, dim=2).permute(0, 2, 1).reshape(N, -1, self.num_classes), dim=-1)  # bs, levels*positions*scales, num_cls
        reg_pred = torch.cat(reg_pred, dim=2).permute(0, 2, 1).reshape(N, -1, 2)   # bs, levels*positions*scales, 2


        candidate_inds = cls_pred > self.pre_nms_thresh
        pre_nms_top_n = candidate_inds.view(N, -1).sum(1)
        pre_nms_top_n = pre_nms_top_n.clamp(max=self.pre_nms_top_n)

        loc_res = []
        score_res = []
        for cls_seq, reg_seq, anchor_seq, pre_nms_top_n_seq, candidate_inds_seq in zip(cls_pred, reg_pred, all_anchors, pre_nms_top_n, candidate_inds):

            loc_pred = self.box_coder.decode(
                reg_seq.view(-1, 2),        # levels*positions*scales, 2
                anchor_seq.view(-1, 2)      # levels*positions*scales, 2
            )

            score_pred = cls_seq

            loc_res.append(loc_pred)
            score_res.append(score_pred)

        return loc_res, score_res


    def cat_boxlist(self, bboxes):
        """
        Concatenates a list of BoxList (having the same image size) into a
        single BoxList

        Arguments:
            bboxes (list[BoxList])
        """
        assert isinstance(bboxes, (list, tuple))
        assert all(isinstance(bbox, dict) for bbox in bboxes)

        res = {}
        res['loc'] = torch.ones([0, 2])
        for bb in bboxes:
            res['loc'] = torch.cat(res['loc'], bb['loc'], dim=0)

        return res

EgoMQ/Models/ActionGenerator.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import torch
import torch.nn as nn
import math

class AnchorGenerator(nn.Module):
    def __init__(self, opt):
        super(AnchorGenerator, self).__init__()
        self.num_levels = opt['num_levels']
        self.tscale = opt["temporal_scale"]
        self.scales = opt['anchor_scale']  #8  # 1
        self.base_stride = opt['base_stride']
        self.strides = []            # distance between anchors with respect to the sequence

        self.base_anchors = []
        for l in range(self.num_levels):
            stride = self.base_stride * pow(2, l)
            self.strides.append(stride)
            self.base_anchors.append( self.get_base_anchors(stride, self.scales))

        self.anchors = self.gen_anchors()


    def gen_anchors(self):
        feat_sizes = [ math.ceil(self.tscale / self.strides[l]) for l in range(self.num_levels)]
        anchors = []
        for size, stride, base_anchors in zip(feat_sizes, self.strides, self.base_anchors):
            shifts = torch.arange(0, size * stride, step=stride, dtype=torch.float32)[:, None].repeat(1, 2)
            anchors.append((shifts.view(-1, 1, 2) + base_anchors.view(1, -1, 2)).reshape(-1, 2))

        return anchors

    def get_base_anchors(self, stride, scales):
        anchors = torch.tensor([1, stride], dtype=torch.float) - 0.5
        anchors = self._scale_enum(anchors, scales)
        return anchors

    def _scale_enum(self, anchor, scales):
        """Enumerate a set of anchors for each scale wrt an anchor."""
        length, center = self._whctrs(anchor)
        ws = length * torch.tensor(scales)
        anchors = self._mkanchors(ws, center)
        return anchors

    def _mkanchors(self, ws, ctr):
        anchors = torch.stack(
            (
                ctr.unsqueeze(0)- 0.5 * (ws.to(dtype=torch.float32) - 1),
                ctr.unsqueeze(0)+ 0.5 * (ws.to(dtype=torch.float32) - 1)
            )
        ).transpose(0,1)

        return anchors

    def _whctrs(self, anchor):

        length = anchor[1] - anchor[0] + 1
        center = anchor[0] + 0.5 * (length - 1)
        return length, center


EgoMQ/Models/AnchorGenerator.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import torch
import torch.nn as nn
from .Loss import Loss_loc_cls
import pdb

class BoundaryAdjust(nn.Module):

    def __init__(self, opt):
        super(BoundaryAdjust, self).__init__()
        self.tscale = opt['temporal_scale']
        self.base_stride = opt['base_stride']
        bb_hidden_dim = opt['bb_hidden_dim']

        self.Loss = Loss_loc_cls(opt)

        self.start_conv = nn.Sequential(
            nn.Conv1d(in_channels=bb_hidden_dim, out_channels=bb_hidden_dim,kernel_size=3,stride=2,padding=0,groups=1),
            nn.ReLU(inplace=True),
            nn.Conv1d(in_channels=bb_hidden_dim, out_channels=1, kernel_size=1),
        )

        self.end_conv = nn.Sequential(
            nn.Conv1d(in_channels=bb_hidden_dim, out_channels=bb_hidden_dim,kernel_size=3,stride=2,padding=0,groups=1),
            nn.ReLU(inplace=True),
            nn.Conv1d(in_channels=bb_hidden_dim, out_channels=1, kernel_size=1),
        )

    def forward(self, loc_box, feat_frmlvl):
        beta = 8.
        bs, C, _ = feat_frmlvl.shape

        if isinstance(loc_box, list):
            loc_box = torch.cat(loc_box, dim=1)
        loc_box = torch.clamp(loc_box, min=0., max=(self.tscale)-1)

        boundary_length = (loc_box[:,:,1] - loc_box[:,:,0] + 1) / beta

        # Starts
        starts = torch.clamp((loc_box[:,:,0] / self.base_stride).to(dtype=torch.long), min=0, max=(self.tscale / self.base_stride)-1).long()
        starts_left = torch.clamp(((loc_box[:,:,0] - boundary_length) / self.base_stride).to(dtype=torch.long), min=0, max=(self.tscale / self.base_stride)-1).long()
        starts_right = torch.clamp(((loc_box[:,:,0] + boundary_length) / self.base_stride).to(dtype=torch.long), min=0, max=(self.tscale / self.base_stride)-1).long()

        #pdb.set_trace()

        start_feats_center = torch.stack([feat_frmlvl[i, :, starts[i, :]] for i in range(bs)]).permute(0, 2, 1).reshape(-1, C)
        start_feats_left = torch.stack([feat_frmlvl[i, :, starts_left[i, :]] for i in range(bs)]).permute(0, 2, 1).reshape(-1, C)
        start_feats_right = torch.stack([feat_frmlvl[i, :, starts_right[i, :]] for i in range(bs)]).permute(0, 2, 1).reshape(-1, C)

        start_feats = torch.stack((start_feats_left, start_feats_center, start_feats_right), dim=-1)

        start_offsets = self.start_conv(start_feats).squeeze().view(bs,-1)

        # Ends
        ends = torch.clamp((loc_box[:,:,1] / self.base_stride).to(dtype=torch.long), min=0, max=(self.tscale / self.base_stride)-1).long()
        ends_left = torch.clamp(((loc_box[:,:,1] - boundary_length) / self.base_stride).to(dtype=torch.long), min=0, max=(self.tscale / self.base_stride)-1).long()
        ends_right = torch.clamp(((loc_box[:,:,1] + boundary_length) / self.base_stride).to(dtype=torch.long), min=0, max=(self.tscale / self.base_stride)-1).long()

        #pdb.set_trace()

        end_feats_center = torch.stack([feat_frmlvl[i, :, ends[i, :]] for i in range(bs)]).permute(0, 2, 1).reshape(-1, C)
        end_feats_left = torch.stack([feat_frmlvl[i, :, ends_left[i, :]] for i in range(bs)]).permute(0, 2, 1).reshape(-1, C)
        end_feats_right = torch.stack([feat_frmlvl[i, :, ends_right[i, :]] for i in range(bs)]).permute(0, 2, 1).reshape(-1, C)

        end_feats = torch.stack((end_feats_left, end_feats_center, end_feats_right), dim=-1)

        end_offsets = self.end_conv(end_feats).squeeze().view(bs,-1)

        return start_offsets, end_offsets

    def cal_loss(self, start_offsets, end_offsets, anchors, gt_bbox, num_gt):

        box_pred = self.update_bd(anchors, start_offsets, end_offsets).view(-1, 2)

        cls_labels, reg_targets = self.Loss.prepare_targets(gt_bbox, num_gt, anchors, stage=2)

        cls_labels = torch.cat(cls_labels, dim=0)  # bs*levels*positions*scales, num_cls
        reg_targets = torch.cat(reg_targets, dim=0)  # bs*levels*positions*scales, left-right

        all_anchors = torch.cat(anchors, dim=1).view(-1, 2) # bs*levels*positions*scales, 2

        pos_inds = torch.nonzero(cls_labels > 0).squeeze(1)
        reg_loss = self.Loss.reg_loss_func(None, reg_targets[pos_inds], all_anchors[pos_inds], box_pred[pos_inds]) / pos_inds.numel()

        return reg_loss

    def update_bd(self, anchors, start_offsets, end_offsets):

        if isinstance(anchors, list):
            anchors = torch.cat(anchors, dim=1)
        loc_st2 = anchors.clone()
        loc_st2[:,:,0] = anchors[:,:,0] + start_offsets
        loc_st2[:,:,1] = anchors[:,:,1] + end_offsets

        return loc_st2

EgoMQ/Models/BoundaryAdjust.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import torch
import math

class BoxCoder(object):

    def __init__(self, opt):
        self.cfg = opt

    def encode(self, gt_boxes, anchors):
        if False: #self.cfg.MODEL.ATSS.REGRESSION_TYPE == 'POINT':
            TO_REMOVE = 1  # TODO remove
            anchors_w = anchors[:, 2] - anchors[:, 0] + TO_REMOVE
            anchors_h = anchors[:, 3] - anchors[:, 1] + TO_REMOVE
            anchors_cx = (anchors[:, 2] + anchors[:, 0]) / 2
            anchors_cy = (anchors[:, 3] + anchors[:, 1]) / 2

            w = self.cfg.MODEL.ATSS.ANCHOR_SIZES[0] / self.cfg.MODEL.ATSS.ANCHOR_STRIDES[0]
            l = w * (anchors_cx - gt_boxes[:, 0]) / anchors_w
            t = w * (anchors_cy - gt_boxes[:, 1]) / anchors_h
            r = w * (gt_boxes[:, 2] - anchors_cx) / anchors_w
            b = w * (gt_boxes[:, 3] - anchors_cy) / anchors_h
            targets = torch.stack([l, t, r, b], dim=1)
        elif True: #self.cfg.MODEL.ATSS.REGRESSION_TYPE == 'BOX':
            TO_REMOVE = 1  # TODO remove
            ex_length = anchors[:, 1] - anchors[:, 0] + TO_REMOVE
            ex_center = (anchors[:, 1] + anchors[:, 0]) / 2

            gt_length = gt_boxes[:, 1] - gt_boxes[:, 0] + TO_REMOVE
            gt_center = (gt_boxes[:, 1] + gt_boxes[:, 0]) / 2

            wx, ww = (10., 5.)
            targets_dx = wx * (gt_center - ex_center) / ex_length
            targets_dw = ww * torch.log(gt_length / ex_length)
            targets = torch.stack((targets_dx, targets_dw), dim=1)

        return targets

    def decode(self, preds, anchors):
        if False: #self.cfg.MODEL.ATSS.REGRESSION_TYPE == 'POINT':
            TO_REMOVE = 1  # TODO remove
            anchors_w = anchors[:, 2] - anchors[:, 0] + TO_REMOVE
            anchors_h = anchors[:, 3] - anchors[:, 1] + TO_REMOVE
            anchors_cx = (anchors[:, 2] + anchors[:, 0]) / 2
            anchors_cy = (anchors[:, 3] + anchors[:, 1]) / 2

            w = self.cfg.MODEL.ATSS.ANCHOR_SIZES[0] / self.cfg.MODEL.ATSS.ANCHOR_STRIDES[0]
            x1 = anchors_cx - preds[:, 0] / w * anchors_w
            y1 = anchors_cy - preds[:, 1] / w * anchors_h
            x2 = anchors_cx + preds[:, 2] / w * anchors_w
            y2 = anchors_cy + preds[:, 3] / w * anchors_h
            pred_boxes = torch.stack([x1, y1, x2, y2], dim=1)
        elif True: #self.cfg.MODEL.ATSS.REGRESSION_TYPE == 'BOX':
            anchors = anchors.to(preds.dtype)

            TO_REMOVE = 1  # TODO remove
            ex_length = anchors[:, 1] - anchors[:, 0] + TO_REMOVE
            ex_center = (anchors[:, 1] + anchors[:, 0]) / 2

            wx, ww = (10, 5.)
            dx = preds[:, 0] / wx
            dw = preds[:, 1] / ww

            # Prevent sending too large values into torch.exp()
            dw = torch.clamp(dw, max=math.log(1000. / 16))

            pred_ctr_x = (dx * ex_length + ex_center)
            pred_w = (torch.exp(dw) * ex_length)

            pred_boxes = torch.zeros_like(preds)
            pred_boxes[:, 0] = pred_ctr_x - 0.5 * (pred_w - 1)
            pred_boxes[:, 1] = pred_ctr_x + 0.5 * (pred_w - 1)
        return pred_boxes

EgoMQ/Models/BoxCoder.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import torch
import torch.nn as nn


# dynamic graph from knn
def knn(x, num_frms, opt, y=None, k=10):
    bs, _, length = x.shape
    if y is None:
        y = x

    # Original neighbors
    dif = torch.sum((x.unsqueeze(2) - y.unsqueeze(3))** 2, dim=1)
    idx_org = dif.topk(k=k, dim=-1, largest=False)[1]

    if not opt['use_VSS']:
        return idx_org
    else:
        idx_new = idx_org.clone()
        max_dif = torch.max(dif)
        ratio = opt['temporal_scale'] / length
        half1_k = int(k / 2)
        half2_k = k - half1_k
        for i in range(bs):
            if num_frms[i] <= (opt['short_ratio'] * opt['temporal_scale']):
                thr = int((num_frms[i] + opt['stitch_gap']) / ratio)
                dif[i, thr:, thr:] = max_dif + 1

                loc1 = torch.tensor(range(length), dtype=torch.long, device=x.device)[:, None].repeat(1, half1_k).view(-1)
                loc2 = idx_org[i, :, :half1_k].reshape(-1)
                dif[i, loc1, loc2] = max_dif + 1

                idx_new[i, :, half1_k:] = dif[i].topk(k=half2_k, dim=-1, largest=False)[1]

    return idx_new


def get_neigh_idx_semantic(x, n_neigh, num_frms, opt):

    B, _, num_prop_v = x.shape
    neigh_idx = knn(x, num_frms, opt, k=n_neigh).to(dtype=torch.float32)
    shift = (torch.tensor(range(B), dtype=torch.float32, device=x.device) * num_prop_v)[:, None, None].repeat(1, num_prop_v, n_neigh)
    neigh_idx = (neigh_idx + shift).view(-1)
    return neigh_idx


class NeighConv(nn.Module):
    def __init__(self, in_features, out_features, opt):
        super(NeighConv, self).__init__()
        self.num_neigh = opt['num_neigh']
        self.nfeat_mode = opt['nfeat_mode']
        self.agg_type = opt['agg_type']
        self.edge_weight = opt['edge_weight']

        self.mlp = nn.Linear(in_features*2, out_features)

    def forward(self, feat_prop, neigh_idx):

        feat_neigh = feat_prop[neigh_idx.to(torch.long)]
        f_neigh_temp = feat_neigh.view(-1, self.num_neigh, feat_neigh.shape[-1])

        if self.nfeat_mode == 'feat_ctr':
            feat_neigh = torch.cat((feat_neigh.view(-1, self.num_neigh, feat_prop.size(-1)), feat_prop.view(-1, 1, feat_prop.size(-1)).repeat(1, self.num_neigh,1)), dim=-1)
        elif self.nfeat_mode == 'dif_ctr':
            feat_prop = feat_prop.view(-1, 1, feat_prop.size(-1)).repeat(1, self.num_neigh,1)
            diff = feat_neigh.view(-1, self.num_neigh, feat_prop.size(-1)) - feat_prop
            feat_neigh = torch.cat((diff, feat_prop), dim=-1)
        elif self.nfeat_mode == 'feat':
            feat_neigh = feat_neigh.view(-1, self.num_neigh, feat_prop.size(-1))

        feat_neigh_out = self.mlp(feat_neigh)
        if self.edge_weight == 'true':

            weight = torch.matmul(f_neigh_temp, feat_prop.unsqueeze(2))
            weight_denom1 = torch.sqrt(torch.sum(f_neigh_temp*f_neigh_temp, dim=2, keepdim=True))
            weight_denom2 = torch.sqrt(torch.sum(feat_prop.unsqueeze(2)*feat_prop.unsqueeze(2), dim=1, keepdim=True))
            weight = (weight / torch.matmul(weight_denom1, weight_denom2)).squeeze(2)
            feat_neigh_out = feat_neigh_out * weight.unsqueeze(2)

        if self.agg_type == 'max':
            feat_neigh_out = feat_neigh_out.max(dim=1, keepdim=False)[0]
        elif self.agg_type == 'mean':
            feat_neigh_out = feat_neigh_out.mean(dim=1, keepdim=False)
        return feat_neigh_out

class xGN(nn.Module):
    def __init__(self,  opt, in_channels, out_channels, stride = 2, bias=True):
        super(xGN, self).__init__()
        self.gcn_insert = opt['gcn_insert']
        self.n_neigh = opt['num_neigh']
        self.opt = opt
        self.tconv1 = nn.Conv1d(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1, groups=1)
        self.nconv1 = NeighConv(in_channels, out_channels, opt)
        self.stride = stride

        self.relu = nn.ReLU(inplace=True)
        if self.stride == 2:
            self.maxpool = nn.MaxPool1d(kernel_size=2, stride=2)

    def forward(self, x, num_frms):
        bs, C, num_frm = x.shape

        # CNN
        c_out = self.tconv1(x)

        # GCN
        if self.gcn_insert == 'par':
            neigh_idx = get_neigh_idx_semantic(x, self.n_neigh, num_frms, self.opt)
            g_out = self.nconv1(x.permute(0, 2, 1).reshape(-1, C), neigh_idx)
            g_out = g_out.view(bs, num_frm, -1).permute(0, 2, 1)
            out = c_out + g_out

        elif self.gcn_insert == 'seq':
            neigh_idx = get_neigh_idx_semantic(c_out, self.n_neigh, num_frms, self.opt)
            g_out = self.nconv1(c_out.permute(0, 2, 1).reshape(-1, C), neigh_idx)
            g_out = g_out.view(bs, num_frm, -1).permute(0, 2, 1)
            out = g_out

        out = self.relu(out)

        if self.stride == 2:
            out = self.maxpool(out)

        return out

EgoMQ/Models/GCNs.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import torch.nn as nn
import torch.nn.functional as F


class Head(nn.Module):
    def __init__(self, opt):
        super(Head, self).__init__()

        self.num_convs_head = opt['num_head_layers']  # 1
        bb_hidden_dim = opt['bb_hidden_dim']
        num_anchors = len(opt['anchor_scale'])
        num_classes = 1 if opt['dataset'] == 'activitynet' else opt['decoder_num_classes']

        cls_tower = []
        bbox_tower = []
        conv_func = nn.Conv1d
        for i in range(self.num_convs_head):

            cls_tower.append(
                conv_func(
                    bb_hidden_dim,
                    bb_hidden_dim,
                    kernel_size=3,
                    stride=1,
                    padding=1,
                    bias=True
                )
            )
            cls_tower.append(nn.GroupNorm(32, bb_hidden_dim))
            cls_tower.append(nn.ReLU())

            bbox_tower.append(
                conv_func(
                    bb_hidden_dim,
                    bb_hidden_dim,
                    kernel_size=3,
                    stride=1,
                    padding=1,
                    bias=True
                )
            )
            bbox_tower.append(nn.GroupNorm(32, bb_hidden_dim))
            bbox_tower.append(nn.ReLU())

        self.add_module('cls_tower', nn.Sequential(*cls_tower))
        self.add_module('bbox_tower', nn.Sequential(*bbox_tower))
        self.cls_logits = nn.Conv1d(
            bb_hidden_dim, num_anchors * num_classes, kernel_size=3, stride=1,
            padding=1
        )
        self.bbox_pred = nn.Conv1d(
            bb_hidden_dim, num_anchors * 2, kernel_size=3, stride=1,
            padding=1
        )

    def forward(self, x):
        logits = []
        bbox_reg = []
        for l, feature in enumerate(x):
            cls_tower = self.cls_tower(feature)
            box_tower = self.bbox_tower(feature)

            logits.append(self.cls_logits(cls_tower))

            bbox_pred = self.bbox_pred(box_tower)
            # if self.cfg.MODEL.ATSS.REGRESSION_TYPE == 'POINT':
            #     bbox_pred = F.relu(bbox_pred)
            bbox_reg.append(bbox_pred)

        return logits, bbox_reg

EgoMQ/Models/Head.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import pdb
import torch
import torch.nn.functional as F
from .matcher import Matcher

from .BoxCoder import  BoxCoder


class Loss_loc_cls(object):

    def __init__(self, opt):

        self.tscale = opt["temporal_scale"]
        self.gamma = 2.0
        self.alpha = opt['focal_alpha']
        self.num_classes = opt['decoder_num_classes']
        self.iou_thresholds = opt['iou_thr']
        self.matcher = Matcher(True)

        self.box_coder = BoxCoder(opt)


    def _iou_anchors_gts(self, anchor, gt):

        anchors_min = anchor[:, 0]
        anchors_max = anchor[:, 1]
        box_min = gt[:, 0]
        box_max = gt[:, 1]
        len_anchors = anchors_max - anchors_min + 1
        int_xmin = torch.max(anchors_min[:, None], box_min)
        int_xmax = torch.min(anchors_max[:, None], box_max)
        inter_len = torch.clamp(int_xmax - int_xmin, min=0)
        union_len = torch.clamp(len_anchors[:, None] + box_max - box_min - inter_len, min=0)
        jaccard = inter_len / union_len
        return jaccard

    def __call__(self, cls_pred_dec, reg_pred_dec, anchors, gt_bbox, num_gt):
        bs = cls_pred_dec[0].shape[0]
        anchors = [anchor.unsqueeze(0).repeat(bs, 1, 1).to(device=gt_bbox.device) for anchor in anchors]
        cls_pred_dec = [cls_pred_dec[i] for i in range(len(cls_pred_dec)-1, -1, -1)]
        reg_pred_dec = [reg_pred_dec[i] for i in range(len(reg_pred_dec)-1, -1, -1)]

        loc_dec = []
        for pred, anchor in zip(reg_pred_dec, anchors):
            pred = pred.permute(0, 2, 1).reshape(-1, 2)
            anchor = anchor.view(-1, 2)
            loc_dec.append(self.box_coder.decode(pred, anchor).view(bs, -1,2))

        _, reg_loss1 = self._loss_one_stage(cls_pred_dec, reg_pred_dec, gt_bbox, num_gt, anchors, stage=1)
        cls_loss1, _ = self._loss_one_stage(cls_pred_dec, reg_pred_dec, gt_bbox, num_gt, loc_dec, stage=1)

        losses = {
            "loss_cls_dec": cls_loss1,
            "loss_reg_dec": reg_loss1,
        }
        return losses, loc_dec


    def _loss_one_stage(self, cls_pred, reg_pred, gt_bbox, num_gt, anchors, stage=0):

        num_cls = self.num_classes

        cls_labels, reg_targets = self.prepare_targets(gt_bbox, num_gt, anchors, stage=stage) # bs, levels*positions*scales, num_cls/left-right

        cls_pred = torch.cat(cls_pred, dim=2).permute(0, 2, 1).reshape(-1, num_cls) # bs, levels*positions, scales*cls --> bs*levels*positions*scales, num_cls
        reg_pred = torch.cat(reg_pred, dim=2).permute(0, 2, 1).reshape(-1, 2) # bs, levels*positions, scales*left-right --> bs*levels*positions*scales, left-right

        cls_labels = torch.cat(cls_labels, dim=0)  # bs*levels*positions*scales, num_cls
        reg_targets = torch.cat(reg_targets, dim=0)  # bs*levels*positions*scales, left-right

        all_anchors = torch.cat(anchors, dim=1).view(-1, 2) # bs*levels*positions*scales, 2

        pos_inds = torch.nonzero(cls_labels > 0).squeeze(1)
        cls_loss = self.cls_loss_func(cls_pred, cls_labels)
        reg_loss = self.reg_loss_func(reg_pred[pos_inds], reg_targets[pos_inds], all_anchors[pos_inds]) / pos_inds.numel()

        return  cls_loss, reg_loss

    def cls_loss_func(self, cls_pred, cls_labels):

        pmask = (cls_labels>0).float()
        nmask = (cls_labels==0).float()
        num_pos = torch.sum(pmask)
        num_neg = torch.sum(nmask)

        CE_loss = torch.nn.CrossEntropyLoss(reduction='none')
        
        #pdb.set_trace()

        loss = CE_loss(cls_pred, cls_labels.to(torch.long))

        pos_loss = torch.sum(loss * pmask) / num_pos
        neg_loss = torch.sum(loss * nmask) / num_neg

        total_loss = pos_loss + neg_loss

        return total_loss

    def reg_loss_func(self, pred, target, anchor, pred_boxes=None, weight=None):
        if type(pred_boxes) == type(None):
            pred_boxes = self.box_coder.decode(pred, anchor)
        pred_x1 = torch.min(pred_boxes[:, 0], pred_boxes[:, 1])
        pred_x2 = torch.max(pred_boxes[:, 0], pred_boxes[:, 1])
        pred_area = (pred_x2 - pred_x1)

        gt_boxes = self.box_coder.decode(target, anchor)
        target_x1 = gt_boxes[:, 0]
        target_x2 = gt_boxes[:, 1]
        target_area = (target_x2 - target_x1)

        x1_intersect = torch.max(pred_x1, target_x1)
        x2_intersect = torch.min(pred_x2, target_x2)
        #pdb.set_trace()
        area_intersect = torch.zeros(pred_x1.size()).to(target)
        mask = (x2_intersect > x1_intersect)
        area_intersect[mask] = (x2_intersect[mask] - x1_intersect[mask])

        x1_enclosing = torch.min(pred_x1, target_x1)
        x2_enclosing = torch.max(pred_x2, target_x2)
        area_enclosing = (x2_enclosing - x1_enclosing)  + 1e-7

        area_union = pred_area + target_area - area_intersect + 1e-7
        ious = area_intersect / area_union
        gious = ious - (area_enclosing - area_union) / area_enclosing

        losses = 1 - gious

        if weight is not None and weight.sum() > 0:
            return (losses * weight).sum()
        else:
            assert losses.numel() != 0
            return losses.sum()


    def prepare_targets(self, gt_bbox, num_gt, anchors, stage=0):

        cls_targets = []
        reg_targets = []
        all_anchors = torch.cat(anchors, dim=1) # bs, levels, positions, scales, left-right

        for i in range(len(gt_bbox)):

            gt_cur_im = gt_bbox[i, :num_gt[i], :-1] * self.tscale
            gt_label = gt_bbox[i, :num_gt[i], -1]
            anchor_cur_im = all_anchors[i]
            iou_matrix = self._iou_anchors_gts(anchor_cur_im, gt_cur_im)

            # Find the corresponding gt for each pred
            matched_idxs = self.matcher(iou_matrix.transpose(0, 1), self.iou_thresholds[stage])

            # Use the label of the corresponding gt as the classification target for the pred
            cls_labels_cur_im = torch.zeros_like(matched_idxs)

            cls_labels_cur_im[:] = gt_label[matched_idxs]

            cls_labels_cur_im[matched_idxs <0] = 0

            # Record the boundary offset as the regression target
            matched_gts = gt_cur_im[matched_idxs.clamp(min=0)]
            reg_targets_cur_im = self.box_coder.encode(matched_gts, anchor_cur_im)

            cls_targets.append(cls_labels_cur_im.to(dtype=torch.int32))
            reg_targets.append(reg_targets_cur_im)

        return cls_targets, reg_targets


def bi_loss(prediction, groundtruth, reduction='mean'):
    gt = groundtruth.view(-1)
    pred = prediction.contiguous().view(-1)

    pmask = (gt>0.5).float()
    num_positive = torch.sum(pmask)
    num_entries = len(gt)
    ratio=num_entries/num_positive

    coef_0=0.5*(ratio)/(ratio-1)
    coef_1=coef_0*(ratio-1)
    loss = coef_1*pmask*torch.log(pred+0.00001) + coef_0*(1.0-pmask)*torch.log(1.0-pred+0.00001)

    if reduction == 'mean':
        loss = -torch.mean(loss)
    elif reduction == 'none':
        loss = -torch.mean(loss.view(groundtruth.shape), dim=1)
    return loss



def get_loss_supplement(pred_action, gt_action, pred_start, gt_start, pred_end, gt_end):
    loss_action = bi_loss(pred_action, gt_action)
    loss_start = bi_loss(pred_start, gt_start)
    loss_end = bi_loss(pred_end, gt_end)

    return loss_action, loss_start, loss_end

EgoMQ/Models/Loss.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import torch.nn as nn
from .XGPN import XGPN
import torch.nn.functional as F
from .Head import Head
from .AnchorGenerator import AnchorGenerator
from .Loss import Loss_loc_cls, get_loss_supplement
from .ActionGenerator import Pred_loc_cls
from .BoundaryAdjust import BoundaryAdjust

class VSGN(nn.Module):
    def __init__(self, opt):
        super(VSGN, self).__init__()

        self.bb_hidden_dim = opt['bb_hidden_dim']
        self.bs = opt["batch_size"]
        self.is_train = opt['is_train']
        self.tem_best_loss = 10000000
        self.hidden_dim_2d = 128
        self.hidden_dim_3d = 512
        self.input_feat_dim = opt['input_feat_dim']

        self.xGPN = XGPN(opt)

        # self.head_enc = Head(opt)
        self.head_dec = Head(opt)

        self.anchors = AnchorGenerator(opt).anchors
        self.loss_loc_and_cls = Loss_loc_cls(opt)
        self.bd_adjust = BoundaryAdjust(opt)

        self.pred_loc_and_cls = Pred_loc_cls(opt)


        # Generate action/start/end scores
        self.head_actionness = nn.Sequential(
            nn.Conv1d(self.bb_hidden_dim, self.bb_hidden_dim, kernel_size=3, padding=1, groups=1),
            # nn.GroupNorm(32, self.bb_hidden_dim),
            nn.ReLU(inplace=True),
            nn.Conv1d(self.bb_hidden_dim, 1, kernel_size=1),
            nn.Sigmoid()
        )
        self.head_startness = nn.Sequential(
            nn.Conv1d(self.bb_hidden_dim, self.bb_hidden_dim, kernel_size=3, padding=1, groups=1),
            # nn.GroupNorm(32, self.bb_hidden_dim),
            nn.ReLU(inplace=True),
            nn.Conv1d(self.bb_hidden_dim, 1, kernel_size=1),
            nn.Sigmoid()
        )
        self.head_endness = nn.Sequential(
            nn.Conv1d(self.bb_hidden_dim, self.bb_hidden_dim, kernel_size=3, padding=1, groups=1),
            # nn.GroupNorm(32, self.bb_hidden_dim),
            nn.ReLU(inplace=True),
            nn.Conv1d(self.bb_hidden_dim, 1, kernel_size=1),
            nn.Sigmoid()
        )

    def _forward_test(self, cls_pred_enc, reg_pred_enc, cls_pred_dec, reg_pred_dec):

        loc_enc, score_enc, loc_dec, score_dec = self.gen_predictions(cls_pred_enc, reg_pred_enc, cls_pred_dec, reg_pred_dec, self.anchors)

        return loc_enc, score_enc, loc_dec, score_dec



    def forward(self, input, num_frms, gt_action= None, gt_start= None, gt_end= None, gt_bbox = None, num_gt = None):
        # Cross-scale graph pyramid network
        feats_enc, feats_dec = self.xGPN(input, num_frms)

        # Scoring and localization
        cls_pred_dec, reg_pred_dec = self.head_dec(feats_dec)
        if self.is_train == 'true':
            losses, loc_dec = self.loss_loc_and_cls(cls_pred_dec, reg_pred_dec, self.anchors, gt_bbox, num_gt)
        else:
            score_dec, loc_dec = self.pred_loc_and_cls(cls_pred_dec, reg_pred_dec, self.anchors)

        # Supplementary scores
        actionness = self.head_actionness(feats_dec[-1])
        startness = self.head_startness(feats_dec[-1])
        endness = self.head_endness(feats_dec[-1])

        actionness = F.interpolate(actionness, size=input.size()[2:], mode='linear', align_corners=True).squeeze(1)
        startness = F.interpolate(startness, size=input.size()[2:], mode='linear', align_corners=True).squeeze(1)
        endness = F.interpolate(endness, size=input.size()[2:], mode='linear', align_corners=True).squeeze(1)

        if self.is_train == 'true':
            loss_action, loss_start, loss_end = get_loss_supplement(actionness, gt_action, startness, gt_start, endness, gt_end)
            losses['loss_action'] = loss_action
            losses['loss_start'] = loss_start
            losses['loss_end'] = loss_end

        # Boundary adjustment
        start_offsets, end_offsets = self.bd_adjust(loc_dec, feats_dec[-1])

        if self.is_train == 'true':
            loss_reg_st2 = self.bd_adjust.cal_loss(start_offsets, end_offsets, loc_dec, gt_bbox, num_gt)
            losses['loss_bd_adjust'] = loss_reg_st2

            return losses, actionness, startness, endness
        else:
            loc_adjusted = self.bd_adjust.update_bd(loc_dec, start_offsets, end_offsets)
            return loc_dec, score_dec, loc_adjusted, actionness, startness, endness







EgoMQ/Models/VSGN.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import torch.nn as nn
from .GCNs import xGN


class XGPN(nn.Module):
    def __init__(self, opt):
        super(XGPN, self).__init__()
        self.input_feat_dim = opt["input_feat_dim"]
        self.bb_hidden_dim = opt['bb_hidden_dim']
        self.batch_size = opt["batch_size"]
        self.tem_best_loss = 10000000
        self.num_levels = opt['num_levels']
        self.use_xGPN = opt['use_xGPN']

        self.conv0 = nn.Sequential(
            nn.Conv1d(in_channels=self.input_feat_dim, out_channels=self.bb_hidden_dim,kernel_size=3,stride=1,padding=1,groups=1),
            nn.ReLU(inplace=True),
        )

        self.levels_enc = nn.ModuleList()
        for i in range(self.num_levels):
            if i == 0:
                stride = 1
            else:
                stride = 2
            self.levels_enc.append(self._make_levels_enc(opt, in_channels=self.bb_hidden_dim, out_channels=self.bb_hidden_dim, stride = stride))

        self.levels_dec = nn.ModuleList()
        for i in range(self.num_levels - 1):
            output_padding = 1
            self.levels_dec.append(self._make_levels_dec(in_channels=self.bb_hidden_dim, out_channels=self.bb_hidden_dim, output_padding = output_padding))

        self.levels1 = nn.ModuleList()
        for i in range(self.num_levels):
            self.levels1.append(self._make_levels(in_channels=self.bb_hidden_dim, out_channels=self.bb_hidden_dim))

        self.levels2 = nn.ModuleList()
        for i in range(self.num_levels - 1):
            self.levels2.append(self._make_levels(in_channels=self.bb_hidden_dim, out_channels=self.bb_hidden_dim))


    def _make_levels_enc(self, opt, in_channels, out_channels, stride = 2):
        if self.use_xGPN:
            return xGN(opt, in_channels=in_channels, out_channels=out_channels, stride = stride)
        else:
            return nn.Sequential(
                nn.Conv1d(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=stride, padding=1,
                          groups=1),
                nn.ReLU(inplace=True)
            )

    def _make_levels_dec(self, in_channels, out_channels, output_padding = 1):

        return nn.Sequential(
            nn.ConvTranspose1d(in_channels=in_channels, out_channels=out_channels,kernel_size=3,stride=2,padding=1, output_padding=output_padding, groups=1),
            nn.ReLU(inplace=True),
        )

    def _make_levels(self, in_channels, out_channels):

        return nn.Sequential(
            nn.Conv1d(in_channels=in_channels, out_channels=out_channels,kernel_size=3,stride=1,padding=1,groups=1),
            nn.ReLU(inplace=True),
        )

    def _encoder(self, input, num_frms):

        feats = []
        x = self.conv0(input)
        for i in range(0, self.num_levels):
            if self.use_xGPN:
                x = self.levels_enc[i](x, num_frms)
            else:
                x = self.levels_enc[i](x)
            feats.append(x)

        return feats

    def _decoder(self, input):

        feats = []
        x = self.levels1[0](input[self.num_levels - 1])
        feats.append(x)

        for i in range(self.num_levels - 1):
            ii = self.num_levels - i - 2
            feat_enc = self.levels2[i](input[ii])
            feat_dec = self.levels_dec[i](x)
            x = self.levels1[i+1](feat_enc + feat_dec)
            feats.append(x)



        return feats

    def forward(self, input, num_frms):

        feats_enc = self._encoder(input, num_frms)
        feats_dec = self._decoder(feats_enc)

        return feats_enc, feats_dec

EgoMQ/Models/XGPN.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import torch

class Matcher(object):
    """
    This class assigns to each predicted "element" (e.g., a box) a ground-truth
    element. Each predicted element will have exactly zero or one matches; each
    ground-truth element may be assigned to zero or more predicted elements.

    Matching is based on the MxN match_quality_matrix, that characterizes how well
    each (ground-truth, predicted)-pair match. For example, if the elements are
    boxes, the matrix may contain box IoU overlap values.

    The matcher returns a tensor of size N containing the index of the ground-truth
    element m that matches to prediction n. If there is no match, a negative value
    is returned.
    """

    BELOW_LOW_THRESHOLD = -1
    BETWEEN_THRESHOLDS = -2

    def __init__(self, allow_low_quality_matches=False):
        """
        Args:
            high_threshold (float): quality values greater than or equal to
                this value are candidate matches.
            low_threshold (float): a lower quality threshold used to stratify
                matches into three levels:
                1) matches >= high_threshold
                2) BETWEEN_THRESHOLDS matches in [low_threshold, high_threshold)
                3) BELOW_LOW_THRESHOLD matches in [0, low_threshold)
            allow_low_quality_matches (bool): if True, produce additional matches
                for predictions that have only low-quality match candidates. See
                set_low_quality_matches_ for more details.
        """
        # assert low_threshold <= high_threshold
        # self.high_threshold = high_threshold
        # self.low_threshold = low_threshold
        self.allow_low_quality_matches = allow_low_quality_matches

    def __call__(self, match_quality_matrix, iou_thr = 0.5):
        """
        Args:
            match_quality_matrix (Tensor[float]): an MxN tensor, containing the
            pairwise quality between M ground-truth elements and N predicted elements.

        Returns:
            matches (Tensor[int64]): an N tensor where N[i] is a matched gt in
            [0, M - 1] or a negative value indicating that prediction i could not
            be matched.
        """
        if match_quality_matrix.numel() == 0:
            # empty targets or proposals not supported during training
            if match_quality_matrix.shape[0] == 0:
                raise ValueError(
                    "No ground-truth boxes available for one of the images "
                    "during training")
            else:
                raise ValueError(
                    "No proposal boxes available for one of the images "
                    "during training")


        matched_vals, matches = match_quality_matrix.max(dim=0)
        if self.allow_low_quality_matches:
            all_matches = matches.clone()

        # Assign candidate matches with low quality to negative (unassigned) values
        below_threshold = matched_vals < iou_thr
        matches[below_threshold] = Matcher.BELOW_LOW_THRESHOLD


        if self.allow_low_quality_matches:
            self.set_low_quality_matches_(matches, all_matches, match_quality_matrix)

        return matches

    def set_low_quality_matches_(self, matches, all_matches, match_quality_matrix):
        """
        Produce additional matches for predictions that have only low-quality matches.
        Specifically, for each ground-truth find the set of predictions that have
        maximum overlap with it (including ties); for each prediction in that set, if
        it is unmatched, then match it to the ground-truth with which it has the highest
        quality value.
        """
        # For each gt, find the prediction with which it has highest quality
        highest_quality_foreach_gt, _ = match_quality_matrix.max(dim=1)
        # Find highest quality match available, even if it is low, including ties
        gt_pred_pairs_of_highest_quality = torch.nonzero(
            match_quality_matrix == highest_quality_foreach_gt[:, None]
        )

        pred_inds_to_update = gt_pred_pairs_of_highest_quality[:, 1]
        matches[pred_inds_to_update] = all_matches[pred_inds_to_update]

EgoMQ/Models/matcher.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import sys
sys.dont_write_bytecode = True
import os
import torch
import torch.nn.parallel
import torch.optim as optim
import numpy as np
from tensorboardX import SummaryWriter
import Utils.opts as opts
from Utils.dataset import VideoDataSet
from Models.VSGN import VSGN
import datetime
from collections import defaultdict

torch.manual_seed(21)

def Train_VSGN(opt):
    path_appendix = '_'.join(string for string in opt['checkpoint_path'].split('_')[1:])
    writer = SummaryWriter(logdir='runs/' + path_appendix)
    model = VSGN(opt)
    device = "cuda"
    model = torch.nn.DataParallel(model)
    model.to(device)

    optimizer = optim.Adam(model.parameters(), lr=opt["train_lr"], weight_decay=opt["weight_decay"])

    start_epoch = 0

    kwargs = {'num_workers': 16, 'pin_memory': True, 'drop_last': True}

    train_loader = torch.utils.data.DataLoader(VideoDataSet(opt, subset="train"),
                                               batch_size=opt["batch_size"], shuffle=True,
                                               **kwargs)

    test_loader = torch.utils.data.DataLoader(VideoDataSet(opt, subset="validation"),
                                              batch_size=opt["batch_size"], shuffle=False,
                                              **kwargs)

    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=opt["step_size"], gamma=opt["step_gamma"])

    for epoch in range(start_epoch, opt["num_epoch"]):

        train_VSGN_epoch(train_loader, model, optimizer, epoch, writer, opt)
        epoch_loss = test_VSGN_epoch(test_loader, model, epoch, writer, opt)

        print((datetime.datetime.now()))
        state = {'epoch': epoch + 1,
                 'state_dict': model.state_dict(),
                 'optimizer': optimizer.state_dict()}
        torch.save(state, opt["checkpoint_path"] + "/checkpoint.pth.tar")
        if epoch_loss < model.module.tem_best_loss:
            print((datetime.datetime.now()))
            print('The best model up to now is from Epoch {}'.format(epoch))
            model.module.tem_best_loss = np.mean(epoch_loss)
            torch.save(state, opt["checkpoint_path"] + "/best.pth.tar")

        scheduler.step()
    writer.close()


def train_VSGN_epoch(data_loader, model, optimizer, epoch, writer, opt, is_train=True):

    if is_train:
        model.train()
    else:
        model.eval()

    epoch_losses = defaultdict(float)
    for n_iter, (input_data, gt_action, gt_start, gt_end, gt_bbox, num_gt, num_frms) in enumerate(data_loader):
        with torch.set_grad_enabled(is_train):
            losses, pred_action, pred_start, pred_end = model(input_data, num_frms, gt_action, gt_start, gt_end,  gt_bbox, num_gt)


        # Overall loss
        loss_cls_dec = torch.mean(losses['loss_cls_dec'])
        loss_reg_dec = torch.mean(losses['loss_reg_dec'])
        loss_action = torch.mean(losses['loss_action'])
        loss_start = torch.mean(losses['loss_start'])
        loss_end = torch.mean(losses['loss_end'])
        loss_bd_adjust = torch.mean(losses['loss_bd_adjust'])

        loss = loss_cls_dec + loss_reg_dec + 0.2*loss_action + 0.2*loss_start +0.2*loss_end + loss_bd_adjust

        if is_train:
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        epoch_losses['loss'] += loss.cpu().detach().numpy()
        epoch_losses['loss_cls_dec'] += loss_cls_dec.cpu().detach().numpy()
        epoch_losses['loss_reg_dec'] += loss_reg_dec.cpu().detach().numpy()
        epoch_losses['loss_action'] += loss_action.cpu().detach().numpy()
        epoch_losses['loss_start'] += loss_start.cpu().detach().numpy()
        epoch_losses['loss_end'] += loss_end.cpu().detach().numpy()
        epoch_losses['loss_bd_adjust'] += loss_bd_adjust.cpu().detach().numpy()

    for k, v in epoch_losses.items():
        epoch_losses[k] = v / (n_iter + 1)

    to_print = ["%s loss (epoch %d): " % ('Train' if is_train else 'Val', epoch)]
    for k, v in epoch_losses.items():
        writer.add_scalar('%s/%s' % ('train' if is_train else 'val', k), v, epoch)
        writer.flush()
        to_print.append('%s: %.04f' % (k, v))
    print(' '.join(to_print))

    return epoch_losses['loss']



def test_VSGN_epoch(data_loader, model, epoch, writer, opt):
    return train_VSGN_epoch(data_loader, model, None, epoch, writer, opt, is_train=False)




if __name__ == '__main__':

    print(datetime.datetime.now())

    opt = opts.parse_opt()
    opt = vars(opt)

    if not os.path.exists(opt["checkpoint_path"]):
        os.makedirs(opt["checkpoint_path"])

    print(opt)

    print("---------------------------------------------------------------------------------------------")
    print("Training starts!")
    print("---------------------------------------------------------------------------------------------")
    Train_VSGN(opt)
    print("Training finishes!")

    print(datetime.datetime.now())


EgoMQ/Train.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import numpy as np
import pandas as pd
import pandas
import numpy
import json
import torch.utils.data as data
import os
import torch
import h5py
import pickle
import torch.nn.functional as F
from scipy.io import loadmat

def load_json(file):
    with open(file) as json_file:
        data = json.load(json_file)
        return data


class VideoDataSet(data.Dataset):
    def __init__(self, opt, subset="train", mode="train"):
        self.temporal_scale = opt["temporal_scale"]
        self.input_feat_dim = opt['input_feat_dim']
        self.temporal_gap = 1. / self.temporal_scale
        self.subset = subset
        self.mode = mode
        self.feature_path = opt["feature_path"]
        self.gap = opt['stitch_gap']
        self.clip_anno = opt['clip_anno']
        self.moment_classes = opt["moment_classes"]
        self.fps = {}


        self._getDatasetDict()

        self.anchor_xmin = [self.temporal_gap * i for i in range(self.temporal_scale)]
        self.anchor_xmax = [self.temporal_gap * i for i in range(1, self.temporal_scale + 1)]

    def _getDatasetDict(self):

        anno_database = load_json(self.clip_anno)

        self.clip_dict = {}
        class_list = []

        for clip_name, clip_info in anno_database.items():

            clip_subset = clip_info['subset']

            if clip_subset in self.subset:
                self.clip_dict[clip_name] = clip_info

                for item in clip_info['annotations']:
                    class_list.append(item['label'])

        self.clip_list = list(self.clip_dict.keys())

        if os.path.exists(self.moment_classes):
            with open(self.moment_classes, 'r') as f:
                self.classes = json.load(f)
        else:
            class_list = list(set(class_list))
            class_list = sorted(class_list)
            print(f'The total number of classes is {len(class_list) + 1}')
            self.classes = {'Background': 0}
            for i,cls in enumerate(class_list):
                self.classes[cls] = i + 1
            with open(self.moment_classes, 'w') as f:
                f.write(json.dumps(self.classes))

    def __getitem__(self, index):
        if self.mode == "train":
            video_data, match_score_action, match_score_start, match_score_end, gt_bbox, num_gt, num_frms = self._get_video_data(index)
            return video_data, match_score_action, match_score_start, match_score_end, gt_bbox, num_gt, num_frms
        else:
            video_data, num_frms = self._get_video_data(index)
            return index, video_data, num_frms

    def _get_video_data(self, index):
        # General data
        clip_name = self.clip_list[index]
        clip_info = self.clip_dict[clip_name]
        # video_name = clip_info['video_id']
        video_name = clip_info['clip_id']   

        # Get video features
        v_data = torch.load(os.path.join(self.feature_path, video_name + '.pt'))
        v_data = torch.transpose(v_data, 0, 1)
        num_frms_v = v_data.shape[-1]
        #fps_v = num_frms_v / clip_info['v_duration']
        fps_v = num_frms_v / (clip_info['parent_end_sec'] - clip_info['parent_start_sec'])  
        self.clip_dict[clip_name]['fps'] = fps_v

        #clip_start = int(clip_info['parent_start_sec'] * fps_v)
        clip_start = 0
        # clip_end = min(int(clip_info['parent_end_sec'] * fps_v), num_frms_v-1)
        clip_end = num_frms_v-1

        video_data = torch.zeros(self.input_feat_dim, self.temporal_scale)
        win_data = v_data[:, clip_start: clip_end+1]
        num_frms = min(win_data.shape[-1], self.temporal_scale)
        video_data[:, :num_frms] = win_data[:, :num_frms]
        
        if self.mode == 'train':
            match_score_action, match_score_start, match_score_end, gt_bbox_padding, num_gt, num_frms = \
                self._get_train_data_label_org(num_frms, clip_name, fps_v)
            return video_data, match_score_action, match_score_start, match_score_end, gt_bbox_padding, num_gt, num_frms
        else:
            return video_data, num_frms

    def _get_train_data_label_org(self, num_frms, clip_name,  fps):
        # Get annotations
        clip_info = self.clip_dict[clip_name]
        clip_labels = clip_info['annotations']

        # Get gt_iou_map
        gt_bbox = []
        for j in range(len(clip_labels)):
            tmp_info = clip_labels[j]

            tmp_start_f = max(min(num_frms-1, tmp_info['start_time']*fps), 0)
            tmp_end_f = max(min(num_frms-1, tmp_info['end_time']*fps), 0)

            tmp_start = tmp_start_f / self.temporal_scale
            tmp_end = tmp_end_f / self.temporal_scale
            tmp_class = self.classes[tmp_info['label']]
            gt_bbox.append([tmp_start, tmp_end, tmp_class])


        # Get actionness scores
        match_score_action = [0] * self.temporal_scale
        for bbox in gt_bbox:
            left_frm = max(round(bbox[0] * self.temporal_scale), 0)
            right_frm = min(round(bbox[1] * self.temporal_scale), self.temporal_scale-1)
            match_score_action[left_frm:right_frm+1] = [bbox[2]] * (right_frm + 1 - left_frm)

        match_score_action = torch.Tensor(match_score_action)

        ####################################################################################################
        # generate R_s and R_e
        gt_bbox = np.array(gt_bbox)
        if gt_bbox.shape[0] == 0:
            print(gt_bbox.shape)

        gt_xmins = gt_bbox[:, 0]
        gt_xmaxs = gt_bbox[:, 1]
        gt_len_small = 3 * self.temporal_gap
        gt_start_bboxs = np.stack((gt_xmins - gt_len_small / 2, gt_xmins + gt_len_small / 2), axis=1)
        gt_end_bboxs = np.stack((gt_xmaxs - gt_len_small / 2, gt_xmaxs + gt_len_small / 2), axis=1)
        #####################################################################################################

        ##########################################################################################################
        # calculate the ioa for all timestamp
        match_score_start = []
        for jdx in range(len(self.anchor_xmin)):
            match_score_start.append(np.max(
                self._ioa_with_anchors(self.anchor_xmin[jdx], self.anchor_xmax[jdx], gt_start_bboxs[:, 0], gt_start_bboxs[:, 1])))
        match_score_end = []
        for jdx in range(len(self.anchor_xmin)):
            match_score_end.append(np.max(
                self._ioa_with_anchors(self.anchor_xmin[jdx], self.anchor_xmax[jdx], gt_end_bboxs[:, 0], gt_end_bboxs[:, 1])))
        match_score_start = torch.Tensor(match_score_start)
        match_score_end = torch.Tensor(match_score_end)
        ############################################################################################################

        max_num_box = 50
        gt_bbox = torch.tensor(gt_bbox, dtype=torch.float32)
        gt_bbox_padding = gt_bbox.new(max_num_box, gt_bbox.size(1)).zero_()
        num_gt = min(gt_bbox.size(0), max_num_box)
        gt_bbox_padding[:num_gt, :] = gt_bbox[:num_gt]
        # labels = BoxList(torch.Tensor(gt_bbox))

        return  match_score_action, match_score_start, match_score_end, gt_bbox_padding, num_gt, num_frms

    def _ioa_with_anchors(self, anchors_min, anchors_max, box_min, box_max):
        len_anchors = anchors_max - anchors_min
        int_xmin = np.maximum(anchors_min, box_min)
        int_xmax = np.minimum(anchors_max, box_max)
        inter_len = np.maximum(int_xmax - int_xmin, 0.)
        scores = np.divide(inter_len, len_anchors)
        return scores

    def __len__(self):
        return len(self.clip_list)


def iou_with_anchors(anchors_min, anchors_max, box_min, box_max):
    """Compute jaccard score between a box and the anchors.
    """
    len_anchors = anchors_max - anchors_min
    int_xmin = np.maximum(anchors_min, box_min)
    int_xmax = np.minimum(anchors_max, box_max)
    inter_len = np.maximum(int_xmax - int_xmin, 0.)
    union_len = len_anchors - inter_len + box_max - box_min
    # print inter_len,union_len
    jaccard = np.divide(inter_len, union_len)
    return jaccard



EgoMQ/Utils/dataset.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import argparse

def parse_opt():

    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--dataset',
        default='ego4d',
        type=str,
        choices=['ego4d', 'ego4d'])
    parser.add_argument(
        '--is_train',
        default='true',
        type=str,
        choices=['true', 'false'])
    parser.add_argument(
        '--out_prop_map',
        default='true',
        type=str,
        choices=['true', 'false'])

    # Dataset and annotation paths
    parser.add_argument(
        '--feature_path',
        type=str,
        default="/mnt/sdb1/Datasets/Ego4d/action_feature_canonical")
    parser.add_argument(
        '--clip_anno',
        type=str,
        default="Evaluation/ego4d/annot/clip_annotations.json")
    parser.add_argument(
        '--moment_classes',
        type=str,
        default="Evaluation/ego4d/annot/moment_classes_idx.json")

    # Output paths
    parser.add_argument(
        '--checkpoint_path',
        type=str,
        default='checkpoint')
    parser.add_argument(
        '--output_path',
        type=str,
        default='output')
    parser.add_argument(
        '--prop_path',
        type=str,
        default='proposals')
    parser.add_argument(
        '--prop_result_file',
        type=str,
        default="proposals_postNMS.json")
    parser.add_argument(
        '--detect_result_file',
        type=str,
        default="detections_postNMS.json")
    parser.add_argument(
        '--retrieval_result_file',
        type=str,
        default="retreival_postNMS.json")
    parser.add_argument(
        '--detad_sensitivity_file',
        type=str,
        default="detad_sensitivity")

    # Training hyper-parameters
    parser.add_argument(
        '--batch_size',
        type=int,
        default=32)
    parser.add_argument(
        '--train_lr',
        type=float,
        default=0.00005)
    parser.add_argument(
        '--weight_decay',
        type=float,
        default=0.0001)
    parser.add_argument(
        '--num_epoch',
        type=int,
        default=50)
    parser.add_argument(
        '--step_size',
        type=int,
        default=15)
    parser.add_argument(
        '--step_gamma',
        type=float,
        default=0.1)
    parser.add_argument(
        '--focal_alpha',
        type=float,
        default=0.01)

    # Post processing
    parser.add_argument(
        '--nms_alpha_detect',
        type=float,
        default=0.46)
    parser.add_argument(
        '--nms_alpha_prop',
        type=float,
        default=0.75)
    parser.add_argument(
        '--nms_thr',
        type=float,
        default=0.4)

    # Model architecture settings
    parser.add_argument(
        '--temporal_scale',
        type=int,
        default=928)
    parser.add_argument(
        '--input_feat_dim',
        type=int,
        default=2304)
    parser.add_argument(
        '--bb_hidden_dim',
        type=int,
        default=256)
    parser.add_argument(
        '--decoder_num_classes',
        type=int,
        default=111)
    parser.add_argument(
        '--num_levels',
        type=int,
        default=5)  # 5
    parser.add_argument(
        '--num_head_layers',
        type=int,
        default=4)

    # Graph network hyper-parameters
    parser.add_argument(
        '--nfeat_mode',
        default='feat_ctr',
        type=str,
        choices=['feat_ctr', 'dif_ctr', 'feat'])
    parser.add_argument(
        '--num_neigh',
        type=int,
        default=12)
    parser.add_argument(
        '--edge_weight',
        default='false',
        type=str,
        choices=['true', 'false'])
    parser.add_argument(
        '--agg_type',
        default='max',
        type=str,
        choices=['max', 'mean'])
    parser.add_argument(
        '--gcn_insert',
        default='par',
        type=str,
        choices=['seq', 'par'])

    # Detection hyper-parameters
    parser.add_argument(
        '--iou_thr',
        nargs='+',
        type=float,
        default=[0.5, 0.5, 0.7])
    parser.add_argument(
        '--anchor_scale',
        nargs='+',
        type=float,
        default=[1, 10])  # 4, 6; 8, 12; 16, 24; 32, 48; 64, 96
    parser.add_argument(
        '--base_stride',
        type=int,
        default=1)


    # VSS hyper-parameters
    parser.add_argument(
        '--stitch_gap',
        type=int,
        default=30)
    parser.add_argument(
        '--short_ratio',
        type=float,
        default=0.4)
    parser.add_argument(
        '--clip_win_size',
        type=float,
        default=0.38)

    # Baselines
    parser.add_argument(
        '--use_xGPN',
        default=False,
        action='store_true')
    parser.add_argument(
        '--use_VSS',
        default=False,
        action='store_true')

    parser.add_argument(
        '--num_props',
        type=int,
        default=200)

    parser.add_argument(
        '--tIoU_thr',
        nargs='+',
        type=float,
        default=[0.1, 0.2, 0.3, 0.4, 0.5])

    parser.add_argument(
        '--eval_stage',
        default='process_eval_prop',
        type=str,
        choices=['process_eval_prop', 'eval_detection', 'eval_retrieval', 'detad', 'all'])

    parser.add_argument(
        '--infer_datasplit',
        default='val',
        type=str,
        choices=['test', 'val'])
    args = parser.parse_args()

    return args

EgoMQ/Utils/opts.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import glob
import argparse

parser = argparse.ArgumentParser(description='Find Best HyperParameters')
parser.add_argument('--root_folder', type=str, help='Root output folder')

args = parser.parse_args()

file_list = glob.glob(str(args.root_folder) + '/*.txt')

best_map = 0

for txt_file in file_list:
    file1 = open(txt_file, 'r')
    Lines = file1.readlines()

    for line in Lines:
        if 'Average-mAP:' in line.strip():
            if float(line.split( )[-1]) > best_map:
                best_map = float(line.split( )[-1])
                best_file_name = txt_file

print("File with the best results: ", best_file_name)

EgoMQ/find_best_parameters.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import os
import pdb
import sys
import json
import pandas as pd
import numpy as np
import torch
import transformers
from tqdm import tqdm

from base.base_dataset import TextVideoDataset
from base.transforms import init_transform_dict, init_video_transform_dict


class NaturalLanguageQueries(TextVideoDataset):

    def __init__(self, dataset_name, text_params, video_params, config='/cis/home/shraman/works_meta_2022/NLQ/Ego4D_episodic_memory-main/NLQ/VSLNet_unified/configs/nlq.json', split='train', tsfms=None, reader='decord'):
        
        f = open(config)
        config = json.load(f)
        meta_dir = config['data_loader']['args']['meta_dir']
        data_dir = config['data_loader']['args']['data_dir']
        self.num_frame = config['data_loader']['args']['video_params']['num_frames']
        #self.args = args
        # Build tokenizer
        self.tokenizer = transformers.AutoTokenizer.from_pretrained(config['arch']['args']['text_params']['model'])
        
        super().__init__(dataset_name=dataset_name, text_params=text_params, video_params=video_params, data_dir=data_dir, meta_dir=meta_dir, split=split, tsfms=tsfms, cut=None, subsample=1, sliding_window_stride=-1, reader=reader, neg_param=None)

    def _load_metadata(self):
        split_files = {
            'train': 'nlq_train.json',
            'val': 'nlq_val.json',            # there is no test
            'test': 'nlq_test_unannotated.json'
        }
        # target_split_fp = split_files[self.split]

        self.metadata = pd.DataFrame(columns=['video_uid', 'clip_uid',
                                              'video_start_sec', 'video_end_sec',
                                              'query'])

        # for split in ['train', 'val']:
        target_split_fp = split_files[split]
        ann_file = os.path.join(self.meta_dir, target_split_fp)
        with open(ann_file) as f:
            anno_json = json.load(f)


        # forward video frames and text
        for anno_video in anno_json["videos"]:
            for anno_clip in anno_video["clips"]:
                clip_times = float(anno_clip["video_start_sec"]), float(anno_clip["video_end_sec"])
                for anno in anno_clip['annotations']:
                    for query in anno["language_queries"]:
                        clip_duration = clip_times[1] - clip_times[0]
                        if 'query' not in query.keys():
                            continue
                        if query['query'] is None:
                            continue
                        new = pd.DataFrame({
                            'video_uid': anno_video['video_uid'],
                            'clip_uid': anno_clip['clip_uid'],
                            'video_start_sec': clip_times[0],
                            'video_end_sec': clip_times[1],
                            'query': query["query"],}, index=[1])
                        self.metadata = self.metadata.append(new, ignore_index=True)

        self.transforms = init_video_transform_dict()['test']

    def _get_video_path(self, sample):
        rel_video_fp = sample[0]
        full_video_fp = os.path.join(self.data_dir, rel_video_fp + '.mp4')
        return full_video_fp, rel_video_fp

    def _get_caption(self, sample):
        caption = sample['query']
        return caption

    def visual_feature_sampling(self, visual_feature, max_num_clips):
        num_clips = visual_feature.shape[0]
        if num_clips <= max_num_clips:
            return visual_feature
        idxs = np.arange(0, max_num_clips + 1, 1.0) / max_num_clips * num_clips
        idxs = np.round(idxs).astype(np.int32)
        idxs[idxs > num_clips - 1] = num_clips - 1
        new_visual_feature = []
        for i in range(max_num_clips):
            s_idx, e_idx = idxs[i], idxs[i + 1]
            if s_idx < e_idx:
                new_visual_feature.append(np.mean(visual_feature[s_idx:e_idx].numpy(), axis=0))
            else:
                new_visual_feature.append(visual_feature[s_idx])
        new_visual_feature = np.asarray(new_visual_feature)
        return torch.from_numpy(new_visual_feature)

    
    def _get_video_text_feats(self, item):
        sample = self.metadata.iloc [item]
        video_fp, rel_fp = self._get_video_path(sample)
       
        
        fps = 1.87
        try:
            imgs, idxs = self.video_reader(video_fp, sample[2]*30, sample[3]*30,
                                               (sample[3]-sample[2]) * fps * self.video_params['num_frames'])
        except:
            print(f"Warning: missing video file {video_fp}.")

        if self.transforms is not None:
            imgs = imgs.transpose(0, 1)  # [T, C, H, W] ---> [C, T, H, W]
            imgs = self.transforms(imgs)
            imgs = imgs.transpose(0, 1)  # recover

        f, c, h, w = imgs.shape[0], imgs.shape[1], imgs.shape[2], imgs.shape[3]
        imgs = imgs[:(f // self.num_frame * self.num_frame), ]
        imgs = imgs.reshape(-1, self.num_frame, c, h, w)

        imgs = self.visual_feature_sampling(imgs, max_num_clips=256) ## max_num_clips = args.max_pos_len
        
        text = self._get_caption(sample)
        text = self.tokenizer(text, return_tensors='pt', padding=True, truncation=True)
        text['input_ids'] = text['input_ids'].repeat(imgs.shape[0], 1)
        text['attention_mask'] = text['attention_mask'].repeat(imgs.shape[0], 1)
        

        meta_arr = {'video_uid': sample[0], 'clip_uid': sample[1], 'video_start': sample[2], 'video_end': sample[3], 'data': video_fp, 'dataset': self.dataset_name}
        data = {'video': imgs, 'text': text, 'meta' : meta_arr}
        #data = {'meta' : meta_arr}
        
        return data

    def __getitem__(self, item):
        return self._get_video_text_feats(item)

if __name__ == "__main__":
    split = 'val'
    kwargs = dict(
        dataset_name="Ego4d_NLQ",
        text_params={
            "input": "text"
        },
        video_params={
        "input_res": 224,
        "num_frames": 4,
        "loading": "lax"
        },
        tsfms=init_video_transform_dict()['test'],
        reader='decord_start_end',
        split=split,
    )
    dataset = NaturalLanguageQueries(**kwargs)
    print(len(dataset))
    
    test_dict = {}
    c = 0
    for i, _ in tqdm(enumerate(range(len(dataset)))):
        item = dataset[i]
        if item['meta']['clip_uid'] in test_dict:
            if test_dict[item['meta']['clip_uid']][0] != item['meta']['video_start'] or test_dict[item['meta']['clip_uid']][1] != item['meta']['video_end']:
                c += 1
        test_dict.update({item['meta']['clip_uid']: (item['meta']['video_start'], item['meta']['video_end'])})
    print(c)


EgoNLQ/Ego4D_NLQ_dataset.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

from .base_dataset import *
from .base_model import *

EgoNLQ/base/__init__.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import os
import pdb

import tqdm
import random
from abc import abstractmethod

import av
import cv2
import decord
import ffmpeg
import numpy as np
import torch
from PIL import Image
from torch.utils.data import Dataset, get_worker_info
from torchvision import transforms


class TextVideoDataset(Dataset):
    def __init__(self,
                 dataset_name,
                 text_params,
                 video_params,
                 data_dir,
                 meta_dir=None,
                 split='train',
                 tsfms=None,
                 cut=None,
                 subsample=1,
                 sliding_window_stride=-1,
                 reader='decord',
                 neg_param=None,
                 ):
        self.dataset_name = dataset_name
        self.text_params = text_params
        self.video_params = video_params
        # check for environment variables
        self.data_dir = os.path.expandvars(data_dir)
        if meta_dir is not None:
            self.meta_dir = os.path.expandvars(meta_dir)
        else:
            self.meta_dir = self.data_dir
        self.split = split
        self.transforms = tsfms
        self.cut = cut
        self.subsample = subsample
        self.sliding_window_stride = sliding_window_stride
        self.video_reader = video_reader[reader]
        self.label_type = 'caption'
        self.neg_param = neg_param
        self._load_metadata()
        if self.sliding_window_stride != -1:
            if self.split != 'test':
                raise ValueError('Fixing frame sampling is for test time only. can remove but...')
            self._fix_temporal_samples()

    @abstractmethod
    def _load_metadata(self):
        raise NotImplementedError("Metadata loading must be implemented by subclass")

    @abstractmethod
    def _get_video_path(self, sample):
        raise NotImplementedError("Get video path function must be implemented by subclass")

    def _get_caption(self, sample):
        raise NotImplementedError("Get caption function must be implemented by subclass")

    def _get_video_lens(self):
        vlen_li = []
        for idx, row in self.metadata.iterrows():
            video_path = self._get_video_path(row)[0]
            vlen_li.append(get_video_len(video_path))

        return vlen_li

    def _fix_temporal_samples(self):
        self.metadata['vlen'] = self._get_video_lens()
        self.metadata['frame_intervals'] = self.metadata['vlen'].apply(
            lambda x: np.linspace(start=0, stop=x, num=min(x, self.video_params['num_frames']) + 1).astype(int))
        self.metadata['fix_start'] = self.metadata['frame_intervals'].apply(
            lambda x: np.arange(0, int(x[-1] / len(x - 1)), self.sliding_window_stride)
        )
        self.metadata = self.metadata.explode('fix_start')

    def __len__(self):
        return len(self.metadata)

    def __getitem__(self, item):
        item = item % len(self.metadata)
        sample = self.metadata.iloc[item]
        video_fp, rel_fp = self._get_video_path(sample)
        caption = self._get_caption(sample)

        video_loading = self.video_params.get('loading', 'strict')
        frame_sample = 'rand'
        fix_start = None
        if self.split == 'test':
            frame_sample = 'uniform'
        if self.sliding_window_stride != -1:
            fix_start = sample['fix_start']

        try:
            if os.path.isfile(video_fp):
                imgs, idxs = self.video_reader(video_fp, self.video_params['num_frames'], frame_sample,
                                               fix_start=fix_start)
            else:
                print(f"Warning: missing video file {video_fp}.")
                assert False
        except Exception as e:
            if video_loading == 'strict':
                raise ValueError(
                    f'Video loading failed for {video_fp}, video loading for this dataset is strict.') from e
            else:
                imgs = Image.new('RGB', (self.video_params['input_res'], self.video_params['input_res']), (0, 0, 0))
                imgs = transforms.ToTensor()(imgs).unsqueeze(0)

        # if self.transforms is not None:
        #     imgs = self.transforms(imgs)

        if self.transforms is not None:
            if self.video_params['num_frames'] > 1:
                imgs = imgs.transpose(0, 1)  # [T, C, H, W] ---> [C, T, H, W]
                imgs = self.transforms(imgs)
                imgs = imgs.transpose(0, 1)  # recover
            else:
                imgs = self.transforms(imgs)

        # if self.transforms is not None:
        #     # for video aug if T>1 else just do image aug
        #     if imgs.size(0) > 1:
        #         # added by Mr. YAN
        #         imgs = imgs.transpose(0, 1) # [T, C, H, W] ---> [C, T, H, W]
        #         imgs = self.transforms(imgs)
        #         imgs = imgs.transpose(0, 1) # recover
        #     else:
        #         imgs = self.transforms(imgs)    # t

        final = torch.zeros([self.video_params['num_frames'], 3, self.video_params['input_res'],
                             self.video_params['input_res']])
        final[:imgs.shape[0]] = imgs

        meta_arr = {'raw_captions': caption, 'paths': rel_fp, 'dataset': self.dataset_name}
        data = {'video': final, 'text': caption, 'meta': meta_arr}
        return data


class TextImageDataset(TextVideoDataset):

    def __getitem__(self, item):
        item = item % len(self.metadata)
        sample = self.metadata.iloc[item]
        video_fp, rel_fp = self._get_video_path(sample)
        caption = self._get_caption(sample)

        video_loading = self.video_params.get('loading', 'strict')

        try:
            img = Image.open(video_fp).convert("RGB")
        except:
            if video_loading == 'strict':
                raise ValueError(f'Image loading failed for {video_fp}, image loading for this dataset is strict.')
            else:
                img = Image.new('RGB', (self.video_params['input_res'], self.video_params['input_res']), (0, 0, 0))

        # convert to tensor because video transforms don't, expand such that its a 1-frame video.
        img = transforms.ToTensor()(img).unsqueeze(0)
        if self.transforms is not None:
            img = self.transforms(img)
        meta_arr = {'raw_captions': caption, 'paths': rel_fp, 'dataset': self.dataset_name}
        data = {'video': img, 'text': caption, 'meta': meta_arr}
        return data


def sample_frames(num_frames, vlen, sample='rand', fix_start=None):
    acc_samples = min(num_frames, vlen)
    intervals = np.linspace(start=0, stop=vlen, num=acc_samples + 1).astype(int)
    ranges = []
    for idx, interv in enumerate(intervals[:-1]):
        ranges.append((interv, intervals[idx + 1] - 1))
    if sample == 'rand':
        frame_idxs = [random.choice(range(x[0], x[1])) for x in ranges]
    elif fix_start is not None:
        frame_idxs = [x[0] + fix_start for x in ranges]
    elif sample == 'uniform':
        frame_idxs = [(x[0] + x[1]) // 2 for x in ranges]
    else:
        raise NotImplementedError

    return frame_idxs

def sample_frames_clips(start, end, vlen, acc_samples):
    start = max(0, start)
    end = min(vlen, end)

    intervals = np.linspace(start=start, stop=end, num=int(acc_samples) + 1).astype(int)
    ranges = []
    for idx, interv in enumerate(intervals[:-1]):
        ranges.append((interv, intervals[idx + 1] - 1))
        frame_idxs = [(x[0] + x[1]) // 2 for x in ranges
                      ]
    return frame_idxs

def sample_frames_start_end(num_frames, start, end, sample='rand', fix_start=None):
    acc_samples = min(num_frames, end)
    intervals = np.linspace(start=start, stop=end, num=acc_samples + 1).astype(int)
    ranges = []
    for idx, interv in enumerate(intervals[:-1]):
        ranges.append((interv, intervals[idx + 1] - 1))
    if sample == 'rand':
        frame_idxs = [random.choice(range(x[0], x[1])) for x in ranges]
    elif fix_start is not None:
        frame_idxs = [x[0] + fix_start for x in ranges]
    elif sample == 'uniform':
        frame_idxs = [(x[0] + x[1]) // 2 for x in ranges]
    else:
        raise NotImplementedError

    return frame_idxs

def read_frames_cv2(video_path, num_frames, sample='rand', fix_start=None):
    cap = cv2.VideoCapture(video_path)
    assert (cap.isOpened())
    vlen = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    # get indexes of sampled frames
    frame_idxs = sample_frames(num_frames, vlen, sample=sample, fix_start=fix_start)
    frames = []
    success_idxs = []
    for index in frame_idxs:
        cap.set(cv2.CAP_PROP_POS_FRAMES, index - 1)
        ret, frame = cap.read()
        if ret:
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frame = torch.from_numpy(frame)
            # (H x W x C) to (C x H x W)
            frame = frame.permute(2, 0, 1)
            frames.append(frame)
            success_idxs.append(index)
        else:
            pass
            # print(frame_idxs, ' fail ', index, f'  (vlen {vlen})')

    frames = torch.stack(frames).float() / 255
    cap.release()
    return frames, success_idxs

def read_frames_cv2_egoclip(video_path_1, video_path_2, num_frames, sample,
                            start_sec, end_sec, bound_sec):
    if video_path_1 == video_path_2:
        cap1 = cv2.VideoCapture(video_path_1)
        cap2 = cap1
        vlen1 = int(cap1.get(cv2.CAP_PROP_FRAME_COUNT))
        vlen2 = vlen1
        assert (cap1.isOpened())
    else:   # some clips may span two segments.
        cap1 = cv2.VideoCapture(video_path_1)
        cap2 = cv2.VideoCapture(video_path_2)
        vlen1 = int(cap1.get(cv2.CAP_PROP_FRAME_COUNT))
        vlen2 = int(cap2.get(cv2.CAP_PROP_FRAME_COUNT))
        assert (cap1.isOpened())
        assert (cap2.isOpened())

    # get indexes of sampled frames
    start_f = max(0, int(start_sec * 30))
    end_f = max(0, int(end_sec * 30))
    bound_f = int(bound_sec * 30)
    frame_idxs = sample_frames_start_end(num_frames, start_f, end_f, sample=sample)

    frames = []
    success_idxs = []
    for index in frame_idxs:
        _index = index % (600 * 30)
        if index > bound_f: # frame from the last video
            _index = min(_index, vlen2)
            cap2.set(cv2.CAP_PROP_POS_FRAMES, _index - 1)
            ret, frame = cap2.read()
        else:   # frame from the first video
            _index = min(_index, vlen1)
            cap1.set(cv2.CAP_PROP_POS_FRAMES, _index - 1)
            ret, frame = cap1.read()

        if ret:
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frame = torch.from_numpy(frame)
            # (H x W x C) to (C x H x W)
            frame = frame.permute(2, 0, 1)
            frames.append(frame)
            success_idxs.append(index)
        else:
            pass

    while len(frames) < num_frames: # complete the frame
        frames.append(frames[-1])

    frames = torch.stack(frames).float() / 255
    cap1.release()
    cap2.release()
    return frames, success_idxs

def read_frames_cv2_epic(video_path, start_frame, stop_frame, num_frames, sample='rand', fix_start=None):
    # get indexes of sampled frames
    frame_idxs = sample_frames_start_end(num_frames, start_frame, stop_frame, sample=sample, fix_start=fix_start)
    frames = []
    success_idxs = []
    for index in frame_idxs:
        img_name = 'frame_' + str(index).zfill(10) + '.jpg'
        frame = cv2.imread(os.path.join(video_path, img_name),cv2.COLOR_BGR2RGB)

        frame = torch.from_numpy(frame)
        # (H x W x C) to (C x H x W)
        frame = frame.permute(2, 0, 1)
        frames.append(frame)
        success_idxs.append(index)

    frames = torch.stack(frames).float() / 255
    return frames, success_idxs

def read_frames_cv2_charades(video_path, num_frames, sample, start_sec=None, end_sec=None):
    cap = cv2.VideoCapture(video_path)
    assert (cap.isOpened())
    vlen = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = cap.get(5)

    # get indexes of sampled frames
    if not start_sec and not end_sec:
        frame_idxs = sample_frames(num_frames, vlen, sample=sample)
    else:
        start_f = max(0, int(start_sec * fps))
        end_f = min(int(end_sec * fps), vlen)
        frame_idxs = sample_frames_start_end(num_frames, start_f, end_f, sample=sample)

    frames = []
    success_idxs = []
    for index in frame_idxs:
        cap.set(cv2.CAP_PROP_POS_FRAMES, index - 1)
        ret, frame = cap.read()
        if ret:
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frame = torch.from_numpy(frame)
            # (H x W x C) to (C x H x W)
            frame = frame.permute(2, 0, 1)
            frames.append(frame)
            success_idxs.append(index)
        else:
            pass

    frames = torch.stack(frames).float() / 255
    cap.release()
    return frames, success_idxs

def read_frames_av(video_path, num_frames, sample='rand', fix_start=None):
    reader = av.open(video_path)
    try:
        frames = []
        frames = [torch.from_numpy(f.to_rgb().to_ndarray()) for f in reader.decode(video=0)]
    except (RuntimeError, ZeroDivisionError) as exception:
        print('{}: WEBM reader cannot open {}. Empty '
              'list returned.'.format(type(exception).__name__, video_path))
    vlen = len(frames)
    frame_idxs = sample_frames(num_frames, vlen, sample=sample, fix_start=fix_start)
    frames = torch.stack([frames[idx] for idx in frame_idxs]).float() / 255
    frames = frames.permute(0, 3, 1, 2)
    return frames, frame_idxs

decord.bridge.set_bridge("torch")

def read_frames_decord(video_path, num_frames, sample='rand', fix_start=None):
    video_reader = decord.VideoReader(video_path, num_threads=1)
    vlen = len(video_reader)
    frame_idxs = sample_frames(num_frames, vlen, sample=sample, fix_start=fix_start)
    video_reader.skip_frames(1)
    frames = video_reader.get_batch(frame_idxs)

    frames = frames.float() / 255
    frames = frames.permute(0, 3, 1, 2)
    return frames, frame_idxs

def read_frames_decord_start_end(video_path, start, end, num_frames):
    video_reader = decord.VideoReader(video_path, num_threads=1)
    vlen = len(video_reader)
    frame_idxs = sample_frames_clips(start, end, vlen, num_frames + 1)
    video_reader.skip_frames(1)
    frames = video_reader.get_batch(frame_idxs)

    frames = frames.float() / 255
    frames = frames.permute(0, 3, 1, 2)
    return frames, frame_idxs

def get_video_len(video_path):
    cap = cv2.VideoCapture(video_path)
    if not (cap.isOpened()):
        return False
    vlen = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    cap.release()
    return vlen

video_reader = {
    'av': read_frames_av,
    'cv2': read_frames_cv2,
    'cv2_epic': read_frames_cv2_epic,
    'cv2_charades': read_frames_cv2_charades,
    'cv2_egoclip': read_frames_cv2_egoclip,
    'decord': read_frames_decord,
    'decord_start_end': read_frames_decord_start_end,
}

EgoNLQ/base/base_dataset.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import torch.nn as nn
import numpy as np
from abc import abstractmethod


class BaseModel(nn.Module):
    """
    Base class for all models
    """
    @abstractmethod
    def forward(self, *inputs):
        """
        Forward pass logic

        :return: Model output
        """
        raise NotImplementedError

    def __str__(self):
        """
        Model prints with number of trainable parameters
        """
        model_parameters = filter(lambda p: p.requires_grad, self.parameters())
        params = sum(np.prod(p.size()) for p in model_parameters)
        return super().__str__() + '\nTrainable parameters: {}'.format(params)

EgoNLQ/base/base_model.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


from torchvision import transforms
from torchvision.transforms._transforms_video import RandomCropVideo, RandomResizedCropVideo,CenterCropVideo, NormalizeVideo,ToTensorVideo,RandomHorizontalFlipVideo


def init_transform_dict(input_res=224,
                        center_crop=256,
                        randcrop_scale=(0.5, 1.0),
                        color_jitter=(0, 0, 0),
                        norm_mean=(0.485, 0.456, 0.406),
                        norm_std=(0.229, 0.224, 0.225)):
    normalize = transforms.Normalize(mean=norm_mean, std=norm_std)
    tsfm_dict = {
        'train': transforms.Compose([
            transforms.RandomResizedCrop(input_res, scale=randcrop_scale),
            transforms.RandomHorizontalFlip(),
            transforms.ColorJitter(brightness=color_jitter[0], saturation=color_jitter[1], hue=color_jitter[2]),
            normalize,
        ]),
        'val': transforms.Compose([
            transforms.Resize(center_crop),
            transforms.CenterCrop(center_crop),
            transforms.Resize(input_res),
            normalize,
        ]),
        'test': transforms.Compose([
            transforms.Resize(center_crop),
            transforms.CenterCrop(center_crop),
            transforms.Resize(input_res),
            normalize,
        ])
    }
    return tsfm_dict

def init_video_transform_dict(input_res=224,
                        center_crop=256,
                        randcrop_scale=(0.5, 1.0),
                        color_jitter=(0, 0, 0),
                        norm_mean=(0.485, 0.456, 0.406),
                        norm_std=(0.229, 0.224, 0.225)):
    print('Video Transform is used!')
    normalize = NormalizeVideo(mean=norm_mean, std=norm_std)
    tsfm_dict = {
        'train': transforms.Compose([
            RandomResizedCropVideo(input_res, scale=randcrop_scale),
            RandomHorizontalFlipVideo(),
            transforms.ColorJitter(brightness=color_jitter[0], saturation=color_jitter[1], hue=color_jitter[2]),
            normalize,
        ]),
        'val': transforms.Compose([
            transforms.Resize(center_crop),
            transforms.CenterCrop(center_crop),
            transforms.Resize(input_res),
            normalize,
        ]),
        'test': transforms.Compose([
            transforms.Resize(center_crop),
            transforms.CenterCrop(center_crop),
            transforms.Resize(input_res),
            normalize,
        ])
    }
    return tsfm_dict

EgoNLQ/base/transforms.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
import torchvision
import os
from torchvision import transforms,models
import torch.nn as nn
from PIL import Image
import clip
import pdb
import argparse
import glob
import h5py
import json
from tqdm import tqdm
#from parse_config import ConfigParser
import model.model_video_unfused as module_arch
from pathlib import Path
from collections import OrderedDict
from segment import cpd_auto

max_segment_num = 20
max_frame_num = 200


import warnings
warnings.filterwarnings("ignore")

def read_json(fname):
    with fname.open('rt') as handle:
        return json.load(handle, object_hook=OrderedDict)

config = read_json(Path('egoclip.json'))

class Image_Dataset(Dataset):
    def __init__(self, image_list, transform=None):

        self.image_list = image_list
        self.transform = transform

    def __len__(self):
        return len(self.image_list)

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()
            
        img_name = self.image_list[idx]
        
        image = Image.open(img_name).convert('RGB')

        if self.transform:
            sample = self.transform(image)
          
        return sample



def extract_video_features(args, video_num):
    
    '''
    parser = argparse.ArgumentParser(description='Extract CLIP features')
    #parser.add_argument('--model_name', default='ViT-B32', type=str, help='name of the model')
    #parser.add_argument('--root_dir', type=str, help='name of root feature dir')
    parser.add_argument('--cuda_base', help="in form cuda:x")
    parser.add_argument('--device_ids', help='delimited list input', type=lambda s: [int(item) for item in s.split(',')])
    #parser.add_argument('--h5_path', type=str, help="path to store the extracted features")
    #parser.add_argument('--config', default='egoclip.json', type=str, help='config file path (default: None)')
    #parser.add_argument('--resume', default=None, type=str, help='path to latest checkpoint (default: None)')

    args = parser.parse_args()
    '''

    args.root_dir = '../Datasets/QFVS/data/video_frames/P0' + str(video_num) + '-moviepy/'

    image_list = sorted(filter(os.path.isfile,
                        glob.glob(os.path.join(args.root_dir,'*.jpg'))))


    transform = transforms.Compose([
                transforms.RandomResizedCrop(224, interpolation=transforms.InterpolationMode.BICUBIC),
                transforms.ToTensor(),
                transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])


    #config = ConfigParser(parser)
    #model = config.initialize('arch', module_arch)

    model = getattr(module_arch, config['arch']['type'])(config['arch']['args']['video_params'], config['arch']['args']['text_params'], 
                                 projection=config['arch']['args']["projection"], load_checkpoint=config['arch']['args']["load_checkpoint"])

    data = Image_Dataset(image_list, transform=transform)
    dataloader = DataLoader(data, batch_size=5, drop_last=True, shuffle=False, num_workers=4)

    device = torch.device(args.cuda_base if torch.cuda.is_available() else 'cpu')
    model = model.to(device)
    model = nn.DataParallel(model, device_ids = args.device_ids)

    feature = []

    with torch.no_grad():
        for i, image in enumerate(tqdm(dataloader)):
            image = image.to(device)
            b, c, h, w = image.shape
            image_reshape = image.reshape(int(b/5), 5, c, h, w)
            output = model(image_reshape)
            output = output.detach().to('cpu').numpy()
            feature.append(output)
            
        feature = np.concatenate(feature, axis=0)
        frame_num=feature.shape[0]
        print("Number of Frames: ", frame_num)

        K=feature[:,0,:]
        K=np.dot(K,K.T)

        cps,_=cpd_auto(K,max_segment_num-1,1,desc_rate=1,verbose=False,lmax=max_frame_num-1) #int(K.shape[0]/25)
        seg_num=len(cps)+1

        assert seg_num<=max_segment_num

        seg_points=cps
        seg_points=np.insert(seg_points,0,0)
        seg_points=np.append(seg_points,frame_num)

        segments=[]
        for i in range(seg_num):
            segments.append(np.arange(seg_points[i],seg_points[i+1],1,dtype=np.int32))

        assert len(segments)<=max_segment_num

        for seg in segments:
            assert len(seg)<=max_frame_num

        seg_len=np.zeros((max_segment_num),dtype=np.int32)
        for index,seg in enumerate(segments):
            seg_len[index]=len(seg)

        # features

        for seg in segments:
            for frame in seg:
                assert frame<frame_num

        feature_dim=768

        features=torch.zeros((max_segment_num, max_frame_num, feature.shape[1], feature_dim))
        for seg_index,seg in enumerate(segments):
            for frame_index,frame in enumerate(seg):
                features[seg_index,frame_index]=torch.tensor(feature[frame])
                # features[seg_index,frame_index]=F.avg_pool1d(t.tensor(feature[frame]).unsqueeze(0).unsqueeze(0),kernel_size=2,stride=2)
        
        
        print("Feature Extraction Done!")

        return features, seg_len


EgoNLQ/extract_features.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

from .logger import *

EgoNLQ/logger/__init__.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import logging
import logging.config
from pathlib import Path

from utils.util import read_json


def setup_logging(save_dir, log_config='/logger/logger_config.json', default_level=logging.INFO):
    """
    Setup logging configuration
    """
    log_config = Path(log_config)
    if log_config.is_file():
        config = read_json(log_config)
        # modify logging paths based on run config
        for _, handler in config['handlers'].items():
            if 'filename' in handler:
                handler['filename'] = str(save_dir / handler['filename'])

        logging.config.dictConfig(config)
    else:
        print("Warning: logging configuration file is not found in {}.".format(log_config))
        logging.basicConfig(level=default_level)

EgoNLQ/logger/logger.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import importlib

from utils import Timer


class TensorboardWriter:
    def __init__(self, log_dir, logger, enabled):
        self.writer = None
        self.selected_module = ""

        if enabled:
            log_dir = str(log_dir)

            # Retrieve visualization writer.
            for module in ["torch.utils.tensorboard", "tensorboardX"]:
                try:
                    self.writer = importlib.import_module(module).SummaryWriter(log_dir)
                    succeeded = True
                    break
                except ImportError:
                    succeeded = False
                self.selected_module = module

            if not succeeded:
                message = "Warning: visualization (Tensorboard) is configured to use, but currently not installed on " \
                    "this machine. Please install either TensorboardX with 'pip install tensorboardx', upgrade " \
                    "PyTorch to version >= 1.1 for using 'torch.utils.tensorboard' or turn off the option in " \
                    "the 'config.json' file."
                logger.warning(message)

        self.step = 0
        self.mode = ''

        self.tb_writer_ftns = {
            'add_scalar', 'add_scalars', 'add_image', 'add_images', 'add_audio',
            'add_text', 'add_histogram', 'add_pr_curve', 'add_embedding'
        }
        self.tag_mode_exceptions = {'add_histogram', 'add_embedding'}
            
        self.timer = Timer()

    def set_step(self, step, mode='train'):
        self.mode = mode
        self.step = step
        if step == 0:
            self.timer.reset()
        else:
            duration = self.timer.check()
            self.add_scalar('steps_per_sec', 1 / duration)

    def __getattr__(self, name):
        """
        If visualization is configured to use:
            return add_data() methods of tensorboard with additional information (step, tag) added.
        Otherwise:
            return a blank function handle that does nothing
        """
        if name in self.tb_writer_ftns:
            add_data = getattr(self.writer, name, None)

            def wrapper(tag, data, *args, **kwargs):
                if add_data is not None:
                    # add mode(train/valid) tag
                    if name not in self.tag_mode_exceptions:
                        tag = '{}/{}'.format(tag, self.mode)
                    add_data(tag, data, self.step, *args, **kwargs)
            return wrapper
        else:
            # default action for returning methods defined in this class, set_step() for instance.
            try:
                attr = object.__getattr__(name)
            except AttributeError:
                raise AttributeError("type object '{}' has no attribute '{}'".format(self.selected_module, name))
            return attr


class SacredNeptuneWriter:
    def __init__(self):
        raise NotImplementedError

EgoNLQ/logger/visualization.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


"""Main script to train/test models for Ego4D NLQ dataset.
"""
import argparse
import os
import pdb

import numpy as np
import options
import torch
import torch.nn as nn
from model.VSLNet import build_optimizer_and_scheduler, VSLNet
from tqdm.auto import tqdm
from utils.data_gen import gen_or_load_dataset, visual_feature_sampling
from utils.data_util import load_json, load_video_features, save_json
from utils.runner_utils import (
    convert_length_to_mask,
    eval_test,
    filter_checkpoints,
    get_last_checkpoint,
    set_th_config,
)
import json
import tqdm
from glob import glob
import torch
import torch.utils.data
from model.model import FrozenInTime
from utils.data_loader import get_test_loader, get_train_loader



class Dataset(torch.utils.data.Dataset):
    def __init__(self, dataset):
        super(Dataset, self).__init__()
        self.dataset = dataset

    def __getitem__(self, index):
        record = self.dataset[index]
        video_feature = torch.load(record["video_frame_path"])
        #video_feature = visual_feature_sampling(video_feature, max_num_clips=self.args.max_pos_len)
        text_tokens = record["text_tokens"]

        return video_feature, text_tokens, record["sample_id"], record["vid"], record["s_time"], record["e_time"], record["duration"], record["query"], record["s_ind"], record["e_ind"], record["v_len"], record["annotation_uid"], record["query_idx"]
        

    def __len__(self):
        return len(self.dataset)



def fused_feature_extract(args, dataset, feat_path, split='train_set'):
    
    train_set = Dataset(dataset = dataset[split])
    train_loader = torch.utils.data.DataLoader(
        dataset=train_set,
        batch_size=1,
        shuffle=False,
    )

    f = open('configs/nlq.json')
    config = json.load(f)

    video_params = {"model": config['arch']['args']['video_params']['model'], "arch_config": config['arch']['args']['video_params']['arch_config'], "num_frames": config['arch']['args']['video_params']['num_frames'], "pretrained": True, "time_init":  config['arch']['args']['video_params']['time_init']}
    text_params = {"model": config['arch']['args']['text_params']['model'], "pretrained": True, "input": config['arch']['args']['text_params']['input']}
    projection_dim=config['arch']['args']["projection_dim"]
    load_checkpoint=args.model_name
    projection='minimal'
    load_temporal_fix='bilinear'
    task_names = 'EgoNCE_ITM_MLM'
    norm_layer = None
    embed_dim=768

    model = FrozenInTime(video_params, text_params, projection_dim=projection_dim, load_checkpoint=load_checkpoint,
                            projection=projection, load_temporal_fix=load_temporal_fix,
                            task_names = task_names, norm_layer = norm_layer, embed_dim=embed_dim)

    device = torch.device(args.cuda_base if torch.cuda.is_available() else 'cpu')
    model = model.to(device)
    model = nn.DataParallel(model, device_ids = args.device_ids)
    model.eval()

    feat_path = os.path.join(feat_path, split.split('_')[0])
    os.mkdir(feat_path)

    with torch.no_grad():
        for idx, batch in enumerate(train_loader):
        
            data = {"video": batch[0][0].to(device), "text": {"input_ids": batch[1]["input_ids"][0].to(device), "attention_mask": batch[1]["attention_mask"][0].to(device)}}
            #fused_video_features = model.forward(data)
            #fused_video_features = model.module.compute_video(data['video'])

            projection_dim = 768
            fused_video_features = torch.zeros(data['video'].shape[0], projection_dim)
            b_s = 256 ## Batch size inside the next loop, as (897, 4, 3, 224, 224) goes out of memory
            
            if data['video'].shape[0] % b_s == 0:
                times = data['video'].shape[0] // b_s 
            else:
                times = 1 + data['video'].shape[0] // b_s  ##  '1 +' is added by us
            
            for j in range(times):
                start = j*b_s
                if (j+1) * b_s > data['video'].shape[0]:
                    end = data['video'].shape[0]
                else:
                    end = (j+1)*b_s

                video_length = end - start

                #fused_video_features[start:end,] = model.module.compute_video(data['video'][start:end,])
                data_batch = {"video": data['video'][start:end,], "text": {"input_ids": data['text']['input_ids'][0:video_length,], "attention_mask": data['text']['attention_mask'][0:video_length,]}}
                fused_video_features[start:end,] = model.forward(data_batch)
            

            #fused_video_features = visual_feature_sampling(fused_video_features, max_num_clips=args.max_pos_len)
            
            dual_text_features = model.module.compute_text_tokens({"input_ids": batch[1]["input_ids"][0][0:1].to(device), "attention_mask": batch[1]["attention_mask"][0][0:1].to(device)}, is_proj=False)

            #saved_feature_dict = {"fused_video_features": fused_video_features.detach().cpu(), "dual_text_features": dual_text_features.detach().cpu(),
            #                      "sample_id": batch[2][0], "vid": batch[3][0], "s_time": batch[4][0], "e_time": batch[5][0],
            #                      "duration": batch[6][0], "query": batch[7][0], "s_ind": batch[8][0], "e_ind": batch[9][0],
            #                      "v_len": batch[10][0], "annotation_uid": batch[11][0], "query_idx": batch[12][0], "text_token": {"input_ids": batch[1]["input_ids"][0][0:1].to(device), "attention_mask": batch[1]["attention_mask"][0][0:1].to(device)}}  

            saved_feature_dict = {"fused_video_features": fused_video_features.detach().cpu(), "dual_text_features": dual_text_features.detach().cpu(),
                                  "sample_id": batch[2][0], "vid": batch[3][0], "s_time": batch[4][0], "e_time": batch[5][0],
                                  "duration": batch[6][0], "query": batch[7][0], "s_ind": batch[8][0], "e_ind": batch[9][0],
                                  "v_len": batch[10][0], "annotation_uid": batch[11][0], "query_idx": batch[12][0], "text_token": {"attention_mask": batch[1]["attention_mask"][0][0:1]}}

            torch.save(saved_feature_dict, os.path.join(feat_path, 'feat_' + str(idx) + '.pt'))


def main(args, parser):

    # set tensorflow configs
    # set_th_config(args.seed)

    feat_path = os.path.join(str(args.feature_dir), str(args.model_name).split('/')[-1].split('.')[0])

    if not os.path.exists(feat_path):
        os.mkdir(feat_path)
        print("Feature is being extracted")
        # prepare or load dataset
        dataset = gen_or_load_dataset(args)
        args.char_size = dataset.get("n_chars", -1)
        args.word_size = dataset.get("n_words", -1)
        fused_feature_extract(args, dataset, feat_path, split='train_set')
        fused_feature_extract(args, dataset, feat_path, split='val_set')
        fused_feature_extract(args, dataset, feat_path, split='test_set')


    train_loader = get_train_loader(
        glob(os.path.join(feat_path, "train") + "/*.pt"), args=args
    )
    train_eval_loader = get_test_loader(
        glob(os.path.join(feat_path, "train") + "/*.pt"), args=args
    )
    val_loader = (get_test_loader(glob(os.path.join(feat_path, "val") + "/*.pt"), args=args))

    test_loader = get_test_loader(
        glob(os.path.join(feat_path, "train") + "/*.pt"), args=args
    )

    args.num_train_steps = len(train_loader) * args.epochs
    num_train_batches = len(train_loader)
    num_val_batches = 0 if val_loader is None else len(val_loader)
    num_test_batches = len(test_loader)

    # Device configuration
    device = torch.device(args.cuda_base if torch.cuda.is_available() else 'cpu')

    # create model dir
    #home_dir = os.path.join(
    #    args.model_dir,
    #    "_".join(
    #        [
    #           args.model_name,
    #            args.task,
    #            args.fv,
    #            str(args.max_pos_len),
    #            args.predictor,
    #        ]
    #    ),
    #)
    #if args.suffix is not None:
    #    home_dir = home_dir + "_" + args.suffix
    #model_dir = os.path.join(home_dir, "model")

    # train and test

    if args.mode.lower() == "train":
        eval_period = num_train_batches // 2
        #pdb.set_trace()
        # build model
        model = VSLNet(configs=args).to(device)

        optimizer, scheduler = build_optimizer_and_scheduler(model, configs=args)
        # start training
        best_metric = -1.0

        new_save_dir = './saved_nlq_results/' + str(args.model_name).split('/')[-1].split('.')[0]
        if not os.path.exists(new_save_dir):
            os.mkdir(new_save_dir)
        score_writer = open(
            os.path.join(new_save_dir, "eval_results_" + str(args.max_pos_len) + "_" + str(args.batch_size) + "_" + str(args.init_lr) + ".txt"), mode="w", encoding="utf-8"
        )
        print("start training...", flush=True)
        
        global_step = 0
        best_r1_iou03 = 0

        for epoch in range(args.epochs):
            model.train()
            for _, data in enumerate(train_loader):
                global_step += 1
                (
                    _,
                    vfeats,
                    vfeat_lens,
                    dual_text_feature,
                    s_labels,
                    e_labels,
                    h_labels,
                    query_attention_mask
                ) = data
                # prepare features
                vfeats, vfeat_lens = vfeats.to(device), vfeat_lens.to(device)

                #pdb.set_trace()
                dual_text_feature = dual_text_feature.to(device)
                s_labels, e_labels, h_labels = (
                    s_labels.to(device),
                    e_labels.to(device),
                    h_labels.to(device),
                )
                query_attention_mask = query_attention_mask.to(device)

                if args.predictor == 'EgoVLP':
                    # generate mask
                    query_mask = query_attention_mask    

                else:
                    raise NotImplementedError("predictor should be EgoVLP")
                    
                # generate mask
                video_mask = convert_length_to_mask(vfeat_lens).to(device)
                # compute logits
                h_score, start_logits, end_logits = model(
                    vfeats, video_mask, query_mask, dual_text_feature
                )

                # compute loss
                highlight_loss = model.compute_highlight_loss(
                    h_score, h_labels, video_mask
                )
                loc_loss = model.compute_loss(
                    start_logits, end_logits, s_labels, e_labels
                )
                total_loss = loc_loss + args.highlight_lambda * highlight_loss
                
                #print(total_loss)
                #total_loss = loc_loss
                # compute and apply gradients
                optimizer.zero_grad()
                total_loss.backward()
                nn.utils.clip_grad_norm_(
                    model.parameters(), args.clip_norm
                )  # clip gradient
                
                optimizer.step()
                scheduler.step()
     
                # evaluate
                if (
                    global_step % eval_period == 0
                    or global_step % num_train_batches == 0
                ):
                    model.eval()
                    print(
                        f"\nEpoch: {epoch + 1:2d} | Step: {global_step:5d}", flush=True
                    )
                    #result_save_path = os.path.join(
                    #    model_dir,
                    #    f"{args.model_name}_{epoch}_{global_step}_preds.json",
                    #)
                    # Evaluate on val, keep the top 3 checkpoints.
                    results, mIoU, score_str = eval_test(
                        model=model,
                        data_loader=val_loader,
                        device=device,
                        mode="val",
                        epoch=epoch + 1,
                        global_step=global_step,
                        gt_json_path=args.eval_gt_json,
                        configs=args,
                    )

                    if best_r1_iou03 < results[0][0]:
                        best_r1_iou03 = results[0][0]
                        best_score_str = score_str
                        epoch_step = "Best Epoch: " + str(epoch+1) + ", Best Step: " + str(global_step)

                    print(score_str, flush=True)
                    score_writer.write(score_str)
                    score_writer.flush()
                    # Recall@1, 0.3 IoU overlap --> best metric.
                    if results[0][0] >= best_metric:
                        best_metric = results[0][0]
                    #    torch.save(
                    #        model.state_dict(),
                    #        os.path.join(
                    #            model_dir,
                    #            "{}_{}.t7".format(configs.model_name, global_step),
                    #        ),
                    #    )
                    #    # only keep the top-3 model checkpoints
                    #    filter_checkpoints(model_dir, suffix="t7", max_to_keep=3)
                    model.train()

        print("=" * 50)
        print(epoch_step + "\n" + best_score_str, flush=True)
        score_writer.write(epoch_step + "\n" + best_score_str)
        score_writer.flush()
        score_writer.close()

    
    '''

    elif configs.mode.lower() == "test":
        if not os.path.exists(model_dir):
            raise ValueError("No pre-trained weights exist")
        # load previous configs
        pre_configs = load_json(os.path.join(model_dir, "configs.json"))
        parser.set_defaults(**pre_configs)
        configs = parser.parse_args()
        # build model
        model = VSLNet(
            configs=configs, word_vectors=dataset.get("word_vector", None)
        ).to(device)

        # get last checkpoint file
        filename = get_last_checkpoint(model_dir, suffix="t7")
        model.load_state_dict(torch.load(filename))
        model.eval()
        result_save_path = filename.replace(".t7", "_test_result.json")
        results, mIoU, score_str = eval_test(
            model=model,
            data_loader=test_loader,
            device=device,
            mode="test",
            result_save_path=result_save_path,
            configs=configs,
        )
        print(score_str, flush=True)
    '''

if __name__ == "__main__":
    configs, parser = options.read_command_line()
    main(configs, parser)

EgoNLQ/main.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


"""VSLNet Baseline for Ego4D Episodic Memory -- Natural Language Queries.
"""
import torch
import torch.nn as nn
from transformers import AdamW, get_linear_schedule_with_warmup

from model.layers import (
    Embedding,
    VisualProjection,
    FeatureEncoder,
    CQAttention,
    CQConcatenate,
    ConditionedPredictor,
    HighLightLayer,
    BertEmbedding,
)


def build_optimizer_and_scheduler(model, configs):
    no_decay = [
        "bias",
        "layer_norm",
        "LayerNorm",
    ]  # no decay for parameters of layer norm and bias
    optimizer_grouped_parameters = [
        {
            "params": [
                p
                for n, p in model.named_parameters()
                if not any(nd in n for nd in no_decay)
            ],
            "weight_decay": 0.01,
        },
        {
            "params": [
                p
                for n, p in model.named_parameters()
                if any(nd in n for nd in no_decay)
            ],
            "weight_decay": 0.0,
        },
    ]
    optimizer = AdamW(optimizer_grouped_parameters, lr=configs.init_lr)
    scheduler = get_linear_schedule_with_warmup(
        optimizer,
        configs.num_train_steps * configs.warmup_proportion,
        configs.num_train_steps,
    )
    return optimizer, scheduler


class VSLNet(nn.Module):
    def __init__(self, configs):
        super(VSLNet, self).__init__()
        self.configs = configs
        self.video_affine = VisualProjection(
            visual_dim=configs.video_feature_dim,
            dim=configs.dim,
            drop_rate=configs.drop_rate,
        )
        self.feature_encoder = FeatureEncoder(
            dim=configs.dim,
            num_heads=configs.num_heads,
            kernel_size=7,
            num_layers=4,
            max_pos_len=configs.max_pos_len,
            drop_rate=configs.drop_rate,
        )
        # video and query fusion
        self.cq_attention = CQAttention(dim=configs.dim, drop_rate=configs.drop_rate)
        self.cq_concat = CQConcatenate(dim=configs.dim)
        # query-guided highlighting
        self.highlight_layer = HighLightLayer(dim=configs.dim)
        # conditioned predictor
        self.predictor = ConditionedPredictor(
            dim=configs.dim,
            num_heads=configs.num_heads,
            drop_rate=configs.drop_rate,
            max_pos_len=configs.max_pos_len,
            predictor=configs.predictor,
        )

        # If pretrained transformer, initialize_parameters and load.
        if self.configs.predictor == 'EgoVLP':
            # Project back from BERT to dim.
            self.query_affine = nn.Linear(768, configs.dim)
            self.init_parameters()
        else:
            raise NotImplementedError("Preditor should be EgoVLP")

    def init_parameters(self):
        def init_weights(m):
            if (
                isinstance(m, nn.Conv2d)
                or isinstance(m, nn.Conv1d)
                or isinstance(m, nn.Linear)
            ):
                torch.nn.init.xavier_uniform_(m.weight)
                if m.bias is not None:
                    torch.nn.init.zeros_(m.bias)
            elif isinstance(m, nn.LSTM):
                m.reset_parameters()

        self.apply(init_weights)

    def forward(self, video_features, v_mask, q_mask, dual_text_feature):
        video_features = self.video_affine(video_features)
        if self.configs.predictor == 'EgoVLP':
            query_features = self.query_affine(dual_text_feature)
        else:
            raise NotImplementedError("Preditor should be EgoVLP")

        query_features = self.feature_encoder(query_features, mask=q_mask)
        video_features = self.feature_encoder(video_features, mask=v_mask)
        features = self.cq_attention(video_features, query_features, v_mask, q_mask)
        features = self.cq_concat(features, query_features, q_mask)
        h_score = self.highlight_layer(features, v_mask)
        features = features * h_score.unsqueeze(2)
        start_logits, end_logits = self.predictor(features, mask=v_mask)
        return h_score, start_logits, end_logits

    def extract_index(self, start_logits, end_logits):
        return self.predictor.extract_index(
            start_logits=start_logits, end_logits=end_logits
        )

    def compute_highlight_loss(self, scores, labels, mask):
        return self.highlight_layer.compute_loss(
            scores=scores, labels=labels, mask=mask
        )

    def compute_loss(self, start_logits, end_logits, start_labels, end_labels):
        return self.predictor.compute_cross_entropy_loss(
            start_logits=start_logits,
            end_logits=end_logits,
            start_labels=start_labels,
            end_labels=end_labels,
        )

EgoNLQ/model/VSLNet.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

"""VSLNet Baseline for Ego4D Episodic Memory -- Natural Language Queries.
"""
import torch
import torch.nn as nn
from transformers import AdamW, get_linear_schedule_with_warmup
import pdb 

from model.layers import (
    Embedding,
    VisualProjection,
    FeatureEncoder,
    CQAttention,
    CQConcatenate,
    ConditionedPredictor,
    HighLightLayer,
    BertEmbedding,
)


def build_optimizer_and_scheduler(model, configs):
    no_decay = [
        "bias",
        "layer_norm",
        "LayerNorm",
    ]  # no decay for parameters of layer norm and bias
    optimizer_grouped_parameters = [
        {
            "params": [
                p
                for n, p in model.named_parameters()
                if not any(nd in n for nd in no_decay)
            ],
            "weight_decay": 0.01,
        },
        {
            "params": [
                p
                for n, p in model.named_parameters()
                if any(nd in n for nd in no_decay)
            ],
            "weight_decay": 0.0,
        },
    ]
    optimizer = AdamW(optimizer_grouped_parameters, lr=configs.init_lr)
    scheduler = get_linear_schedule_with_warmup(
        optimizer,
        configs.num_train_steps * configs.warmup_proportion,
        configs.num_train_steps,
    )
    return optimizer, scheduler


class VSLNet(nn.Module):
    def __init__(self, configs):
        super(VSLNet, self).__init__()
        
        self.configs = configs
        self.video_affine = VisualProjection(
            visual_dim=configs.video_feature_dim,
            dim=configs.dim,
            drop_rate=configs.drop_rate,
        )

        self.feature_encoder = FeatureEncoder(
                    dim=configs.dim,
                    num_heads=configs.num_heads,
                    kernel_size=7,
                    num_layers=4,
                    max_pos_len=configs.max_pos_len,
                    drop_rate=configs.drop_rate,
                )

        # video and query fusion
        self.cq_concat = CQConcatenate(dim=configs.dim)
        # query-guided highlighting
        self.highlight_layer = HighLightLayer(dim=configs.dim)
        # conditioned predictor
        self.predictor = ConditionedPredictor(
            dim=configs.dim,
            num_heads=configs.num_heads,
            drop_rate=configs.drop_rate,
            max_pos_len=configs.max_pos_len,
            predictor=configs.predictor,
        )

        # If pretrained transformer, initialize_parameters and load.
        if self.configs.predictor == 'EgoVLP':
            # Project back from BERT to dim.
            self.query_affine = nn.Linear(768, configs.dim)
            self.init_parameters()
        else:
            raise NotImplementedError("predictor should be EgoVLP")


    def init_parameters(self):
        def init_weights(m):
            if (
                isinstance(m, nn.Conv2d)
                or isinstance(m, nn.Conv1d)
                or isinstance(m, nn.Linear)
            ):
                torch.nn.init.xavier_uniform_(m.weight)
                if m.bias is not None:
                    torch.nn.init.zeros_(m.bias)
            elif isinstance(m, nn.LSTM):
                m.reset_parameters()

        self.apply(init_weights)

    def forward(self, video_features, v_mask, q_mask, dual_text_feature):
        
        video_features = self.video_affine(video_features)
        
        if self.configs.predictor == 'EgoVLP':
            query_features = self.query_affine(dual_text_feature)
        else:
            raise NotImplementedError("predictor should be EgoVLP")

        query_features = self.feature_encoder(query_features, mask=q_mask)
        video_features = self.feature_encoder(video_features, mask=v_mask)
        
        features = self.cq_concat(video_features, query_features, q_mask)
        h_score = self.highlight_layer(features, v_mask)
        #h_score = 0
        features = features * h_score.unsqueeze(2)
        start_logits, end_logits = self.predictor(features, mask=v_mask)
        start_logits, end_logits = self.predictor(video_features, mask=v_mask)
        return h_score, start_logits, end_logits

    def extract_index(self, start_logits, end_logits):
        return self.predictor.extract_index(
            start_logits=start_logits, end_logits=end_logits
        )

    def compute_highlight_loss(self, scores, labels, mask):
        return self.highlight_layer.compute_loss(
            scores=scores, labels=labels, mask=mask
        )

    def compute_loss(self, start_logits, end_logits, start_labels, end_labels):
        return self.predictor.compute_cross_entropy_loss(
            start_logits=start_logits,
            end_logits=end_logits,
            start_labels=start_labels,
            end_labels=end_labels,
        )

EgoNLQ/model/VSLNet_small.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import torch
import torch.nn as nn
import torch.nn.functional as F

from transformers.models.bert.modeling_bert import BertPredictionHeadTransform


class Pooler(nn.Module):
    def __init__(self, hidden_size):
        super().__init__()
        self.dense = nn.Linear(hidden_size, hidden_size)
        self.activation = nn.Tanh()

    def forward(self, hidden_states):
        first_token_tensor = hidden_states #[:, 0]
        pooled_output = self.dense(first_token_tensor)
        pooled_output = self.activation(pooled_output)
        return pooled_output


class ITMHead(nn.Module):
    def __init__(self, hidden_size):
        super().__init__()
        self.fc = nn.Linear(hidden_size, 2)

    def forward(self, x):
        x = self.fc(x)
        return x


class MLMHead(nn.Module):
    def __init__(self, config, weight=None):
        super().__init__()
        self.transform = BertPredictionHeadTransform(config)
        self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.bias = nn.Parameter(torch.zeros(config.vocab_size))
        if weight is not None:
            self.decoder.weight = weight

    def forward(self, x):
        x = self.transform(x)
        x = self.decoder(x) + self.bias
        return x

EgoNLQ/model/heads.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


"""
Layers to construct the VSLNet model.
"""
import math

import torch
import torch.nn as nn
import torch.nn.functional as F


def mask_logits(inputs, mask, mask_value=-1e30):
    mask = mask.type(torch.float32)
    return inputs + (1.0 - mask) * mask_value


class Conv1D(nn.Module):
    def __init__(self, in_dim, out_dim, kernel_size=1, stride=1, padding=0, bias=True):
        super(Conv1D, self).__init__()
        self.conv1d = nn.Conv1d(
            in_channels=in_dim,
            out_channels=out_dim,
            kernel_size=kernel_size,
            padding=padding,
            stride=stride,
            bias=bias,
        )

    def forward(self, x):
        # suppose all the input with shape (batch_size, seq_len, dim)
        x = x.transpose(1, 2)  # (batch_size, dim, seq_len)
        x = self.conv1d(x)
        return x.transpose(1, 2)  # (batch_size, seq_len, dim)


class WordEmbedding(nn.Module):
    def __init__(self, num_words, word_dim, drop_rate, word_vectors=None):
        super(WordEmbedding, self).__init__()
        self.is_pretrained = False if word_vectors is None else True
        if self.is_pretrained:
            self.pad_vec = nn.Parameter(
                torch.zeros(size=(1, word_dim), dtype=torch.float32),
                requires_grad=False,
            )
            unk_vec = torch.empty(
                size=(1, word_dim), requires_grad=True, dtype=torch.float32
            )
            nn.init.xavier_uniform_(unk_vec)
            self.unk_vec = nn.Parameter(unk_vec, requires_grad=True)
            self.glove_vec = nn.Parameter(
                torch.tensor(word_vectors, dtype=torch.float32), requires_grad=False
            )
        else:
            self.word_emb = nn.Embedding(num_words, word_dim, padding_idx=0)
        self.dropout = nn.Dropout(p=drop_rate)

    def forward(self, word_ids):
        if self.is_pretrained:
            word_emb = F.embedding(
                word_ids,
                torch.cat([self.pad_vec, self.unk_vec, self.glove_vec], dim=0),
                padding_idx=0,
            )
        else:
            word_emb = self.word_emb(word_ids)
        return self.dropout(word_emb)


class CharacterEmbedding(nn.Module):
    def __init__(self, num_chars, char_dim, drop_rate):
        super(CharacterEmbedding, self).__init__()
        self.char_emb = nn.Embedding(num_chars, char_dim, padding_idx=0)
        kernels, channels = [1, 2, 3, 4], [10, 20, 30, 40]
        self.char_convs = nn.ModuleList(
            [
                nn.Sequential(
                    nn.Conv2d(
                        in_channels=char_dim,
                        out_channels=channel,
                        kernel_size=(1, kernel),
                        stride=(1, 1),
                        padding=0,
                        bias=True,
                    ),
                    nn.ReLU(),
                )
                for kernel, channel in zip(kernels, channels)
            ]
        )
        self.dropout = nn.Dropout(p=drop_rate)

    def forward(self, char_ids):
        char_emb = self.char_emb(
            char_ids
        )  # (batch_size, w_seq_len, c_seq_len, char_dim)
        char_emb = self.dropout(char_emb)
        char_emb = char_emb.permute(
            0, 3, 1, 2
        )  # (batch_size, char_dim, w_seq_len, c_seq_len)
        char_outputs = []
        for conv_layer in self.char_convs:
            output = conv_layer(char_emb)
            output, _ = torch.max(
                output, dim=3, keepdim=False
            )  # reduce max (batch_size, channel, w_seq_len)
            char_outputs.append(output)
        char_output = torch.cat(
            char_outputs, dim=1
        )  # (batch_size, sum(channels), w_seq_len)
        return char_output.permute(0, 2, 1)  # (batch_size, w_seq_len, sum(channels))


class Embedding(nn.Module):
    def __init__(
        self,
        num_words,
        num_chars,
        word_dim,
        char_dim,
        drop_rate,
        out_dim,
        word_vectors=None,
    ):
        super(Embedding, self).__init__()
        self.word_emb = WordEmbedding(
            num_words, word_dim, drop_rate, word_vectors=word_vectors
        )
        self.char_emb = CharacterEmbedding(num_chars, char_dim, drop_rate)
        # output linear layer
        self.linear = Conv1D(
            in_dim=word_dim + 100,
            out_dim=out_dim,
            kernel_size=1,
            stride=1,
            padding=0,
            bias=True,
        )

    def forward(self, word_ids, char_ids):
        word_emb = self.word_emb(word_ids)  # (batch_size, w_seq_len, word_dim)
        char_emb = self.char_emb(char_ids)  # (batch_size, w_seq_len, 100)
        emb = torch.cat(
            [word_emb, char_emb], dim=2
        )  # (batch_size, w_seq_len, word_dim + 100)
        emb = self.linear(emb)  # (batch_size, w_seq_len, dim)
        return emb


class BertEmbedding(nn.Module):
    def __init__(self, text_agnostic=False):
        super(BertEmbedding, self).__init__()
        from transformers import BertModel, BertConfig

        # Load pretrained BERT model if not text_agnostic, else random BERT.
        if text_agnostic:
            self.embedder = BertModel(BertConfig())
        else:
            self.embedder = BertModel.from_pretrained("bert-base-uncased")
        # Freeze the model.
        for param in self.embedder.parameters():
            param.requires_grad = False

    def forward(self, word_ids):
        outputs = self.embedder(**word_ids)
        return outputs["last_hidden_state"].detach()


class PositionalEmbedding(nn.Module):
    """Construct the embeddings from word, position and token_type embeddings."""

    def __init__(self, num_embeddings, embedding_dim):
        super(PositionalEmbedding, self).__init__()
        self.position_embeddings = nn.Embedding(num_embeddings, embedding_dim)

    def forward(self, inputs):
        bsz, seq_length = inputs.shape[:2]
        position_ids = torch.arange(seq_length, dtype=torch.long, device=inputs.device)
        position_ids = position_ids.unsqueeze(0).repeat(bsz, 1)  # (N, L)
        position_embeddings = self.position_embeddings(position_ids)
        return position_embeddings


class VisualProjection(nn.Module):
    def __init__(self, visual_dim, dim, drop_rate=0.0):
        super(VisualProjection, self).__init__()
        self.drop = nn.Dropout(p=drop_rate)
        self.linear = Conv1D(
            in_dim=visual_dim,
            out_dim=dim,
            kernel_size=1,
            stride=1,
            bias=True,
            padding=0,
        )

    def forward(self, visual_features):
        # the input visual feature with shape (batch_size, seq_len, visual_dim)
        visual_features = self.drop(visual_features)
        output = self.linear(visual_features)  # (batch_size, seq_len, dim)
        return output


class DepthwiseSeparableConvBlock(nn.Module):
    def __init__(self, dim, kernel_size, drop_rate, num_layers=4):
        super(DepthwiseSeparableConvBlock, self).__init__()
        self.depthwise_separable_conv = nn.ModuleList(
            [
                nn.Sequential(
                    nn.Conv1d(
                        in_channels=dim,
                        out_channels=dim,
                        kernel_size=kernel_size,
                        groups=dim,
                        padding=kernel_size // 2,
                        bias=False,
                    ),
                    nn.Conv1d(
                        in_channels=dim,
                        out_channels=dim,
                        kernel_size=1,
                        padding=0,
                        bias=True,
                    ),
                    nn.ReLU(),
                )
                for _ in range(num_layers)
            ]
        )
        self.layer_norms = nn.ModuleList(
            [nn.LayerNorm(dim, eps=1e-6) for _ in range(num_layers)]
        )
        self.dropout = nn.Dropout(p=drop_rate)

    def forward(self, x):
        output = x  # (batch_size, seq_len, dim)
        for idx, conv_layer in enumerate(self.depthwise_separable_conv):
            residual = output
            output = self.layer_norms[idx](output)  # (batch_size, seq_len, dim)
            output = output.transpose(1, 2)  # (batch_size, dim, seq_len)
            output = conv_layer(output)
            output = self.dropout(output)
            output = output.transpose(1, 2) + residual  # (batch_size, seq_len, dim)
        return output


class MultiHeadAttentionBlock(nn.Module):
    def __init__(self, dim, num_heads, drop_rate):
        super(MultiHeadAttentionBlock, self).__init__()
        assert (
            dim % num_heads == 0
        ), "The channels (%d) is not a multiple of attention heads (%d)" % (
            dim,
            num_heads,
        )
        self.head_size, self.num_heads, self.dim = int(dim / num_heads), num_heads, dim
        self.dropout = nn.Dropout(p=drop_rate)
        self.query = Conv1D(
            in_dim=dim, out_dim=dim, kernel_size=1, stride=1, padding=0, bias=True
        )
        self.key = Conv1D(
            in_dim=dim, out_dim=dim, kernel_size=1, stride=1, padding=0, bias=True
        )
        self.value = Conv1D(
            in_dim=dim, out_dim=dim, kernel_size=1, stride=1, padding=0, bias=True
        )
        self.layer_norm1 = nn.LayerNorm(dim, eps=1e-6)
        self.layer_norm2 = nn.LayerNorm(dim, eps=1e-6)
        self.out_layer = Conv1D(
            in_dim=dim, out_dim=dim, kernel_size=1, stride=1, padding=0, bias=True
        )

    def transpose_for_scores(self, x):
        new_x_shape = x.size()[:-1] + (self.num_heads, self.head_size)
        x = x.view(*new_x_shape)
        return x.permute(0, 2, 1, 3)  # (batch_size, num_heads, w_seq_len, head_size)

    @staticmethod
    def combine_last_two_dim(x):
        old_shape = list(x.size())
        new_shape = old_shape[:-2] + [old_shape[-2] * old_shape[-1]]
        return x.reshape(shape=new_shape)

    def forward(self, x, mask=None):
        output = self.layer_norm1(x)  # (batch_size, seq_len, dim)
        output = self.dropout(output)
        # multi-head attention layer
        query = self.transpose_for_scores(
            self.query(output)
        )  # (batch_size, num_heads, seq_len, head_size)
        key = self.transpose_for_scores(self.key(output))
        value = self.transpose_for_scores(self.value(output))
        attention_scores = torch.matmul(
            query, key.transpose(-1, -2)
        )  # (batch_size, num_heads, seq_len, seq_len)
        attention_scores = attention_scores / math.sqrt(self.head_size)
        if mask is not None:  # masking
            mask = mask.unsqueeze(1).unsqueeze(2)  # (batch_size, 1, 1, seq_len)
            attention_scores = mask_logits(attention_scores, mask)
        attention_probs = nn.Softmax(dim=-1)(
            attention_scores
        )  # (batch_size, num_heads, seq_len, seq_len)
        attention_probs = self.dropout(attention_probs)
        value = torch.matmul(
            attention_probs, value
        )  # (batch_size, num_heads, seq_len, head_size)
        value = self.combine_last_two_dim(
            value.permute(0, 2, 1, 3)
        )  # (batch_size, seq_len, dim)
        # intermediate layer
        output = self.dropout(value)
        residual = output + x
        output = self.layer_norm2(residual)
        output = self.dropout(output)
        output = self.out_layer(output)
        output = self.dropout(output) + residual
        return output


class FeatureEncoder(nn.Module):
    def __init__(
        self, dim, num_heads, max_pos_len, kernel_size=7, num_layers=4, drop_rate=0.0
    ):
        super(FeatureEncoder, self).__init__()
        self.pos_embedding = PositionalEmbedding(
            num_embeddings=max_pos_len, embedding_dim=dim
        )
        self.conv_block = DepthwiseSeparableConvBlock(
            dim=dim, kernel_size=kernel_size, drop_rate=drop_rate, num_layers=num_layers
        )
        self.attention_block = MultiHeadAttentionBlock(
            dim=dim, num_heads=num_heads, drop_rate=drop_rate
        )

    def forward(self, x, mask=None):
        features = x + self.pos_embedding(x)  # (batch_size, seq_len, dim)
        features = self.conv_block(features)  # (batch_size, seq_len, dim)
        features = self.attention_block(
            features, mask=mask
        )  # (batch_size, seq_len, dim)
        return features


class CQAttention(nn.Module):
    def __init__(self, dim, drop_rate=0.0):
        super(CQAttention, self).__init__()
        w4C = torch.empty(dim, 1)
        w4Q = torch.empty(dim, 1)
        w4mlu = torch.empty(1, 1, dim)
        nn.init.xavier_uniform_(w4C)
        nn.init.xavier_uniform_(w4Q)
        nn.init.xavier_uniform_(w4mlu)
        self.w4C = nn.Parameter(w4C, requires_grad=True)
        self.w4Q = nn.Parameter(w4Q, requires_grad=True)
        self.w4mlu = nn.Parameter(w4mlu, requires_grad=True)
        self.dropout = nn.Dropout(p=drop_rate)
        self.cqa_linear = Conv1D(
            in_dim=4 * dim, out_dim=dim, kernel_size=1, stride=1, padding=0, bias=True
        )

    def forward(self, context, query, c_mask, q_mask):
        score = self.trilinear_attention(
            context, query
        )  # (batch_size, c_seq_len, q_seq_len)
        score_ = nn.Softmax(dim=2)(
            mask_logits(score, q_mask.unsqueeze(1))
        )  # (batch_size, c_seq_len, q_seq_len)
        score_t = nn.Softmax(dim=1)(
            mask_logits(score, c_mask.unsqueeze(2))
        )  # (batch_size, c_seq_len, q_seq_len)
        score_t = score_t.transpose(1, 2)  # (batch_size, q_seq_len, c_seq_len)
        c2q = torch.matmul(score_, query)  # (batch_size, c_seq_len, dim)
        q2c = torch.matmul(
            torch.matmul(score_, score_t), context
        )  # (batch_size, c_seq_len, dim)
        output = torch.cat(
            [context, c2q, torch.mul(context, c2q), torch.mul(context, q2c)], dim=2
        )
        output = self.cqa_linear(output)  # (batch_size, c_seq_len, dim)
        return output

    def trilinear_attention(self, context, query):
        batch_size, c_seq_len, dim = context.shape
        batch_size, q_seq_len, dim = query.shape
        context = self.dropout(context)
        query = self.dropout(query)
        subres0 = torch.matmul(context, self.w4C).expand(
            [-1, -1, q_seq_len]
        )  # (batch_size, c_seq_len, q_seq_len)
        subres1 = (
            torch.matmul(query, self.w4Q).transpose(1, 2).expand([-1, c_seq_len, -1])
        )
        subres2 = torch.matmul(context * self.w4mlu, query.transpose(1, 2))
        res = subres0 + subres1 + subres2  # (batch_size, c_seq_len, q_seq_len)
        return res


class WeightedPool(nn.Module):
    def __init__(self, dim):
        super(WeightedPool, self).__init__()
        weight = torch.empty(dim, 1)
        nn.init.xavier_uniform_(weight)
        self.weight = nn.Parameter(weight, requires_grad=True)

    def forward(self, x, mask):
        alpha = torch.tensordot(
            x, self.weight, dims=1
        )  # shape = (batch_size, seq_length, 1)
        alpha = mask_logits(alpha, mask=mask.unsqueeze(2))
        alphas = nn.Softmax(dim=1)(alpha)
        pooled_x = torch.matmul(x.transpose(1, 2), alphas)  # (batch_size, dim, 1)
        pooled_x = pooled_x.squeeze(2)
        return pooled_x


class CQConcatenate(nn.Module):
    def __init__(self, dim):
        super(CQConcatenate, self).__init__()
        self.weighted_pool = WeightedPool(dim=dim)
        self.conv1d = Conv1D(
            in_dim=2 * dim, out_dim=dim, kernel_size=1, stride=1, padding=0, bias=True
        )

    def forward(self, context, query, q_mask):
        pooled_query = self.weighted_pool(query, q_mask)  # (batch_size, dim)
        _, c_seq_len, _ = context.shape
        pooled_query = pooled_query.unsqueeze(1).repeat(
            1, c_seq_len, 1
        )  # (batch_size, c_seq_len, dim)
        output = torch.cat(
            [context, pooled_query], dim=2
        )  # (batch_size, c_seq_len, 2*dim)
        output = self.conv1d(output)
        return output


class HighLightLayer(nn.Module):
    def __init__(self, dim):
        super(HighLightLayer, self).__init__()
        self.conv1d = Conv1D(
            in_dim=dim, out_dim=1, kernel_size=1, stride=1, padding=0, bias=True
        )

    def forward(self, x, mask):
        # compute logits
        logits = self.conv1d(x)
        logits = logits.squeeze(2)
        logits = mask_logits(logits, mask)
        # compute score
        scores = nn.Sigmoid()(logits)
        return scores

    @staticmethod
    def compute_loss(scores, labels, mask, epsilon=1e-12):
        labels = labels.type(torch.float32)
        weights = torch.where(labels == 0.0, labels + 1.0, 2.0 * labels)
        loss_per_location = nn.BCELoss(reduction="none")(scores, labels)
        loss_per_location = loss_per_location * weights
        mask = mask.type(torch.float32)
        loss = torch.sum(loss_per_location * mask) / (torch.sum(mask) + epsilon)
        return loss


class DynamicRNN(nn.Module):
    def __init__(self, dim):
        super(DynamicRNN, self).__init__()
        self.lstm = nn.LSTM(
            input_size=dim,
            hidden_size=dim,
            num_layers=1,
            bias=True,
            batch_first=True,
            bidirectional=False,
        )

    def forward(self, x, mask):
        out, _ = self.lstm(x)  # (bsz, seq_len, dim)
        mask = mask.type(torch.float32)
        mask = mask.unsqueeze(2)
        out = out * mask
        return out


class ConditionedPredictor(nn.Module):
    def __init__(self, dim, num_heads, max_pos_len, drop_rate=0.0, predictor="rnn"):
        super(ConditionedPredictor, self).__init__()
        self.predictor = predictor
        if predictor == "rnn":
            self.start_encoder = DynamicRNN(dim=dim)
            self.end_encoder = DynamicRNN(dim=dim)
        else:
            self.encoder = FeatureEncoder(
                dim=dim,
                num_heads=num_heads,
                kernel_size=7,
                num_layers=4,
                max_pos_len=max_pos_len,
                drop_rate=drop_rate,
            )
            self.start_layer_norm = nn.LayerNorm(dim, eps=1e-6)
            self.end_layer_norm = nn.LayerNorm(dim, eps=1e-6)

        self.start_block = nn.Sequential(
            Conv1D(
                in_dim=2 * dim,
                out_dim=dim,
                kernel_size=1,
                stride=1,
                padding=0,
                bias=True,
            ),
            nn.ReLU(),
            Conv1D(
                in_dim=dim, out_dim=1, kernel_size=1, stride=1, padding=0, bias=True
            ),
        )
        self.end_block = nn.Sequential(
            Conv1D(
                in_dim=2 * dim,
                out_dim=dim,
                kernel_size=1,
                stride=1,
                padding=0,
                bias=True,
            ),
            nn.ReLU(),
            Conv1D(
                in_dim=dim, out_dim=1, kernel_size=1, stride=1, padding=0, bias=True
            ),
        )

    def forward(self, x, mask):
        if self.predictor == "rnn":
            start_features = self.start_encoder(x, mask)  # (batch_size, seq_len, dim)
            end_features = self.end_encoder(start_features, mask)
        else:
            start_features = self.encoder(x, mask)
            end_features = self.encoder(start_features, mask)
            start_features = self.start_layer_norm(start_features)
            end_features = self.end_layer_norm(end_features)
        start_features = self.start_block(
            torch.cat([start_features, x], dim=2)
        )  # (batch_size, seq_len, 1)
        end_features = self.end_block(torch.cat([end_features, x], dim=2))
        start_logits = mask_logits(start_features.squeeze(2), mask=mask)
        end_logits = mask_logits(end_features.squeeze(2), mask=mask)
        return start_logits, end_logits

    @staticmethod
    def extract_index(start_logits, end_logits):
        start_prob = nn.Softmax(dim=1)(start_logits)
        end_prob = nn.Softmax(dim=1)(end_logits)
        outer = torch.matmul(start_prob.unsqueeze(dim=2), end_prob.unsqueeze(dim=1))
        outer = torch.triu(outer, diagonal=0)

        # _, start_index = torch.max(torch.max(outer, dim=2)[0], dim=1)  # (batch_size, )
        # _, end_index = torch.max(torch.max(outer, dim=1)[0], dim=1)  # (batch_size, )
        # return start_index, end_index

        # Get top 5 start and end indices.
        batch_size, height, width = outer.shape
        outer_flat = outer.view(batch_size, -1)
        _, flat_indices = outer_flat.topk(5, dim=-1)
        start_indices = flat_indices // width
        end_indices = flat_indices % width
        return start_indices, end_indices

    @staticmethod
    def compute_cross_entropy_loss(start_logits, end_logits, start_labels, end_labels):
        start_loss = nn.CrossEntropyLoss(reduction="mean")(start_logits, start_labels)
        end_loss = nn.CrossEntropyLoss(reduction="mean")(end_logits, end_labels)
        return start_loss + end_loss

EgoNLQ/model/layers.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import torch

path = '/apdcephfs/private_qinghonglin/video_codebase/frozen-in-time-main/results/EgoClip_M_EgoNCE_N_V_Neg_Seg_60/models/0509_00/checkpoint-epoch3.pth'
# path = '/apdcephfs/private_qinghonglin/video_codebase/frozen-in-time-main/results/EgoClip_SE_1_mid_scale_th_v2/models/0502_01/checkpoint-epoch4.pth'
checkpoint = torch.load(path)
state_dict = checkpoint['state_dict']


EgoNLQ/model/load_checkpoint.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import pdb
import torch
import torch.nn.functional as F
from torch import nn
import pickle

class NormSoftmaxLoss(nn.Module):
    def __init__(self, temperature=0.05):
        super().__init__()

        self.temperature = temperature

    def forward(self, x):
        "Assumes input x is similarity matrix of N x M \in [-1, 1], computed using the cosine similarity between normalised vectors"
        i_logsm = F.log_softmax(x/self.temperature, dim=1)
        j_logsm = F.log_softmax(x.t()/self.temperature, dim=1)

        # sum over positives
        idiag = torch.diag(i_logsm)
        loss_i = idiag.sum() / len(idiag)

        jdiag = torch.diag(j_logsm)
        loss_j = jdiag.sum() / len(jdiag)

        return - loss_i - loss_j


EgoNLQ/model/loss.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


"""Module for computing performance metrics

"""
import math
import numbers
from pathlib import Path
import numpy as np
import torch
import scipy.stats
from sklearn.metrics import average_precision_score
import os
import pandas as pd
import pickle
from utils import nDCG, mAP

def t2v_metrics(sims, query_masks=None):
    """Compute retrieval metrics from a similiarity matrix.

    Args:
        sims (th.Tensor): N x M matrix of similarities between embeddings, where
             x_{i,j} = <text_embd[i], vid_embed[j]>
        query_masks (th.Tensor): mask any missing queries from the dataset (two videos
             in MSRVTT only have 19, rather than 20 captions)

    Returns:
        (dict[str:float]): retrieval metrics
    """
    assert sims.ndim == 2, "expected a matrix"
    num_queries, num_vids = sims.shape
    dists = -sims
    sorted_dists = np.sort(dists, axis=1)

    # The indices are computed such that they slice out the ground truth distances
    # from the psuedo-rectangular dist matrix
    queries_per_video = num_queries // num_vids
    gt_idx = [[np.ravel_multi_index([ii, jj], (num_queries, num_vids))
               for ii in range(jj * queries_per_video, (jj + 1) * queries_per_video)]
              for jj in range(num_vids)]
    gt_idx = np.array(gt_idx)
    gt_dists = dists.reshape(-1)[gt_idx.reshape(-1)]
    gt_dists = gt_dists[:, np.newaxis]
    rows, cols = np.where((sorted_dists - gt_dists) == 0)  # find column position of GT

    # --------------------------------
    # NOTE: Breaking ties
    # --------------------------------
    # We sometimes need to break ties (in general, these should occur extremely rarely,
    # but there are pathological cases when they can distort the scores, such as when
    # the similarity matrix is all zeros). Previous implementations (e.g. the t2i
    # evaluation function used
    # here: https://github.com/niluthpol/multimodal_vtt/blob/master/evaluation.py and
    # here: https://github.com/linxd5/VSE_Pytorch/blob/master/evaluation.py#L87) generally
    # break ties "optimistically".  However, if the similarity matrix is constant this
    # can evaluate to a perfect ranking. A principled option is to average over all
    # possible partial orderings implied by the ties. See # this paper for a discussion:
    #    McSherry, Frank, and Marc Najork,
    #    "Computing information retrieval performance measures efficiently in the presence
    #    of tied scores." European conference on information retrieval. Springer, Berlin,
    #    Heidelberg, 2008.
    # http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.145.8892&rep=rep1&type=pdf

    break_ties = "optimistically"
    #break_ties = "averaging"

    if rows.size > num_queries:
        assert np.unique(rows).size == num_queries, "issue in metric evaluation"
        if break_ties == "optimistically":
            _, idx = np.unique(rows, return_index=True)
            cols = cols[idx]
        elif break_ties == "averaging":
            # fast implementation, based on this code:
            # https://stackoverflow.com/a/49239335
            locs = np.argwhere((sorted_dists - gt_dists) == 0)

            # Find the split indices
            steps = np.diff(locs[:, 0])
            splits = np.nonzero(steps)[0] + 1
            splits = np.insert(splits, 0, 0)

            # Compute the result columns
            summed_cols = np.add.reduceat(locs[:, 1], splits)
            counts = np.diff(np.append(splits, locs.shape[0]))
            avg_cols = summed_cols / counts
            if False:
                print("Running slower code to verify rank averaging across ties")
                # slow, but more interpretable version, used for testing
                avg_cols_slow = [np.mean(cols[rows == idx]) for idx in range(num_queries)]
                assert np.array_equal(avg_cols, avg_cols_slow), "slow vs fast difference"
                print("passed num check")
            cols = avg_cols

    msg = "expected ranks to match queries ({} vs {}) "
    if cols.size != num_queries:
        #import ipdb;
        #ipdb.set_trace()
        pass
    assert cols.size == num_queries, msg

    if False:
        # overload mask to check that we can recover the scores for single-query
        # retrieval
        print("DEBUGGING MODE")
        query_masks = np.zeros_like(query_masks)
        query_masks[:, 0] = 1  # recover single query score

    if query_masks is not None:
        # remove invalid queries
        assert query_masks.size == num_queries, "invalid query mask shape"
        cols = cols[query_masks.reshape(-1).astype(np.bool)]
        assert cols.size == query_masks.sum(), "masking was not applied correctly"
        # update number of queries to account for those that were missing
        num_queries = query_masks.sum()

    if False:
        # sanity check against old logic for square matrices
        gt_dists_old = np.diag(dists)
        gt_dists_old = gt_dists_old[:, np.newaxis]
        _, cols_old = np.where((sorted_dists - gt_dists_old) == 0)
        assert np.array_equal(cols_old, cols), "new metric doesn't match"

    return cols2metrics(cols, num_queries)


def v2t_metrics(sims, query_masks=None):
    """Compute retrieval metrics from a similiarity matrix.

    Args:
        sims (th.Tensor): N x M matrix of similarities between embeddings, where
             x_{i,j} = <text_embd[i], vid_embed[j]>
        query_masks (th.Tensor): mask any missing captions from the dataset

    Returns:
        (dict[str:float]): retrieval metrics

    NOTES: We find the closest "GT caption" in the style of VSE, which corresponds
    to finding the rank of the closest relevant caption in embedding space:
    github.com/ryankiros/visual-semantic-embedding/blob/master/evaluation.py#L52-L56
    """
    # switch axes of text and video
    sims = sims.T

    if False:
        # experiment with toy example
        sims = np.ones((3, 3))
        sims[0, 0] = 2
        sims[1, 1:2] = 2
        sims[2, :] = 2
        query_masks = None

    assert sims.ndim == 2, "expected a matrix"
    num_queries, num_caps = sims.shape
    dists = -sims
    caps_per_video = num_caps // num_queries
    break_ties = "averaging"

    MISSING_VAL = 1E8
    query_ranks = []
    for ii in range(num_queries):
        row_dists = dists[ii, :]
        if query_masks is not None:
            # Set missing queries to have a distance of infinity.  A missing query
            # refers to a query position `n` for a video that had less than `n`
            # captions (for example, a few MSRVTT videos only have 19 queries)
            row_dists[np.logical_not(query_masks.reshape(-1))] = MISSING_VAL

        # NOTE: Using distance subtraction to perform the ranking is easier to make
        # deterministic than using argsort, which suffers from the issue of defining
        # "stability" for equal distances.  Example of distance subtraction code:
        # github.com/antoine77340/Mixture-of-Embedding-Experts/blob/master/train.py
        sorted_dists = np.sort(row_dists)

        min_rank = np.inf
        for jj in range(ii * caps_per_video, (ii + 1) * caps_per_video):
            if row_dists[jj] == MISSING_VAL:
                # skip rankings of missing captions
                continue
            ranks = np.where((sorted_dists - row_dists[jj]) == 0)[0]
            if break_ties == "optimistically":
                rank = ranks[0]
            elif break_ties == "averaging":
                # NOTE: If there is more than one caption per video, its possible for the
                # method to do "worse than chance" in the degenerate case when all
                # similarities are tied.  TODO(Samuel): Address this case.
                rank = ranks.mean()
            if rank < min_rank:
                min_rank = rank
        query_ranks.append(min_rank)
    query_ranks = np.array(query_ranks)

    # sanity check against old version of code
    if False:
        sorted_dists = np.sort(dists, axis=1)
        gt_dists_old = np.diag(dists)
        gt_dists_old = gt_dists_old[:, np.newaxis]
        rows_old, cols_old = np.where((sorted_dists - gt_dists_old) == 0)
        if rows_old.size > num_queries:
            _, idx = np.unique(rows_old, return_index=True)
            cols_old = cols_old[idx]
        num_diffs = (1 - (cols_old == query_ranks)).sum()
        msg = f"new metric doesn't match in {num_diffs} places"
        assert np.array_equal(cols_old, query_ranks), msg

        # visualise the distance matrix
        import sys
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
        sys.path.insert(0, str(Path.home() / "coding/src/zsvision/python"))
        from zsvision.zs_iterm import zs_dispFig  # NOQA
        plt.matshow(dists)
        zs_dispFig()

    return cols2metrics(query_ranks, num_queries)


EgoNLQ/model/metric.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import os
import pdb

import timm
import torch
import yaml
import torch.nn as nn
import numpy as np
import torch.nn.functional as F
from transformers import AutoModel
from einops import rearrange, repeat

from base import BaseModel
from model import video_transformer
from model.video_transformer import SpaceTimeTransformer
from utils.util import state_dict_data_parallel_fix

from model import roberta
from model.roberta import RobertaModel, _prepare_decoder_attention_mask
from model import heads
from transformers import RobertaConfig
from functools import partial
import copy
import torch.distributed as dist

with open('./EgoNCE_MLM_ITM_Config.yml') as f:
    config = yaml.load(f, Loader=yaml.FullLoader)

def init_weights(module):
    if isinstance(module, (nn.Linear, nn.Embedding)):
        module.weight.data.normal_(mean=0.0, std=0.02)
    elif isinstance(module, nn.LayerNorm):
        module.bias.data.zero_()
        module.weight.data.fill_(1.0)

    if isinstance(module, nn.Linear) and module.bias is not None:
        module.bias.data.zero_()


class FrozenInTime(BaseModel):
    def __init__(self,
                 video_params,
                 text_params,
                 projection_dim=4096,
                 load_checkpoint=None,
                 projection='minimal',
                 load_temporal_fix='zeros', #bilinear
                 config = config,
                 task_names = 'EgoNCE_ITM_MLM',
                 norm_layer = None,
                 embed_dim=768):
        super().__init__()

        self.video_params = video_params
        self.text_params = text_params
        self.load_temporal_fix = load_temporal_fix
        self.config = config
        self.task_names = task_names
        if not text_params['pretrained']:
            raise NotImplementedError("Huggingface text models require pretrained init.")

        if self.text_params['model'].startswith('roberta'):
            self.text_model = RobertaModel.from_pretrained("roberta-base")
        self.text_model.train()

        pretrained = video_params['pretrained']
        if video_params['model'] == "SpaceTimeTransformer":
            self.num_frames = video_params['num_frames'] 
            time_init = 'zeros'
            attention_style = 'frozen-in-time'
            arch_config = 'base_patch16_224'
            vit_init = 'imagenet-21k'
            if arch_config == 'base_patch16_224':
                vit_model = torch.load("/cis/home/shraman/works_meta_2022/pre-training/EgoVLP_Fused_HardNegITM_Checkpoint_multinode/frozen-in-time-main/pretrained/jx_vit_base_p16_224-80ecf9dd.pth", map_location="cpu")
                model = SpaceTimeTransformer(num_frames=self.num_frames,
                                            time_init=time_init,
                                            attention_style=attention_style)
            else:
                raise NotImplementedError

            model.head = nn.Identity()
            model.pre_logits = nn.Identity()
            ftr_dim = model.embed_dim
           
            if load_checkpoint in ["", None]:
                vit_checkpoint = vit_model
                new_vit_dict = state_dict_data_parallel_fix(vit_checkpoint, model.state_dict())
                model.load_state_dict(new_vit_dict, strict=False)
            self.video_model = model
        else:
            raise NotImplementedError(f"{video_params['model']} not implemented")

        # for backwards compatibility (old models)
        self.video_model.fc = nn.Identity()

        # Project to a common embedding
        if projection == 'minimal':

            txt_proj = nn.Sequential(
                nn.Linear(self.text_model.config.hidden_size, projection_dim, bias=False),
                nn.ReLU(inplace=True), nn.Linear(projection_dim, projection_dim, bias=True),
                nn.ReLU(inplace=True), nn.Linear(projection_dim, projection_dim, bias=True)
                )

            vid_proj = nn.Sequential(
                nn.Linear(ftr_dim, projection_dim, bias=False),
                nn.ReLU(inplace=True), nn.Linear(projection_dim, projection_dim, bias=True),
                nn.ReLU(inplace=True), nn.Linear(projection_dim, projection_dim, bias=True)
                )
        
        elif projection == '':
            txt_proj = nn.Identity()
            vid_proj = nn.Identity()
        else:
            raise NotImplementedError
        self.txt_proj = txt_proj
        self.vid_proj = vid_proj

        if ('MLM' in self.task_names or 'ITM' in self.task_names):
            # for FIBER-like cross-attention

            bert_config = RobertaConfig(
                vocab_size=self.config["vocab_size"],
                hidden_size=self.config["hidden_size"],
                num_hidden_layers=self.config["num_layers"],
                num_attention_heads=self.config["num_heads"],
                intermediate_size=self.config["hidden_size"] * config["mlp_ratio"],
                #max_position_embeddings=maxlen, [was used in BTGOT script]
                hidden_dropout_prob=self.config["drop_rate"],
                attention_probs_dropout_prob=self.config["drop_rate"],
            )

            self.num_fuse_block=self.config["num_fuse_block"]
            self.num_text_layer=self.config["num_layers"]
            roberta.NUM_FUSE_BLOCK = self.video_model.NUM_FUSE_BLOCK=self.num_fuse_block
            roberta.DIM_IMG=self.config["input_image_embed_size"]
            self.video_model.DIM_TXT=self.config["input_text_embed_size"]

            self.cross_modal_text_transform = nn.Linear(self.config["input_text_embed_size"], self.config["hidden_size"])
            self.cross_modal_text_transform.apply(init_weights)
            self.cross_modal_video_transform = nn.Linear(self.config["input_image_embed_size"], self.config["hidden_size"])
            self.cross_modal_video_transform.apply(init_weights)

            self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))

            self.num_patches = self.video_model.patch_embed.num_patches
            self.patches_per_frame = self.num_patches//self.num_frames
            norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
            self.norm = norm_layer(embed_dim)
            self.pre_logits = nn.Identity()


            self.avgpool = nn.AdaptiveAvgPool1d(1)
            self.cross_modal_video_pooler = heads.Pooler(config["hidden_size"])
            self.cross_modal_video_pooler.apply(init_weights)
            self.cross_modal_text_pooler = heads.Pooler(config["hidden_size"])
            self.cross_modal_text_pooler.apply(init_weights)

            ## einops transformations
            self.einops_from_space = 'b (f n) d'
            self.einops_to_space = '(b f) n d'
            self.einops_from_time = 'b (f n) d'
            self.einops_to_time = '(b n) f d'

        if 'MLM' in self.task_names:
            self.mlm_score = heads.MLMHead(bert_config)
            self.mlm_score.apply(init_weights)

        if 'ITM' in self.task_names:
            self.itm_score = heads.ITMHead(config["hidden_size"] * 2)
            self.itm_score.apply(init_weights)

        if load_checkpoint not in ["", None]:
            checkpoint = torch.load(load_checkpoint, map_location='cpu')
            state_dict = checkpoint['state_dict']
            new_state_dict = state_dict_data_parallel_fix(state_dict, self.state_dict())
            new_state_dict = self._inflate_positional_embeds(new_state_dict)
            self.load_state_dict(new_state_dict, strict=True)
            print("Model is loaded with pre-trained parameters")

    def set_device(self, device):
        self.device = device

    def forward(self, data, video_only=False, return_embeds=True):
        
        text_data = data['text']
        video_data = data['video']


        b, curr_frames, channels, _, _ = video_data.shape
        video_data_itm = self.video_model.patch_embed(video_data)
        video_data_itm = video_data_itm.flatten(2).transpose(2, 1)
        video_data_itm = video_data_itm.reshape(b, -1, self.video_model.patch_embed.embed_dim)

        BF = video_data_itm.shape[0]
        cls_tokens = self.cls_token.expand(BF, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        video_data_itm = torch.cat((cls_tokens, video_data_itm), dim=1)
        # positional embed needs to be tiled for each frame (this does [1,2,3] --> [1,2,3,1,2,3]...)
        cls_embed = self.video_model.pos_embed[:, 0, :].unsqueeze(1)
        tile_pos_embed = self.video_model.pos_embed[:, 1:, :].repeat(1, self.num_frames, 1)
        # temporal embed needs to be repeated within each frame (this does [1,2,3] --> [1,1,1,2,2,2,3,3,3]...)
        tile_temporal_embed = self.video_model.temporal_embed.repeat_interleave(self.patches_per_frame, 1)
        total_pos_embed = tile_pos_embed + tile_temporal_embed
        total_pos_embed = torch.cat([cls_embed, total_pos_embed], dim=1)

        n = self.patches_per_frame
        f = curr_frames

        curr_patches = video_data_itm.shape[1]
        video_data_itm = video_data_itm + total_pos_embed[:, :curr_patches]
        video_data_itm = self.video_model.pos_drop(video_data_itm)

        unfused_blocks = self.num_text_layer - self.num_fuse_block
            

        for blk_i, blk in enumerate(self.video_model.blocks[:unfused_blocks]):
            if config['use_checkpoint']:
                video_data_itm = torch.utils.checkpoint.checkpoint(blk, video_data_itm, self.einops_from_space, self.einops_to_space, self.einops_from_time, self.einops_to_time,
                                n, f)
            else:
                video_data_itm = blk(video_data_itm, self.einops_from_space, self.einops_to_space, self.einops_from_time, self.einops_to_time,
                                time_n=n, space_f=f)
                            
            
        text_embeds = self.text_model.embeddings(input_ids=text_data['input_ids']) # before it was input_ids=text_ids
        device = text_embeds.device
        text_masks = text_data['attention_mask']
        input_shape = text_masks.size()
        extend_text_masks = self.text_model.get_extended_attention_mask(text_masks, input_shape, device)
        for layer_i, layer in enumerate(self.text_model.encoder.layer[:unfused_blocks]):
            if config['use_checkpoint']:
                text_embeds = torch.utils.checkpoint.checkpoint(layer, text_embeds, extend_text_masks)[0]
            else:
                text_embeds = layer(text_embeds, extend_text_masks)[0]


        for blk_i, blk in enumerate(self.video_model.blocks[unfused_blocks:self.num_text_layer]):
            if config['use_checkpoint']:
                    
                fuse_video_data = torch.utils.checkpoint.checkpoint(blk, video_data_itm, self.einops_from_space, self.einops_to_space, self.einops_from_time, self.einops_to_time, 
                                          n, f, text_embeds, extend_text_masks)
                text_embeds = torch.utils.checkpoint.checkpoint(self.text_model.encoder.layer[blk_i + unfused_blocks],
                                          text_embeds, extend_text_masks, None, (video_data_itm), None, None, False, True)[0]
            else:
                fuse_video_data = blk(video_data_itm, self.einops_from_space, self.einops_to_space, self.einops_from_time, self.einops_to_time, 
                                          y=text_embeds, y_mask=extend_text_masks, time_n=n, space_f=f)
                text_embeds = self.text_model.encoder.layer[blk_i + unfused_blocks](text_embeds, extend_text_masks, encoder_hidden_states=(video_data_itm), last_norm=True)[0]
            video_data_itm = fuse_video_data

            
        video_data_itm = self.norm(video_data_itm)[:, 0]
        video_data_itm = self.pre_logits(video_data_itm)

        return video_data_itm

    
    def compute_text(self, text_data):
        if self.text_params['model'].startswith('bert'):
            text_embeddings = self.text_model(text_data['input_ids'], attention_mask=text_data['attention_mask'])[
                'pooler_output']
        elif self.text_params['model'].startswith('distilbert'):
            text_embeddings = self.text_model(**text_data).last_hidden_state[:, 0, :]
        elif self.text_params['model'].startswith('roberta'):
            text_embeddings = self.text_model(**text_data).last_hidden_state[:, 0, :]
        else:
            raise NotImplementedError
        text_embeddings = self.txt_proj(text_embeddings)
        return text_embeddings

    def compute_text_tokens(self, text_data, is_proj=True):
        if self.text_params['model'].startswith('bert'):
            text_embeddings = self.text_model(text_data['input_ids'], attention_mask=text_data['attention_mask'])[
                'pooler_output']    # not implement for bert
        elif self.text_params['model'].startswith('distilbert'):
            text_embeddings = self.text_model(**text_data).last_hidden_state
        elif self.text_params['model'].startswith('roberta'):
            text_embeddings = self.text_model(**text_data).last_hidden_state
        else:
            raise NotImplementedError
        
        if is_proj:
            text_embeddings = self.txt_proj(text_embeddings)
        
        return text_embeddings

    def compute_video(self, video_data):
        video_embeddings = self.video_model(video_data)
        video_embeddings = self.vid_proj(video_embeddings)
        return video_embeddings

    def _inflate_positional_embeds(self, new_state_dict):
        # allow loading of timesformer with fewer num_frames
        curr_keys = list(self.state_dict().keys())
        if 'video_model.temporal_embed' in new_state_dict and 'video_model.temporal_embed' in curr_keys:
            load_temporal_embed = new_state_dict['video_model.temporal_embed']
            load_num_frames = load_temporal_embed.shape[1]
            curr_num_frames = self.video_params['num_frames']
            embed_dim = load_temporal_embed.shape[2]

            if load_num_frames != curr_num_frames:
                if load_num_frames > curr_num_frames:
                    print(f'### loaded {self.video_params["model"]} model has MORE frames than current...'
                          f'### loading weights, filling in the extras via {self.load_temporal_fix}')
                    new_temporal_embed = load_temporal_embed[:, :curr_num_frames, :]
                else:
                    print(f'### loaded {self.video_params["model"]} model has FEWER frames than current...'
                          f'### loading weights, filling in the extras via {self.load_temporal_fix}')
                    if self.load_temporal_fix == 'zeros':
                        new_temporal_embed = torch.zeros([load_temporal_embed.shape[0], curr_num_frames, embed_dim])
                        new_temporal_embed[:, :load_num_frames] = load_temporal_embed
                    elif self.load_temporal_fix in ['interp', 'bilinear']:
                        # interpolate
                        # unsqueeze so pytorch thinks its an image
                        mode = 'nearest'
                        if self.load_temporal_fix == 'bilinear':
                            mode = 'bilinear'
                        load_temporal_embed = load_temporal_embed.unsqueeze(0)
                        new_temporal_embed = F.interpolate(load_temporal_embed,
                                                           (curr_num_frames, embed_dim), mode=mode, align_corners=True).squeeze(0)
                    else:
                        raise NotImplementedError
                new_state_dict['video_model.temporal_embed'] = new_temporal_embed
        # allow loading with smaller spatial patches. assumes custom border crop, to append the
        # border patches to the input sequence
        if 'video_model.pos_embed' in new_state_dict and 'video_model.pos_embed' in curr_keys:
            load_pos_embed = new_state_dict['video_model.pos_embed']
            load_num_patches = load_pos_embed.shape[1]
            curr_pos_embed = self.state_dict()['video_model.pos_embed']
            if load_num_patches != curr_pos_embed.shape[1]:
                raise NotImplementedError(
                    'Loading models with different spatial resolution / patch number not yet implemented, sorry.')

        return new_state_dict

def sim_matrix(a, b, eps=1e-8):
    """
    added eps for numerical stability
    """
    a_n, b_n = a.norm(dim=1)[:, None], b.norm(dim=1)[:, None]
    a_norm = a / torch.max(a_n, eps * torch.ones_like(a_n))
    b_norm = b / torch.max(b_n, eps * torch.ones_like(b_n))
    sim_mt = torch.mm(a_norm, b_norm.transpose(0, 1))
    return sim_mt


def sim_matrix_batch_val(a, b, eps=1e-8):
    """
    added eps for numerical stability
    """
    a_n, b_n = a.norm(dim=-1).unsqueeze(-1), b.norm(dim=-1).unsqueeze(-1)
    a_norm = a / torch.max(a_n, eps * torch.ones_like(a_n))
    b_norm = b / torch.max(b_n, eps * torch.ones_like(b_n))
    sim_mt = torch.bmm(a_norm, b_norm.transpose(1, 2))
    return sim_mt


if __name__ == "__main__":
    video_params = {"model": SpaceTimeTransformer, "arch_config": "base_patch16_224", "num_frames": 4, "pretrained": True, "time_init": "zeros"}
    text_params: {"model": "roberta-base", "pretrained": True, "input": "text"}
    model = FrozenInTime(video_params, text_params)

EgoNLQ/model/model.py
facebookresearch/EgoVLPv2
# coding=utf-8
# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""PyTorch RoBERTa model. """

import yaml

with open('./EgoNCE_MLM_ITM_Config.yml') as f:
    config_yaml = yaml.load(f, Loader=yaml.FullLoader)

NUM_FUSE_BLOCK = config_yaml['num_fuse_block']
DIM_IMG = 768

import math
import inspect
from typing import Callable, List, Optional, Set, Tuple, Union

import torch
import torch.nn as nn
import torch.utils.checkpoint
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from transformers.activations import ACT2FN, gelu
from transformers.file_utils import (
    add_code_sample_docstrings,
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    replace_return_docstrings,
)
from transformers.modeling_outputs import (
    BaseModelOutputWithPastAndCrossAttentions,
    BaseModelOutputWithPoolingAndCrossAttentions,
    CausalLMOutputWithCrossAttentions,
    MaskedLMOutput,
    MultipleChoiceModelOutput,
    QuestionAnsweringModelOutput,
    SequenceClassifierOutput,
    TokenClassifierOutput,
)
from transformers.modeling_utils import (
    PreTrainedModel,
    find_pruneable_heads_and_indices,
    prune_linear_layer,
)
from transformers.utils import logging
from transformers.models.roberta.configuration_roberta import RobertaConfig


logger = logging.get_logger(__name__)

_CHECKPOINT_FOR_DOC = "roberta-base"
_CONFIG_FOR_DOC = "RobertaConfig"
_TOKENIZER_FOR_DOC = "RobertaTokenizer"

ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST = [
    "roberta-base",
    "roberta-large",
    "roberta-large-mnli",
    "distilroberta-base",
    "roberta-base-openai-detector",
    "roberta-large-openai-detector",
    # See all RoBERTa models at https://huggingface.co/models?filter=roberta
]


def apply_chunking_to_forward(
    forward_fn: Callable[..., torch.Tensor], chunk_size: int, chunk_dim: int, *input_tensors, last_norm=True
) -> torch.Tensor:
    """
    This function chunks the `input_tensors` into smaller input tensor parts of size `chunk_size` over the dimension
    `chunk_dim`. It then applies a layer `forward_fn` to each chunk independently to save memory.
    If the `forward_fn` is independent across the `chunk_dim` this function will yield the same result as directly
    applying `forward_fn` to `input_tensors`.
    Args:
        forward_fn (`Callable[..., torch.Tensor]`):
            The forward function of the model.
        chunk_size (`int`):
            The chunk size of a chunked tensor: `num_chunks = len(input_tensors[0]) / chunk_size`.
        chunk_dim (`int`):
            The dimension over which the `input_tensors` should be chunked.
        input_tensors (`Tuple[torch.Tensor]`):
            The input tensors of `forward_fn` which will be chunked
    Returns:
        `torch.Tensor`: A tensor with the same shape as the `forward_fn` would have given if applied`.
    Examples:
    ```python
    # rename the usual forward() fn to forward_chunk()
    def forward_chunk(self, hidden_states):
        hidden_states = self.decoder(hidden_states)
        return hidden_states
    # implement a chunked forward function
    def forward(self, hidden_states):
        return apply_chunking_to_forward(self.forward_chunk, self.chunk_size_lm_head, self.seq_len_dim, hidden_states)
    ```"""

    assert len(input_tensors) > 0, f"{input_tensors} has to be a tuple/list of tensors"

    # inspect.signature exist since python 3.5 and is a python method -> no problem with backward compatibility
    num_args_in_forward_chunk_fn = len(inspect.signature(forward_fn).parameters)
    # if num_args_in_forward_chunk_fn != len(input_tensors):
    #    raise ValueError(
    #        f"forward_chunk_fn expects {num_args_in_forward_chunk_fn} arguments, but only {len(input_tensors)} input "
    #        "tensors are given"
    #    )

    if chunk_size > 0:
        tensor_shape = input_tensors[0].shape[chunk_dim]
        for input_tensor in input_tensors:
            if input_tensor.shape[chunk_dim] != tensor_shape:
                raise ValueError(
                    f"All input tenors have to be of the same shape: {tensor_shape}, "
                    f"found shape {input_tensor.shape[chunk_dim]}"
                )

        if input_tensors[0].shape[chunk_dim] % chunk_size != 0:
            raise ValueError(
                f"The dimension to be chunked {input_tensors[0].shape[chunk_dim]} has to be a multiple of the chunk "
                f"size {chunk_size}"
            )

        num_chunks = input_tensors[0].shape[chunk_dim] // chunk_size

        # chunk input tensor into tuples
        input_tensors_chunks = tuple(input_tensor.chunk(num_chunks, dim=chunk_dim) for input_tensor in input_tensors)
        # apply forward fn to every tuple
        output_chunks = tuple(
            forward_fn(*input_tensors_chunk, last_norm=last_norm) for input_tensors_chunk in zip(*input_tensors_chunks)
        )
        # concatenate output at same dimension
        return torch.cat(output_chunks, dim=chunk_dim)

    return forward_fn(*input_tensors, last_norm=last_norm)


class RobertaEmbeddings(nn.Module):
    """
    Same as BertEmbeddings with a tiny tweak for positional embeddings indexing.
    """

    # Copied from transformers.models.bert.modeling_bert.BertEmbeddings.__init__
    def __init__(self, config):
        super().__init__()
        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
        self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)

        # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
        # any TensorFlow checkpoint file
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

        # position_ids (1, len position emb) is contiguous in memory and exported when serialized
        self.register_buffer("position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)))
        self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")

        # End copy
        self.padding_idx = config.pad_token_id
        self.position_embeddings = nn.Embedding(
            config.max_position_embeddings, config.hidden_size, padding_idx=self.padding_idx
        )

    def forward(
        self, input_ids=None, token_type_ids=None, position_ids=None, inputs_embeds=None, past_key_values_length=0
    ):
        if position_ids is None:
            if input_ids is not None:
                # Create the position ids from the input token ids. Any padded tokens remain padded.
                position_ids = create_position_ids_from_input_ids(
                    input_ids, self.padding_idx, past_key_values_length
                ).to(input_ids.device)
            else:
                position_ids = self.create_position_ids_from_inputs_embeds(inputs_embeds)

        if input_ids is not None:
            input_shape = input_ids.size()
        else:
            input_shape = inputs_embeds.size()[:-1]

        if token_type_ids is None:
            token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)

        if inputs_embeds is None:
            inputs_embeds = self.word_embeddings(input_ids)
        token_type_embeddings = self.token_type_embeddings(token_type_ids)

        embeddings = inputs_embeds + token_type_embeddings
        if self.position_embedding_type == "absolute":
            position_embeddings = self.position_embeddings(position_ids)
            embeddings += position_embeddings
        embeddings = self.LayerNorm(embeddings)
        embeddings = self.dropout(embeddings)
        return embeddings

    def create_position_ids_from_inputs_embeds(self, inputs_embeds):
        """
        We are provided embeddings directly. We cannot infer which are padded so just generate sequential position ids.
        Args:
            inputs_embeds: torch.Tensor
        Returns: torch.Tensor
        """
        input_shape = inputs_embeds.size()[:-1]
        sequence_length = input_shape[1]

        position_ids = torch.arange(
            self.padding_idx + 1, sequence_length + self.padding_idx + 1, dtype=torch.long, device=inputs_embeds.device
        )
        return position_ids.unsqueeze(0).expand(input_shape)


# Copied from transformers.models.bert.modeling_bert.BertSelfAttention with Bert->Roberta
class RobertaSelfAttention(nn.Module):
    def __init__(self, config, layer_index=None):
        super().__init__()
        if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
            raise ValueError(
                f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
                f"heads ({config.num_attention_heads})"
            )

        self.num_attention_heads = config.num_attention_heads
        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
        self.all_head_size = self.num_attention_heads * self.attention_head_size

        self.query = nn.Linear(config.hidden_size, self.all_head_size)
        if layer_index is None:
            self.key = nn.Linear(config.hidden_size, self.all_head_size)
            self.value = nn.Linear(config.hidden_size, self.all_head_size)
        else:
            self.key = nn.Linear(DIM_IMG, self.all_head_size)
            self.value = nn.Linear(DIM_IMG, self.all_head_size)

        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
        self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")
        if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
            self.max_position_embeddings = config.max_position_embeddings
            self.distance_embedding = nn.Embedding(2 * config.max_position_embeddings - 1, self.attention_head_size)

        self.is_decoder = config.is_decoder

    def transpose_for_scores(self, x):
        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
        x = x.view(*new_x_shape)
        return x.permute(0, 2, 1, 3)

    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        past_key_value=None,
        output_attentions=False,
    ):
        mixed_query_layer = self.query(hidden_states)

        # If this is instantiated as a cross-attention module, the keys
        # and values come from an encoder; the attention mask needs to be
        # such that the encoder's padding tokens are not attended to.
        is_cross_attention = encoder_hidden_states is not None

        if encoder_hidden_states is not None:
            key_layer = self.transpose_for_scores(self.key(encoder_hidden_states))
            value_layer = self.transpose_for_scores(self.value(encoder_hidden_states))
            attention_mask = encoder_attention_mask
        else:
            key_layer = self.transpose_for_scores(self.key(hidden_states))
            value_layer = self.transpose_for_scores(self.value(hidden_states))

        query_layer = self.transpose_for_scores(mixed_query_layer)

        # Take the dot product between "query" and "key" to get the raw attention scores.
        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))

        if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
            seq_length = hidden_states.size()[1]
            position_ids_l = torch.arange(seq_length, dtype=torch.long, device=hidden_states.device).view(-1, 1)
            position_ids_r = torch.arange(seq_length, dtype=torch.long, device=hidden_states.device).view(1, -1)
            distance = position_ids_l - position_ids_r
            positional_embedding = self.distance_embedding(distance + self.max_position_embeddings - 1)
            positional_embedding = positional_embedding.to(dtype=query_layer.dtype)  # fp16 compatibility

            if self.position_embedding_type == "relative_key":
                relative_position_scores = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
                attention_scores = attention_scores + relative_position_scores
            elif self.position_embedding_type == "relative_key_query":
                relative_position_scores_query = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
                relative_position_scores_key = torch.einsum("bhrd,lrd->bhlr", key_layer, positional_embedding)
                attention_scores = attention_scores + relative_position_scores_query + relative_position_scores_key

        attention_scores = attention_scores / math.sqrt(self.attention_head_size)
        if attention_mask is not None:
            # Apply the attention mask is (precomputed for all layers in RobertaModel forward() function)
            attention_scores = attention_scores + attention_mask

        # Normalize the attention scores to probabilities.
        attention_probs = nn.Softmax(dim=-1)(attention_scores)

        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.
        attention_probs = self.dropout(attention_probs)

        # Mask heads if we want to
        if head_mask is not None:
            attention_probs = attention_probs * head_mask

        context_layer = torch.matmul(attention_probs, value_layer)

        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
        context_layer = context_layer.view(*new_context_layer_shape)

        outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)

        return outputs


# Copied from transformers.models.bert.modeling_bert.BertSelfOutput
class RobertaSelfOutput(nn.Module):
    def __init__(self, config, layer_index=None):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        if layer_index is None:
            self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, hidden_states):
        #print("Shape of hidden_states: ", hidden_states.shape)
        hidden_states = self.dense(hidden_states)
        hidden_states = self.dropout(hidden_states)
        return hidden_states


# Copied from transformers.models.bert.modeling_bert.BertAttention with Bert->Roberta
class RobertaAttention(nn.Module):
    def __init__(self, config, layer_index=None):
        super().__init__()
        self.self = RobertaSelfAttention(config, layer_index=layer_index)
        self.output = RobertaSelfOutput(config, layer_index=layer_index)
        self.pruned_heads = set()

    def prune_heads(self, heads):
        if len(heads) == 0:
            return
        heads, index = find_pruneable_heads_and_indices(
            heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads
        )

        # Prune linear layers
        self.self.query = prune_linear_layer(self.self.query, index)
        self.self.key = prune_linear_layer(self.self.key, index)
        self.self.value = prune_linear_layer(self.self.value, index)
        self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)

        # Update hyper params and store pruned heads
        self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
        self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads
        self.pruned_heads = self.pruned_heads.union(heads)

    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        past_key_value=None,
        output_attentions=False,
    ):
        self_outputs = self.self(
            hidden_states,
            attention_mask,
            head_mask,
            encoder_hidden_states,
            encoder_attention_mask,
            past_key_value,
            output_attentions,
        )
        attention_output = self.output(self_outputs[0])
        outputs = (attention_output,) + self_outputs[1:]  # add attentions if we output them
        return outputs


# Copied from transformers.models.bert.modeling_bert.BertIntermediate
class RobertaIntermediate(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
        if isinstance(config.hidden_act, str):
            self.intermediate_act_fn = ACT2FN[config.hidden_act]
        else:
            self.intermediate_act_fn = config.hidden_act

    def forward(self, hidden_states):
        hidden_states = self.dense(hidden_states)
        hidden_states = self.intermediate_act_fn(hidden_states)
        return hidden_states


# Copied from transformers.models.bert.modeling_bert.BertOutput
class RobertaOutput(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, hidden_states, input_tensor, last_norm=True):
        hidden_states = self.dense(hidden_states)
        hidden_states = self.dropout(hidden_states)
        hidden_states = hidden_states + input_tensor
        if last_norm:
            hidden_states = self.LayerNorm(hidden_states)
        return hidden_states


# Copied from transformers.models.bert.modeling_bert.BertLayer with Bert->Roberta
class RobertaLayer(nn.Module):
    def __init__(self, config, layer_index=None):
        super().__init__()
        self.chunk_size_feed_forward = config.chunk_size_feed_forward
        self.seq_len_dim = 1
        self.attention = RobertaAttention(config)
        self.is_decoder = config.is_decoder
        self.add_cross_attention = config.add_cross_attention
        if layer_index >= 12 - NUM_FUSE_BLOCK:
            self.crossattention_t2i = RobertaAttention(config, layer_index=layer_index)
            self.alpha_t2i = nn.Parameter(torch.Tensor([0]))
        self.intermediate = RobertaIntermediate(config)
        self.output = RobertaOutput(config)

    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        past_key_value=None,
        output_attentions=False,
        last_norm=True,
    ):
        # decoder uni-directional self-attention cached key/values tuple is at positions 1,2
        self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
        self_attention_outputs = self.attention(
            hidden_states,
            attention_mask,
            head_mask,
            output_attentions=output_attentions,
            past_key_value=self_attn_past_key_value,
        )
        attention_output = self_attention_outputs[0]

        # if decoder, the last output is tuple of self-attn cache
        outputs = self_attention_outputs[1:]  # add self attentions if we output attention weights

        cross_attn_present_key_value = None
        if encoder_hidden_states is not None:
            assert hasattr(
                self, "crossattention_t2i"
            ), f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers by setting `config.add_cross_attention=True`"

            # cross_attn cached key/values tuple is at positions 3,4 of past_key_value tuple
            cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None
            cross_attention_outputs = self.crossattention_t2i(
                attention_output,
                attention_mask,
                head_mask,
                encoder_hidden_states,
                encoder_attention_mask,
                cross_attn_past_key_value,
                output_attentions,
            )
            attention_output = self.alpha_t2i * cross_attention_outputs[0] + attention_output

        attention_output = self.attention.output.LayerNorm(attention_output + hidden_states)

        layer_output = apply_chunking_to_forward(
            self.feed_forward_chunk,
            self.chunk_size_feed_forward,
            self.seq_len_dim,
            attention_output,
            last_norm=last_norm,
        )
        outputs = (layer_output,) + outputs

        # if decoder, return the attn key/values as the last output
        return outputs

    def feed_forward_chunk(self, attention_output, last_norm=True):
        intermediate_output = self.intermediate(attention_output)
        layer_output = self.output(intermediate_output, attention_output, last_norm=last_norm)
        return layer_output


# Copied from transformers.models.bert.modeling_bert.BertEncoder with Bert->Roberta
class RobertaEncoder(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.config.gradient_checkpointing = config_yaml['use_checkpoint']
        self.layer = nn.ModuleList(
            [RobertaLayer(config, layer_index=layer_i) for layer_i in range(config.num_hidden_layers)]
        )

    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        past_key_values=None,
        use_cache=None,
        output_attentions=False,
        output_hidden_states=False,
        return_dict=True,
    ):
        all_hidden_states = () if output_hidden_states else None
        all_self_attentions = () if output_attentions else None
        all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None

        next_decoder_cache = () if use_cache else None
        for i, layer_module in enumerate(self.layer):
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

            layer_head_mask = head_mask[i] if head_mask is not None else None
            past_key_value = past_key_values[i] if past_key_values is not None else None

            if getattr(self.config, "gradient_checkpointing", False) and self.training:

                if use_cache:
                    logger.warning(
                        "`use_cache=True` is incompatible with `config.gradient_checkpointing=True`. Setting "
                        "`use_cache=False`..."
                    )
                    use_cache = False

                #def create_custom_forward(module):
                    #def custom_forward(*inputs):
                        #return module(*inputs, past_key_value, output_attentions)
                    #return custom_forward

                layer_outputs = torch.utils.checkpoint.checkpoint(
                    layer_module,
                    hidden_states,
                    attention_mask,
                    layer_head_mask,
                    encoder_hidden_states,
                    encoder_attention_mask,
                    past_key_value, 
                    output_attentions
                )
            else:
                layer_outputs = layer_module(
                    hidden_states,
                    attention_mask,
                    layer_head_mask,
                    encoder_hidden_states,
                    encoder_attention_mask,
                    past_key_value,
                    output_attentions,
                )


            hidden_states = layer_outputs[0]
            if use_cache:
                next_decoder_cache += (layer_outputs[-1],)
            if output_attentions:
                all_self_attentions = all_self_attentions + (layer_outputs[1],)
                if self.config.add_cross_attention:
                    all_cross_attentions = all_cross_attentions + (layer_outputs[2],)

        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        if not return_dict:
            return tuple(
                v
                for v in [
                    hidden_states,
                    next_decoder_cache,
                    all_hidden_states,
                    all_self_attentions,
                    all_cross_attentions,
                ]
                if v is not None
            )
        return BaseModelOutputWithPastAndCrossAttentions(
            last_hidden_state=hidden_states,
            past_key_values=next_decoder_cache,
            hidden_states=all_hidden_states,
            attentions=all_self_attentions,
            cross_attentions=all_cross_attentions,
        )

class RobertaPooler(nn.Module):
    def __init__(self, config):
        super().__init__()
        #self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        self.activation = nn.Tanh()

    def forward(self, hidden_states):
        # We "pool" the model by simply taking the hidden state corresponding
        # to the first token.
        first_token_tensor = hidden_states[:, 0]
        #pooled_output = self.dense(first_token_tensor)
        pooled_output = self.activation(first_token_tensor)
        return pooled_output


class RobertaPreTrainedModel(PreTrainedModel):
    """
    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
    models.
    """

    config_class = RobertaConfig
    base_model_prefix = "roberta"

    # Copied from transformers.models.bert.modeling_bert.BertPreTrainedModel._init_weights
    def _init_weights(self, module):
        """Initialize the weights"""
        if isinstance(module, nn.Linear):
            # Slightly different from the TF version which uses truncated_normal for initialization
            # cf https://github.com/pytorch/pytorch/pull/5617
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)


ROBERTA_START_DOCSTRING = r"""
    This model inherits from :class:`~transformers.PreTrainedModel`. Check the superclass documentation for the generic
    methods the library implements for all its model (such as downloading or saving, resizing the input embeddings,
    pruning heads etc.)
    This model is also a PyTorch `torch.nn.Module <https://pytorch.org/docs/stable/nn.html#torch.nn.Module>`__
    subclass. Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to
    general usage and behavior.
    Parameters:
        config (:class:`~transformers.RobertaConfig`): Model configuration class with all the parameters of the
            model. Initializing with a config file does not load the weights associated with the model, only the
            configuration. Check out the :meth:`~transformers.PreTrainedModel.from_pretrained` method to load the model
            weights.
"""

ROBERTA_INPUTS_DOCSTRING = r"""
    Args:
        input_ids (:obj:`torch.LongTensor` of shape :obj:`({0})`):
            Indices of input sequence tokens in the vocabulary.
            Indices can be obtained using :class:`~transformers.RobertaTokenizer`. See
            :meth:`transformers.PreTrainedTokenizer.encode` and :meth:`transformers.PreTrainedTokenizer.__call__` for
            details.
            `What are input IDs? <../glossary.html#input-ids>`__
        attention_mask (:obj:`torch.FloatTensor` of shape :obj:`({0})`, `optional`):
            Mask to avoid performing attention on padding token indices. Mask values selected in ``[0, 1]``:
            - 1 for tokens that are **not masked**,
            - 0 for tokens that are **masked**.
            `What are attention masks? <../glossary.html#attention-mask>`__
        token_type_ids (:obj:`torch.LongTensor` of shape :obj:`({0})`, `optional`):
            Segment token indices to indicate first and second portions of the inputs. Indices are selected in ``[0,
            1]``:
            - 0 corresponds to a `sentence A` token,
            - 1 corresponds to a `sentence B` token.
            `What are token type IDs? <../glossary.html#token-type-ids>`_
        position_ids (:obj:`torch.LongTensor` of shape :obj:`({0})`, `optional`):
            Indices of positions of each input sequence tokens in the position embeddings. Selected in the range ``[0,
            config.max_position_embeddings - 1]``.
            `What are position IDs? <../glossary.html#position-ids>`_
        head_mask (:obj:`torch.FloatTensor` of shape :obj:`(num_heads,)` or :obj:`(num_layers, num_heads)`, `optional`):
            Mask to nullify selected heads of the self-attention modules. Mask values selected in ``[0, 1]``:
            - 1 indicates the head is **not masked**,
            - 0 indicates the head is **masked**.
        inputs_embeds (:obj:`torch.FloatTensor` of shape :obj:`({0}, hidden_size)`, `optional`):
            Optionally, instead of passing :obj:`input_ids` you can choose to directly pass an embedded representation.
            This is useful if you want more control over how to convert :obj:`input_ids` indices into associated
            vectors than the model's internal embedding lookup matrix.
        output_attentions (:obj:`bool`, `optional`):
            Whether or not to return the attentions tensors of all attention layers. See ``attentions`` under returned
            tensors for more detail.
        output_hidden_states (:obj:`bool`, `optional`):
            Whether or not to return the hidden states of all layers. See ``hidden_states`` under returned tensors for
            more detail.
        return_dict (:obj:`bool`, `optional`):
            Whether or not to return a :class:`~transformers.file_utils.ModelOutput` instead of a plain tuple.
"""


@add_start_docstrings(
    "The bare RoBERTa Model transformer outputting raw hidden-states without any specific head on top.",
    ROBERTA_START_DOCSTRING,
)
class RobertaModel(RobertaPreTrainedModel):
    """
    The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of
    cross-attention is added between the self-attention layers, following the architecture described in `Attention is
    all you need`_ by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz
    Kaiser and Illia Polosukhin.
    To behave as an decoder the model needs to be initialized with the :obj:`is_decoder` argument of the configuration
    set to :obj:`True`. To be used in a Seq2Seq model, the model needs to initialized with both :obj:`is_decoder`
    argument and :obj:`add_cross_attention` set to :obj:`True`; an :obj:`encoder_hidden_states` is then expected as an
    input to the forward pass.
    .. _`Attention is all you need`: https://arxiv.org/abs/1706.03762
    """

    _keys_to_ignore_on_load_missing = [r"position_ids"]

    # Copied from transformers.models.bert.modeling_bert.BertModel.__init__ with Bert->Roberta
    def __init__(self, config, add_pooling_layer=True):
        super().__init__(config)
        self.config = config

        self.embeddings = RobertaEmbeddings(config)
        self.encoder = RobertaEncoder(config)

        self.pooler = RobertaPooler(config) if add_pooling_layer else None

        self.init_weights()

    def get_input_embeddings(self):
        return self.embeddings.word_embeddings

    def set_input_embeddings(self, value):
        self.embeddings.word_embeddings = value

    def _prune_heads(self, heads_to_prune):
        """
        Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
        class PreTrainedModel
        """
        for layer, heads in heads_to_prune.items():
            self.encoder.layer[layer].attention.prune_heads(heads)

    @add_start_docstrings_to_model_forward(ROBERTA_INPUTS_DOCSTRING.format("(batch_size, sequence_length)"))
    @add_code_sample_docstrings(
        processor_class=_TOKENIZER_FOR_DOC,
        checkpoint=_CHECKPOINT_FOR_DOC,
        output_type=BaseModelOutputWithPoolingAndCrossAttentions,
        config_class=_CONFIG_FOR_DOC,
    )
    # Copied from transformers.models.bert.modeling_bert.BertModel.forward
    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        past_key_values=None,
        use_cache=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
    ):
        r"""
        encoder_hidden_states  (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`):
            Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
            the model is configured as a decoder.
        encoder_attention_mask (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
            Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
            the cross-attention if the model is configured as a decoder. Mask values selected in ``[0, 1]``:
            - 1 for tokens that are **not masked**,
            - 0 for tokens that are **masked**.
        past_key_values (:obj:`tuple(tuple(torch.FloatTensor))` of length :obj:`config.n_layers` with each tuple having 4 tensors of shape :obj:`(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
            Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
            If :obj:`past_key_values` are used, the user can optionally input only the last :obj:`decoder_input_ids`
            (those that don't have their past key value states given to this model) of shape :obj:`(batch_size, 1)`
            instead of all :obj:`decoder_input_ids` of shape :obj:`(batch_size, sequence_length)`.
        use_cache (:obj:`bool`, `optional`):
            If set to :obj:`True`, :obj:`past_key_values` key value states are returned and can be used to speed up
            decoding (see :obj:`past_key_values`).
        """
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        use_cache = False

        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot secify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
            input_shape = input_ids.size()
            batch_size, seq_length = input_shape
        elif inputs_embeds is not None:
            input_shape = inputs_embeds.size()[:-1]
            batch_size, seq_length = input_shape
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        device = input_ids.device if input_ids is not None else inputs_embeds.device

        # past_key_values_length
        past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0

        if attention_mask is None:
            attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length)), device=device)
        if token_type_ids is None:
            token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)

        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
        # ourselves in which case we just need to make it broadcastable to all heads.
        extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape, device)

        # If a 2D or 3D attention mask is provided for the cross-attention
        # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
        if encoder_hidden_states is not None:
            encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
            encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
            if encoder_attention_mask is None:
                encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
            encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
        else:
            encoder_extended_attention_mask = None

        # Prepare head mask if needed
        # 1.0 in head_mask indicate we keep the head
        # attention_probs has shape bsz x n_heads x N x N
        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)

        embedding_output = self.embeddings(
            input_ids=input_ids,
            position_ids=position_ids,
            token_type_ids=token_type_ids,
            inputs_embeds=inputs_embeds,
            past_key_values_length=past_key_values_length,
        )
        encoder_outputs = self.encoder(
            embedding_output,
            attention_mask=extended_attention_mask,
            head_mask=head_mask,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_extended_attention_mask,
            past_key_values=past_key_values,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        sequence_output = encoder_outputs[0]
        pooled_output = self.pooler(sequence_output) if self.pooler is not None else None

        if not return_dict:
            return (sequence_output, pooled_output) + encoder_outputs[1:]

        return BaseModelOutputWithPoolingAndCrossAttentions(
            last_hidden_state=sequence_output,
            pooler_output=pooled_output,
            past_key_values=encoder_outputs.past_key_values,
            hidden_states=encoder_outputs.hidden_states,
            attentions=encoder_outputs.attentions,
            cross_attentions=encoder_outputs.cross_attentions,
        )


def create_position_ids_from_input_ids(input_ids, padding_idx, past_key_values_length=0):
    """
    Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding symbols
    are ignored. This is modified from fairseq's `utils.make_positions`.
    Args:
        x: torch.Tensor x:
    Returns: torch.Tensor
    """
    # The series of casts and type-conversions here are carefully balanced to both work with ONNX export and XLA.
    mask = input_ids.ne(padding_idx).int()
    incremental_indices = (torch.cumsum(mask, dim=1).type_as(mask) + past_key_values_length) * mask
    return incremental_indices.long() + padding_idx


def _make_causal_mask(input_ids_shape: torch.Size, dtype: torch.dtype, past_key_values_length: int = 0):
    """
    Make causal mask used for bi-directional self-attention.
    """
    bsz, tgt_len = input_ids_shape
    mask = torch.full((tgt_len, tgt_len), float("-inf"))
    mask_cond = torch.arange(mask.size(-1))
    mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
    mask = mask.to(dtype)

    if past_key_values_length > 0:
        mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype), mask], dim=-1)
    return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)


def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len=None):
    """
    Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
    """
    bsz, src_len = mask.size()
    tgt_len = tgt_len if tgt_len is not None else src_len

    expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)

    inverted_mask = 1.0 - expanded_mask

    return inverted_mask.masked_fill(inverted_mask.bool(), torch.finfo(dtype).min)


def _prepare_decoder_attention_mask(
    attention_mask, input_shape, inputs_embeds, device, past_key_values_length: int = 0
):
    # create causal mask
    # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
    combined_attention_mask = None
    if input_shape[-1] > 1:
        combined_attention_mask = _make_causal_mask(
            input_shape, inputs_embeds.dtype, past_key_values_length=past_key_values_length
        ).to(device)

    if attention_mask is not None:
        # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
        expanded_attn_mask = _expand_mask(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1])
        combined_attention_mask = (
            expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask
        )

    return combined_attention_mask

EgoNLQ/model/roberta.py
facebookresearch/EgoVLPv2
"""
Implementations of Video Transformers in PyTorch

A PyTorch implementation of space-time transformer as described in
'Frozen in Time: A Joint Image and Video Encoder for End-to-End Retrieval' - https://arxiv.org/abs/2104.00650

A PyTorch implementation of timesformer as described in
'Is Space-Time Attention All You Need for Video Understanding?' - https://arxiv.org/abs/2102.05095

Acknowledgments:
- This code builds on Ross Wightman's vision_transformer code in pytorch-image-models:
https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py

- It is also inspired by lucidrains timesformer implementation:
https://github.com/lucidrains/TimeSformer-pytorch

Hacked together by Max Bain
"""

from collections import OrderedDict
from functools import partial
import yaml

import torch
from einops import rearrange, repeat
from timm.models.layers import DropPath, to_2tuple, trunc_normal_
from torch import einsum, nn

with open('./EgoNCE_MLM_ITM_Config.yml') as f:
    config_yaml = yaml.load(f, Loader=yaml.FullLoader)

NUM_FUSE_BLOCK = config_yaml['num_fuse_block']
DIM_TEXT = 768

def attn(q, k, v):
    sim = einsum('b i d, b j d -> b i j', q, k)
    attn = sim.softmax(dim=-1)
    out = einsum('b i j, b j d -> b i d', attn, v)
    return out


class Mlp(nn.Module):
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x


class VideoPatchEmbed(nn.Module):
    """ Video to Patch Embedding
    """

    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768,
                 num_frames=8):
        super().__init__()
        img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)
        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0]) * num_frames
        self.img_size = img_size
        self.patch_size = patch_size
        self.num_patches = num_patches
        self.num_frames = num_frames
        self.embed_dim = embed_dim
        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)

    def forward(self, x):
        B, F, C, H, W = x.shape
        assert F == self.num_frames, print(F, self.num_frames)
        x = x.view(-1, C, H, W)
        x = self.proj(x)
        return x


class VarAttention(nn.Module):
    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.,
                 initialize='random', dim_text=None, norm_layer=nn.LayerNorm, space_attn=True):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        # NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights
        self.scale = qk_scale or head_dim ** -0.5
        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.proj = nn.Linear(dim, dim)
        if initialize == 'zeros':
            self.qkv.weight.data.fill_(0)
            self.qkv.bias.data.fill_(0)
            # fill proj weight with 1 here to improve training dynamics. Otherwise temporal attention inputs
            # are multiplied by 0*0, which is hard for the model to move out of.
            self.proj.weight.data.fill_(1)
            self.proj.bias.data.fill_(0)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj_drop = nn.Dropout(proj_drop)

        self.softmax = nn.Softmax(dim=-1)

        if dim_text is not None and space_attn:
            self.qkv_text_i2t = nn.Linear(dim_text, dim * 2, bias=qkv_bias)
            self.qkv_i2t = nn.Linear(dim, dim, bias=qkv_bias)
            self.attn_drop_i2t = nn.Dropout(attn_drop)
            self.proj_i2t = nn.Linear(dim, dim)
            self.proj_drop_i2t = nn.Dropout(proj_drop)
            self.alpha_i2t = nn.Parameter(torch.Tensor([0]))
            self.norm_i2t_i = norm_layer(dim)

    def forward(self, x, einops_from, einops_to, y=None, y_mask=None, **einops_dims):
        h = self.num_heads
        # project x to q, k, v vaalues
        q, k, v = self.qkv(x).chunk(3, dim=-1)
        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))

        q = q*self.scale

        # splice out CLS token at index 1
        (cls_q, q_), (cls_k, k_), (cls_v, v_) = map(lambda t: (t[:, 0:1], t[:, 1:]), (q, k, v))

        # let CLS token attend to key / values of all patches across time and space
        cls_out = attn(cls_q, k, v)
        # rearrange across time or space
        q_, k_, v_ = map(lambda t: rearrange(t, f'{einops_from} -> {einops_to}', **einops_dims), (q_, k_, v_))

        # expand cls token keys and values across time or space and concat
        r = q_.shape[0] // cls_k.shape[0]
        cls_k, cls_v = map(lambda t: repeat(t, 'b () d -> (b r) () d', r=r), (cls_k, cls_v))

        k_ = torch.cat((cls_k, k_), dim=1)
        v_ = torch.cat((cls_v, v_), dim=1)

        # attention
        out = attn(q_, k_, v_)

        # merge back time or space
        out = rearrange(out, f'{einops_to} -> {einops_from}', **einops_dims)

        # concat back the cls token
        out = torch.cat((cls_out, out), dim=1)

        # merge back the heads
        out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
        ## to out
        x = self.proj(out)
        x = self.proj_drop(x)

        if y is not None:
            B_, N, C = x.shape
            B_text, N_text, C_text = y.shape

            kv_text = (
                self.qkv_text_i2t(y)
                .reshape(B_text, N_text, 2, self.num_heads, C // self.num_heads)
                .permute(2, 0, 3, 1, 4)
            )
            k_text, v_text = kv_text[0], kv_text[1]

            q_i2t = self.qkv_i2t(self.norm_i2t_i(x))
            q_i2t = q_i2t.reshape(B_, N, 1, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
            q_i2t = q_i2t[0]

            # image to text attention
            text_scale = k_text.size(-1) ** -0.5
            q_i2t = q_i2t * text_scale
            attn_i2t = q_i2t @ k_text.transpose(-2, -1)  # B_, nH, N, N_text

            # add image to text bias and text_mask
            if y_mask is not None:
                mask_and_i2t_bias = y_mask.view(B_text, 1, 1, N_text)
                attn_i2t = attn_i2t + mask_and_i2t_bias

            attn_i2t = self.softmax(attn_i2t)
            attn_i2t = self.attn_drop_i2t(attn_i2t)
            y = (attn_i2t @ v_text).transpose(1, 2).reshape(B_, N, C)
            y = self.proj_i2t(y)
            y = self.proj_drop_i2t(y)
            x = x + self.alpha_i2t * y

        return x


class SpaceTimeBlock(nn.Module):

    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, time_init='zeros',
                 attention_style='frozen-in-time', dim_text=None):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.attn = VarAttention(
            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop, dim_text=dim_text, 
            norm_layer=norm_layer, space_attn=True)

        self.timeattn = VarAttention(
            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop,
            initialize=time_init, dim_text=dim_text, norm_layer=norm_layer, space_attn=False)

        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
        self.norm3 = norm_layer(dim)

        self.attention_style = attention_style

    def forward(self, x, einops_from_space, einops_to_space, einops_from_time, einops_to_time,
                time_n, space_f, y=None, y_mask=None):

        time_output = self.timeattn(self.norm3(x), einops_from_time, einops_to_time, n=time_n, y=None, y_mask=None)
        time_residual = x + time_output
        space_output = self.attn(self.norm1(time_residual), einops_from_space,
                                 einops_to_space, f=space_f, y=y, y_mask=y_mask)
        if self.attention_style == 'frozen-in-time':
            space_residual = x + self.drop_path(space_output)
        else:
            raise NotImplementedError

        x = space_residual + self.drop_path(self.mlp(self.norm2(space_residual)))

        return x


class SpaceTimeTransformer(nn.Module):
    """ Vision Transformer

    A PyTorch impl of : `Space-Time Transformer` from Frozen-in-time  - by Max Bain.
        https://arxiv.org/abs/2104.00650

    Based off:
     - ViT implementation from the timm library [https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py]
    lucidrains timesformer implementation [https://github.com/lucidrains/TimeSformer-pytorch].

    Notable differences:
     - allows for variable length input frames (<= num_frames)
     - allows for variable length input resolution  (<= (img_size, img_size)) [UNTESTED]
     - different attention block mechanism
    """

    def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12,
                 num_heads=12, mlp_ratio=4., qkv_bias=True, qk_scale=None, representation_size=None,
                 drop_rate=0., attn_drop_rate=0., drop_path_rate=0., hybrid_backbone=None,
                 num_frames=8, time_init='rand', attention_style='frozen-in-time', norm_layer=nn.LayerNorm, dim_text=None):
        """
        Args:
            img_size (int, tuple): input image size
            patch_size (int, tuple): patch size
            in_chans (int): number of input channels
            num_classes (int): number of classes for classification head
            embed_dim (int): embedding dimension
            depth (int): depth of transformer
            num_heads (int): number of attention heads
            mlp_ratio (int): ratio of mlp hidden dim to embedding dim
            qkv_bias (bool): enable bias for qkv if True
            qk_scale (float): override default qk scale of head_dim ** -0.5 if set
            representation_size (Optional[int]): enable and set representation layer (pre-logits) to this value if set
            drop_rate (float): dropout rate
            attn_drop_rate (float): attention dropout rate
            drop_path_rate (float): stochastic depth rate
            hybrid_backbone (nn.Module): CNN backbone to use in-place of PatchEmbed module
            norm_layer: (nn.Module): normalization layer
            num_frames: (int) maximum number of frames expected as input
            time_init: (str) how to initialise the time attention layer, 'zeros' allows for the timesformer to start off
                        as ViT.
            attention_style: (str) how to attend to space and time.
        """
        super().__init__()
        self.num_classes = num_classes
        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
        self.num_frames = num_frames
        self.embed_dim = embed_dim
        norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
        print("######USING ATTENTION STYLE: ", attention_style)
        if hybrid_backbone is not None:
            raise NotImplementedError('hybrid backbone not implemented')
        else:
            self.patch_embed = VideoPatchEmbed(
                img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim, num_frames=num_frames)
        num_patches = self.patch_embed.num_patches
        self.patches_per_frame = num_patches // num_frames

        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.pos_embed = nn.Parameter(
            torch.zeros(1, self.patches_per_frame + 1,
                        embed_dim))  # remember to take pos_embed[1:] for tiling over time
        self.temporal_embed = nn.Parameter(torch.zeros(1, num_frames, embed_dim))

        self.pos_drop = nn.Dropout(p=drop_rate)

        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
        self.blocks = nn.ModuleList([
            SpaceTimeBlock(
                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
                drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer, time_init=time_init,
                attention_style=attention_style, dim_text=None if i < 6 else DIM_TEXT)
            for i in range(depth)])
        self.norm = norm_layer(embed_dim)

        # Representation layer
        if representation_size:
            self.num_features = representation_size
            self.pre_logits = nn.Sequential(OrderedDict([
                ('fc', nn.Linear(embed_dim, representation_size)),
                ('act', nn.Tanh())
            ]))
        else:
            self.pre_logits = nn.Identity()

        # Classifier head
        self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()

        trunc_normal_(self.pos_embed, std=.02)
        trunc_normal_(self.cls_token, std=.02)

        # if num_frames > 1, then we perform ViT inflation and initialise time attention to zero so not necessary.
        if num_frames == 1:
            self.apply(self._init_weights)

        ## einops transformations
        self.einops_from_space = 'b (f n) d'
        self.einops_to_space = '(b f) n d'
        self.einops_from_time = 'b (f n) d'
        self.einops_to_time = '(b n) f d'


    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    @torch.jit.ignore
    def no_weight_decay(self):
        return {'pos_embed', 'cls_token'}

    def get_classifier(self):
        return self.head

    def reset_classifier(self, num_classes, global_pool=''):
        self.num_classes = num_classes
        self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()

    def forward_features(self, x):
        b, curr_frames, channels, _, _ = x.shape
        x = self.patch_embed(x)
        x = x.flatten(2).transpose(2, 1)
        x = x.reshape(b, -1, self.patch_embed.embed_dim)

        BF = x.shape[0]
        cls_tokens = self.cls_token.expand(BF, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        # positional embed needs to be tiled for each frame (this does [1,2,3] --> [1,2,3,1,2,3]...)
        cls_embed = self.pos_embed[:, 0, :].unsqueeze(1)
        tile_pos_embed = self.pos_embed[:, 1:, :].repeat(1, self.num_frames, 1)
        # temporal embed needs to be repeated within each frame (this does [1,2,3] --> [1,1,1,2,2,2,3,3,3]...)
        tile_temporal_embed = self.temporal_embed.repeat_interleave(self.patches_per_frame, 1)
        total_pos_embed = tile_pos_embed + tile_temporal_embed
        total_pos_embed = torch.cat([cls_embed, total_pos_embed], dim=1)

        curr_patches = x.shape[1]
        x = x + total_pos_embed[:, :curr_patches]
        x = self.pos_drop(x)
        n = self.patches_per_frame
        f = curr_frames

        for blk in self.blocks:
            if config_yaml['use_checkpoint']:

                def create_custom_forward(module):
                    def custom_forward(*inputs):
                        return module(*inputs, time_n=n, space_f=f)

                    return custom_forward

                x = torch.utils.checkpoint.checkpoint(create_custom_forward(blk), x, self.einops_from_space, self.einops_to_space, self.einops_from_time,
                    self.einops_to_time)
            else:
                x = blk(x, self.einops_from_space, self.einops_to_space, self.einops_from_time,
                    self.einops_to_time, time_n=n, space_f=f)

        x = self.norm(x)[:, 0]
        x = self.pre_logits(x)

        return x

    def forward(self, x):
        x = self.forward_features(x)
        x = self.head(x)
        return x

if __name__ == "__main__":
    network = SpaceTimeTransformer(num_frames=4)
    data = torch.rand((3, 4, 3, 224, 224))
    network(data)

EgoNLQ/model/video_transformer.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


"""
Reading command line options.
"""

from __future__ import absolute_import, division, print_function, unicode_literals

import argparse


def read_command_line():
    parser = argparse.ArgumentParser()
     #data parameters
    parser.add_argument(
        "--save_dir",
        type=str,
        default="datasets",
        help="path to save processed dataset",
    )
    parser.add_argument("--task", type=str, default="charades", help="target task")
    parser.add_argument(
        "--eval_gt_json",
        type=str,
        default=None,
        help="Provide GT JSON to evaluate while training"
    )
    parser.add_argument(
        "--max_pos_len",
        type=int,
        default=128,
        help="maximal position sequence length allowed",
    )
    parser.add_argument(
        "--num_workers",
        type=int,
        default=1,
        help="Number of CPU workers to process the data",
    )
    # model parameters
    parser.add_argument("--word_size", type=int, default=None, help="number of words")
    parser.add_argument(
        "--char_size", type=int, default=None, help="number of characters"
    )
    parser.add_argument(
        "--word_dim", type=int, default=300, help="word embedding dimension"
    )
    parser.add_argument(
        "--video_feature_dim",
        type=int,
        default=768,
        help="video feature input dimension",
    )
    parser.add_argument(
        "--char_dim",
        type=int,
        default=50,
        help="character dimension, set to 100 for activitynet",
    )
    parser.add_argument("--dim", type=int, default=128, help="hidden size")
    parser.add_argument(
        "--highlight_lambda",
        type=float,
        default=5.0,
        help="lambda for highlight region",
    )
    parser.add_argument("--num_heads", type=int, default=8, help="number of heads")
    parser.add_argument("--drop_rate", type=float, default=0.2, help="dropout rate")
    parser.add_argument(
        "--predictor", type=str, default="rnn", help="[rnn | transformer]"
    )
    # training/evaluation parameters
    parser.add_argument("--gpu_idx", type=str, default="0", help="GPU index")
    parser.add_argument("--seed", type=int, default=12345, help="random seed")
    parser.add_argument("--mode", type=str, default="train", help="[train | test]")
    parser.add_argument("--epochs", type=int, default=100, help="number of epochs")
    parser.add_argument("--batch_size", type=int, default=32, help="batch size")
    parser.add_argument(
        "--num_train_steps", type=int, default=None, help="number of training steps"
    )
    parser.add_argument(
        "--init_lr", type=float, default=0.0001, help="initial learning rate"
    )
    parser.add_argument(
        "--clip_norm", type=float, default=1.0, help="gradient clip norm"
    )
    parser.add_argument(
        "--warmup_proportion", type=float, default=0.0, help="warmup proportion"
    )
    parser.add_argument(
        "--extend", type=float, default=0.1, help="highlight region extension"
    )
    parser.add_argument(
        "--period", type=int, default=100, help="training loss print period"
    )
    parser.add_argument(
        "--text_agnostic",
        dest="text_agnostic",
        action="store_true",
        default=False,
        help="Text agnostic model; random text features",
    )
    parser.add_argument(
        "--video_agnostic",
        dest="video_agnostic",
        action="store_true",
        default=False,
        help="Video agnostic model; random video features",
    )
    parser.add_argument("--model_name", type=str, help="model name")
    parser.add_argument(
        "--suffix",
        type=str,
        default=None,
        help="set to the last `_xxx` in ckpt repo to eval results",
    )
    parser.add_argument('--cuda_base', help="in form cuda:x")
    parser.add_argument('--device_ids', help='delimited list input', type=lambda s: [int(item) for item in s.split(',')])
    parser.add_argument("--feature_dir", type=str, default="./saved_features", help="model name")
    parser.add_argument("--query_max_len", type=int, default=30, help="max_len of query")

    configs = parser.parse_args()
    return configs, parser

EgoNLQ/options.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import os
import logging
import pdb
from pathlib import Path
from functools import reduce
from operator import getitem
from datetime import datetime
# import pdb; pdb.set_trace()
from logger import setup_logging
from utils.util import read_json, write_json
import time
import inspect


class ConfigParser:
    def __init__(self, args, options='', timestamp=True, test=False, eval_mode=None):
        # parse default and custom cli options
        for opt in options:
            args.add_argument(*opt.flags, default=None, type=opt.type)
        args = args.parse_args()
        self.args = args
        if args.device:
            os.environ["CUDA_VISIBLE_DEVICES"] = args.device
        if args.resume is None:
            msg_no_cfg = "Configuration file need to be specified. Add '-c config.json', for example."
            assert args.config is not None, msg_no_cfg
            self.cfg_fname = Path(args.config)
            config = read_json(self.cfg_fname)
            self.resume = None
        else:
            self.resume = Path(args.resume)
            resume_cfg_fname = self.resume.parent / 'config.json'
            if eval_mode == "epic":
                resume_cfg_fname = Path('configs/epic.json')
            if eval_mode == "charades":
                resume_cfg_fname = Path('configs/charades.json')
            if eval_mode == "nlq":
                resume_cfg_fname = Path('configs/nlq.json')
            if eval_mode == "mq":
                resume_cfg_fname = Path('configs/mq.json')

            config = read_json(resume_cfg_fname)
            if args.config is not None:
                config.update(read_json(Path(args.config)))

        # load config file and apply custom cli options
        self._config = _update_config(config, options, args)

        # set save_dir where trained model and log will be saved.
        #save_dir = Path(self.config['trainer']['save_dir'])
        save_dir = Path(args.save_dir)
        # timestamp = datetime.now().strftime(r'%m%d_%H%M%S') if timestamp else ''
        timestamp = datetime.now().strftime(r'%m%d_%H') if timestamp else ''

        exper_name = self.config['name']

        self._save_dir = save_dir / 'models' /  timestamp
        self._web_log_dir = save_dir / 'web' /  timestamp
        self._log_dir = save_dir / 'log' /  timestamp
        self._tf_dir = save_dir / 'tf' / timestamp

        if not test:
            self.save_dir.mkdir(parents=True, exist_ok=True)
            self.log_dir.mkdir(parents=True, exist_ok=True)
            self._tf_dir.mkdir(parents=True, exist_ok=True)

        # if set, remove all previous experiments with the current config
        if vars(args).get("purge_exp_dir", False):
            for dirpath in (self._save_dir, self._log_dir, self._web_log_dir):
                config_dir = dirpath.parent
                existing = list(config_dir.glob("*"))
                print(f"purging {len(existing)} directories from config_dir...")
                tic = time.time()
                os.system(f"rm -rf {config_dir}")
                print(f"Finished purge in {time.time() - tic:.3f}s")

        # save updated config file to the checkpoint dir
        if not test:
            write_json(self.config, self.save_dir / 'config.json')

            # configure logging module
            setup_logging(self.log_dir)
            self.log_levels = {
                0: logging.WARNING,
                1: logging.INFO,
                2: logging.DEBUG
            }

    def initialize(self, name, module, *args, index=None, **kwargs):
        """
        finds a function handle with the name given as 'type' in config, and returns the
        instance initialized with corresponding keyword args given as 'args'.
        """
        if index is None:
            module_name = self[name]['type']
            module_args = dict(self[name]['args'])
            assert all([k not in module_args for k in kwargs]), 'Overwriting kwargs given in config file is not allowed'
            module_args.update(kwargs)
        else:
            module_name = self[name][index]['type']
            module_args = dict(self[name][index]['args'])

        # if parameter not in config subdict, then check if it's in global config.
        signature = inspect.signature(getattr(module, module_name).__init__)
        print(module_name)
        for param in signature.parameters.keys():
            if param not in module_args and param in self.config:
                module_args[param] = self[param]
            if module_name == 'FrozenInTime' and param == 'args':
                module_args[param] = self.args
            if module_name == 'MultiDistTextVideoDataLoader' and param == 'args':
                module_args[param] = self.args

        return getattr(module, module_name)(*args, **module_args)

    def __getitem__(self, name):
        return self.config[name]

    def get_logger(self, name, verbosity=2):
        msg_verbosity = 'verbosity option {} is invalid. Valid options are {}.'.format(verbosity,
                                                                                       self.log_levels.keys())
        assert verbosity in self.log_levels, msg_verbosity
        logger = logging.getLogger(name)
        logger.setLevel(self.log_levels[verbosity])
        return logger

    # setting read-only attributes
    @property
    def config(self):
        return self._config

    @property
    def save_dir(self):
        return self._save_dir

    @property
    def log_dir(self):
        return self._log_dir

    @property
    def tf_dir(self):
        return self._tf_dir

# helper functions used to update config dict with custom cli options
def _update_config(config, options, args):
    for opt in options:
        value = getattr(args, _get_opt_name(opt.flags))
        if value is not None:
            _set_by_path(config, opt.target, value)
    return config


def _get_opt_name(flags):
    for flg in flags:
        if flg.startswith('--'):
            return flg.replace('--', '')
    return flags[0].replace('--', '')


def _set_by_path(tree, keys, value):
    """Set a value in a nested object in tree by sequence of keys."""
    _get_by_path(tree, keys[:-1])[keys[-1]] = value


def _get_by_path(tree, keys):
    """Access a nested object in tree by sequence of keys."""
    return reduce(getitem, keys, tree)

EgoNLQ/parse_config.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

from .util import *

EgoNLQ/utils/__init__.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import numbers
from typing import List, Tuple

import torch
from torch import Tensor
from torchvision.transforms import functional_pil as F_pil, functional_tensor as F_t
from torchvision.transforms.functional import center_crop, crop


def _get_image_size(img: Tensor) -> List[int]:
    """Returns image size as [w, h]
    """
    if isinstance(img, torch.Tensor):
        return F_t._get_image_size(img)

    return F_pil._get_image_size(img)

def center_plus_four_crops(img: Tensor, size: List[int],
                              margin_h: int, margin_w: int) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
    """Crop the given image into four tiled borders and the central crop.
    """

    if isinstance(size, numbers.Number):
        size = (int(size), int(size))
    elif isinstance(size, (tuple, list)) and len(size) == 1:
        size = (size[0], size[0])

    if len(size) != 2:
        raise ValueError("Please provide only two dimensions (h, w) for size.")

    image_width, image_height = _get_image_size(img)

    crop_height, crop_width = size

    if crop_width > image_width or crop_height > image_height:
        msg = "Requested crop size {} is bigger than input size {}"
        raise ValueError(msg.format(size, (image_height, image_width)))

    if crop_width + margin_w > image_width:
        msg = "Requested margin size {} + input {} is bigger than input size {}"
        raise ValueError(msg.format((margin_h, margin_w), size, (image_height, image_width)))

    #vertical_border_height = image_height - crop_height
    #horizontal_border_height = image_width - crop_width

    #x1 = horizontal_border_height // 2
    x11 = (image_width - crop_width - 2 * margin_w) // 2
    x12 = x11 + margin_w
    x21 = x12 + crop_width
    x22 = x21 + margin_w

    y11 = (image_height - crop_height - 2 * margin_h) // 2
    y12 = y11 + margin_h
    y21 = y12 + crop_height
    y22 = y21 + margin_h

    tl = crop(img, y11, x11, margin_h, margin_w + crop_width)
    tr = crop(img, y11, x21, margin_h + crop_height, margin_w)
    bl = crop(img, y12, x11, margin_h + crop_height, margin_w)
    br = crop(img, y21, x12, margin_h, margin_w + crop_width)
    center = center_crop(img, [crop_height, crop_width])

    return tl, tr, bl, br, center



def center_plus_twohori_crops(img: Tensor, size: List[int],
                              margin_w: int) -> Tuple[Tensor, Tensor, Tensor]:
    """Crop the given image into four tiled borders and the central crop.
    """

    if isinstance(size, numbers.Number):
        size = (int(size), int(size))
    elif isinstance(size, (tuple, list)) and len(size) == 1:
        size = (size[0], size[0])

    if len(size) != 2:
        raise ValueError("Please provide only two dimensions (h, w) for size.")

    image_width, image_height = _get_image_size(img)

    crop_height, crop_width = size

    if crop_width > image_width or crop_height > image_height:
        msg = "Requested crop size {} is bigger than input size {}"
        raise ValueError(msg.format(size, (image_height, image_width)))

    if crop_width + margin_w > image_width :
        msg = "Requested margin size {} + input {} is bigger than input size {}"
        raise ValueError(msg.format((0, margin_w), size, (image_height, image_width)))

    # vertical_border_height = image_height - crop_height
    # horizontal_border_height = image_width - crop_width

    # x1 = horizontal_border_height // 2
    x11 = (image_width - crop_width - 2 * margin_w) // 2
    x12 = x11 + margin_w
    x21 = x12 + crop_width

    y11 = (image_height - crop_height) // 2

    left = crop(img, y11, x11, crop_height, margin_w)
    right = crop(img, y11, x21, crop_height, margin_w)
    center = center_crop(img, [crop_height, crop_width])

    return left, right, center

from torch import nn
class TwoHoriCrop(nn.Module):
    def __init__(self, size, margin_w):
        super().__init__()
        self.size = size
        self.margin_w = margin_w

    def forward(self, x):
        return center_plus_twohori_crops(x, self.size, self.margin_w)

if __name__ == "__main__":
    from PIL import Image

    img = Image.open('visualisations/guitar.png')
    crops = center_plus_four_crops(img, [336, 336], 112, 112)
    order = ['tl', 'tr', 'bl', 'br', 'center']

    for idx, subimg in zip(order, crops):
        subimg.save(f'visualisations/guitar_{idx}.png')

    crops = center_plus_twohori_crops(img, [448, 448], 112)
    order = ['left', 'right', 'center2']

    for idx, subimg in zip(order, crops):
        subimg.save(f'visualisations/guitar_{idx}.png')

EgoNLQ/utils/custom_transforms.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import codecs
import multiprocessing
import os
from collections import Counter

import numpy as np
from nltk.tokenize import word_tokenize
from tqdm import tqdm
import json
import pdb

from utils.data_util import (
    load_json,
    load_lines,
    load_pickle,
    save_pickle,
    time_to_index,
)
import transformers
from base.transforms import init_transform_dict, init_video_transform_dict
import decord
import torch

PAD, UNK = "<PAD>", "<UNK>"

# from utils.frozen_ego_minimum import parse_config


class EpisodicNLQProcessor:
    def __init__(self):
        super(EpisodicNLQProcessor, self).__init__()
        self.idx_counter = 0
        
        config = '/cis/home/shraman/works_meta_2022/NLQ/Ego4D_episodic_memory-main/NLQ/VSLNet_unified/configs/nlq.json'
        f = open(config)
        self.config = json.load(f)
        self.meta_dir = self.config['data_loader']['args']['meta_dir']
        self.data_dir = self.config['data_loader']['args']['data_dir']
        self.num_frame = self.config['data_loader']['args']['video_params']['num_frames']
        #self.args = args
        # Build tokenizer
        self.tokenizer = transformers.AutoTokenizer.from_pretrained(self.config['arch']['args']['text_params']['model'])
        self.transforms = init_video_transform_dict()['test']

    def reset_idx_counter(self):
        self.idx_counter = 0


    def create_meta_dict(self, split):
        
        split_files = {
            'train': 'nlq_train.json',
            'val': 'nlq_val.json',            # there is no test
            'test': 'nlq_test_unannotated.json'
        }
        # target_split_fp = split_files[self.split]

        # for split in ['train', 'val']:
        target_split_fp = split_files[split]
        ann_file = os.path.join(self.meta_dir, target_split_fp)
        with open(ann_file) as f:
            anno_json = json.load(f)

        meta_dict = dict()

        # forward video frames and text
        for anno_video in anno_json["videos"]:
            for anno_clip in anno_video["clips"]:
                clip_times = float(anno_clip["video_start_sec"]), float(anno_clip["video_end_sec"])
                meta_dict.update({anno_clip["clip_uid"]: {"video_uid": anno_video["video_uid"], "video_start": clip_times[0], "video_end": clip_times[1]}})

        return meta_dict

    def process_data_tan(self, data, scope, args):
        results = []
        meta_dict_scope = self.create_meta_dict(scope)

        for vid, data_item in tqdm(data.items(), total=len(data), desc=f"process episodic nlq {scope}"):
            fps = float(data_item["fps"])
            duration = float(data_item["num_frames"]) / fps
            zipper = zip(
                data_item["timestamps"],
                data_item["exact_times"],
                data_item["sentences"],
                data_item["annotation_uids"],
                data_item["query_idx"],
            )
            for timestamp, exact_time, sentence, ann_uid, query_idx in zipper:
                start_time = max(0.0, float(timestamp[0]) / fps)
                end_time = min(float(timestamp[1]) / fps, duration)

                if self._predictor == "EgoVLP":
                    text = self.tokenizer(sentence.strip().lower(), return_tensors='pt', padding='max_length', max_length=args.query_max_len, truncation=True)
                    text['input_ids'] = text['input_ids'].repeat(256, 1) ## This 256 is hardcoded to be the same as b_s in main.py
                    text['attention_mask'] = text['attention_mask'].repeat(256, 1) ## This 256 is hardcoded to be the same as b_s in main.py
                
                else:
                    raise NotImplementedError("args.predictor should be EgoVLP")
                
                record = {
                    "sample_id": self.idx_counter,
                    "vid": str(vid),
                    "num_frame": self.num_frame,
                    "transforms": self.transforms,
                    "data_dir": self.data_dir,
                    "s_time": start_time,
                    "e_time": end_time,
                    "exact_s_time": exact_time[0],
                    "exact_e_time": exact_time[1],
                    "duration": duration,
                    "query": sentence.strip().lower(),
                    "annotation_uid": ann_uid,
                    "query_idx": query_idx,
                    "query_feats":text,
                    "meta_dict_scope": meta_dict_scope,
                }
                results.append(record)
                self.idx_counter += 1


        return results

    def convert(self, args):
        predictor = args.predictor
        self._predictor = predictor
        self.reset_idx_counter()
        
        # load raw data
        train_data = load_json(os.path.join('jsons', "train.json"))
        val_data = load_json(os.path.join('jsons', "val.json"))
        test_data = load_json(os.path.join('jsons', "test.json"))

        # process data
        train_set = self.process_data_tan(train_data, scope="train", args=args)
        val_set = self.process_data_tan(val_data, scope="val", args=args)
        test_set = self.process_data_tan(test_data, scope="test", args=args)
        return train_set, val_set, test_set


def sample_frames_clips(start, end, vlen, acc_samples):
    start = max(0, start)
    end = min(vlen, end)

    intervals = np.linspace(start=start, stop=end, num=int(acc_samples) + 1).astype(int)
    ranges = []
    for idx, interv in enumerate(intervals[:-1]):
        ranges.append((interv, intervals[idx + 1] - 1))
        frame_idxs = [(x[0] + x[1]) // 2 for x in ranges
                      ]
    return frame_idxs

def read_frames_decord_start_end(video_path, start, end, num_frames):
    video_reader = decord.VideoReader(video_path, num_threads=8)
    vlen = len(video_reader)
    frame_idxs = sample_frames_clips(start, end, vlen, num_frames + 1)
    video_reader.skip_frames(1)
    frames = video_reader.get_batch(frame_idxs)

    frames = frames.float() / 255
    frames = frames.permute(0, 3, 1, 2)
    return frames, frame_idxs

def _get_video_path(sample, data_dir):
    rel_video_fp = sample
    full_video_fp = os.path.join(data_dir, rel_video_fp + '.mp4')
    return full_video_fp, rel_video_fp

def visual_feature_sampling(visual_feature, max_num_clips):
    num_clips = visual_feature.shape[0]
    if num_clips <= max_num_clips:
        return visual_feature
    idxs = np.arange(0, max_num_clips + 1, 1.0) / max_num_clips * num_clips
    idxs = np.round(idxs).astype(np.int32)
    idxs[idxs > num_clips - 1] = num_clips - 1
    new_visual_feature = []
    for i in range(max_num_clips):
        s_idx, e_idx = idxs[i], idxs[i + 1]
        if s_idx < e_idx:
            new_visual_feature.append(np.mean(visual_feature[s_idx:e_idx].numpy(), axis=0))
        else:
            new_visual_feature.append(visual_feature[s_idx].numpy())
    new_visual_feature = np.asarray(new_visual_feature)
    #pdb.set_trace()
    return torch.from_numpy(new_visual_feature)
    
def _get_video_text_feats(meta_dict_scope, clip_uid, num_frames, args, transforms, data_dir):
    
    video_fp, _ = _get_video_path(meta_dict_scope[clip_uid]["video_uid"], data_dir)

    fps = 1.87
    try:
        imgs, idxs = read_frames_decord_start_end(video_fp, meta_dict_scope[clip_uid]["video_start"]*30, meta_dict_scope[clip_uid]["video_end"]*30,
                                               (meta_dict_scope[clip_uid]["video_end"]-meta_dict_scope[clip_uid]["video_start"]) * fps * num_frames)
    except:
        print(f"Warning: missing video file {video_fp}.")

    if transforms is not None:
        imgs = imgs.transpose(0, 1)  # [T, C, H, W] ---> [C, T, H, W]
        imgs = transforms(imgs)
        imgs = imgs.transpose(0, 1)  # recover

    f, c, h, w = imgs.shape[0], imgs.shape[1], imgs.shape[2], imgs.shape[3]
    imgs = imgs[:(f // num_frames * num_frames), ]
    imgs = imgs.reshape(-1, num_frames, c, h, w)

    #imgs = visual_feature_sampling(imgs, max_num_clips=args.max_pos_len) ## max_num_clips = args.max_pos_len

    return imgs


def dataset_gen_EgoVLP(data, vfeat_lens, tokenizer, max_pos_len, scope, args, num_workers=1):
    dataset = []
    if not os.path.exists('./saved_video_frames/' + str(scope)):
        os.mkdir('./saved_video_frames/' + str(scope))

    print("DATASET INITIALIZED!!!!!!!")
    for indx, record in tqdm(enumerate(data), desc="dataset_gen_EgoVLP"):
        vid = record["vid"]
        if vid not in vfeat_lens:
            continue
        s_ind, e_ind, _ = time_to_index(
            record["s_time"], record["e_time"], vfeat_lens[vid], record["duration"]
        )
        
        video_frame_path = './saved_video_frames/' + str(scope) + '/' + str(record["vid"]) + '_frames.pt'

        if 'video_frame_path' in record['meta_dict_scope'][record['vid']]:
            video_path = record['meta_dict_scope'][record['vid']]['video_frame_path']
        else:
            video = _get_video_text_feats(record["meta_dict_scope"], record["vid"], record["num_frame"], args, record["transforms"], record["data_dir"])
            record['meta_dict_scope'][record['vid']].update({'video_frame_path':video_frame_path})
            video_path = video_frame_path
            torch.save(video, video_frame_path)

        result = {
            "sample_id": record["sample_id"],
            "vid": record["vid"],
            "video_frame_path": video_path,
            "s_time": record["s_time"],
            "e_time": record["e_time"],
            "duration": record["duration"],
            "query": record["query"],
            "s_ind": int(s_ind),
            "e_ind": int(e_ind),
            "v_len": vfeat_lens[vid],
            "text_tokens": record["query_feats"],
            "annotation_uid": record["annotation_uid"],
            "query_idx": record["query_idx"],
        } 
        #pdb.set_trace()
        dataset.append(result)
        
        #if indx == 40:
        #   print(len(dataset))
        #   break

    return dataset


def gen_or_load_dataset(args):
    if not os.path.exists(args.save_dir):
        os.makedirs(args.save_dir)
    data_dir = os.path.join("data", "dataset", args.task)
    #feature_dir = os.path.join("data", "features", args.task, args.fv)
    #if args.suffix is None:
    #    save_path = os.path.join(
    #        args.save_dir,
    #        "_".join(
    #            [args.task, args.fv, str(args.max_pos_len)]
    #        )
    #        + ".pkl",
    #    )
    #else:
    #    save_path = os.path.join(
    #        args.save_dir,
    #        "_".join(
    #            [args.task, args.fv, str(args.max_pos_len), args.suffix]
    #        )
    #        + ".pkl",
    #    )
    
    #if os.path.exists(save_path): ### Use this block with care, preferable to compute text fetaures before running
    #    dataset = load_pickle(save_path)
    #    return dataset
    

    feat_len_path = os.path.join("feature_shapes.json")
    #emb_path = os.path.join("data", "features", "glove.840B.300d.txt")
    # load video feature length
    vfeat_lens = load_json(feat_len_path)
    for vid, vfeat_len in vfeat_lens.items():
        vfeat_lens[vid] = min(args.max_pos_len, vfeat_len)
    # load data
    processor = EpisodicNLQProcessor()

    train_data, val_data, test_data = processor.convert(args)
    # generate dataset
    data_list = (
        [train_data, test_data]
        if val_data is None
        else [train_data, val_data, test_data]
    )
    
    if args.predictor == "EgoVLP":

        tokenizer = None
        train_set = dataset_gen_EgoVLP(
            train_data,
            vfeat_lens,
            tokenizer,
            args.max_pos_len,
            "train",
            args,
            num_workers=args.num_workers,
        )
        if val_data:
            val_set = dataset_gen_EgoVLP(
                val_data,
                vfeat_lens,
                tokenizer,
                args.max_pos_len,
                "val",
                args,
                num_workers=args.num_workers,
            )
        else:
            val_set = None
        test_set = dataset_gen_EgoVLP(
            test_data,
            vfeat_lens,
            tokenizer,
            args.max_pos_len,
            "test",
            args,
            num_workers=args.num_workers,
        )
        n_val = 0 if val_set is None else len(val_set)
        dataset = {
            "train_set": train_set,
            "val_set": val_set,
            "test_set": test_set,
            "n_train": len(train_set),
            "n_val": n_val,
            "n_test": len(test_set),
        }
    else:
        raise NotImplementedError("args.predictor should be EgoVLP")

    #save_pickle(dataset, save_path)
    return dataset

EgoNLQ/utils/data_gen.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import numpy as np
import torch
import torch.utils.data
from torch.nn.utils.rnn import pad_sequence
import pdb

from utils.data_util import pad_seq, pad_char_seq, pad_video_seq
from utils.data_gen import visual_feature_sampling

class Dataset(torch.utils.data.Dataset):
    def __init__(self, feature_root_list, args):
        super(Dataset, self).__init__()
        self.feature_root_list = feature_root_list 
        self.args = args

    def __getitem__(self, index):

        feature = torch.load(self.feature_root_list[index])
        video_feature = feature["fused_video_features"]
        video_feature = visual_feature_sampling(video_feature, max_num_clips=self.args.max_pos_len)
        dual_text_feature = feature["dual_text_features"][0]
        query_attention_mask = feature["text_token"]["attention_mask"][0]
        s_ind, e_ind = int(feature["s_ind"].item()), int(feature["e_ind"].item())
        return feature, video_feature, dual_text_feature, s_ind, e_ind, query_attention_mask

    def __len__(self):
        return len(self.feature_root_list)


def train_collate_fn(data):

    feature, video_features, dual_text_feature, s_inds, e_inds, query_attention_mask = zip(*data)

    # process video features
    vfeats, vfeat_lens = pad_video_seq(video_features)

    #for _idx in range(len(vfeats)):
    #    print(vfeats[_idx].shape)

    #pdb.set_trace()

    vfeats = np.asarray([vf.numpy() for vf in vfeats], dtype=np.float32)  # (batch_size, v_seq_len, v_dim)
    vfeat_lens = np.asarray(vfeat_lens, dtype=np.int32)  # (batch_size, )
    # process labels
    max_len = np.max(vfeat_lens)
    batch_size = vfeat_lens.shape[0]
    s_labels = np.asarray(s_inds, dtype=np.int64)
    e_labels = np.asarray(e_inds, dtype=np.int64)
    h_labels = np.zeros(shape=[batch_size, max_len], dtype=np.int32)
    extend = 0.1
    for idx in range(batch_size):
        st, et = s_inds[idx], e_inds[idx]
        cur_max_len = vfeat_lens[idx]
        extend_len = round(extend * float(et - st + 1))
        if extend_len > 0:
            st_ = max(0, st - extend_len)
            et_ = min(et + extend_len, cur_max_len - 1)
            h_labels[idx][st_ : (et_ + 1)] = 1
        else:
            h_labels[idx][st : (et + 1)] = 1
    # convert to torch tensor
    vfeats = torch.tensor(vfeats, dtype=torch.float32)
    vfeat_lens = torch.tensor(vfeat_lens, dtype=torch.int64)
    dual_text_feature = torch.stack(list(dual_text_feature), dim=0)
    query_attention_mask = torch.stack(list(query_attention_mask), dim=0)
    s_labels = torch.tensor(s_labels, dtype=torch.int64)
    e_labels = torch.tensor(e_labels, dtype=torch.int64)
    h_labels = torch.tensor(h_labels, dtype=torch.int64)
    return feature, vfeats, vfeat_lens, dual_text_feature, s_labels, e_labels, h_labels, query_attention_mask


def test_collate_fn(data):

    feature, video_features, dual_text_feature, _, _, query_attention_mask = zip(*data)

    # process video features
    vfeats, vfeat_lens = pad_video_seq(video_features)
    vfeats = np.asarray([vf.numpy() for vf in vfeats], dtype=np.float32)  # (batch_size, v_seq_len, v_dim)
    vfeat_lens = np.asarray(vfeat_lens, dtype=np.int32)  # (batch_size, )
    # convert to torch tensoet number

    vfeats = torch.tensor(vfeats, dtype=torch.float32)
    vfeat_lens = torch.tensor(vfeat_lens, dtype=torch.int64)
    dual_text_feature = torch.stack(list(dual_text_feature), dim=0)
    query_attention_mask = torch.stack(list(query_attention_mask), dim=0)
    return feature, vfeats, vfeat_lens, dual_text_feature, query_attention_mask


def get_train_loader(feature_root_list, args):
    train_set = Dataset(feature_root_list=feature_root_list, args=args)
    train_loader = torch.utils.data.DataLoader(
        dataset=train_set,
        batch_size=args.batch_size,
        shuffle=True,
        collate_fn=train_collate_fn,
    )
    return train_loader


def get_test_loader(feature_root_list, args):
    test_set = Dataset(feature_root_list=feature_root_list, args=args)
    test_loader = torch.utils.data.DataLoader(
        dataset=test_set,
        batch_size=args.batch_size,
        shuffle=False,
        collate_fn=test_collate_fn,
    )
    return test_loader

EgoNLQ/utils/data_loader.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import glob
import json
import os
import pickle
import pdb

import numpy as np
import torch
from tqdm import tqdm


def load_json(filename):
    with open(filename, mode="r", encoding="utf-8") as f:
        data = json.load(f)
    return data


def save_json(data, filename, save_pretty=False, sort_keys=False):
    with open(filename, mode="w", encoding="utf-8") as f:
        if save_pretty:
            f.write(json.dumps(data, indent=4, sort_keys=sort_keys))
        else:
            json.dump(data, f)


def load_lines(filename):
    with open(filename, mode="r", encoding="utf-8") as f:
        return [e.strip("\n") for e in f.readlines()]


def save_lines(data, filename):
    with open(filename, mode="w", encoding="utf-8") as f:
        f.write("\n".join(data))


def load_pickle(filename):
    with open(filename, mode="rb") as handle:
        data = pickle.load(handle)
        return data


def save_pickle(data, filename):
    with open(filename, mode="wb") as handle:
        pickle.dump(data, handle, protocol=pickle.HIGHEST_PROTOCOL)


def load_video_features(root, max_position_length):
    video_features = dict()
    extension = "*.pt"

    filenames = glob.glob(os.path.join(root, extension))
    for filename in tqdm(filenames, total=len(filenames), desc="load video features"):
        video_id = filename.split("/")[-1].split(".")[0]
        feature = torch.load(filename).numpy()
        if max_position_length is None:
            video_features[video_id] = feature
        else:
            new_feature = visual_feature_sampling(
                feature, max_num_clips=max_position_length
            )
            video_features[video_id] = new_feature
            # if new_feature.shape[0] != feature.shape[0]:
            #    print(f"Reduced: {feature.shape[0]} --> {new_feature.shape[0]}")
    return video_features


def visual_feature_sampling(visual_feature, max_num_clips):
    num_clips = visual_feature.shape[0]
    if num_clips <= max_num_clips:
        return visual_feature
    idxs = np.arange(0, max_num_clips + 1, 1.0) / max_num_clips * num_clips
    idxs = np.round(idxs).astype(np.int32)
    idxs[idxs > num_clips - 1] = num_clips - 1
    new_visual_feature = []
    for i in range(max_num_clips):
        s_idx, e_idx = idxs[i], idxs[i + 1]
        if s_idx < e_idx:
            new_visual_feature.append(np.mean(visual_feature[s_idx:e_idx], axis=0))
        else:
            new_visual_feature.append(visual_feature[s_idx])
    new_visual_feature = np.asarray(new_visual_feature)
    return new_visual_feature


def compute_overlap(pred, gt):
    # check format
    assert isinstance(pred, list) and isinstance(gt, list)
    pred_is_list = isinstance(pred[0], list)
    gt_is_list = isinstance(gt[0], list)
    pred = pred if pred_is_list else [pred]
    gt = gt if gt_is_list else [gt]
    # compute overlap
    pred, gt = np.array(pred), np.array(gt)
    inter_left = np.maximum(pred[:, 0, None], gt[None, :, 0])
    inter_right = np.minimum(pred[:, 1, None], gt[None, :, 1])
    inter = np.maximum(0.0, inter_right - inter_left)
    union_left = np.minimum(pred[:, 0, None], gt[None, :, 0])
    union_right = np.maximum(pred[:, 1, None], gt[None, :, 1])
    union = np.maximum(1e-12, union_right - union_left)
    overlap = 1.0 * inter / union
    # reformat output
    overlap = overlap if gt_is_list else overlap[:, 0]
    overlap = overlap if pred_is_list else overlap[0]
    return overlap


def time_to_index(start_time, end_time, num_units, duration):
    s_times = np.arange(0, num_units).astype(np.float32) / float(num_units) * duration
    e_times = (
        np.arange(1, num_units + 1).astype(np.float32) / float(num_units) * duration
    )
    candidates = np.stack(
        [
            np.repeat(s_times[:, None], repeats=num_units, axis=1),
            np.repeat(e_times[None, :], repeats=num_units, axis=0),
        ],
        axis=2,
    ).reshape((-1, 2))
    overlaps = compute_overlap(candidates.tolist(), [start_time, end_time]).reshape(
        num_units, num_units
    )
    start_index = np.argmax(overlaps) // num_units
    end_index = np.argmax(overlaps) % num_units
    return start_index, end_index, overlaps


def index_to_time(start_index, end_index, num_units, duration):
    s_times = np.arange(0, num_units).astype(np.float32) * duration / float(num_units)
    e_times = (
        np.arange(1, num_units + 1).astype(np.float32) * duration / float(num_units)
    )
    start_time = s_times[start_index]
    end_time = e_times[end_index]
    return start_time, end_time


def pad_seq(sequences, pad_tok=None, max_length=None):
    if pad_tok is None:
        pad_tok = 0  # 0: "PAD" for words and chars, "PAD" for tags
    if max_length is None:
        max_length = max([len(seq) for seq in sequences])
    sequence_padded, sequence_length = [], []
    for seq in sequences:
        seq_ = seq[:max_length] + [pad_tok] * max(max_length - len(seq), 0)
        sequence_padded.append(seq_)
        sequence_length.append(min(len(seq), max_length))
    return sequence_padded, sequence_length


def pad_char_seq(sequences, max_length=None, max_length_2=None):
    sequence_padded, sequence_length = [], []
    if max_length is None:
        max_length = max(map(lambda x: len(x), sequences))
    if max_length_2 is None:
        max_length_2 = max([max(map(lambda x: len(x), seq)) for seq in sequences])
    for seq in sequences:
        sp, sl = pad_seq(seq, max_length=max_length_2)
        sequence_padded.append(sp)
        sequence_length.append(sl)
    sequence_padded, _ = pad_seq(
        sequence_padded, pad_tok=[0] * max_length_2, max_length=max_length
    )
    sequence_length, _ = pad_seq(sequence_length, max_length=max_length)
    return sequence_padded, sequence_length


def pad_video_seq(sequences, max_length=None):
    if max_length is None:
        max_length = max([vfeat.shape[0] for vfeat in sequences])
    feature_length = sequences[0].shape[1]
    sequence_padded, sequence_length = [], []
    for seq in sequences:
        add_length = max_length - seq.shape[0]
        sequence_length.append(seq.shape[0])
        if add_length > 0:
            add_feature = np.zeros(shape=[add_length, feature_length], dtype=np.float32)
            seq_ = np.concatenate([seq, add_feature], axis=0)
            seq_ = torch.from_numpy(seq_)
        else:
            seq_ = seq
        sequence_padded.append(seq_)
    return sequence_padded, sequence_length

EgoNLQ/utils/data_util.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


"""
Script to evaluate performance of any model for Ego4d Episodic Memory.

Natural Language Queries (NLQ)
"""

from __future__ import absolute_import, division, print_function, unicode_literals

import argparse
import json

import numpy as np
import terminaltables


def display_results(results, mIoU, thresholds, topK, title=None):
    display_data = [
        [f"Rank@{ii}\nmIoU@{jj}" for ii in topK for jj in thresholds] + ["mIoU"]
    ]
    results *= 100
    mIoU *= 100
    display_data.append(
        [
            f"{results[jj][ii]:.02f}"
            for ii in range(len(topK))
            for jj in range(len(thresholds))
        ]
        + [f"{mIoU:.02f}"]
    )
    table = terminaltables.AsciiTable(display_data, title)
    for ii in range(len(thresholds) * len(topK)):
        table.justify_columns[ii] = "center"
    return table.table


def compute_IoU(pred, gt):
    """Compute the IoU given predicted and ground truth windows."""
    assert isinstance(pred, list) and isinstance(gt, list)
    pred_is_list = isinstance(pred[0], list)
    gt_is_list = isinstance(gt[0], list)
    if not pred_is_list:
        pred = [pred]
    if not gt_is_list:
        gt = [gt]
    pred, gt = np.array(pred), np.array(gt)
    inter_left = np.maximum(pred[:, 0, None], gt[None, :, 0])
    inter_right = np.minimum(pred[:, 1, None], gt[None, :, 1])
    inter = np.maximum(0.0, inter_right - inter_left)
    union_left = np.minimum(pred[:, 0, None], gt[None, :, 0])
    union_right = np.maximum(pred[:, 1, None], gt[None, :, 1])
    union = np.maximum(0.0, union_right - union_left)
    overlap = 1.0 * inter / union
    if not gt_is_list:
        overlap = overlap[:, 0]
    if not pred_is_list:
        overlap = overlap[0]
    return overlap


def evaluate_nlq_performance(
    predictions, ground_truth, thresholds, topK, per_instance=False
):
    """Evalutes the performances."""
    gt_dict = {}
    num_gt_queries = 0

    for video_datum in ground_truth["videos"]:
        for clip_datum in video_datum["clips"]:
            clip_uid = clip_datum["clip_uid"]
            for ann_datum in clip_datum["annotations"]:
                key = (clip_uid, ann_datum["annotation_uid"])
                gt_dict[key] = ann_datum
                num_gt_queries += len(ann_datum["language_queries"])

    results = [[[] for _ in topK] for _ in thresholds]
    average_IoU = []
    num_instances = 0
    pred_datum_list = []
    gt_query_datum_list = []

    for pred_datum in predictions:
        key = (pred_datum["clip_uid"], pred_datum["annotation_uid"])
        assert key in gt_dict, "Instance not present!"
        query_id = pred_datum["query_idx"]
        gt_datum = gt_dict[key]
        gt_query_datum = gt_datum["language_queries"][query_id]

        # Compute overlap and recalls.
        overlap = compute_IoU(
            pred_datum["predicted_times"],
            [[gt_query_datum["clip_start_sec"], gt_query_datum["clip_end_sec"]]],
        )
        pred_datum_list.append(pred_datum)
        gt_query_datum_list.append(gt_query_datum)

        average_IoU.append(np.mean(np.sort(overlap[0])[-3:]))
        for tt, threshold in enumerate(thresholds):
            for rr, KK in enumerate(topK):
                results[tt][rr].append((overlap > threshold)[:KK].any())
        num_instances += 1

    mean_results = np.array(results).mean(axis=-1)
    mIoU = np.mean(average_IoU)
    print(f"Evaluated: {num_instances} / {num_gt_queries} instances")
    if per_instance:
        per_instance_results = {
            "overlap": overlap,
            "average_IoU": average_IoU,
            "results": results,
            "pred_datum_list": pred_datum_list,
            "gt_query_datum_list": gt_query_datum_list
        }
        return mean_results, mIoU, per_instance_results
    else:
        return mean_results, mIoU


def main(args):
    print(f"""Reading predictions: {args["model_prediction_json"]}""")
    with open(args["model_prediction_json"], "r") as file_id:
        predictions = json.load(file_id)

    print(f"""Reading gt: {args["ground_truth_json"]}""")
    with open(args["ground_truth_json"], "r") as file_id:
        ground_truth = json.load(file_id)

    assert predictions.get("version", None) == "1.0", "Ego4D version does not match!"
    assert predictions.get("challenge", None) == "ego4d_nlq_challenge", (
        "Ego4D challenge does not match!"
    )
    results, mIoU = evaluate_nlq_performance(
        predictions["results"], ground_truth, args["thresholds"], args["topK"]
    )
    print(display_results(results, mIoU, args["thresholds"], args["topK"]))


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--ground_truth_json",
        required=True,
        help="Ground truth temporal windows",
    )
    parser.add_argument(
        "--model_prediction_json",
        required=True,
        help="Model predicted temporal windows",
    )
    parser.add_argument(
        "--thresholds",
        required=True,
        nargs="+",
        type=float,
        help="Thresholds for IoU computation",
    )
    parser.add_argument(
        "--topK",
        required=True,
        nargs="+",
        type=int,
        help="Top K for computing recall@k",
    )

    try:
        parsed_args = vars(parser.parse_args())
    except (IOError) as msg:
        parser.error(str(msg))
    main(parsed_args)

EgoNLQ/utils/evaluate_ego4d_nlq.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import os
import dominate
from dominate.tags import a, attr, br, h3, img, meta, p, source, span, table, td, tr, video


class HTML:
    """This HTML class allows us to save images and write texts into a single HTML file.

     It consists of functions such as <add_header> (add a text header to the HTML file),
     <add_images> (add a row of images to the HTML file), and <save> (save the HTML to the disk).
     It is based on Python library 'dominate', a Python library for creating and
     manipulating HTML documents using a DOM API.
    """

    def __init__(self, web_dir, title, refresh=0):
        """Initialize the HTML classes

        Parameters:
            web_dir (str) -- a directory that stores the webpage. HTML file will be
            created at <web_dir>/index.html; images will be saved at <web_dir/images/
            title (str)   -- the webpage name
            reflect (int) -- how often the website refresh itself; if 0; no refreshing
        """
        self.title = title
        self.web_dir = web_dir
        self.img_dir = os.path.join(self.web_dir, "images")
        if not os.path.exists(self.web_dir):
            os.makedirs(self.web_dir)
        if not os.path.exists(self.img_dir):
            os.makedirs(self.img_dir)

        self.doc = dominate.document(title=title)
        if refresh > 0:
            with self.doc.head:
                meta(http_equiv="refresh", content=str(refresh))

    def get_image_dir(self):
        """Return the directory that stores images"""
        return self.img_dir

    def add_header(self, text):
        """Insert a header to the HTML file

        Parameters:
            text (str) -- the header text
        """
        with self.doc:
            h3(text)

    def add_videos(self, vids, txts, links, width=400, hidden_tag="hidden"):
        """add images to the HTML file

        Parameters:
            vids (str list)   -- a list of image paths
            txts (str list)  -- a list of image names shown on the website
            links (str list) --  a list of hyperref links; when you click an image,
            it will redirect you to a new page
        """
        self.t = table(border=1, style="table-layout: fixed;")  # Insert a table
        self.doc.add(self.t)
        colors = ["red", "blue", "gold", "salman"]
        with self.t:
            with tr():
                for vid, txt, link in zip(vids, txts, links):
                    td_style = "word-wrap: break-word; width:{}px".format(width)
                    with td(style=td_style, halign="center", valign="top"):
                        with p():
                            vid_path = str(vid)
                            if vid_path == hidden_tag:
                                p_style = "font-weight: bold; width:{}px;"
                                p_style = p_style.format(width * 3)
                                p("hidden video", style=p_style)
                            else:
                                with a(href=str(link)):
                                    with video():
                                        attr(controls="controls")
                                        source(src=vid_path, type="video/mp4")
                            br()
                            rows = txt.split("<br>")
                            for idx, row in enumerate(rows):
                                color = colors[idx % len(colors)]
                                bold_tag = "<b>"
                                if not row.startswith(bold_tag):
                                    s_style = "color:{};".format(color)
                                else:
                                    s_style = "color:black; font-weight: bold;"
                                    row = row[len(bold_tag):]
                                span(row, style=s_style)

    def add_images(self, ims, txts, links, width=400):
        """add images to the HTML file

        Parameters:
            ims (str list)   -- a list of image paths
            txts (str list)  -- a list of image names shown on the website
            links (str list) --  a list of hyperref links; when you click an image,
            it will redirect you to a new page
        """
        self.t = table(border=1, style="table-layout: fixed;")  # Insert a table
        self.doc.add(self.t)
        colors = ["red", "blue", "gold", "salman"]
        with self.t:
            with tr():
                for im, txt, link in zip(ims, txts, links):
                    td_style = "word-wrap: break-word;"
                    with td(style=td_style, halign="center", valign="top"):
                        with p():
                            with a(href=link):
                                img(
                                    style="width:%dpx" % width,
                                    src=im,
                                )
                            br()
                            rows = txt.split("<br>")
                            for idx, row in enumerate(rows):
                                color = colors[idx % len(colors)]
                                bold_tag = "<b>"
                                if not row.startswith(bold_tag):
                                    s_style = "color:{};".format(color)
                                else:
                                    s_style = "color:black; font-weight: bold;"
                                    row = row[len(bold_tag):]
                                span(row, style=s_style)

    def save(self):
        """save the current content to the HMTL file"""
        html_file = "%s/index.html" % self.web_dir
        f = open(html_file, "wt")
        f.write(self.doc.render())
        f.close()


if __name__ == "__main__":  # we show an example usage here.
    html = HTML("web/", "test_html")
    html.add_header("hello world")

    ims, txts, links = [], [], []
    for n in range(4):
        ims.append("image_%d.png" % n)
        txts.append("text_%d" % n)
        links.append("image_%d.png" % n)
    html.add_images(ims, txts, links)
    html.save()

EgoNLQ/utils/html.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import numpy as np

def calculate_mAP(sim_mat, relevancy_matrix):
    """
    Computes the mean average precision according to the following formula of
    average precision:
    \frac{\sum_{k=1}^n p(k) x rel(k)}{num_rel_docs}

    where p(k) is the precision at k, rel(k) is an indicator function
    determining whether the kth returned item is relevant or not and
    num_rel_docs is the number of relevant items to find within the search.

    The mean average precision is the mean of the average precision for each
    query item (i.e row in the matrix)

    This function takes in two parameters:
        - sim_mat: a NxM matrix which represents the similarity between two
        modalities (with modality 1 being of size N and modality 2 of size M).
        - relevancy_matrix: an NxM matrix which represents the relevancy between two
        modalities of items (with modality 1 being of size N and modality 2 of
        size M).
    """
    #Find the order of the items in modality 2 according to modality 1
    ranked_order = (-sim_mat).argsort()
    ranked_sim_mat = sim_mat[np.arange(sim_mat.shape[0])[:, None], ranked_order]
    #re-order the relevancy matrix to accommodate the proposals
    ranked_rel_mat = relevancy_matrix[np.arange(relevancy_matrix.shape[0])[:, None], ranked_order]

    #find the number of relevant items found at each k
    cumulative_rel_mat = np.cumsum(ranked_rel_mat, axis=1)
    #Mask this ensuring that it is non zero if the kth term is 1 (rel(k) above)
    cumulative_rel_mat[ranked_rel_mat != 1] = 0
    #find the divisor for p(k)
    divisor = np.arange(ranked_rel_mat.shape[1]) + 1

    #find the number of relevant docs per query item
    number_rel_docs = np.sum(ranked_rel_mat==1, axis=1)

    #find the average precision per query, within np.sum finds p(k) * rel(k)
    avg_precision = np.sum(cumulative_rel_mat / divisor, axis=1) / number_rel_docs
    mAP = np.mean(avg_precision)
    return mAP

EgoNLQ/utils/mAP.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import numpy as np

def calculate_DCG(similarity_matrix, relevancy_matrix, k_counts):
    """
    Calculates the Discounted Cumulative Gain (DCG) between two modalities for
    the first modality.
    DCG = \sum_{i=1}^k \frac{rel_i}{log_2(i + 1)}
    i.e. the sum of the k relevant retrievals which is calculated as the scaled
    relevancy for the ith item. The scale is designed such that early
    retrievals are more important than later retrievals.
    Params:
        - similarity_matrix: matrix of size n1 x n2 where n1 is the number of
          items in the first modality and n2 is the number of items in the
          second modality. The [ith,jth] element is the predicted similarity
          between the ith item from the first modality and the jth item from
          the second modality.
        - relevancy_matrix: matrix of size n1 x n2 (see similarity_matrix
          above). The [ith, jth] element is the semantic relevancy between the
          ith item from the first modality and the jth item from the second
          modality.
        - k_counts: matrix of size n1 x n2 (see similarity_matrix above) which
          includes information on which items to use to calculate the DCG for
          (see calculate_k_counts for more info on this matrix).
    Returns:
        - The DCG for each item in the first modality, a n1 length vector.
    """
    x_sz, y_sz = similarity_matrix.shape
    ranks = np.argsort(similarity_matrix)[:, ::-1]
    #Create vector of size (n,) where n is the length of the last dimension in
    #similarity matrix
    #This vector is of the form log(i+1)
    logs = np.log2(np.arange(y_sz) + 2)
    #Convert logs into the divisor for the DCG calculation, of size similarity
    #matrix
    divisors = np.repeat(np.expand_dims(logs, axis=0), x_sz, axis=0)

    #mask out the sorted relevancy matrix to only use the first k relevant
    #retrievals for each item.
    columns = np.repeat(np.expand_dims(np.arange(x_sz), axis=1), y_sz, axis=1)
    numerators = relevancy_matrix[columns, ranks] * k_counts
    #Calculate the final DCG score (note that this isn't expected to sum to 1)
    return np.sum(numerators / divisors, axis=1)

def calculate_k_counts(relevancy_matrix):
    """
    Works out the maximum number of allowed retrievals when working out the
    Discounted Cumulative Gain. For each query the DCG only uses the first k
    items retrieved which constitute the k relevant items for that query
    (otherwise the nDCG scores can be deceptively high for bad rankings).
    Params:
        - relevancy_matrix: matrix of size n1 x n2 where n1 is the number of
          items in the first modality and n2 is the number of items in the
          second modality.  The [ith, jth] element is the semantic relevancy
          between the ith item from the first modality and the jth item from
          the second modality.
    Returns:
        - Matrix of size n1 x n2 (see relevancy matrix for more info). This is
          created as a mask such that if the [ith, jth] element is 1 it
          represents a valid item to use for the calculation of DCG for the
          ith item after sorting. For example, if relevancy matrix of:
        [[1, 0.5, 0],
          [0, 0  , 1]]
          is given, then the k_counts matrix will be:
        [[1, 1, 0],
         [1, 0, 0]]
         i.e. the first row has 2 non-zero items, so the first two retrieved
         items should be used in the calculation. In the second row there is
         only 1 relevant item, therefore only the first retrieved item should
         be used for the DCG calculation.
    """
    return (np.sort(relevancy_matrix)[:, ::-1] > 0).astype(int)


def calculate_IDCG(relevancy_matrix, k_counts):
    """
    Calculates the Ideal Discounted Cumulative Gain (IDCG) which is the value
    of the Discounted Cumulative Gain (DCG) for a perfect retrieval, i.e. the
    items in the second modality were retrieved in order of their descending
    relevancy.
    Params:
        - relevancy_matrix: matrix of size n1 x n2 where n1 is the number of
          items in the first modality and n2 is the number of items in the
          second modality. The [ith, jth] element is the semantic relevancy
          between the ith item from the first modality and the jth item from
          the second modality.
        - k_counts: matrix of size n1 x n2 (see similarity_matrix above) which
          includes information on which items to use to calculate the DCG for
          (see calculate_k_counts for more info on this matrix).
    """
    return calculate_DCG(relevancy_matrix, relevancy_matrix, k_counts)

def calculate_nDCG(similarity_matrix, relevancy_matrix, k_counts=None, IDCG=None, reduction='mean'):
    """
    Calculates the normalised Discounted Cumulative Gain (nDCG) between two
    modalities for the first modality using the Discounted Cumulative Gain
    (DCG) and the Ideal Discounted Cumulative Gain (IDCG).

    nDCG = \frac{DCG}{IDCG}

    Params:
        - similarity_matrix: matrix of size n1 x n2 where n1 is the number of
          items in the first modality and n2 is the number of items in the second
          modality. The [ith,jth] element is the predicted similarity between
          the ith item from the first modality and the jth item from the second
          modality.
        - relevancy_matrix: matrix of size n1 x n2 (see similarity_matrix
          above). The [ith, jth] element is the semantic relevancy between the
          ith item from the first modality and the jth item from the second
          modality.
        - k_counts: optional parameter: matrix of size n1 x n2 (see
          similarity_matrix above) which includes information on which items to
          use to calculate the DCG for (see calculate_k_counts for more info on
          this matrix). This will be calculated using calculate_IDCG if not
          present, but should be pre-processed for efficiency.
        - IDCG: Optional parameter which includes the pre-processed Ideal
          Discounted Cumulative Gain (IDCG). This is a vector of size n1 (see
          similarity_matrix above) which contains the IDCG value for each item
          from the first modality. This will be calculated using calculate_IDCG
          if not present, but should be pre-processed for efficiency.
        - reduction: what to use to reduce the different nDCG scores. By
          default this applies np.mean across all different queries.
    Returns:
        - The nDCG values for the first modality.
    """
    if k_counts is None:
        k_counts = calculate_k_counts(relevancy_matrix)
    DCG = calculate_DCG(similarity_matrix, relevancy_matrix, k_counts)
    if IDCG is None:
        IDCG = calculate_IDCG(relevancy_matrix, k_counts)
    if reduction == 'mean':
        return np.mean(DCG / IDCG)
    elif reduction is None:
        return DCG / IDCG


if __name__ == '__main__':
    sim_matrix = np.array([
        [1.0, 0.7, 0.4, 0.0],
        [0.3, 0.9, 0.6, 0.1],
        [0.2, 0.5, 0.8, 0.4]
    ])
    rel_matrix = np.array([
        [1.0, 0.5, 0.25, 0.0],
        [0.0, 1.0, 0.4, 0.0],
        [0.5, 0.3, 1.0, 0.0]
    ])
    k_counts = np.array([
        [1, 1, 1, 0],
        [1, 1, 0, 0],
        [1, 1, 1, 0]
    ])
    assert (k_counts == calculate_k_counts(rel_matrix)).all()
    nDCG = calculate_nDCG(sim_matrix, rel_matrix, k_counts)
    assert nDCG == 0.9371789900735429
    DCG = calculate_DCG(sim_matrix, rel_matrix, k_counts)
    IDCG = calculate_IDCG(rel_matrix, k_counts)

    assert nDCG == np.mean(DCG / IDCG)

    pre_nDCG = calculate_nDCG(sim_matrix, rel_matrix, k_counts, IDCG=IDCG)
    assert pre_nDCG == nDCG

    post_mean_nDCG = calculate_nDCG(sim_matrix, rel_matrix, k_counts, IDCG=IDCG, reduction=None)
    assert np.mean(post_mean_nDCG) == pre_nDCG

EgoNLQ/utils/nDCG.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import copy
import glob
import json
import os
import random
import pdb

import numpy as np
import torch
import torch.backends.cudnn
import torch.utils.data
from tqdm import tqdm

import utils.evaluate_ego4d_nlq as ego4d_eval
from utils.data_util import index_to_time


def set_th_config(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True


def filter_checkpoints(model_dir, suffix="t7", max_to_keep=5):
    model_paths = glob.glob(os.path.join(model_dir, "*.{}".format(suffix)))
    if len(model_paths) > max_to_keep:
        model_file_dict = dict()
        suffix_len = len(suffix) + 1
        for model_path in model_paths:
            step = int(os.path.basename(model_path).split("_")[1][0:-suffix_len])
            model_file_dict[step] = model_path
        sorted_tuples = sorted(model_file_dict.items())
        unused_tuples = sorted_tuples[0:-max_to_keep]
        for _, model_path in unused_tuples:
            os.remove(model_path)


def get_last_checkpoint(model_dir, suffix="t7"):
    model_filenames = glob.glob(os.path.join(model_dir, "*.{}".format(suffix)))
    model_file_dict = dict()
    suffix_len = len(suffix) + 1
    for model_filename in model_filenames:
        step = int(os.path.basename(model_filename).split("_")[1][0:-suffix_len])
        model_file_dict[step] = model_filename
    sorted_tuples = sorted(model_file_dict.items())
    last_checkpoint = sorted_tuples[-1]
    return last_checkpoint[1]


def convert_length_to_mask(lengths):
    max_len = lengths.max().item()
    mask = torch.arange(max_len, device=lengths.device).expand(
        lengths.size()[0], max_len
    ) < lengths.unsqueeze(1)
    mask = mask.float()
    return mask


def eval_test(
    model,
    data_loader,
    device,
    mode="test",
    gt_json_path=None,
    epoch=None,
    global_step=None,
    configs=None,
):
    predictions = []
    with torch.no_grad():
        for idx, (records, vfeats, vfeat_lens, dual_text_feature, query_attention_mask) in enumerate(data_loader):
            # prepare features
            vfeats, vfeat_lens = vfeats.to(device), vfeat_lens.to(device)
            dual_text_feature = dual_text_feature.to(device)
            query_attention_mask = query_attention_mask.to(device)

            if configs.predictor == 'EgoVLP':
                # generate mask
                query_mask = query_attention_mask
            else:
                raise NotImplementedError("predictor should be EgoVLP")

            # generate mask
            video_mask = convert_length_to_mask(vfeat_lens).to(device)
            # compute predicted results
            _, start_logits, end_logits = model(
                vfeats, video_mask, query_mask, dual_text_feature
            )
            start_indices, end_indices = model.extract_index(start_logits, end_logits)
            start_indices = start_indices.cpu().numpy()
            end_indices = end_indices.cpu().numpy()


            # Record output and use standard evalution script for NLQ.
            for record, starts, ends in zip(records, start_indices, end_indices):
                # Convert all indices to times.
                timewindow_predictions = []
                for start, end in zip(starts, ends):
                    start_time, end_time = index_to_time(
                        start, end, record["v_len"].item(), record["duration"].item()
                    )
                    timewindow_predictions.append([float(start_time), float(end_time)])
                new_datum = {
                    "clip_uid": record["vid"],
                    "annotation_uid": record["annotation_uid"],
                    "query_idx": int(record["query_idx"]),
                    "predicted_times": copy.deepcopy(timewindow_predictions),
                }
                predictions.append(new_datum)

    ## Save predictions if path is provided.
    #if result_save_path:
    #    with open(result_save_path, "w") as file_id:
    #        json.dump(
    #            {
    #                "version": "1.0",
    #                "challenge": "ego4d_nlq_challenge",
    #                "results": predictions,
    #            }, file_id
    #        )

    # Evaluate if ground truth JSON file is provided.

    if gt_json_path:
        with open(gt_json_path) as file_id:
            ground_truth = json.load(file_id)
        thresholds = [0.3, 0.5, 0.01]
        topK = [1, 3, 5]
        results, mIoU = ego4d_eval.evaluate_nlq_performance(
            predictions, ground_truth, thresholds, topK
        )
        title = f"Epoch {epoch}, Step {global_step}"
        score_str = ego4d_eval.display_results(
            results, mIoU, thresholds, topK, title=title
        )
    else:
        results = None
        mIoU = None
        score_str = None

    return results, mIoU, score_str

EgoNLQ/utils/runner_utils.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import functools
import json
import os
import socket
import time
from collections import OrderedDict
from datetime import datetime
from itertools import repeat
from pathlib import Path

import humanize
import numpy as np
import psutil


def replace_nested_dict_item(obj, key, replace_value):
    for k, v in obj.items():
        if isinstance(v, dict):
            obj[k] = replace_nested_dict_item(v, key, replace_value)
    if key in obj:
        obj[key] = replace_value
    return obj


def state_dict_data_parallel_fix(load_state_dict, curr_state_dict):
    load_keys = list(load_state_dict.keys())
    curr_keys = list(curr_state_dict.keys())

    redo_dp = False
    undo_dp = False
    if not curr_keys[0].startswith('module.') and load_keys[0].startswith('module.'):   # this
        undo_dp = True
    elif curr_keys[0].startswith('module.') and not load_keys[0].startswith('module.'):
        redo_dp = True

    if undo_dp: # this
        from collections import OrderedDict
        new_state_dict = OrderedDict()
        for k, v in load_state_dict.items():
            name = k[7:]  # remove `module.`
            new_state_dict[name] = v
        # load params
    elif redo_dp:
        from collections import OrderedDict
        new_state_dict = OrderedDict()
        for k, v in load_state_dict.items():
            name = 'module.' + k  # remove `module.`
            new_state_dict[name] = v
    else:
        new_state_dict = load_state_dict
    return new_state_dict

def print_numpy(x, val=True, shp=False):
    """Print the mean, min, max, median, std, and size of a numpy array
    Parameters:
        val (bool) -- if print the values of the numpy array
        shp (bool) -- if print the shape of the numpy array
    """
    x = x.astype(np.float64)
    if shp:
        print('shape,', x.shape)
    if val:
        x = x.flatten()
        print('mean = %3.3f, min = %3.3f, max = %3.3f, median = %3.3f, std=%3.3f' % (
            np.mean(x), np.min(x), np.max(x), np.median(x), np.std(x)))


def mkdirs(paths):
    """create empty directories if they don't exist
    Parameters:
        paths (str list) -- a list of directory paths
    """
    if isinstance(paths, list) and not isinstance(paths, str):
        for path in paths:
            mkdir(path)
    else:
        mkdir(paths)


def mkdir(path):
    """create a single empty directory if it didn't exist
    Parameters:
        path (str) -- a single directory path
    """
    if not os.path.exists(path):
        os.makedirs(path)

def read_json(fname):
    with fname.open('rt') as handle:
        return json.load(handle, object_hook=OrderedDict)

def write_json(content, fname):
    with fname.open('wt') as handle:
        json.dump(content, handle, indent=4, sort_keys=False)

def inf_loop(data_loader):
    ''' wrapper function for endless data loader. '''
    for loader in repeat(data_loader):
        yield from loader

def memory_summary():
    vmem = psutil.virtual_memory()
    msg = (
        f">>> Currently using {vmem.percent}% of system memory "
        f"{humanize.naturalsize(vmem.used)}/{humanize.naturalsize(vmem.available)}"
    )
    print(msg)

@functools.lru_cache(maxsize=64, typed=False)
def memcache(path):
    suffix = Path(path).suffix
    print(f"loading features >>>", end=" ")
    tic = time.time()
    if suffix == ".npy":
        res = np_loader(path)
    else:
        raise ValueError(f"unknown suffix: {suffix} for path {path}")
    print(f"[Total: {time.time() - tic:.1f}s] ({socket.gethostname() + ':' + str(path)})")
    return res

def np_loader(np_path, l2norm=False):
    with open(np_path, "rb") as f:
        data = np.load(f, encoding="latin1", allow_pickle=True)
    if isinstance(data, np.ndarray) and data.size == 1:
        data = data[()]  # handle numpy dict storage convnetion
    if l2norm:
        print("L2 normalizing features")
        if isinstance(data, dict):
            for key in data:
                feats_ = data[key]
                feats_ = feats_ / max(np.linalg.norm(feats_), 1E-6)
                data[key] = feats_
        elif data.ndim == 2:
            data_norm = np.linalg.norm(data, axis=1)
            data = data / np.maximum(data_norm.reshape(-1, 1), 1E-6)
        else:
            raise ValueError("unexpected data format {}".format(type(data)))
    return data


class Timer:
    def __init__(self):
        self.cache = datetime.now()

    def check(self):
        now = datetime.now()
        duration = now - self.cache
        self.cache = now
        return duration.total_seconds()

    def reset(self):
        self.cache = datetime.now()

EgoNLQ/utils/util.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import random
import cv2
import numpy as np
import torch


def load_frames_from_video_path(path, num_frames, sample='rand'):
    cap = cv2.VideoCapture(path)
    assert (cap.isOpened())
    vlen = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    acc_samples = min(num_frames, vlen)
    intervals = np.linspace(start=0, stop=vlen, num=acc_samples + 1).astype(int)
    ranges = []
    for idx, interv in enumerate(intervals[:-1]):
        ranges.append((interv, intervals[idx + 1] - 1))
    if sample == 'rand':
        frame_idxs = [random.choice(range(x[0], x[1])) for x in ranges]
    elif sample == 'uniform':
        frame_idxs = [(x[0] + x[1]) // 2 for x in ranges]
    else:
        raise NotImplementedError

    frames = []
    for index in frame_idxs:
        cap.set(cv2.CAP_PROP_POS_FRAMES, index)
        ret, frame = cap.read()
        if ret:
            cv2.imwrite(f'images/{index}.jpg', frame)
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frame = torch.from_numpy(frame)
            # (H x W x C) to (C x H x W)
            frame = frame.permute(2, 0, 1)
            frames.append(frame)
        else:
            raise ValueError

    frames = torch.stack(frames).float() / 255
    cap.release()
    return frames, frame_idxs

EgoNLQ/utils/video.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import os
import cv2
import shutil
import json
import sys
import multiprocessing
from multiprocessing import Pool
import subprocess
import pandas as pd
from csv import reader, writer

video_dir = './dataset/ego4d_256_multiprocessing'
output_dir = './dataset/ego4d_chunked_multiprocessing'

save_dir = os.path.join(output_dir)
if not os.path.exists(save_dir):
    os.makedirs(save_dir)

dur_limit = 600

def video2segments(infos):
    global count
    index, uid, dur = infos[0], infos[1], infos[2]
    input_path = os.path.join(video_dir, uid + '.mp4')

    output_uid_dir = os.path.join(output_dir, uid)
    if not os.path.exists(output_uid_dir):
        os.makedirs(output_uid_dir)

    #if index % num_partition != partition:
    #    return

    assert os.path.exists(input_path)

    cap = cv2.VideoCapture(input_path)
    rate = cap.get(5)
    frame_num = cap.get(7)
    duration = frame_num / rate

    if duration <= dur_limit:
        shutil.copyfile(input_path, os.path.join(output_uid_dir, '0.mp4'))
        return

    num_seg = duration // dur_limit

    s1_time = 0;
    s2_time = dur_limit
    num_finished = 0
    while num_finished <= num_seg:
        output_mp4_path = os.path.join(output_uid_dir, str(num_finished) + '.mp4')

        cmd = 'ffmpeg -y -i {} -ss {} -to {} -async 1 {}'.format(input_path, s1_time, s2_time,
                                                                                 output_mp4_path)
        # print(cmd)
        subprocess.call(cmd, shell=True)

        # Update for next steps
        s1_time = s2_time
        s2_time = min(s1_time + dur_limit, duration)
        num_finished += 1
    return

if __name__ == "__main__":
    with open('./dataset/ego4d/manifest.csv', 'r') as csv:
        csv_reader = list(reader(csv))[1:]

    downloaded = os.listdir('./dataset/ego4d')

    uid_list = []
    infos_list = []
    num_valid = 0
    for i_item, item in enumerate(csv_reader):
        uid, dur = item[0], float(item[2])
        existed = uid + '.mp4' in downloaded

        if not existed:
            continue

        uid_list.append(uid)
        infos_list.append([num_valid, uid, dur])
        num_valid += 1

    # for infos in infos_list:
    #     video2segments(infos)

    pool = Pool(multiprocessing.cpu_count())
    pool.map(video2segments, tuple(infos_list))

EgoNLQ/utils/video_chunk.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import os
import time
import sys
import subprocess
import multiprocessing
from multiprocessing import Pool, Value

folder_path = './dataset/ego4d'
output_path = './dataset/ego4d_256_multiprocessing'

def videos_resize(videoinfos):
    global count

    videoid, videoname = videoinfos

    if os.path.exists(os.path.join(output_path, videoname)):
        print(f'{videoname} is resized.')
        return

    inname = folder_path + '/' + videoname
    outname = output_path + '/' + videoname

    cmd = "ffmpeg -y -i {} -filter:v scale=\"trunc(oh*a/2)*2:256\" -c:a copy {}".format( inname, outname)
    subprocess.call(cmd, shell=True)

    return

if __name__ == "__main__":
    file_list = []
    mp4_list = [item for item in os.listdir(folder_path) if item.endswith('.mp4')]
    for id, video in enumerate(mp4_list):
        file_list.append([id, video])

    pool = Pool(multiprocessing.cpu_count()) 
    pool.map(videos_resize, tuple(file_list))

    # for file in file_list:
    #    videos_resize(file)

EgoNLQ/utils/video_resize.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import matplotlib
import numpy as np
import torch

matplotlib.use('Agg')


def visualise_path(pred, target, window):
    """
    :param pred: (P, 2) Tensor where P is the number of predictions, and 2 is the (i,j) coordinate
    :param target: (T, 2) Tensor where T is the number of targets, and 2 is the (i,j) coordinate
    :param dims: (H, W) tup/list the desired height and width of matrix (should be >= to max(i), max(j))
    :param assignment_method: Method of assignment (dtw, minimum etc.)
    :return: image, visualisation of path prediction and target.
    """
    tp = torch.Tensor((64, 191, 64))
    fp = torch.Tensor((191, 64, 64))
    gt = torch.Tensor((102, 153, 255))

    grid = torch.ones_like(window).unsqueeze(0).repeat(3, 1, 1) * 255
    inf = 130 * torch.ones_like(grid)
    grid = torch.where(torch.isnan(window), inf, grid)

    clip_idxs = [t[0] for t in target]
    local_idxs = np.unique(np.array(clip_idxs)).tolist()

    for t in target:
        local_idx = local_idxs.index(t[0])
        grid[:, local_idx, t[1]] = gt

    for p in pred:
        local_idx = local_idxs.index(p[0])
        if (grid[:, local_idx, p[1]] == gt).all():
            grid[:, local_idx, p[1]] = tp
        else:
            grid[:, local_idx, p[1]] = fp

    return grid / 255


def batch_path_vis(pred_dict, target, window):
    grids = []

    window = window.cpu()
    for key, pred in pred_dict.items():
        tmp_window = window
        if key == 'min_dist':
            tmp_window = torch.zeros_like(window)
        grids.append(visualise_path(pred, target, tmp_window))

    return torch.stack(grids)


if __name__ == "__main__":
    pred = [[1, 1], [2, 4]]
    gt = [[1, 1], [3, 4]]
    window = torch.zeros((5, 6))
    visualise_path(pred, gt, window)

EgoNLQ/utils/visualisation.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import os
import sys
import torch
import pandas as pd
import pdb
from base.base_dataset import TextVideoDataset
# from data_loader.transforms import init_transform_dict, init_video_transform_dict
from transforms import init_transform_dict, init_video_transform_dict
import json
import transformers


class EgoTaskQA(TextVideoDataset):
    def _load_metadata(self, args):
        metadata_dir = 'Data/qa/' + args.dataset_split_type
        split_files = {
            'train': 'formatted_train_qas_encode.json',
            'val': 'formatted_val_qas_encode.json',            # there is no test
            'test': 'formatted_test_qas_encode.json',  # there is no test
        }
        target_split_fp = split_files[self.split]

        with open(os.path.join(metadata_dir, target_split_fp),'r') as load_f:
            metadata = json.load(load_f)
        if self.subsample < 1:
            metadata = metadata.sample(frac=self.subsample)
        self.metadata = pd.DataFrame(metadata)

        if self.split in ['train']:
            self.frame_sample = 'rand'
        elif self.split in ['val', 'test']:
            self.frame_sample = 'uniform'

    def _get_video_path(self, sample):
        rel_video_fp = sample['interval'] + '.mp4'
        full_video_fp = os.path.join(self.data_dir, rel_video_fp)
        return full_video_fp

    def _get_video_frames(self, video_fp):
        video_loading = self.video_params.get('loading', 'strict')
        # pdb.set_trace()
        try:
            if os.path.isfile(video_fp):
                imgs, idxs = self.video_reader(video_fp, self.video_params['num_frames'])
            else:
                print(f"Warning: missing video file {video_fp}.")
                assert False
        except Exception as e:
            if video_loading == 'strict':
                raise ValueError(
                    f'Video loading failed for {video_fp}, video loading for this dataset is strict.') from e
            else:
                imgs = Image.new('RGB', (self.video_params['input_res'], self.video_params['input_res']), (0, 0, 0))
                imgs = transforms.ToTensor()(imgs).unsqueeze(0)

        if self.transforms is not None:
            if self.video_params['num_frames'] > 1:
                imgs = imgs.transpose(0, 1)  # [T, C, H, W] ---> [C, T, H, W]
                imgs = self.transforms(imgs)
                imgs = imgs.transpose(0, 1)  # recover
            else:
                imgs = self.transforms(imgs)

        final = torch.zeros([self.video_params['num_frames'], 3, self.video_params['input_res'], self.video_params['input_res']])
        final[:imgs.shape[0]] = imgs
        return final

    def __getitem__(self, item):
        item = item % len(self.metadata)
        sample = self.metadata.iloc[item]
        video_fp  = self._get_video_path(sample)

        final = self._get_video_frames(video_fp)

        #pdb.set_trace()
        meta_arr = {'type': sample['type'], 'category': sample['category'], 'semantic': sample['semantic'], 'reasoning': sample['reasoning_type'].split('$')}
        return {'video': final, 'text': sample['question'], 'meta': meta_arr, 'answer': sample['answer_encode']}

def collate_func(batch):
    
    config = 'configs/egotaskqa.json'
    unique_dict = torch.load('reasoning_unique_cat.pth')
    f = open(config)
    config = json.load(f)
    tokenizer = transformers.AutoTokenizer.from_pretrained(config['arch']['args']['text_params']['model'])
    video_frames_list, text_tokens_list, attention_masks_list, answers_list, reasoning_list = [], [], [], [], []
                            
    for i, item in enumerate(batch):
        video_frames_list.append(item['video'])
        tokenized_text = tokenizer(item['text'].strip().lower(), return_tensors='pt', padding='max_length', max_length=30, truncation=True)
        text_tokens_list.append(tokenized_text['input_ids'])
        attention_masks_list.append(tokenized_text['attention_mask'])
        answers_list.append(torch.tensor(item['answer']))

        index_list = []
        for _idx in range(len(item['meta']['reasoning'])):
            index_list.append(unique_dict[item['meta']['reasoning'][_idx]])
        reasoning_list.append(torch.tensor(index_list))

    video_frames = torch.stack(video_frames_list, dim=0)
    text_tokens = torch.stack(text_tokens_list, dim=0).squeeze(1)
    attention_masks = torch.stack(attention_masks_list, dim=0).squeeze(1)
    answers = torch.stack(answers_list, dim=0)

    # # torch.Size([4, 20, 3, 224, 224]) torch.Size([4, 12]) torch.Size([4]) torch.Size([4, 20, 2048]) torch.Size([4, 25, 15])
    return video_frames, text_tokens, attention_masks, answers, reasoning_list

if __name__ == "__main__":
    kwargs = dict(
        dataset_name="EgoTaskQA",
        text_params={
            "input": "text"
        },
        video_params={
        "input_res": 224,
        "num_frames": 4,
        "loading": "lax"
        },
        data_dir="EgoTaskQA/qa_videos",
        meta_dir="Data/qa/indirect",
        tsfms=init_video_transform_dict()['test'],
        reader='decord',
        split='train',
        neg_param=60
    )
    dataset = EgoTaskQA(**kwargs)
    import tqdm
    max_class = 0
    for i in tqdm.tqdm(range(len(dataset))):
        item = dataset[i]
        pdb.set_trace()
        class_no = item['answer']
        if class_no > max_class:
            max_class = class_no
        # print(item.keys())
    print(max_class)

EgoTaskQA/EgoTaskQA_dataset.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

from .base_dataset import *
from .base_model import *

EgoTaskQA/base/__init__.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import os
import pdb

import tqdm
import random
from abc import abstractmethod

import av
import cv2
import decord
import ffmpeg
import numpy as np
import torch
from PIL import Image
from torch.utils.data import Dataset, get_worker_info
from torchvision import transforms


class TextVideoDataset(Dataset):
    def __init__(self,
                 dataset_name,
                 text_params,
                 video_params,
                 data_dir,
                 meta_dir=None,
                 split='train',
                 tsfms=None,
                 cut=None,
                 subsample=1,
                 sliding_window_stride=-1,
                 reader='decord',
                 neg_param=None,
                 args=None
                 ):
        self.dataset_name = dataset_name
        self.text_params = text_params
        self.video_params = video_params
        # check for environment variables
        self.data_dir = os.path.expandvars(data_dir)
        if meta_dir is not None:
            self.meta_dir = os.path.expandvars(meta_dir)
        else:
            self.meta_dir = self.data_dir
        self.split = split
        self.transforms = tsfms
        self.cut = cut
        self.subsample = subsample
        self.sliding_window_stride = sliding_window_stride
        self.video_reader = video_reader[reader]
        self.label_type = 'caption'
        self.neg_param = neg_param
        self._load_metadata(args)
        if self.sliding_window_stride != -1:
            if self.split != 'test':
                raise ValueError('Fixing frame sampling is for test time only. can remove but...')
            self._fix_temporal_samples()

    @abstractmethod
    def _load_metadata(self):
        raise NotImplementedError("Metadata loading must be implemented by subclass")

    @abstractmethod
    def _get_video_path(self, sample):
        raise NotImplementedError("Get video path function must be implemented by subclass")

    def _get_caption(self, sample):
        raise NotImplementedError("Get caption function must be implemented by subclass")

    def _get_video_lens(self):
        vlen_li = []
        for idx, row in self.metadata.iterrows():
            video_path = self._get_video_path(row)[0]
            vlen_li.append(get_video_len(video_path))

        return vlen_li

    def _fix_temporal_samples(self):
        self.metadata['vlen'] = self._get_video_lens()
        self.metadata['frame_intervals'] = self.metadata['vlen'].apply(
            lambda x: np.linspace(start=0, stop=x, num=min(x, self.video_params['num_frames']) + 1).astype(int))
        self.metadata['fix_start'] = self.metadata['frame_intervals'].apply(
            lambda x: np.arange(0, int(x[-1] / len(x - 1)), self.sliding_window_stride)
        )
        self.metadata = self.metadata.explode('fix_start')

    def __len__(self):
        return len(self.metadata)

    def __getitem__(self, item):
        item = item % len(self.metadata)
        sample = self.metadata.iloc[item]
        video_fp, rel_fp = self._get_video_path(sample)
        caption = self._get_caption(sample)

        video_loading = self.video_params.get('loading', 'strict')
        frame_sample = 'rand'
        fix_start = None
        if self.split == 'test':
            frame_sample = 'uniform'
        if self.sliding_window_stride != -1:
            fix_start = sample['fix_start']

        try:
            if os.path.isfile(video_fp):
                imgs, idxs = self.video_reader(video_fp, self.video_params['num_frames'], frame_sample,
                                               fix_start=fix_start)
            else:
                print(f"Warning: missing video file {video_fp}.")
                assert False
        except Exception as e:
            if video_loading == 'strict':
                raise ValueError(
                    f'Video loading failed for {video_fp}, video loading for this dataset is strict.') from e
            else:
                imgs = Image.new('RGB', (self.video_params['input_res'], self.video_params['input_res']), (0, 0, 0))
                imgs = transforms.ToTensor()(imgs).unsqueeze(0)

        # if self.transforms is not None:
        #     imgs = self.transforms(imgs)

        if self.transforms is not None:
            if self.video_params['num_frames'] > 1:
                imgs = imgs.transpose(0, 1)  # [T, C, H, W] ---> [C, T, H, W]
                imgs = self.transforms(imgs)
                imgs = imgs.transpose(0, 1)  # recover
            else:
                imgs = self.transforms(imgs)

        # if self.transforms is not None:
        #     # for video aug if T>1 else just do image aug
        #     if imgs.size(0) > 1:
        #         # added by Mr. YAN
        #         imgs = imgs.transpose(0, 1) # [T, C, H, W] ---> [C, T, H, W]
        #         imgs = self.transforms(imgs)
        #         imgs = imgs.transpose(0, 1) # recover
        #     else:
        #         imgs = self.transforms(imgs)    # t

        final = torch.zeros([self.video_params['num_frames'], 3, self.video_params['input_res'],
                             self.video_params['input_res']])
        final[:imgs.shape[0]] = imgs

        meta_arr = {'raw_captions': caption, 'paths': rel_fp, 'dataset': self.dataset_name}
        data = {'video': final, 'text': caption, 'meta': meta_arr}
        return data


class TextImageDataset(TextVideoDataset):

    def __getitem__(self, item):
        item = item % len(self.metadata)
        sample = self.metadata.iloc[item]
        video_fp, rel_fp = self._get_video_path(sample)
        caption = self._get_caption(sample)

        video_loading = self.video_params.get('loading', 'strict')

        try:
            img = Image.open(video_fp).convert("RGB")
        except:
            if video_loading == 'strict':
                raise ValueError(f'Image loading failed for {video_fp}, image loading for this dataset is strict.')
            else:
                img = Image.new('RGB', (self.video_params['input_res'], self.video_params['input_res']), (0, 0, 0))

        # convert to tensor because video transforms don't, expand such that its a 1-frame video.
        img = transforms.ToTensor()(img).unsqueeze(0)
        if self.transforms is not None:
            img = self.transforms(img)
        meta_arr = {'raw_captions': caption, 'paths': rel_fp, 'dataset': self.dataset_name}
        data = {'video': img, 'text': caption, 'meta': meta_arr}
        return data


def sample_frames(num_frames, vlen, sample='rand', fix_start=None):
    acc_samples = min(num_frames, vlen)
    intervals = np.linspace(start=0, stop=vlen, num=acc_samples + 1).astype(int)
    ranges = []
    for idx, interv in enumerate(intervals[:-1]):
        ranges.append((interv, intervals[idx + 1] - 1))
    if sample == 'rand':
        frame_idxs = [random.choice(range(x[0], x[1])) for x in ranges]
    elif fix_start is not None:
        frame_idxs = [x[0] + fix_start for x in ranges]
    elif sample == 'uniform':
        frame_idxs = [(x[0] + x[1]) // 2 for x in ranges]
    else:
        raise NotImplementedError

    return frame_idxs

def sample_frames_clips(start, end, vlen, acc_samples):
    start = max(0, start)
    end = min(vlen, end)

    intervals = np.linspace(start=start, stop=end, num=int(acc_samples) + 1).astype(int)
    ranges = []
    for idx, interv in enumerate(intervals[:-1]):
        ranges.append((interv, intervals[idx + 1] - 1))
        frame_idxs = [(x[0] + x[1]) // 2 for x in ranges
                      ]
    return frame_idxs

def sample_frames_start_end(num_frames, start, end, sample='rand', fix_start=None):
    acc_samples = min(num_frames, end)
    intervals = np.linspace(start=start, stop=end, num=acc_samples + 1).astype(int)
    ranges = []
    for idx, interv in enumerate(intervals[:-1]):
        ranges.append((interv, intervals[idx + 1] - 1))
    if sample == 'rand':
        frame_idxs = [random.choice(range(x[0], x[1])) for x in ranges]
    elif fix_start is not None:
        frame_idxs = [x[0] + fix_start for x in ranges]
    elif sample == 'uniform':
        frame_idxs = [(x[0] + x[1]) // 2 for x in ranges]
    else:
        raise NotImplementedError

    return frame_idxs

def read_frames_cv2(video_path, num_frames, sample='rand', fix_start=None):
    cap = cv2.VideoCapture(video_path)
    assert (cap.isOpened())
    vlen = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    # get indexes of sampled frames
    frame_idxs = sample_frames(num_frames, vlen, sample=sample, fix_start=fix_start)
    frames = []
    success_idxs = []
    for index in frame_idxs:
        cap.set(cv2.CAP_PROP_POS_FRAMES, index - 1)
        ret, frame = cap.read()
        if ret:
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frame = torch.from_numpy(frame)
            # (H x W x C) to (C x H x W)
            frame = frame.permute(2, 0, 1)
            frames.append(frame)
            success_idxs.append(index)
        else:
            pass
            # print(frame_idxs, ' fail ', index, f'  (vlen {vlen})')

    frames = torch.stack(frames).float() / 255
    cap.release()
    return frames, success_idxs

def read_frames_cv2_egoclip(video_path_1, video_path_2, num_frames, sample,
                            start_sec, end_sec, bound_sec):
    if video_path_1 == video_path_2:
        cap1 = cv2.VideoCapture(video_path_1)
        cap2 = cap1
        vlen1 = int(cap1.get(cv2.CAP_PROP_FRAME_COUNT))
        vlen2 = vlen1
        assert (cap1.isOpened())
    else:   # some clips may span two segments.
        cap1 = cv2.VideoCapture(video_path_1)
        cap2 = cv2.VideoCapture(video_path_2)
        vlen1 = int(cap1.get(cv2.CAP_PROP_FRAME_COUNT))
        vlen2 = int(cap2.get(cv2.CAP_PROP_FRAME_COUNT))
        assert (cap1.isOpened())
        assert (cap2.isOpened())

    # get indexes of sampled frames
    start_f = max(0, int(start_sec * 30))
    end_f = max(0, int(end_sec * 30))
    bound_f = int(bound_sec * 30)
    frame_idxs = sample_frames_start_end(num_frames, start_f, end_f, sample=sample)

    frames = []
    success_idxs = []
    for index in frame_idxs:
        _index = index % (600 * 30)
        if index > bound_f: # frame from the last video
            _index = min(_index, vlen2)
            cap2.set(cv2.CAP_PROP_POS_FRAMES, _index - 1)
            ret, frame = cap2.read()
        else:   # frame from the first video
            _index = min(_index, vlen1)
            cap1.set(cv2.CAP_PROP_POS_FRAMES, _index - 1)
            ret, frame = cap1.read()

        if ret:
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frame = torch.from_numpy(frame)
            # (H x W x C) to (C x H x W)
            frame = frame.permute(2, 0, 1)
            frames.append(frame)
            success_idxs.append(index)
        else:
            pass

    while len(frames) < num_frames: # complete the frame
        frames.append(frames[-1])

    frames = torch.stack(frames).float() / 255
    cap1.release()
    cap2.release()
    return frames, success_idxs

def read_frames_cv2_epic(video_path, start_frame, stop_frame, num_frames, sample='rand', fix_start=None):
    # get indexes of sampled frames
    frame_idxs = sample_frames_start_end(num_frames, start_frame, stop_frame, sample=sample, fix_start=fix_start)
    frames = []
    success_idxs = []
    for index in frame_idxs:
        img_name = 'frame_' + str(index).zfill(10) + '.jpg'
        frame = cv2.imread(os.path.join(video_path, img_name),cv2.COLOR_BGR2RGB)

        frame = torch.from_numpy(frame)
        # (H x W x C) to (C x H x W)
        frame = frame.permute(2, 0, 1)
        frames.append(frame)
        success_idxs.append(index)

    frames = torch.stack(frames).float() / 255
    return frames, success_idxs

def read_frames_cv2_charades(video_path, num_frames, sample, start_sec=None, end_sec=None):
    cap = cv2.VideoCapture(video_path)
    assert (cap.isOpened())
    vlen = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = cap.get(5)

    # get indexes of sampled frames
    if not start_sec and not end_sec:
        frame_idxs = sample_frames(num_frames, vlen, sample=sample)
    else:
        start_f = max(0, int(start_sec * fps))
        end_f = min(int(end_sec * fps), vlen)
        frame_idxs = sample_frames_start_end(num_frames, start_f, end_f, sample=sample)

    frames = []
    success_idxs = []
    for index in frame_idxs:
        cap.set(cv2.CAP_PROP_POS_FRAMES, index - 1)
        ret, frame = cap.read()
        if ret:
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frame = torch.from_numpy(frame)
            # (H x W x C) to (C x H x W)
            frame = frame.permute(2, 0, 1)
            frames.append(frame)
            success_idxs.append(index)
        else:
            pass

    frames = torch.stack(frames).float() / 255
    cap.release()
    return frames, success_idxs

def read_frames_av(video_path, num_frames, sample='rand', fix_start=None):
    reader = av.open(video_path)
    try:
        frames = []
        frames = [torch.from_numpy(f.to_rgb().to_ndarray()) for f in reader.decode(video=0)]
    except (RuntimeError, ZeroDivisionError) as exception:
        print('{}: WEBM reader cannot open {}. Empty '
              'list returned.'.format(type(exception).__name__, video_path))
    vlen = len(frames)
    frame_idxs = sample_frames(num_frames, vlen, sample=sample, fix_start=fix_start)
    frames = torch.stack([frames[idx] for idx in frame_idxs]).float() / 255
    frames = frames.permute(0, 3, 1, 2)
    return frames, frame_idxs

decord.bridge.set_bridge("torch")

def read_frames_decord(video_path, num_frames, sample='rand', fix_start=None):
    video_reader = decord.VideoReader(video_path, num_threads=1)
    vlen = len(video_reader)
    frame_idxs = sample_frames(num_frames, vlen, sample=sample, fix_start=fix_start)
    video_reader.skip_frames(1)
    frames = video_reader.get_batch(frame_idxs)

    frames = frames.float() / 255
    frames = frames.permute(0, 3, 1, 2)
    return frames, frame_idxs

def read_frames_decord_start_end(video_path, start, end, num_frames):
    video_reader = decord.VideoReader(video_path, num_threads=1)
    vlen = len(video_reader)
    frame_idxs = sample_frames_clips(start, end, vlen, num_frames + 1)
    video_reader.skip_frames(1)
    frames = video_reader.get_batch(frame_idxs)

    frames = frames.float() / 255
    frames = frames.permute(0, 3, 1, 2)
    return frames, frame_idxs

def get_video_len(video_path):
    cap = cv2.VideoCapture(video_path)
    if not (cap.isOpened()):
        return False
    vlen = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    cap.release()
    return vlen

video_reader = {
    'av': read_frames_av,
    'cv2': read_frames_cv2,
    'cv2_epic': read_frames_cv2_epic,
    'cv2_charades': read_frames_cv2_charades,
    'cv2_egoclip': read_frames_cv2_egoclip,
    'decord': read_frames_decord,
    'decord_start_end': read_frames_decord_start_end,
}

EgoTaskQA/base/base_dataset.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import torch.nn as nn
import numpy as np
from abc import abstractmethod


class BaseModel(nn.Module):
    """
    Base class for all models
    """
    @abstractmethod
    def forward(self, *inputs):
        """
        Forward pass logic

        :return: Model output
        """
        raise NotImplementedError

    def __str__(self):
        """
        Model prints with number of trainable parameters
        """
        model_parameters = filter(lambda p: p.requires_grad, self.parameters())
        params = sum(np.prod(p.size()) for p in model_parameters)
        return super().__str__() + '\nTrainable parameters: {}'.format(params)

EgoTaskQA/base/base_model.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

from .logger import *
from .visualization import *

EgoTaskQA/logger/__init__.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import logging
import logging.config
from pathlib import Path

from utils import read_json


def setup_logging(save_dir, log_config='/cis/home/shraman/works_meta_2022/pre-training/EgoVLP_Fused_HardNegITM_Checkpoint_multinode_lrmodified_FIBER/logger/logger_config.json', default_level=logging.INFO):
    """
    Setup logging configuration
    """
    log_config = Path(log_config)
    if log_config.is_file():
        config = read_json(log_config)
        # modify logging paths based on run config
        for _, handler in config['handlers'].items():
            if 'filename' in handler:
                handler['filename'] = str(save_dir / handler['filename'])

        logging.config.dictConfig(config)
    else:
        print("Warning: logging configuration file is not found in {}.".format(log_config))
        logging.basicConfig(level=default_level)

EgoTaskQA/logger/logger.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import importlib
from utils import Timer


class TensorboardWriter:
    def __init__(self, log_dir, logger, enabled):
        self.writer = None
        self.selected_module = ""

        if enabled:
            log_dir = str(log_dir)

            # Retrieve visualization writer.
            for module in ["torch.utils.tensorboard", "tensorboardX"]:
                try:
                    self.writer = importlib.import_module(module).SummaryWriter(log_dir)
                    succeeded = True
                    break
                except ImportError:
                    succeeded = False
                self.selected_module = module

            if not succeeded:
                message = "Warning: visualization (Tensorboard) is configured to use, but currently not installed on " \
                    "this machine. Please install either TensorboardX with 'pip install tensorboardx', upgrade " \
                    "PyTorch to version >= 1.1 for using 'torch.utils.tensorboard' or turn off the option in " \
                    "the 'config.json' file."
                logger.warning(message)

        self.step = 0
        self.mode = ''

        self.tb_writer_ftns = {
            'add_scalar', 'add_scalars', 'add_image', 'add_images', 'add_audio',
            'add_text', 'add_histogram', 'add_pr_curve', 'add_embedding'
        }
        self.tag_mode_exceptions = {'add_histogram', 'add_embedding'}
            
        self.timer = Timer()

    def set_step(self, step, mode='train'):
        self.mode = mode
        self.step = step
        if step == 0:
            self.timer.reset()
        else:
            duration = self.timer.check()
            self.add_scalar('steps_per_sec', 1 / duration)

    def __getattr__(self, name):
        """
        If visualization is configured to use:
            return add_data() methods of tensorboard with additional information (step, tag) added.
        Otherwise:
            return a blank function handle that does nothing
        """
        if name in self.tb_writer_ftns:
            add_data = getattr(self.writer, name, None)

            def wrapper(tag, data, *args, **kwargs):
                if add_data is not None:
                    # add mode(train/valid) tag
                    if name not in self.tag_mode_exceptions:
                        tag = '{}/{}'.format(tag, self.mode)
                    add_data(tag, data, self.step, *args, **kwargs)
            return wrapper
        else:
            # default action for returning methods defined in this class, set_step() for instance.
            try:
                attr = object.__getattr__(name)
            except AttributeError:
                raise AttributeError("type object '{}' has no attribute '{}'".format(self.selected_module, name))
            return attr


class SacredNeptuneWriter:
    def __init__(self):
        raise NotImplementedError

EgoTaskQA/logger/visualization.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import torch
import numpy as np
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
import sys, os
from tqdm import tqdm
from datetime import datetime
import argparse, time
import pdb
from yaml import parse
from tqdm import tqdm
import json
import glob
import copy
import math
from model.video_qa_model_linear_end2end import FrozenInTime
from transformers import (
    get_polynomial_decay_schedule_with_warmup,
    get_cosine_schedule_with_warmup)

from EgoTaskQA_dataset import EgoTaskQA, collate_func
from baselines.utils.utils import ReasongingTypeAccCalculator
from transforms import init_transform_dict, init_video_transform_dict
import torch.distributed as dist
from utils.util import state_dict_data_parallel_fix


class AllGather_multi(torch.autograd.Function):
    """An autograd function that performs allgather on a tensor."""

    @staticmethod
    def forward(ctx, tensor, n_gpu, args):
        output = [torch.empty_like(tensor) for _ in range(args.world_size)]
        dist.all_gather(output, tensor)
        ctx.rank = args.rank
        ctx.batch_size = tensor.shape[0]
        return torch.cat(output, 0)

    @staticmethod
    def backward(ctx, grad_output):
        return (
            grad_output[ctx.batch_size * ctx.rank : ctx.batch_size * (ctx.rank + 1)],
            None, None,
        )


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--basedir", type=str, default='outputs_end2end/',
                        help='where to store ckpts and logs')

    parser.add_argument("--per_gpu_batch_size", type=int, default=32)

    parser.add_argument("--nepoch", type=int, default=36,  
                        help='num of total epoches')
    parser.add_argument("--lr", type=float, default=5e-5,  
                        help='')
    
    parser.add_argument("--i_val",   type=int, default=20, 
                        help='frequency of console printout and metric loggin')
    parser.add_argument("--i_test",   type=int, default=20, 
                        help='frequency of console printout and metric loggin')
    parser.add_argument("--i_print",   type=int, default=5, 
                        help='frequency of console printout and metric loggin')
    parser.add_argument("--i_weight", type=int, default=20, 
                        help='frequency of weight ckpt saving')
    
    parser.add_argument('--dataset_split_type', type=str, default='direct', help='direct/indirect')
    parser.add_argument('--num_frames_per_video', type=int, default=4)
    parser.add_argument('--frame_resolution', type=int, default=224)
    parser.add_argument('--model_name', type=str, help="Path of the pretrained checkpoint")
    parser.add_argument('--resume_finetune_model_path', type=str, default=None, help ="Path of the previously fine-tuned checkpoint for resume")
    parser.add_argument('--test_only_model_path', type=str, default=None, help ="Path of the final fine-tuned checkpoint only for evaluation")
    parser.add_argument('--base_data_dir', type=str, default='data')
    args = parser.parse_args()
    return args

def main_worker(gpu, args):
    
    args.rank += gpu
    torch.distributed.init_process_group(
        backend='nccl', init_method=args.dist_url,
        world_size=args.world_size, rank=args.rank)

    torch.cuda.set_device(gpu)
    torch.backends.cudnn.benchmark = True
    
    dataset_train = EgoTaskQA(dataset_name="EgoTaskQA", text_params={"input": "text"},
                                  video_params={"input_res": args.frame_resolution,
                                                "num_frames": args.num_frames_per_video,
                                                "loading": "lax"},
                                  data_dir="./qa_videos",
                                  meta_dir="./Data/qa/" + args.dataset_split_type,
                                  tsfms=init_video_transform_dict()['test'], reader='decord', split='train', neg_param=60, args=args)
    
    sampler_train = torch.utils.data.distributed.DistributedSampler(dataset_train)
    dataloader_train = DataLoader(dataset_train, batch_size=args.per_gpu_batch_size, drop_last=True, sampler=sampler_train, pin_memory=True, collate_fn=collate_func)


    dataset_val = EgoTaskQA(dataset_name="EgoTaskQA", text_params={"input": "text"},
                                video_params={"input_res": args.frame_resolution,
                                              "num_frames": args.num_frames_per_video,
                                              "loading": "lax"},
                                data_dir="./qa_videos",
                                meta_dir="./Data/qa/" + args.dataset_split_type,
                                tsfms=init_video_transform_dict()['test'], reader='decord', split='val', neg_param=60, args=args)
    sampler_val = torch.utils.data.distributed.DistributedSampler(dataset_val)
    dataloader_val = DataLoader(dataset_val, batch_size=args.per_gpu_batch_size, drop_last=False, sampler=sampler_val, collate_fn=collate_func)


    dataset_test = EgoTaskQA(dataset_name="EgoTaskQA", text_params={"input": "text"},
                                 video_params={"input_res": args.frame_resolution,
                                               "num_frames": args.num_frames_per_video,
                                               "loading": "lax"},
                                 data_dir="./qa_videos",
                                 meta_dir="./Data/qa/" + args.dataset_split_type,
                                 tsfms=init_video_transform_dict()['test'], reader='decord', split='test', neg_param=60, args=args)
    sampler_test = torch.utils.data.distributed.DistributedSampler(dataset_test)
    dataloader_test = DataLoader(dataset_test, batch_size=args.per_gpu_batch_size, drop_last=False, sampler=sampler_test, collate_fn=collate_func)

    f = open('configs/egotaskqa.json')
    config = json.load(f)

    video_params = {"model": config['arch']['args']['video_params']['model'], "arch_config": config['arch']['args']['video_params']['arch_config'], "num_frames": config['arch']['args']['video_params']['num_frames'], "pretrained": True, "time_init":  config['arch']['args']['video_params']['time_init']}
    text_params = {"model": config['arch']['args']['text_params']['model'], "pretrained": True, "input": config['arch']['args']['text_params']['input']}
    projection_dim=config['arch']['args']["projection_dim"]
    if args.model_name is not None:
        load_checkpoint=args.model_name
    else:
        load_checkpoint=""
    projection='minimal'
    load_temporal_fix='bilinear'
    task_names = 'EgoNCE_ITM_MLM'
    norm_layer = None
    embed_dim=768

    with open('./Data/qa/' + args.dataset_split_type + '/answer_set.txt', 'r') as ansf:
        answers = ansf.readlines()
        args.output_dim = len(answers) # # output_dim == len(answers)

    model = FrozenInTime(video_params, text_params, projection_dim=projection_dim, load_checkpoint=load_checkpoint,
                         projection=projection, load_temporal_fix=load_temporal_fix,
                         task_names = task_names, norm_layer = norm_layer, embed_dim=embed_dim, model_dim=768, output_dim=args.output_dim).cuda(gpu)


    global_step = 0
    start_epoch = 0
    TIMESTAMP = "{0:%Y-%m-%dT%H-%M-%S/}".format(datetime.now())

    criterion = nn.CrossEntropyLoss().cuda(gpu)
    optimizer = optim.AdamW(model.parameters(), lr=args.lr)
    scheduler = get_cosine_schedule_with_warmup(
                         optimizer,
                         num_warmup_steps = int(args.nepoch * len(dataloader_train) * 0.1),
                         num_training_steps = int(args.nepoch * len(dataloader_train)))


    if args.resume_finetune_model_path is not None:
        ckpt = torch.load(args.resume_finetune_model_path, map_location='cpu')
        new_ckpt_model = state_dict_data_parallel_fix(ckpt['model_state_dict'], model.state_dict())
        model.load_state_dict(new_ckpt_model)
        optimizer.load_state_dict(ckpt['optimizer_state_dict'])
        scheduler.load_state_dict(ckpt['scheduler_state_dict'])
        global_step = ckpt['global_step']
        start_epoch = (global_step+1)//len(dataloader_train)

    if args.test_only_model_path is not None:
        ckpt = torch.load(args.test_only_model_path, map_location='cpu')
        new_ckpt_model = state_dict_data_parallel_fix(ckpt['model_state_dict'], model.state_dict())
        model.load_state_dict(new_ckpt_model)


    model = nn.SyncBatchNorm.convert_sync_batchnorm(model)


    scaler = torch.cuda.amp.GradScaler()
    ddp_args={}
    ddp_args['static_graph'] = True
    model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[gpu], find_unused_parameters=True, **ddp_args)

    with open('./Data/qa/' + args.dataset_split_type + '/all_reasoning_types.txt', 'r') as reasonf:
        all_reasoning_types = reasonf.readlines()
        all_reasoning_types = [item.strip() for item in all_reasoning_types]

    train_acc_calculator = ReasongingTypeAccCalculator(reasoning_types=all_reasoning_types)
    test_acc_calculator = ReasongingTypeAccCalculator(reasoning_types=all_reasoning_types)

    if args.test_only_model_path is not None:
        test_acc_calculator.reset()
        test_loss, test_acc = validate(model, dataloader_test, 0, args, gpu, acc_calculator=test_acc_calculator)
        acc_dct = test_acc_calculator.get_acc()
        if args.rank == 0:
            print("Evaluation Accuracy: ", acc_dct)
        return 0

    
    if args.rank == 0:
        args.basedir = os.path.join(args.basedir, args.model_name.split('/')[-1].split('.')[0])
        args.basedir = os.path.join(args.basedir, args.dataset_split_type)
        log_dir = os.path.join(args.basedir, TIMESTAMP, 'logs')
        os.makedirs(log_dir)

        with open(os.path.join(log_dir, 'argument.txt'), 'w') as f:
            for key, value in vars(args).items():
                f.write('%s:%s\n'%(key, value))
                print(key, value)

        log_file = open(os.path.join(log_dir, 'log.txt'), 'w')
        stats_file = open(os.path.join(log_dir, 'stats.txt'), 'a', buffering=1)
        print(' '.join(sys.argv))
        print(' '.join(sys.argv), file=stats_file)
        os.makedirs(os.path.join(args.basedir, TIMESTAMP, 'ckpts'), exist_ok=True)


    best_overall_accuracy = -1
    unique_dict = torch.load('./reasoning_unique_cat.pth')

    for epoch in range(start_epoch, args.nepoch):
        sampler_train.set_epoch(epoch)
        model.train()
        train_acc_calculator.reset()
        pbar = tqdm(total = len(dataloader_train))

        for idx, (video_frames, text_tokens, attention_masks, answers, reasoning_type_lst) in tqdm(enumerate(dataloader_train)):

            optimizer.zero_grad()

            B, num_frame_per_video, C, H, W = video_frames.shape
            video_frames, text_tokens, attention_masks, answers = video_frames.cuda(gpu, non_blocking=True), text_tokens.cuda(gpu, non_blocking=True), attention_masks.cuda(gpu, non_blocking=True), answers.cuda(gpu, non_blocking=True)

            data_batch = {"video": video_frames, "text": {"input_ids": text_tokens, "attention_mask": attention_masks}}
            
            with torch.set_grad_enabled(True):
                with torch.cuda.amp.autocast():
                    
                    logits = model.forward(data_batch)
                    logits = args.allgather(logits, gpu, args)
                    answers = args.allgather(answers, gpu, args)
                    loss = criterion(logits, answers.long())

                    max_len = 7
                    # pad all tensors to have same length
                    reasoning_type_lst = [torch.nn.functional.pad(x, pad=(0, max_len - x.numel()), mode='constant', value= -100) for x in reasoning_type_lst]
                    reasoning_type_lst = torch.stack(reasoning_type_lst).cuda(gpu, non_blocking=True)
                    reasoning_type_lst = args.allgather(reasoning_type_lst, gpu, args)
                    

            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
            scheduler.step()

            reasoning_lst = []

            for _i in range(reasoning_type_lst.shape[0]):
                _list = []
                for _j in range(reasoning_type_lst.shape[1]):
                    if reasoning_type_lst[_i][_j] > -100:
                        _list.append([*unique_dict.keys()][int(reasoning_type_lst[_i][_j])])
                reasoning_lst.append(_list)

            reasoning_type_lst = reasoning_lst
            del reasoning_lst

            pred = torch.argmax(logits, dim=1)
            train_acc = sum(pred == answers) / B*args.world_size

            train_acc_calculator.update(reasoning_type_lst, pred, answers)

            pbar.update(1)

            args.i_test = args.i_val = args.i_weight = len(dataloader_train)
            

            if args.rank == 0:
                if global_step % args.i_print  == 0:
                    print(f"global_step:{global_step}, train_loss:{loss.item()}, train_acc:{train_acc}")
                    stats = dict(epoch=epoch, step=global_step,
                             lr_weights=optimizer.param_groups[0]['lr'],
                             loss=loss.item(), train_acc=train_acc.item())
                    print(json.dumps(stats), file=stats_file)

            
            if (global_step+1) % args.i_val == 0:
                test_acc_calculator.reset()
                val_loss, val_acc = validate(model, dataloader_val, epoch, args, gpu, acc_calculator=test_acc_calculator)
                acc_dct = test_acc_calculator.get_acc()
                stats = dict(epoch=epoch, step=global_step, val_loss=val_loss.item(), val_acc=val_acc.item())
                for key, value in acc_dct.items():
                    stats.update({f'val/reasoning_{key}': value})
                
                if args.rank == 0:
                    print(json.dumps(stats), file=stats_file)
                    log_file.write(f'[VAL]: epoch: {epoch}, global_step: {global_step}\n')
                    log_file.write(f'true count dct: {test_acc_calculator.true_count_dct}\nall count dct: {test_acc_calculator.all_count_dct}\n\n')


            if (global_step+1) % args.i_test == 0:
                test_acc_calculator.reset()
                test_loss, test_acc = validate(model, dataloader_test, epoch, args, gpu, acc_calculator=test_acc_calculator)
                acc_dct = test_acc_calculator.get_acc()
                stats = dict(epoch=epoch, step=global_step, test_loss=test_loss.item(), test_acc=test_acc.item())
                
                for key, value in acc_dct.items():
                    stats.update({f'test/reasoning_{key}': value})
                    if key == args.dataset_split_type:
                        if best_overall_accuracy < value:
                            best_act_dict = copy.deepcopy(acc_dct)
                            best_overall_accuracy = value
                            if args.rank == 0:
                                torch.save({'model_state_dict': model.state_dict(),
                                            'optimizer_state_dict': optimizer.state_dict(),
                                            'scheduler_state_dict': scheduler.state_dict(),
                                            'loss': loss,
                                            'global_step': global_step,}, os.path.join(args.basedir, TIMESTAMP, 'ckpts', f"model_best.tar"))

                
                if args.rank == 0:
                    print(json.dumps(stats), file=stats_file)
                    log_file.write(f'[TEST]: epoch: {epoch}, global_step: {global_step}\n')
                    log_file.write(f'true count dct: {test_acc_calculator.true_count_dct}\nall count dct: {test_acc_calculator.all_count_dct}\n\n')

            
            if (global_step+1) % args.i_weight == 0:
                if args.rank == 0:
                    torch.save({
                        'model_state_dict': model.state_dict(),
                        'optimizer_state_dict': optimizer.state_dict(),
                        'scheduler_state_dict': scheduler.state_dict(),
                        'loss': loss,
                        'global_step': global_step,
                    }, os.path.join(args.basedir, TIMESTAMP, 'ckpts', f"model_{global_step}.tar"))

            global_step += 1


        if args.rank == 0:
            acc_dct = train_acc_calculator.get_acc()
            for key, value in acc_dct.items():
                stats.update({f'train/reasoning_{key}': value})
            print(json.dumps(stats), file=stats_file)

        
            log_file.write(f'[TRAIN]: epoch: {epoch}, global_step: {global_step}\n')
            log_file.write(f'true count dct: {train_acc_calculator.true_count_dct}\nall count dct: {train_acc_calculator.all_count_dct}\n\n')
            log_file.write(f'============================ Printing Best Overall Accuracy till Now ==========================================\n\n')
            log_file.write(f'BEST RESULTS: {best_act_dict}\n\n')
            log_file.write(f'===============================================================================================================\n\n')
            log_file.flush()
        
def validate(model, val_dataloader, epoch, args, gpu, acc_calculator):
    model.eval()
    all_acc = 0
    all_loss = 0
    acc_calculator.reset()
    
    print('validating...')
    unique_dict = torch.load('./reasoning_unique_cat.pth')
    with torch.no_grad():
        starttime = time.time()
        for idx, (video_frames, text_tokens, attention_masks, answers, reasoning_type_lst) in tqdm(enumerate(val_dataloader)):

            B, num_frame_per_video, C, H, W = video_frames.shape
            video_frames, text_tokens, attention_masks, answers = video_frames.cuda(gpu, non_blocking=True), text_tokens.cuda(gpu, non_blocking=True), attention_masks.cuda(gpu, non_blocking=True), answers.cuda(gpu, non_blocking=True)

            data_batch = {"video": video_frames, "text": {"input_ids": text_tokens, "attention_mask": attention_masks}}
            logits = model(data_batch)

            logits_all = [torch.zeros_like(logits) for _ in range(args.world_size)]
            torch.distributed.all_gather(logits_all, logits)
            logits_all = torch.cat(logits_all, dim=0)

            answers_all = [torch.zeros_like(answers) for _ in range(args.world_size)]
            torch.distributed.all_gather(answers_all, answers)
            answers_all = torch.cat(answers_all, dim=0)

            all_loss += nn.CrossEntropyLoss().cuda(gpu)(logits_all, answers_all.long())
            # print('validate finish in', (time.time() - starttime) * (len(val_loader) - i), 's')
            # starttime = time.time()
            pred_all = torch.argmax(logits_all, dim=1)
            test_acc = sum(pred_all == answers_all) / B * args.world_size
            all_acc += test_acc


            max_len = 7
            # pad all tensors to have same length
            reasoning_type_lst = [torch.nn.functional.pad(x, pad=(0, max_len - x.numel()), mode='constant', value= -100) for x in reasoning_type_lst]
            reasoning_type_lst = torch.stack(reasoning_type_lst).cuda(gpu, non_blocking=True)
            reasoning_type_lst = args.allgather(reasoning_type_lst, gpu, args)

            reasoning_lst = []

            for _i in range(reasoning_type_lst.shape[0]):
                _list = []
                for _j in range(reasoning_type_lst.shape[1]):
                    if reasoning_type_lst[_i][_j] > -100:
                        _list.append([*unique_dict.keys()][int(reasoning_type_lst[_i][_j])])
                reasoning_lst.append(_list)

            acc_calculator.update(reasoning_lst, pred_all, answers_all)


    all_loss /= len(val_dataloader)
    all_acc /= len(val_dataloader)
    model.train()
    return all_loss, all_acc


def main():
    args = parse_args()
    args.allgather = AllGather_multi.apply
    args.ngpus_per_node = torch.cuda.device_count()
    if 'SLURM_JOB_ID' in os.environ:
        # single-node and multi-node distributed training on SLURM cluster
        # requeue job on SLURM preemption
        signal.signal(signal.SIGUSR1, handle_sigusr1)
        signal.signal(signal.SIGTERM, handle_sigterm)
        # find a common host name on all nodes
        # assume scontrol returns hosts in the same order on all nodes
        cmd = 'scontrol show hostnames ' + os.getenv('SLURM_JOB_NODELIST')
        stdout = subprocess.check_output(cmd.split())
        host_name = stdout.decode().splitlines()[0]
        args.rank = int(os.getenv('SLURM_NODEID')) * args.ngpus_per_node
        args.world_size = int(os.getenv('SLURM_NNODES')) * args.ngpus_per_node
        args.dist_url = f'tcp://{host_name}:58472'
    else:
        # single-node distributed training
        args.rank = 0
        args.dist_url = 'tcp://localhost:58472'
        args.world_size = args.ngpus_per_node
    torch.multiprocessing.spawn(main_worker, (args,), args.ngpus_per_node)


if __name__ == '__main__':
    main()

EgoTaskQA/main_end2end.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import torch
import torch.nn as nn
import torch.nn.functional as F

from transformers.models.bert.modeling_bert import BertPredictionHeadTransform


class Pooler(nn.Module):
    def __init__(self, hidden_size):
        super().__init__()
        self.dense = nn.Linear(hidden_size, hidden_size)
        self.activation = nn.Tanh()

    def forward(self, hidden_states):
        first_token_tensor = hidden_states #[:, 0]
        pooled_output = self.dense(first_token_tensor)
        pooled_output = self.activation(pooled_output)
        return pooled_output


class ITMHead(nn.Module):
    def __init__(self, hidden_size):
        super().__init__()
        self.fc = nn.Linear(hidden_size, 2)

    def forward(self, x):
        x = self.fc(x)
        return x


class MLMHead(nn.Module):
    def __init__(self, config, weight=None):
        super().__init__()
        self.transform = BertPredictionHeadTransform(config)
        self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.bias = nn.Parameter(torch.zeros(config.vocab_size))
        if weight is not None:
            self.decoder.weight = weight

    def forward(self, x):
        x = self.transform(x)
        x = self.decoder(x) + self.bias
        return x

EgoTaskQA/model/heads.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import os
import pdb

import timm
import torch
import yaml
import torch.nn as nn
import numpy as np
import torch.nn.functional as F
from transformers import AutoModel
from einops import rearrange, repeat

from base import BaseModel
from model import video_transformer
from model.video_transformer import SpaceTimeTransformer
from utils.util import state_dict_data_parallel_fix

from model import roberta
from model.roberta import RobertaModel, _prepare_decoder_attention_mask
from model import heads
from transformers import RobertaConfig
from functools import partial
import copy
import torch.distributed as dist

with open('./EgoNCE_MLM_ITM_Config.yml') as f:
    config = yaml.load(f, Loader=yaml.FullLoader)

def init_weights(module):
    if isinstance(module, (nn.Linear, nn.Embedding)):
        module.weight.data.normal_(mean=0.0, std=0.02)
    elif isinstance(module, nn.LayerNorm):
        module.bias.data.zero_()
        module.weight.data.fill_(1.0)

    if isinstance(module, nn.Linear) and module.bias is not None:
        module.bias.data.zero_()


class FrozenInTime(BaseModel):
    def __init__(self,
                 video_params,
                 text_params,
                 projection_dim=4096,
                 load_checkpoint=None,
                 projection='minimal',
                 load_temporal_fix='zeros', #bilinear
                 config = config,
                 task_names = 'EgoNCE_ITM_MLM',
                 norm_layer = None,
                 embed_dim=768):
                 #egotaskqa_head_output_dim=1010):
        super().__init__()

        self.video_params = video_params
        self.text_params = text_params
        self.load_temporal_fix = load_temporal_fix
        self.config = config
        #self.egotaskqa_head_output_dim = egotaskqa_head_output_dim
        self.task_names = task_names
        if not text_params['pretrained']:
            raise NotImplementedError("Huggingface text models require pretrained init.")

        if self.text_params['model'].startswith('roberta'):
            self.text_model = RobertaModel.from_pretrained("roberta-base")
        self.text_model.train()

        pretrained = video_params['pretrained']
        if video_params['model'] == "SpaceTimeTransformer":
            self.num_frames = 4 #self.config["data_loader"][0]["args"]["video_params"]["num_frames"]
            time_init = 'zeros'
            attention_style = 'frozen-in-time'
            arch_config = 'base_patch16_224'
            vit_init = 'imagenet-21k'
            if arch_config == 'base_patch16_224':
                #vit_model = timm.models.vision_transformer.vit_base_patch16_224(pretrained=pretrained)
                vit_model = torch.load("/cis/home/shraman/works_meta_2022/pre-training/EgoVLP_Fused_HardNegITM_Checkpoint_multinode/frozen-in-time-main/pretrained/jx_vit_base_p16_224-80ecf9dd.pth", map_location="cpu")
                model = SpaceTimeTransformer(num_frames=self.num_frames,
                                            time_init=time_init,
                                            attention_style=attention_style)
            else:
                raise NotImplementedError

            model.head = nn.Identity()
            model.pre_logits = nn.Identity()
            ftr_dim = model.embed_dim
           
            if load_checkpoint in ["", None]:
                # vit_checkpoint = vit_model.state_dict()
                # model.load_state_dict(vit_checkpoint, strict=False)
                vit_checkpoint = vit_model
                new_vit_dict = state_dict_data_parallel_fix(vit_checkpoint, model.state_dict())
                model.load_state_dict(new_vit_dict, strict=False)
            self.video_model = model
        else:
            raise NotImplementedError(f"{video_params['model']} not implemented")

        # for backwards compatibility (old models)
        self.video_model.fc = nn.Identity()

        # Project to a common embedding
        if projection == 'minimal':

            txt_proj = nn.Sequential(
                                    nn.Linear(self.text_model.config.hidden_size, projection_dim, bias=False),
                                    nn.ReLU(inplace=True), nn.Linear(projection_dim, projection_dim, bias=True),
                                    nn.ReLU(inplace=True), nn.Linear(projection_dim, projection_dim, bias=True)
                                    )

            vid_proj = nn.Sequential(
                                    nn.Linear(ftr_dim, projection_dim, bias=False),
                                    nn.ReLU(inplace=True), nn.Linear(projection_dim, projection_dim, bias=True),
                                    nn.ReLU(inplace=True), nn.Linear(projection_dim, projection_dim, bias=True)
                                    )
        
        elif projection == '':
            txt_proj = nn.Identity()
            vid_proj = nn.Identity()
        else:
            raise NotImplementedError
        self.txt_proj = txt_proj
        self.vid_proj = vid_proj


        if ('MLM' in self.task_names or 'ITM' in self.task_names):
            # for FIBER-like cross-attention

            bert_config = RobertaConfig(
                vocab_size=self.config["vocab_size"],
                hidden_size=self.config["hidden_size"],
                num_hidden_layers=self.config["num_layers"],
                num_attention_heads=self.config["num_heads"],
                intermediate_size=self.config["hidden_size"] * config["mlp_ratio"],
                #max_position_embeddings=maxlen, [was used in BTGOT script]
                hidden_dropout_prob=self.config["drop_rate"],
                attention_probs_dropout_prob=self.config["drop_rate"],
            )

            self.num_fuse_block=self.config["num_fuse_block"]
            self.num_text_layer=self.config["num_layers"]
            roberta.NUM_FUSE_BLOCK = self.video_model.NUM_FUSE_BLOCK=self.num_fuse_block
            roberta.DIM_IMG=self.config["input_image_embed_size"]
            self.video_model.DIM_TXT=self.config["input_text_embed_size"]

            self.cross_modal_text_transform = nn.Linear(self.config["input_text_embed_size"], self.config["hidden_size"])
            self.cross_modal_text_transform.apply(init_weights)
            self.cross_modal_video_transform = nn.Linear(self.config["input_image_embed_size"], self.config["hidden_size"])
            self.cross_modal_video_transform.apply(init_weights)

            self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))

            self.num_patches = self.video_model.patch_embed.num_patches
            self.patches_per_frame = self.num_patches//self.num_frames
            norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
            self.norm = norm_layer(embed_dim)
            self.pre_logits = nn.Identity()


            self.avgpool = nn.AdaptiveAvgPool1d(1)
            self.cross_modal_video_pooler = heads.Pooler(config["hidden_size"])
            self.cross_modal_video_pooler.apply(init_weights)
            self.cross_modal_text_pooler = heads.Pooler(config["hidden_size"])
            self.cross_modal_text_pooler.apply(init_weights)

            ## einops transformations
            self.einops_from_space = 'b (f n) d'
            self.einops_to_space = '(b f) n d'
            self.einops_from_time = 'b (f n) d'
            self.einops_to_time = '(b n) f d'

        if 'MLM' in self.task_names:
            self.mlm_score = heads.MLMHead(bert_config)
            self.mlm_score.apply(init_weights)

        if 'ITM' in self.task_names:
            self.itm_score = heads.ITMHead(config["hidden_size"] * 2)
            self.itm_score.apply(init_weights)

        if load_checkpoint not in ["", None]:
            checkpoint = torch.load(load_checkpoint, map_location='cpu')
            state_dict = checkpoint['state_dict']
            new_state_dict = state_dict_data_parallel_fix(state_dict, self.state_dict())
            new_state_dict = self._inflate_positional_embeds(new_state_dict)
            self.load_state_dict(new_state_dict, strict=False)
            print("Model is loaded with pre-trained parameters")

    def set_device(self, device):
        self.device = device

    def forward(self, data, video_only=False, return_embeds=True):
        

        text_data = data['text']
        video_data = data['video']


        b, curr_frames, channels, _, _ = video_data.shape
        video_data_itm = self.video_model.patch_embed(video_data)
        video_data_itm = video_data_itm.flatten(2).transpose(2, 1)
        video_data_itm = video_data_itm.reshape(b, -1, self.video_model.patch_embed.embed_dim)

        BF = video_data_itm.shape[0]
        cls_tokens = self.cls_token.expand(BF, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        video_data_itm = torch.cat((cls_tokens, video_data_itm), dim=1)
        # positional embed needs to be tiled for each frame (this does [1,2,3] --> [1,2,3,1,2,3]...)
        cls_embed = self.video_model.pos_embed[:, 0, :].unsqueeze(1)
        tile_pos_embed = self.video_model.pos_embed[:, 1:, :].repeat(1, self.num_frames, 1)
        # temporal embed needs to be repeated within each frame (this does [1,2,3] --> [1,1,1,2,2,2,3,3,3]...)
        tile_temporal_embed = self.video_model.temporal_embed.repeat_interleave(self.patches_per_frame, 1)
        total_pos_embed = tile_pos_embed + tile_temporal_embed
        total_pos_embed = torch.cat([cls_embed, total_pos_embed], dim=1)

        n = self.patches_per_frame
        f = curr_frames

        curr_patches = video_data_itm.shape[1]
        video_data_itm = video_data_itm + total_pos_embed[:, :curr_patches]
        video_data_itm = self.video_model.pos_drop(video_data_itm)

        unfused_blocks = self.num_text_layer - self.num_fuse_block
            

        for blk_i, blk in enumerate(self.video_model.blocks[:unfused_blocks]):
            if config['use_checkpoint']:
                video_data_itm = torch.utils.checkpoint.checkpoint(blk, video_data_itm, self.einops_from_space, self.einops_to_space, self.einops_from_time, self.einops_to_time,
                                n, f)
            else:
                video_data_itm = blk(video_data_itm, self.einops_from_space, self.einops_to_space, self.einops_from_time, self.einops_to_time,
                                time_n=n, space_f=f)
                            
            
        text_embeds = self.text_model.embeddings(input_ids=text_data['input_ids']) # before it was input_ids=text_ids
        device = text_embeds.device
        #print(text_data.keys())
        text_masks = text_data['attention_mask']
        input_shape = text_masks.size()
        extend_text_masks = self.text_model.get_extended_attention_mask(text_masks, input_shape, device)
        for layer_i, layer in enumerate(self.text_model.encoder.layer[:unfused_blocks]):
            if config['use_checkpoint']:
                text_embeds = torch.utils.checkpoint.checkpoint(layer, text_embeds, extend_text_masks)[0]
            else:
                text_embeds = layer(text_embeds, extend_text_masks)[0]


        for blk_i, blk in enumerate(self.video_model.blocks[unfused_blocks:self.num_text_layer]):
            if config['use_checkpoint']:
                    
                fuse_video_data = torch.utils.checkpoint.checkpoint(blk, video_data_itm, self.einops_from_space, self.einops_to_space, self.einops_from_time, self.einops_to_time, 
                                          n, f, text_embeds, extend_text_masks)
                text_embeds = torch.utils.checkpoint.checkpoint(self.text_model.encoder.layer[blk_i + unfused_blocks],
                                          text_embeds, extend_text_masks, None, (video_data_itm), None, None, False, True)[0]
            else:
                fuse_video_data = blk(video_data_itm, self.einops_from_space, self.einops_to_space, self.einops_from_time, self.einops_to_time, 
                                          y=text_embeds, y_mask=extend_text_masks, time_n=n, space_f=f)
                text_embeds = self.text_model.encoder.layer[blk_i + unfused_blocks](text_embeds, extend_text_masks, encoder_hidden_states=(video_data_itm), last_norm=True)[0]
            video_data_itm = fuse_video_data

            
        #print("Shape of model output", video_data.shape)
        video_data_itm = self.norm(video_data_itm)[:, :]
        video_data_itm = self.pre_logits(video_data_itm)

        #return self.egotaskqa_head(video_data_itm)
        return video_data_itm
        #return sim_matrix(text_embeddings, video_embeddings)

    
    def compute_text(self, text_data):
        if self.text_params['model'].startswith('bert'):
            text_embeddings = self.text_model(text_data['input_ids'], attention_mask=text_data['attention_mask'])[
                'pooler_output']
        elif self.text_params['model'].startswith('distilbert'):
            text_embeddings = self.text_model(**text_data).last_hidden_state[:, 0, :]
        elif self.text_params['model'].startswith('roberta'):
            text_embeddings = self.text_model(**text_data).last_hidden_state[:, 0, :]
        else:
            raise NotImplementedError
        text_embeddings = self.txt_proj(text_embeddings)
        return text_embeddings

    def compute_text_tokens(self, text_data, is_proj=True):
        if self.text_params['model'].startswith('bert'):
            text_embeddings = self.text_model(text_data['input_ids'], attention_mask=text_data['attention_mask'])[
                'pooler_output']    # not implement for bert
        elif self.text_params['model'].startswith('distilbert'):
            text_embeddings = self.text_model(**text_data).last_hidden_state
        elif self.text_params['model'].startswith('roberta'):
            text_embeddings = self.text_model(**text_data).last_hidden_state
        else:
            raise NotImplementedError
        
        if is_proj:
            text_embeddings = self.txt_proj(text_embeddings)
        
        return text_embeddings

    def compute_video(self, video_data):
        video_embeddings = self.video_model(video_data)
        video_embeddings = self.vid_proj(video_embeddings)
        return video_embeddings

    def _inflate_positional_embeds(self, new_state_dict):
        # allow loading of timesformer with fewer num_frames
        curr_keys = list(self.state_dict().keys())
        if 'video_model.temporal_embed' in new_state_dict and 'video_model.temporal_embed' in curr_keys:
            load_temporal_embed = new_state_dict['video_model.temporal_embed']
            load_num_frames = load_temporal_embed.shape[1]
            curr_num_frames = self.video_params['num_frames']
            embed_dim = load_temporal_embed.shape[2]

            if load_num_frames != curr_num_frames:
                if load_num_frames > curr_num_frames:
                    print(f'### loaded {self.video_params["model"]} model has MORE frames than current...'
                          f'### loading weights, filling in the extras via {self.load_temporal_fix}')
                    new_temporal_embed = load_temporal_embed[:, :curr_num_frames, :]
                else:
                    print(f'### loaded {self.video_params["model"]} model has FEWER frames than current...'
                          f'### loading weights, filling in the extras via {self.load_temporal_fix}')
                    if self.load_temporal_fix == 'zeros':
                        new_temporal_embed = torch.zeros([load_temporal_embed.shape[0], curr_num_frames, embed_dim])
                        new_temporal_embed[:, :load_num_frames] = load_temporal_embed
                    elif self.load_temporal_fix in ['interp', 'bilinear']:
                        # interpolate
                        # unsqueeze so pytorch thinks its an image
                        mode = 'nearest'
                        if self.load_temporal_fix == 'bilinear':
                            mode = 'bilinear'
                        load_temporal_embed = load_temporal_embed.unsqueeze(0)
                        new_temporal_embed = F.interpolate(load_temporal_embed,
                                                           (curr_num_frames, embed_dim), mode=mode, align_corners=True).squeeze(0)
                    else:
                        raise NotImplementedError
                new_state_dict['video_model.temporal_embed'] = new_temporal_embed
        # allow loading with smaller spatial patches. assumes custom border crop, to append the
        # border patches to the input sequence
        if 'video_model.pos_embed' in new_state_dict and 'video_model.pos_embed' in curr_keys:
            load_pos_embed = new_state_dict['video_model.pos_embed']
            load_num_patches = load_pos_embed.shape[1]
            curr_pos_embed = self.state_dict()['video_model.pos_embed']
            if load_num_patches != curr_pos_embed.shape[1]:
                raise NotImplementedError(
                    'Loading models with different spatial resolution / patch number not yet implemented, sorry.')

        return new_state_dict

def sim_matrix(a, b, eps=1e-8):
    """
    added eps for numerical stability
    """
    a_n, b_n = a.norm(dim=1)[:, None], b.norm(dim=1)[:, None]
    a_norm = a / torch.max(a_n, eps * torch.ones_like(a_n))
    b_norm = b / torch.max(b_n, eps * torch.ones_like(b_n))
    sim_mt = torch.mm(a_norm, b_norm.transpose(0, 1))
    return sim_mt


def sim_matrix_batch_val(a, b, eps=1e-8):
    """
    added eps for numerical stability
    """
    a_n, b_n = a.norm(dim=-1).unsqueeze(-1), b.norm(dim=-1).unsqueeze(-1)
    a_norm = a / torch.max(a_n, eps * torch.ones_like(a_n))
    b_norm = b / torch.max(b_n, eps * torch.ones_like(b_n))
    sim_mt = torch.bmm(a_norm, b_norm.transpose(1, 2))
    return sim_mt


if __name__ == "__main__":
    video_params = {"model": SpaceTimeTransformer, "arch_config": "base_patch16_224", "num_frames": 4, "pretrained": True, "time_init": "zeros"}
    text_params: {"model": "roberta-base", "pretrained": True, "input": "text"}
    model = FrozenInTime(video_params, text_params)

EgoTaskQA/model/model.py
facebookresearch/EgoVLPv2
# coding=utf-8
# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""PyTorch RoBERTa model. """

import yaml

with open('./EgoNCE_MLM_ITM_Config.yml') as f:
    config_yaml = yaml.load(f, Loader=yaml.FullLoader)

NUM_FUSE_BLOCK = config_yaml['num_fuse_block']
DIM_IMG = 768

import math
import inspect
from typing import Callable, List, Optional, Set, Tuple, Union

import torch
import torch.nn as nn
import torch.utils.checkpoint
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from transformers.activations import ACT2FN, gelu
from transformers.file_utils import (
    add_code_sample_docstrings,
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    replace_return_docstrings,
)
from transformers.modeling_outputs import (
    BaseModelOutputWithPastAndCrossAttentions,
    BaseModelOutputWithPoolingAndCrossAttentions,
    CausalLMOutputWithCrossAttentions,
    MaskedLMOutput,
    MultipleChoiceModelOutput,
    QuestionAnsweringModelOutput,
    SequenceClassifierOutput,
    TokenClassifierOutput,
)
from transformers.modeling_utils import (
    PreTrainedModel,
    find_pruneable_heads_and_indices,
    prune_linear_layer,
)
from transformers.utils import logging
from transformers.models.roberta.configuration_roberta import RobertaConfig


logger = logging.get_logger(__name__)

_CHECKPOINT_FOR_DOC = "roberta-base"
_CONFIG_FOR_DOC = "RobertaConfig"
_TOKENIZER_FOR_DOC = "RobertaTokenizer"

ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST = [
    "roberta-base",
    "roberta-large",
    "roberta-large-mnli",
    "distilroberta-base",
    "roberta-base-openai-detector",
    "roberta-large-openai-detector",
    # See all RoBERTa models at https://huggingface.co/models?filter=roberta
]


def apply_chunking_to_forward(
    forward_fn: Callable[..., torch.Tensor], chunk_size: int, chunk_dim: int, *input_tensors, last_norm=True
) -> torch.Tensor:
    """
    This function chunks the `input_tensors` into smaller input tensor parts of size `chunk_size` over the dimension
    `chunk_dim`. It then applies a layer `forward_fn` to each chunk independently to save memory.
    If the `forward_fn` is independent across the `chunk_dim` this function will yield the same result as directly
    applying `forward_fn` to `input_tensors`.
    Args:
        forward_fn (`Callable[..., torch.Tensor]`):
            The forward function of the model.
        chunk_size (`int`):
            The chunk size of a chunked tensor: `num_chunks = len(input_tensors[0]) / chunk_size`.
        chunk_dim (`int`):
            The dimension over which the `input_tensors` should be chunked.
        input_tensors (`Tuple[torch.Tensor]`):
            The input tensors of `forward_fn` which will be chunked
    Returns:
        `torch.Tensor`: A tensor with the same shape as the `forward_fn` would have given if applied`.
    Examples:
    ```python
    # rename the usual forward() fn to forward_chunk()
    def forward_chunk(self, hidden_states):
        hidden_states = self.decoder(hidden_states)
        return hidden_states
    # implement a chunked forward function
    def forward(self, hidden_states):
        return apply_chunking_to_forward(self.forward_chunk, self.chunk_size_lm_head, self.seq_len_dim, hidden_states)
    ```"""

    assert len(input_tensors) > 0, f"{input_tensors} has to be a tuple/list of tensors"

    # inspect.signature exist since python 3.5 and is a python method -> no problem with backward compatibility
    num_args_in_forward_chunk_fn = len(inspect.signature(forward_fn).parameters)
    # if num_args_in_forward_chunk_fn != len(input_tensors):
    #    raise ValueError(
    #        f"forward_chunk_fn expects {num_args_in_forward_chunk_fn} arguments, but only {len(input_tensors)} input "
    #        "tensors are given"
    #    )

    if chunk_size > 0:
        tensor_shape = input_tensors[0].shape[chunk_dim]
        for input_tensor in input_tensors:
            if input_tensor.shape[chunk_dim] != tensor_shape:
                raise ValueError(
                    f"All input tenors have to be of the same shape: {tensor_shape}, "
                    f"found shape {input_tensor.shape[chunk_dim]}"
                )

        if input_tensors[0].shape[chunk_dim] % chunk_size != 0:
            raise ValueError(
                f"The dimension to be chunked {input_tensors[0].shape[chunk_dim]} has to be a multiple of the chunk "
                f"size {chunk_size}"
            )

        num_chunks = input_tensors[0].shape[chunk_dim] // chunk_size

        # chunk input tensor into tuples
        input_tensors_chunks = tuple(input_tensor.chunk(num_chunks, dim=chunk_dim) for input_tensor in input_tensors)
        # apply forward fn to every tuple
        output_chunks = tuple(
            forward_fn(*input_tensors_chunk, last_norm=last_norm) for input_tensors_chunk in zip(*input_tensors_chunks)
        )
        # concatenate output at same dimension
        return torch.cat(output_chunks, dim=chunk_dim)

    return forward_fn(*input_tensors, last_norm=last_norm)


class RobertaEmbeddings(nn.Module):
    """
    Same as BertEmbeddings with a tiny tweak for positional embeddings indexing.
    """

    # Copied from transformers.models.bert.modeling_bert.BertEmbeddings.__init__
    def __init__(self, config):
        super().__init__()
        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
        self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)

        # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
        # any TensorFlow checkpoint file
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

        # position_ids (1, len position emb) is contiguous in memory and exported when serialized
        self.register_buffer("position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)))
        self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")

        # End copy
        self.padding_idx = config.pad_token_id
        self.position_embeddings = nn.Embedding(
            config.max_position_embeddings, config.hidden_size, padding_idx=self.padding_idx
        )

    def forward(
        self, input_ids=None, token_type_ids=None, position_ids=None, inputs_embeds=None, past_key_values_length=0
    ):
        if position_ids is None:
            if input_ids is not None:
                # Create the position ids from the input token ids. Any padded tokens remain padded.
                position_ids = create_position_ids_from_input_ids(
                    input_ids, self.padding_idx, past_key_values_length
                ).to(input_ids.device)
            else:
                position_ids = self.create_position_ids_from_inputs_embeds(inputs_embeds)

        if input_ids is not None:
            input_shape = input_ids.size()
        else:
            input_shape = inputs_embeds.size()[:-1]

        if token_type_ids is None:
            token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)

        if inputs_embeds is None:
            inputs_embeds = self.word_embeddings(input_ids)
        token_type_embeddings = self.token_type_embeddings(token_type_ids)

        embeddings = inputs_embeds + token_type_embeddings
        if self.position_embedding_type == "absolute":
            position_embeddings = self.position_embeddings(position_ids)
            embeddings += position_embeddings
        embeddings = self.LayerNorm(embeddings)
        embeddings = self.dropout(embeddings)
        return embeddings

    def create_position_ids_from_inputs_embeds(self, inputs_embeds):
        """
        We are provided embeddings directly. We cannot infer which are padded so just generate sequential position ids.
        Args:
            inputs_embeds: torch.Tensor
        Returns: torch.Tensor
        """
        input_shape = inputs_embeds.size()[:-1]
        sequence_length = input_shape[1]

        position_ids = torch.arange(
            self.padding_idx + 1, sequence_length + self.padding_idx + 1, dtype=torch.long, device=inputs_embeds.device
        )
        return position_ids.unsqueeze(0).expand(input_shape)


# Copied from transformers.models.bert.modeling_bert.BertSelfAttention with Bert->Roberta
class RobertaSelfAttention(nn.Module):
    def __init__(self, config, layer_index=None):
        super().__init__()
        if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
            raise ValueError(
                f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
                f"heads ({config.num_attention_heads})"
            )

        self.num_attention_heads = config.num_attention_heads
        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
        self.all_head_size = self.num_attention_heads * self.attention_head_size

        self.query = nn.Linear(config.hidden_size, self.all_head_size)
        if layer_index is None:
            self.key = nn.Linear(config.hidden_size, self.all_head_size)
            self.value = nn.Linear(config.hidden_size, self.all_head_size)
        else:
            self.key = nn.Linear(DIM_IMG, self.all_head_size)
            self.value = nn.Linear(DIM_IMG, self.all_head_size)

        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
        self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")
        if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
            self.max_position_embeddings = config.max_position_embeddings
            self.distance_embedding = nn.Embedding(2 * config.max_position_embeddings - 1, self.attention_head_size)

        self.is_decoder = config.is_decoder

    def transpose_for_scores(self, x):
        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
        x = x.view(*new_x_shape)
        return x.permute(0, 2, 1, 3)

    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        past_key_value=None,
        output_attentions=False,
    ):
        mixed_query_layer = self.query(hidden_states)

        # If this is instantiated as a cross-attention module, the keys
        # and values come from an encoder; the attention mask needs to be
        # such that the encoder's padding tokens are not attended to.
        is_cross_attention = encoder_hidden_states is not None

        if encoder_hidden_states is not None:
            key_layer = self.transpose_for_scores(self.key(encoder_hidden_states))
            value_layer = self.transpose_for_scores(self.value(encoder_hidden_states))
            attention_mask = encoder_attention_mask
        else:
            key_layer = self.transpose_for_scores(self.key(hidden_states))
            value_layer = self.transpose_for_scores(self.value(hidden_states))

        query_layer = self.transpose_for_scores(mixed_query_layer)

        # Take the dot product between "query" and "key" to get the raw attention scores.
        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))

        if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
            seq_length = hidden_states.size()[1]
            position_ids_l = torch.arange(seq_length, dtype=torch.long, device=hidden_states.device).view(-1, 1)
            position_ids_r = torch.arange(seq_length, dtype=torch.long, device=hidden_states.device).view(1, -1)
            distance = position_ids_l - position_ids_r
            positional_embedding = self.distance_embedding(distance + self.max_position_embeddings - 1)
            positional_embedding = positional_embedding.to(dtype=query_layer.dtype)  # fp16 compatibility

            if self.position_embedding_type == "relative_key":
                relative_position_scores = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
                attention_scores = attention_scores + relative_position_scores
            elif self.position_embedding_type == "relative_key_query":
                relative_position_scores_query = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
                relative_position_scores_key = torch.einsum("bhrd,lrd->bhlr", key_layer, positional_embedding)
                attention_scores = attention_scores + relative_position_scores_query + relative_position_scores_key

        attention_scores = attention_scores / math.sqrt(self.attention_head_size)
        if attention_mask is not None:
            # Apply the attention mask is (precomputed for all layers in RobertaModel forward() function)
            attention_scores = attention_scores + attention_mask

        # Normalize the attention scores to probabilities.
        attention_probs = nn.Softmax(dim=-1)(attention_scores)

        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.
        attention_probs = self.dropout(attention_probs)

        # Mask heads if we want to
        if head_mask is not None:
            attention_probs = attention_probs * head_mask

        context_layer = torch.matmul(attention_probs, value_layer)

        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
        context_layer = context_layer.view(*new_context_layer_shape)

        outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)

        return outputs


# Copied from transformers.models.bert.modeling_bert.BertSelfOutput
class RobertaSelfOutput(nn.Module):
    def __init__(self, config, layer_index=None):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        if layer_index is None:
            self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, hidden_states):
        #print("Shape of hidden_states: ", hidden_states.shape)
        hidden_states = self.dense(hidden_states)
        hidden_states = self.dropout(hidden_states)
        return hidden_states


# Copied from transformers.models.bert.modeling_bert.BertAttention with Bert->Roberta
class RobertaAttention(nn.Module):
    def __init__(self, config, layer_index=None):
        super().__init__()
        self.self = RobertaSelfAttention(config, layer_index=layer_index)
        self.output = RobertaSelfOutput(config, layer_index=layer_index)
        self.pruned_heads = set()

    def prune_heads(self, heads):
        if len(heads) == 0:
            return
        heads, index = find_pruneable_heads_and_indices(
            heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads
        )

        # Prune linear layers
        self.self.query = prune_linear_layer(self.self.query, index)
        self.self.key = prune_linear_layer(self.self.key, index)
        self.self.value = prune_linear_layer(self.self.value, index)
        self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)

        # Update hyper params and store pruned heads
        self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
        self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads
        self.pruned_heads = self.pruned_heads.union(heads)

    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        past_key_value=None,
        output_attentions=False,
    ):
        self_outputs = self.self(
            hidden_states,
            attention_mask,
            head_mask,
            encoder_hidden_states,
            encoder_attention_mask,
            past_key_value,
            output_attentions,
        )
        attention_output = self.output(self_outputs[0])
        outputs = (attention_output,) + self_outputs[1:]  # add attentions if we output them
        return outputs


# Copied from transformers.models.bert.modeling_bert.BertIntermediate
class RobertaIntermediate(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
        if isinstance(config.hidden_act, str):
            self.intermediate_act_fn = ACT2FN[config.hidden_act]
        else:
            self.intermediate_act_fn = config.hidden_act

    def forward(self, hidden_states):
        hidden_states = self.dense(hidden_states)
        hidden_states = self.intermediate_act_fn(hidden_states)
        return hidden_states


# Copied from transformers.models.bert.modeling_bert.BertOutput
class RobertaOutput(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, hidden_states, input_tensor, last_norm=True):
        hidden_states = self.dense(hidden_states)
        hidden_states = self.dropout(hidden_states)
        hidden_states = hidden_states + input_tensor
        if last_norm:
            hidden_states = self.LayerNorm(hidden_states)
        return hidden_states


# Copied from transformers.models.bert.modeling_bert.BertLayer with Bert->Roberta
class RobertaLayer(nn.Module):
    def __init__(self, config, layer_index=None):
        super().__init__()
        self.chunk_size_feed_forward = config.chunk_size_feed_forward
        self.seq_len_dim = 1
        self.attention = RobertaAttention(config)
        self.is_decoder = config.is_decoder
        self.add_cross_attention = config.add_cross_attention
        if layer_index >= 12 - NUM_FUSE_BLOCK:
            self.crossattention_t2i = RobertaAttention(config, layer_index=layer_index)
            self.alpha_t2i = nn.Parameter(torch.Tensor([0]))
        self.intermediate = RobertaIntermediate(config)
        self.output = RobertaOutput(config)

    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        past_key_value=None,
        output_attentions=False,
        last_norm=True,
    ):
        # decoder uni-directional self-attention cached key/values tuple is at positions 1,2
        self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
        self_attention_outputs = self.attention(
            hidden_states,
            attention_mask,
            head_mask,
            output_attentions=output_attentions,
            past_key_value=self_attn_past_key_value,
        )
        attention_output = self_attention_outputs[0]

        # if decoder, the last output is tuple of self-attn cache
        outputs = self_attention_outputs[1:]  # add self attentions if we output attention weights

        cross_attn_present_key_value = None
        if encoder_hidden_states is not None:
            assert hasattr(
                self, "crossattention_t2i"
            ), f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers by setting `config.add_cross_attention=True`"

            # cross_attn cached key/values tuple is at positions 3,4 of past_key_value tuple
            cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None
            cross_attention_outputs = self.crossattention_t2i(
                attention_output,
                attention_mask,
                head_mask,
                encoder_hidden_states,
                encoder_attention_mask,
                cross_attn_past_key_value,
                output_attentions,
            )
            attention_output = self.alpha_t2i * cross_attention_outputs[0] + attention_output

        attention_output = self.attention.output.LayerNorm(attention_output + hidden_states)

        layer_output = apply_chunking_to_forward(
            self.feed_forward_chunk,
            self.chunk_size_feed_forward,
            self.seq_len_dim,
            attention_output,
            last_norm=last_norm,
        )
        outputs = (layer_output,) + outputs

        # if decoder, return the attn key/values as the last output
        return outputs

    def feed_forward_chunk(self, attention_output, last_norm=True):
        intermediate_output = self.intermediate(attention_output)
        layer_output = self.output(intermediate_output, attention_output, last_norm=last_norm)
        return layer_output


# Copied from transformers.models.bert.modeling_bert.BertEncoder with Bert->Roberta
class RobertaEncoder(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.config.gradient_checkpointing = config_yaml['use_checkpoint']
        self.layer = nn.ModuleList(
            [RobertaLayer(config, layer_index=layer_i) for layer_i in range(config.num_hidden_layers)]
        )

    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        past_key_values=None,
        use_cache=None,
        output_attentions=False,
        output_hidden_states=False,
        return_dict=True,
    ):
        all_hidden_states = () if output_hidden_states else None
        all_self_attentions = () if output_attentions else None
        all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None

        next_decoder_cache = () if use_cache else None
        for i, layer_module in enumerate(self.layer):
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

            layer_head_mask = head_mask[i] if head_mask is not None else None
            past_key_value = past_key_values[i] if past_key_values is not None else None

            if getattr(self.config, "gradient_checkpointing", False) and self.training:

                if use_cache:
                    logger.warning(
                        "`use_cache=True` is incompatible with `config.gradient_checkpointing=True`. Setting "
                        "`use_cache=False`..."
                    )
                    use_cache = False

                #def create_custom_forward(module):
                    #def custom_forward(*inputs):
                        #return module(*inputs, past_key_value, output_attentions)
                    #return custom_forward

                layer_outputs = torch.utils.checkpoint.checkpoint(
                    layer_module,
                    hidden_states,
                    attention_mask,
                    layer_head_mask,
                    encoder_hidden_states,
                    encoder_attention_mask,
                    past_key_value, 
                    output_attentions
                )
            else:
                layer_outputs = layer_module(
                    hidden_states,
                    attention_mask,
                    layer_head_mask,
                    encoder_hidden_states,
                    encoder_attention_mask,
                    past_key_value,
                    output_attentions,
                )


            hidden_states = layer_outputs[0]
            if use_cache:
                next_decoder_cache += (layer_outputs[-1],)
            if output_attentions:
                all_self_attentions = all_self_attentions + (layer_outputs[1],)
                if self.config.add_cross_attention:
                    all_cross_attentions = all_cross_attentions + (layer_outputs[2],)

        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        if not return_dict:
            return tuple(
                v
                for v in [
                    hidden_states,
                    next_decoder_cache,
                    all_hidden_states,
                    all_self_attentions,
                    all_cross_attentions,
                ]
                if v is not None
            )
        return BaseModelOutputWithPastAndCrossAttentions(
            last_hidden_state=hidden_states,
            past_key_values=next_decoder_cache,
            hidden_states=all_hidden_states,
            attentions=all_self_attentions,
            cross_attentions=all_cross_attentions,
        )

class RobertaPooler(nn.Module):
    def __init__(self, config):
        super().__init__()
        #self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        self.activation = nn.Tanh()

    def forward(self, hidden_states):
        # We "pool" the model by simply taking the hidden state corresponding
        # to the first token.
        first_token_tensor = hidden_states[:, 0]
        #pooled_output = self.dense(first_token_tensor)
        pooled_output = self.activation(first_token_tensor)
        return pooled_output


class RobertaPreTrainedModel(PreTrainedModel):
    """
    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
    models.
    """

    config_class = RobertaConfig
    base_model_prefix = "roberta"

    # Copied from transformers.models.bert.modeling_bert.BertPreTrainedModel._init_weights
    def _init_weights(self, module):
        """Initialize the weights"""
        if isinstance(module, nn.Linear):
            # Slightly different from the TF version which uses truncated_normal for initialization
            # cf https://github.com/pytorch/pytorch/pull/5617
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)


ROBERTA_START_DOCSTRING = r"""
    This model inherits from :class:`~transformers.PreTrainedModel`. Check the superclass documentation for the generic
    methods the library implements for all its model (such as downloading or saving, resizing the input embeddings,
    pruning heads etc.)
    This model is also a PyTorch `torch.nn.Module <https://pytorch.org/docs/stable/nn.html#torch.nn.Module>`__
    subclass. Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to
    general usage and behavior.
    Parameters:
        config (:class:`~transformers.RobertaConfig`): Model configuration class with all the parameters of the
            model. Initializing with a config file does not load the weights associated with the model, only the
            configuration. Check out the :meth:`~transformers.PreTrainedModel.from_pretrained` method to load the model
            weights.
"""

ROBERTA_INPUTS_DOCSTRING = r"""
    Args:
        input_ids (:obj:`torch.LongTensor` of shape :obj:`({0})`):
            Indices of input sequence tokens in the vocabulary.
            Indices can be obtained using :class:`~transformers.RobertaTokenizer`. See
            :meth:`transformers.PreTrainedTokenizer.encode` and :meth:`transformers.PreTrainedTokenizer.__call__` for
            details.
            `What are input IDs? <../glossary.html#input-ids>`__
        attention_mask (:obj:`torch.FloatTensor` of shape :obj:`({0})`, `optional`):
            Mask to avoid performing attention on padding token indices. Mask values selected in ``[0, 1]``:
            - 1 for tokens that are **not masked**,
            - 0 for tokens that are **masked**.
            `What are attention masks? <../glossary.html#attention-mask>`__
        token_type_ids (:obj:`torch.LongTensor` of shape :obj:`({0})`, `optional`):
            Segment token indices to indicate first and second portions of the inputs. Indices are selected in ``[0,
            1]``:
            - 0 corresponds to a `sentence A` token,
            - 1 corresponds to a `sentence B` token.
            `What are token type IDs? <../glossary.html#token-type-ids>`_
        position_ids (:obj:`torch.LongTensor` of shape :obj:`({0})`, `optional`):
            Indices of positions of each input sequence tokens in the position embeddings. Selected in the range ``[0,
            config.max_position_embeddings - 1]``.
            `What are position IDs? <../glossary.html#position-ids>`_
        head_mask (:obj:`torch.FloatTensor` of shape :obj:`(num_heads,)` or :obj:`(num_layers, num_heads)`, `optional`):
            Mask to nullify selected heads of the self-attention modules. Mask values selected in ``[0, 1]``:
            - 1 indicates the head is **not masked**,
            - 0 indicates the head is **masked**.
        inputs_embeds (:obj:`torch.FloatTensor` of shape :obj:`({0}, hidden_size)`, `optional`):
            Optionally, instead of passing :obj:`input_ids` you can choose to directly pass an embedded representation.
            This is useful if you want more control over how to convert :obj:`input_ids` indices into associated
            vectors than the model's internal embedding lookup matrix.
        output_attentions (:obj:`bool`, `optional`):
            Whether or not to return the attentions tensors of all attention layers. See ``attentions`` under returned
            tensors for more detail.
        output_hidden_states (:obj:`bool`, `optional`):
            Whether or not to return the hidden states of all layers. See ``hidden_states`` under returned tensors for
            more detail.
        return_dict (:obj:`bool`, `optional`):
            Whether or not to return a :class:`~transformers.file_utils.ModelOutput` instead of a plain tuple.
"""


@add_start_docstrings(
    "The bare RoBERTa Model transformer outputting raw hidden-states without any specific head on top.",
    ROBERTA_START_DOCSTRING,
)
class RobertaModel(RobertaPreTrainedModel):
    """
    The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of
    cross-attention is added between the self-attention layers, following the architecture described in `Attention is
    all you need`_ by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz
    Kaiser and Illia Polosukhin.
    To behave as an decoder the model needs to be initialized with the :obj:`is_decoder` argument of the configuration
    set to :obj:`True`. To be used in a Seq2Seq model, the model needs to initialized with both :obj:`is_decoder`
    argument and :obj:`add_cross_attention` set to :obj:`True`; an :obj:`encoder_hidden_states` is then expected as an
    input to the forward pass.
    .. _`Attention is all you need`: https://arxiv.org/abs/1706.03762
    """

    _keys_to_ignore_on_load_missing = [r"position_ids"]

    # Copied from transformers.models.bert.modeling_bert.BertModel.__init__ with Bert->Roberta
    def __init__(self, config, add_pooling_layer=True):
        super().__init__(config)
        self.config = config

        self.embeddings = RobertaEmbeddings(config)
        self.encoder = RobertaEncoder(config)

        self.pooler = RobertaPooler(config) if add_pooling_layer else None

        self.init_weights()

    def get_input_embeddings(self):
        return self.embeddings.word_embeddings

    def set_input_embeddings(self, value):
        self.embeddings.word_embeddings = value

    def _prune_heads(self, heads_to_prune):
        """
        Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
        class PreTrainedModel
        """
        for layer, heads in heads_to_prune.items():
            self.encoder.layer[layer].attention.prune_heads(heads)

    @add_start_docstrings_to_model_forward(ROBERTA_INPUTS_DOCSTRING.format("(batch_size, sequence_length)"))
    @add_code_sample_docstrings(
        processor_class=_TOKENIZER_FOR_DOC,
        checkpoint=_CHECKPOINT_FOR_DOC,
        output_type=BaseModelOutputWithPoolingAndCrossAttentions,
        config_class=_CONFIG_FOR_DOC,
    )
    # Copied from transformers.models.bert.modeling_bert.BertModel.forward
    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        past_key_values=None,
        use_cache=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
    ):
        r"""
        encoder_hidden_states  (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`):
            Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
            the model is configured as a decoder.
        encoder_attention_mask (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
            Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
            the cross-attention if the model is configured as a decoder. Mask values selected in ``[0, 1]``:
            - 1 for tokens that are **not masked**,
            - 0 for tokens that are **masked**.
        past_key_values (:obj:`tuple(tuple(torch.FloatTensor))` of length :obj:`config.n_layers` with each tuple having 4 tensors of shape :obj:`(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
            Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
            If :obj:`past_key_values` are used, the user can optionally input only the last :obj:`decoder_input_ids`
            (those that don't have their past key value states given to this model) of shape :obj:`(batch_size, 1)`
            instead of all :obj:`decoder_input_ids` of shape :obj:`(batch_size, sequence_length)`.
        use_cache (:obj:`bool`, `optional`):
            If set to :obj:`True`, :obj:`past_key_values` key value states are returned and can be used to speed up
            decoding (see :obj:`past_key_values`).
        """
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        use_cache = False

        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot secify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
            input_shape = input_ids.size()
            batch_size, seq_length = input_shape
        elif inputs_embeds is not None:
            input_shape = inputs_embeds.size()[:-1]
            batch_size, seq_length = input_shape
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        device = input_ids.device if input_ids is not None else inputs_embeds.device

        # past_key_values_length
        past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0

        if attention_mask is None:
            attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length)), device=device)
        if token_type_ids is None:
            token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)

        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
        # ourselves in which case we just need to make it broadcastable to all heads.
        extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape, device)

        # If a 2D or 3D attention mask is provided for the cross-attention
        # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
        if encoder_hidden_states is not None:
            encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
            encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
            if encoder_attention_mask is None:
                encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
            encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
        else:
            encoder_extended_attention_mask = None

        # Prepare head mask if needed
        # 1.0 in head_mask indicate we keep the head
        # attention_probs has shape bsz x n_heads x N x N
        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)

        embedding_output = self.embeddings(
            input_ids=input_ids,
            position_ids=position_ids,
            token_type_ids=token_type_ids,
            inputs_embeds=inputs_embeds,
            past_key_values_length=past_key_values_length,
        )
        encoder_outputs = self.encoder(
            embedding_output,
            attention_mask=extended_attention_mask,
            head_mask=head_mask,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_extended_attention_mask,
            past_key_values=past_key_values,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        sequence_output = encoder_outputs[0]
        pooled_output = self.pooler(sequence_output) if self.pooler is not None else None

        if not return_dict:
            return (sequence_output, pooled_output) + encoder_outputs[1:]

        return BaseModelOutputWithPoolingAndCrossAttentions(
            last_hidden_state=sequence_output,
            pooler_output=pooled_output,
            past_key_values=encoder_outputs.past_key_values,
            hidden_states=encoder_outputs.hidden_states,
            attentions=encoder_outputs.attentions,
            cross_attentions=encoder_outputs.cross_attentions,
        )


def create_position_ids_from_input_ids(input_ids, padding_idx, past_key_values_length=0):
    """
    Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding symbols
    are ignored. This is modified from fairseq's `utils.make_positions`.
    Args:
        x: torch.Tensor x:
    Returns: torch.Tensor
    """
    # The series of casts and type-conversions here are carefully balanced to both work with ONNX export and XLA.
    mask = input_ids.ne(padding_idx).int()
    incremental_indices = (torch.cumsum(mask, dim=1).type_as(mask) + past_key_values_length) * mask
    return incremental_indices.long() + padding_idx


def _make_causal_mask(input_ids_shape: torch.Size, dtype: torch.dtype, past_key_values_length: int = 0):
    """
    Make causal mask used for bi-directional self-attention.
    """
    bsz, tgt_len = input_ids_shape
    mask = torch.full((tgt_len, tgt_len), float("-inf"))
    mask_cond = torch.arange(mask.size(-1))
    mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
    mask = mask.to(dtype)

    if past_key_values_length > 0:
        mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype), mask], dim=-1)
    return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)


def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len=None):
    """
    Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
    """
    bsz, src_len = mask.size()
    tgt_len = tgt_len if tgt_len is not None else src_len

    expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)

    inverted_mask = 1.0 - expanded_mask

    return inverted_mask.masked_fill(inverted_mask.bool(), torch.finfo(dtype).min)


def _prepare_decoder_attention_mask(
    attention_mask, input_shape, inputs_embeds, device, past_key_values_length: int = 0
):
    # create causal mask
    # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
    combined_attention_mask = None
    if input_shape[-1] > 1:
        combined_attention_mask = _make_causal_mask(
            input_shape, inputs_embeds.dtype, past_key_values_length=past_key_values_length
        ).to(device)

    if attention_mask is not None:
        # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
        expanded_attn_mask = _expand_mask(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1])
        combined_attention_mask = (
            expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask
        )

    return combined_attention_mask

EgoTaskQA/model/roberta.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import os
import pdb

import timm
import torch
import yaml
import torch.nn as nn
import numpy as np
import torch.nn.functional as F
from transformers import AutoModel
from einops import rearrange, repeat

from base import BaseModel
from model import video_transformer
from model.video_transformer import SpaceTimeTransformer
from utils.util import state_dict_data_parallel_fix

from model import roberta
from model.roberta import RobertaModel, _prepare_decoder_attention_mask
from model import heads
from transformers import RobertaConfig
from functools import partial
import copy
import torch.distributed as dist

with open('./EgoNCE_MLM_ITM_Config.yml') as f:
    config = yaml.load(f, Loader=yaml.FullLoader)

def init_weights(module):
    if isinstance(module, (nn.Linear, nn.Embedding)):
        module.weight.data.normal_(mean=0.0, std=0.02)
    elif isinstance(module, nn.LayerNorm):
        module.bias.data.zero_()
        module.weight.data.fill_(1.0)

    if isinstance(module, nn.Linear) and module.bias is not None:
        module.bias.data.zero_()


class FrozenInTime(BaseModel):
    def __init__(self,
                 video_params,
                 text_params,
                 projection_dim=4096,
                 load_checkpoint=None,
                 projection='minimal',
                 load_temporal_fix='bilinear', #bilinear
                 config = config,
                 task_names = 'EgoNCE_ITM_MLM',
                 norm_layer = None,
                 embed_dim=768,
                 model_dim=768,
                 output_dim=1010):
                 #egotaskqa_head_output_dim=1010):
        super().__init__()

        self.video_params = video_params
        self.text_params = text_params
        self.load_temporal_fix = load_temporal_fix
        self.config = config
        #self.egotaskqa_head_output_dim = egotaskqa_head_output_dim
        self.task_names = task_names
        if not text_params['pretrained']:
            raise NotImplementedError("Huggingface text models require pretrained init.")

        if self.text_params['model'].startswith('roberta'):
            self.text_model = RobertaModel.from_pretrained("roberta-base")
        self.text_model.train()

        pretrained = video_params['pretrained']
        if video_params['model'] == "SpaceTimeTransformer":
            self.num_frames = 16 #self.config["data_loader"][0]["args"]["video_params"]["num_frames"]
            time_init = 'zeros'
            attention_style = 'frozen-in-time'
            arch_config = 'base_patch16_224'
            vit_init = 'imagenet-21k'
            if arch_config == 'base_patch16_224':
                #vit_model = timm.models.vision_transformer.vit_base_patch16_224(pretrained=pretrained)
                vit_model = torch.load("/cis/home/shraman/works_meta_2022/pre-training/EgoVLP_Fused_HardNegITM_Checkpoint_multinode/frozen-in-time-main/pretrained/jx_vit_base_p16_224-80ecf9dd.pth", map_location="cpu")
                model = SpaceTimeTransformer(num_frames=self.num_frames,
                                            time_init=time_init,
                                            attention_style=attention_style)
            else:
                raise NotImplementedError

            model.head = nn.Identity()
            model.pre_logits = nn.Identity()
            ftr_dim = model.embed_dim
           
            if load_checkpoint in ["", None]:
                vit_checkpoint = vit_model
                new_vit_dict = state_dict_data_parallel_fix(vit_checkpoint, model.state_dict())
                model.load_state_dict(new_vit_dict, strict=False)
            self.video_model = model
        else:
            raise NotImplementedError(f"{video_params['model']} not implemented")

        # for backwards compatibility (old models)
        self.video_model.fc = nn.Identity()

        # Project to a common embedding
        if projection == 'minimal':

            txt_proj = nn.Sequential(
                                    nn.Linear(self.text_model.config.hidden_size, projection_dim, bias=False),
                                    nn.ReLU(inplace=True), nn.Linear(projection_dim, projection_dim, bias=True),
                                    nn.ReLU(inplace=True), nn.Linear(projection_dim, projection_dim, bias=True)
                                    )

            vid_proj = nn.Sequential(
                                    nn.Linear(ftr_dim, projection_dim, bias=False),
                                    nn.ReLU(inplace=True), nn.Linear(projection_dim, projection_dim, bias=True),
                                    nn.ReLU(inplace=True), nn.Linear(projection_dim, projection_dim, bias=True)
                                    )
        
        elif projection == '':
            txt_proj = nn.Identity()
            vid_proj = nn.Identity()
        else:
            raise NotImplementedError
        self.txt_proj = txt_proj
        self.vid_proj = vid_proj


        if ('MLM' in self.task_names or 'ITM' in self.task_names):
            # for FIBER-like cross-attention

            bert_config = RobertaConfig(
                vocab_size=self.config["vocab_size"],
                hidden_size=self.config["hidden_size"],
                num_hidden_layers=self.config["num_layers"],
                num_attention_heads=self.config["num_heads"],
                intermediate_size=self.config["hidden_size"] * config["mlp_ratio"],
                #max_position_embeddings=maxlen, [was used in BTGOT script]
                hidden_dropout_prob=self.config["drop_rate"],
                attention_probs_dropout_prob=self.config["drop_rate"],
            )

            self.num_fuse_block=self.config["num_fuse_block"]
            self.num_text_layer=self.config["num_layers"]
            roberta.NUM_FUSE_BLOCK = self.video_model.NUM_FUSE_BLOCK=self.num_fuse_block
            roberta.DIM_IMG=self.config["input_image_embed_size"]
            self.video_model.DIM_TXT=self.config["input_text_embed_size"]

            self.cross_modal_text_transform = nn.Linear(self.config["input_text_embed_size"], self.config["hidden_size"])
            self.cross_modal_text_transform.apply(init_weights)
            self.cross_modal_video_transform = nn.Linear(self.config["input_image_embed_size"], self.config["hidden_size"])
            self.cross_modal_video_transform.apply(init_weights)

            self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))

            self.num_patches = self.video_model.patch_embed.num_patches
            self.patches_per_frame = self.num_patches//self.num_frames
            norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
            self.norm = norm_layer(embed_dim)
            self.pre_logits = nn.Identity()


            self.avgpool = nn.AdaptiveAvgPool1d(1)
            self.cross_modal_video_pooler = heads.Pooler(config["hidden_size"])
            self.cross_modal_video_pooler.apply(init_weights)
            self.cross_modal_text_pooler = heads.Pooler(config["hidden_size"])
            self.cross_modal_text_pooler.apply(init_weights)

            self.dropout = nn.Dropout(p=0.2)
            self.relu = nn.ReLU()
            self.projector_1 = nn.Linear(model_dim, output_dim)
            self.projector_2 = nn.Linear(output_dim, output_dim)

            ## einops transformations
            self.einops_from_space = 'b (f n) d'
            self.einops_to_space = '(b f) n d'
            self.einops_from_time = 'b (f n) d'
            self.einops_to_time = '(b n) f d'

        if 'MLM' in self.task_names:
            self.mlm_score = heads.MLMHead(bert_config)
            self.mlm_score.apply(init_weights)

        if 'ITM' in self.task_names:
            self.itm_score = heads.ITMHead(config["hidden_size"] * 2)
            self.itm_score.apply(init_weights)

        if load_checkpoint not in ["", None]:
            checkpoint = torch.load(load_checkpoint, map_location='cpu')
            print("This loop is being executed")
            state_dict = checkpoint['state_dict']
            new_state_dict = state_dict_data_parallel_fix(state_dict, self.state_dict())
            new_state_dict = self._inflate_positional_embeds(new_state_dict)
            self.load_state_dict(new_state_dict, strict=False)
            print("Model is loaded with pre-trained parameters")

    def set_device(self, device):
        self.device = device

    def forward(self, data, video_only=False, return_embeds=True):
        

        text_data = data['text']
        video_data = data['video']


        b, curr_frames, channels, _, _ = video_data.shape
        video_data_itm = self.video_model.patch_embed(video_data)
        video_data_itm = video_data_itm.flatten(2).transpose(2, 1)
        video_data_itm = video_data_itm.reshape(b, -1, self.video_model.patch_embed.embed_dim)

        BF = video_data_itm.shape[0]
        cls_tokens = self.cls_token.expand(BF, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        video_data_itm = torch.cat((cls_tokens, video_data_itm), dim=1)
        # positional embed needs to be tiled for each frame (this does [1,2,3] --> [1,2,3,1,2,3]...)
        cls_embed = self.video_model.pos_embed[:, 0, :].unsqueeze(1)
        tile_pos_embed = self.video_model.pos_embed[:, 1:, :].repeat(1, self.num_frames, 1)
        # temporal embed needs to be repeated within each frame (this does [1,2,3] --> [1,1,1,2,2,2,3,3,3]...)
        tile_temporal_embed = self.video_model.temporal_embed.repeat_interleave(self.patches_per_frame, 1)
        total_pos_embed = tile_pos_embed + tile_temporal_embed
        total_pos_embed = torch.cat([cls_embed, total_pos_embed], dim=1)

        n = self.patches_per_frame
        f = curr_frames

        curr_patches = video_data_itm.shape[1]
        video_data_itm = video_data_itm + total_pos_embed[:, :curr_patches]
        video_data_itm = self.video_model.pos_drop(video_data_itm)

        unfused_blocks = self.num_text_layer - self.num_fuse_block
            

        for blk_i, blk in enumerate(self.video_model.blocks[:unfused_blocks]):
            if config['use_checkpoint']:
                video_data_itm = torch.utils.checkpoint.checkpoint(blk, video_data_itm, self.einops_from_space, self.einops_to_space, self.einops_from_time, self.einops_to_time,
                                n, f)
            else:
                video_data_itm = blk(video_data_itm, self.einops_from_space, self.einops_to_space, self.einops_from_time, self.einops_to_time,
                                time_n=n, space_f=f)
                            
            
        text_embeds = self.text_model.embeddings(input_ids=text_data['input_ids']) # before it was input_ids=text_ids
        device = text_embeds.device
        #print(text_data.keys())
        text_masks = text_data['attention_mask']
        input_shape = text_masks.size()
        extend_text_masks = self.text_model.get_extended_attention_mask(text_masks, input_shape, device)
        for layer_i, layer in enumerate(self.text_model.encoder.layer[:unfused_blocks]):
            if config['use_checkpoint']:
                text_embeds = torch.utils.checkpoint.checkpoint(layer, text_embeds, extend_text_masks)[0]
            else:
                text_embeds = layer(text_embeds, extend_text_masks)[0]


        for blk_i, blk in enumerate(self.video_model.blocks[unfused_blocks:self.num_text_layer]):
            if config['use_checkpoint']:
                    
                fuse_video_data = torch.utils.checkpoint.checkpoint(blk, video_data_itm, self.einops_from_space, self.einops_to_space, self.einops_from_time, self.einops_to_time, 
                                          n, f, text_embeds, extend_text_masks)
                text_embeds = torch.utils.checkpoint.checkpoint(self.text_model.encoder.layer[blk_i + unfused_blocks],
                                          text_embeds, extend_text_masks, None, (video_data_itm), None, None, False, True)[0]
            else:
                fuse_video_data = blk(video_data_itm, self.einops_from_space, self.einops_to_space, self.einops_from_time, self.einops_to_time, 
                                          y=text_embeds, y_mask=extend_text_masks, time_n=n, space_f=f)
                text_embeds = self.text_model.encoder.layer[blk_i + unfused_blocks](text_embeds, extend_text_masks, encoder_hidden_states=(video_data_itm), last_norm=True)[0]
            video_data_itm = fuse_video_data

            
        #print("Shape of model output", video_data.shape)
        video_data_itm = self.norm(video_data_itm)[:, :]
        video_data_itm = self.pre_logits(video_data_itm)

        video_cls = video_data_itm[:,0,:]
        batch_out = self.projector_2(self.dropout(self.relu(self.projector_1(video_cls))))


        return batch_out

    
    def compute_text(self, text_data):
        if self.text_params['model'].startswith('bert'):
            text_embeddings = self.text_model(text_data['input_ids'], attention_mask=text_data['attention_mask'])[
                'pooler_output']
        elif self.text_params['model'].startswith('distilbert'):
            text_embeddings = self.text_model(**text_data).last_hidden_state[:, 0, :]
        elif self.text_params['model'].startswith('roberta'):
            text_embeddings = self.text_model(**text_data).last_hidden_state[:, 0, :]
        else:
            raise NotImplementedError
        text_embeddings = self.txt_proj(text_embeddings)
        return text_embeddings

    def compute_text_tokens(self, text_data, is_proj=True):
        if self.text_params['model'].startswith('bert'):
            text_embeddings = self.text_model(text_data['input_ids'], attention_mask=text_data['attention_mask'])[
                'pooler_output']    # not implement for bert
        elif self.text_params['model'].startswith('distilbert'):
            text_embeddings = self.text_model(**text_data).last_hidden_state
        elif self.text_params['model'].startswith('roberta'):
            text_embeddings = self.text_model(**text_data).last_hidden_state
        else:
            raise NotImplementedError
        
        if is_proj:
            text_embeddings = self.txt_proj(text_embeddings)
        
        return text_embeddings

    def compute_video(self, video_data):
        video_embeddings = self.video_model(video_data)
        video_embeddings = self.vid_proj(video_embeddings)
        return video_embeddings

    def _inflate_positional_embeds(self, new_state_dict):
        # allow loading of timesformer with fewer num_frames
        curr_keys = list(self.state_dict().keys())
        if 'video_model.temporal_embed' in new_state_dict and 'video_model.temporal_embed' in curr_keys:
            load_temporal_embed = new_state_dict['video_model.temporal_embed']
            load_num_frames = load_temporal_embed.shape[1]
            curr_num_frames = self.video_params['num_frames']
            embed_dim = load_temporal_embed.shape[2]

            if load_num_frames != curr_num_frames:
                if load_num_frames > curr_num_frames:
                    print(f'### loaded {self.video_params["model"]} model has MORE frames than current...'
                          f'### loading weights, filling in the extras via {self.load_temporal_fix}')
                    new_temporal_embed = load_temporal_embed[:, :curr_num_frames, :]
                else:
                    print(f'### loaded {self.video_params["model"]} model has FEWER frames than current...'
                          f'### loading weights, filling in the extras via {self.load_temporal_fix}')
                    if self.load_temporal_fix == 'zeros':
                        new_temporal_embed = torch.zeros([load_temporal_embed.shape[0], curr_num_frames, embed_dim])
                        new_temporal_embed[:, :load_num_frames] = load_temporal_embed
                    elif self.load_temporal_fix in ['interp', 'bilinear']:
                        # interpolate
                        # unsqueeze so pytorch thinks its an image
                        mode = 'nearest'
                        if self.load_temporal_fix == 'bilinear':
                            mode = 'bilinear'
                        load_temporal_embed = load_temporal_embed.unsqueeze(0)
                        new_temporal_embed = F.interpolate(load_temporal_embed,
                                                           (curr_num_frames, embed_dim), mode=mode, align_corners=True).squeeze(0)
                    else:
                        raise NotImplementedError
                new_state_dict['video_model.temporal_embed'] = new_temporal_embed
        # allow loading with smaller spatial patches. assumes custom border crop, to append the
        # border patches to the input sequence
        if 'video_model.pos_embed' in new_state_dict and 'video_model.pos_embed' in curr_keys:
            load_pos_embed = new_state_dict['video_model.pos_embed']
            load_num_patches = load_pos_embed.shape[1]
            curr_pos_embed = self.state_dict()['video_model.pos_embed']
            if load_num_patches != curr_pos_embed.shape[1]:
                raise NotImplementedError(
                    'Loading models with different spatial resolution / patch number not yet implemented, sorry.')

        return new_state_dict

def sim_matrix(a, b, eps=1e-8):
    """
    added eps for numerical stability
    """
    a_n, b_n = a.norm(dim=1)[:, None], b.norm(dim=1)[:, None]
    a_norm = a / torch.max(a_n, eps * torch.ones_like(a_n))
    b_norm = b / torch.max(b_n, eps * torch.ones_like(b_n))
    sim_mt = torch.mm(a_norm, b_norm.transpose(0, 1))
    return sim_mt


def sim_matrix_batch_val(a, b, eps=1e-8):
    """
    added eps for numerical stability
    """
    a_n, b_n = a.norm(dim=-1).unsqueeze(-1), b.norm(dim=-1).unsqueeze(-1)
    a_norm = a / torch.max(a_n, eps * torch.ones_like(a_n))
    b_norm = b / torch.max(b_n, eps * torch.ones_like(b_n))
    sim_mt = torch.bmm(a_norm, b_norm.transpose(1, 2))
    return sim_mt


if __name__ == "__main__":
    video_params = {"model": SpaceTimeTransformer, "arch_config": "base_patch16_224", "num_frames": 4, "pretrained": True, "time_init": "zeros"}
    text_params: {"model": "roberta-base", "pretrained": True, "input": "text"}
    model = FrozenInTime(video_params, text_params)

EgoTaskQA/model/video_qa_model_linear_end2end.py
facebookresearch/EgoVLPv2
"""
Implementations of Video Transformers in PyTorch

A PyTorch implementation of space-time transformer as described in
'Frozen in Time: A Joint Image and Video Encoder for End-to-End Retrieval' - https://arxiv.org/abs/2104.00650

A PyTorch implementation of timesformer as described in
'Is Space-Time Attention All You Need for Video Understanding?' - https://arxiv.org/abs/2102.05095

Acknowledgments:
- This code builds on Ross Wightman's vision_transformer code in pytorch-image-models:
https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py

- It is also inspired by lucidrains timesformer implementation:
https://github.com/lucidrains/TimeSformer-pytorch

Hacked together by Max Bain
"""

from collections import OrderedDict
from functools import partial
import yaml

import torch
from einops import rearrange, repeat
from timm.models.layers import DropPath, to_2tuple, trunc_normal_
from torch import einsum, nn

with open('./EgoNCE_MLM_ITM_Config.yml') as f:
    config_yaml = yaml.load(f, Loader=yaml.FullLoader)

NUM_FUSE_BLOCK = config_yaml['num_fuse_block']
DIM_TEXT = 768

def attn(q, k, v):
    sim = einsum('b i d, b j d -> b i j', q, k)
    attn = sim.softmax(dim=-1)
    out = einsum('b i j, b j d -> b i d', attn, v)
    return out


class Mlp(nn.Module):
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x


class VideoPatchEmbed(nn.Module):
    """ Video to Patch Embedding
    """

    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768,
                 num_frames=8):
        super().__init__()
        img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)
        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0]) * num_frames
        self.img_size = img_size
        self.patch_size = patch_size
        self.num_patches = num_patches
        self.num_frames = num_frames
        self.embed_dim = embed_dim
        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)

    def forward(self, x):
        B, F, C, H, W = x.shape
        assert F == self.num_frames, print(F, self.num_frames)
        x = x.view(-1, C, H, W)
        x = self.proj(x)
        return x


class VarAttention(nn.Module):
    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.,
                 initialize='random', dim_text=None, norm_layer=nn.LayerNorm, space_attn=True):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        # NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights
        self.scale = qk_scale or head_dim ** -0.5
        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.proj = nn.Linear(dim, dim)
        if initialize == 'zeros':
            self.qkv.weight.data.fill_(0)
            self.qkv.bias.data.fill_(0)
            # fill proj weight with 1 here to improve training dynamics. Otherwise temporal attention inputs
            # are multiplied by 0*0, which is hard for the model to move out of.
            self.proj.weight.data.fill_(1)
            self.proj.bias.data.fill_(0)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj_drop = nn.Dropout(proj_drop)

        self.softmax = nn.Softmax(dim=-1)

        if dim_text is not None and space_attn:
            self.qkv_text_i2t = nn.Linear(dim_text, dim * 2, bias=qkv_bias)
            self.qkv_i2t = nn.Linear(dim, dim, bias=qkv_bias)
            self.attn_drop_i2t = nn.Dropout(attn_drop)
            self.proj_i2t = nn.Linear(dim, dim)
            self.proj_drop_i2t = nn.Dropout(proj_drop)
            self.alpha_i2t = nn.Parameter(torch.Tensor([0]))
            self.norm_i2t_i = norm_layer(dim)

    def forward(self, x, einops_from, einops_to, y=None, y_mask=None, **einops_dims):
        h = self.num_heads
        # project x to q, k, v vaalues
        q, k, v = self.qkv(x).chunk(3, dim=-1)
        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))

        q = q*self.scale

        # splice out CLS token at index 1
        (cls_q, q_), (cls_k, k_), (cls_v, v_) = map(lambda t: (t[:, 0:1], t[:, 1:]), (q, k, v))

        # let CLS token attend to key / values of all patches across time and space
        cls_out = attn(cls_q, k, v)
        # rearrange across time or space
        q_, k_, v_ = map(lambda t: rearrange(t, f'{einops_from} -> {einops_to}', **einops_dims), (q_, k_, v_))

        # expand cls token keys and values across time or space and concat
        r = q_.shape[0] // cls_k.shape[0]
        cls_k, cls_v = map(lambda t: repeat(t, 'b () d -> (b r) () d', r=r), (cls_k, cls_v))

        k_ = torch.cat((cls_k, k_), dim=1)
        v_ = torch.cat((cls_v, v_), dim=1)

        # attention
        out = attn(q_, k_, v_)

        # merge back time or space
        out = rearrange(out, f'{einops_to} -> {einops_from}', **einops_dims)

        # concat back the cls token
        out = torch.cat((cls_out, out), dim=1)

        # merge back the heads
        out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
        ## to out
        x = self.proj(out)
        x = self.proj_drop(x)

        if y is not None:
            B_, N, C = x.shape
            B_text, N_text, C_text = y.shape

            kv_text = (
                self.qkv_text_i2t(y)
                .reshape(B_text, N_text, 2, self.num_heads, C // self.num_heads)
                .permute(2, 0, 3, 1, 4)
            )
            k_text, v_text = kv_text[0], kv_text[1]

            q_i2t = self.qkv_i2t(self.norm_i2t_i(x))
            q_i2t = q_i2t.reshape(B_, N, 1, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
            q_i2t = q_i2t[0]

            # image to text attention
            text_scale = k_text.size(-1) ** -0.5
            q_i2t = q_i2t * text_scale
            attn_i2t = q_i2t @ k_text.transpose(-2, -1)  # B_, nH, N, N_text

            # add image to text bias and text_mask
            if y_mask is not None:
                mask_and_i2t_bias = y_mask.view(B_text, 1, 1, N_text)
                attn_i2t = attn_i2t + mask_and_i2t_bias

            attn_i2t = self.softmax(attn_i2t)
            attn_i2t = self.attn_drop_i2t(attn_i2t)
            y = (attn_i2t @ v_text).transpose(1, 2).reshape(B_, N, C)
            y = self.proj_i2t(y)
            y = self.proj_drop_i2t(y)
            x = x + self.alpha_i2t * y

        return x


class SpaceTimeBlock(nn.Module):

    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, time_init='zeros',
                 attention_style='frozen-in-time', dim_text=None):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.attn = VarAttention(
            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop, dim_text=dim_text, 
            norm_layer=norm_layer, space_attn=True)

        self.timeattn = VarAttention(
            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop,
            initialize=time_init, dim_text=dim_text, norm_layer=norm_layer, space_attn=False)

        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
        self.norm3 = norm_layer(dim)

        self.attention_style = attention_style

    def forward(self, x, einops_from_space, einops_to_space, einops_from_time, einops_to_time,
                time_n, space_f, y=None, y_mask=None):

        time_output = self.timeattn(self.norm3(x), einops_from_time, einops_to_time, n=time_n, y=None, y_mask=None)
        time_residual = x + time_output
        space_output = self.attn(self.norm1(time_residual), einops_from_space,
                                 einops_to_space, f=space_f, y=y, y_mask=y_mask)
        if self.attention_style == 'frozen-in-time':
            space_residual = x + self.drop_path(space_output)
        else:
            raise NotImplementedError

        x = space_residual + self.drop_path(self.mlp(self.norm2(space_residual)))

        return x


class SpaceTimeTransformer(nn.Module):
    """ Vision Transformer

    A PyTorch impl of : `Space-Time Transformer` from Frozen-in-time  - by Max Bain.
        https://arxiv.org/abs/2104.00650

    Based off:
     - ViT implementation from the timm library [https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py]
    lucidrains timesformer implementation [https://github.com/lucidrains/TimeSformer-pytorch].

    Notable differences:
     - allows for variable length input frames (<= num_frames)
     - allows for variable length input resolution  (<= (img_size, img_size)) [UNTESTED]
     - different attention block mechanism
    """

    def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12,
                 num_heads=12, mlp_ratio=4., qkv_bias=True, qk_scale=None, representation_size=None,
                 drop_rate=0., attn_drop_rate=0., drop_path_rate=0., hybrid_backbone=None,
                 num_frames=8, time_init='rand', attention_style='frozen-in-time', norm_layer=nn.LayerNorm, dim_text=None):
        """
        Args:
            img_size (int, tuple): input image size
            patch_size (int, tuple): patch size
            in_chans (int): number of input channels
            num_classes (int): number of classes for classification head
            embed_dim (int): embedding dimension
            depth (int): depth of transformer
            num_heads (int): number of attention heads
            mlp_ratio (int): ratio of mlp hidden dim to embedding dim
            qkv_bias (bool): enable bias for qkv if True
            qk_scale (float): override default qk scale of head_dim ** -0.5 if set
            representation_size (Optional[int]): enable and set representation layer (pre-logits) to this value if set
            drop_rate (float): dropout rate
            attn_drop_rate (float): attention dropout rate
            drop_path_rate (float): stochastic depth rate
            hybrid_backbone (nn.Module): CNN backbone to use in-place of PatchEmbed module
            norm_layer: (nn.Module): normalization layer
            num_frames: (int) maximum number of frames expected as input
            time_init: (str) how to initialise the time attention layer, 'zeros' allows for the timesformer to start off
                        as ViT.
            attention_style: (str) how to attend to space and time.
        """
        super().__init__()
        self.num_classes = num_classes
        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
        self.num_frames = num_frames
        self.embed_dim = embed_dim
        norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
        print("######USING ATTENTION STYLE: ", attention_style)
        if hybrid_backbone is not None:
            raise NotImplementedError('hybrid backbone not implemented')
        else:
            self.patch_embed = VideoPatchEmbed(
                img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim, num_frames=num_frames)
        num_patches = self.patch_embed.num_patches
        self.patches_per_frame = num_patches // num_frames

        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.pos_embed = nn.Parameter(
            torch.zeros(1, self.patches_per_frame + 1,
                        embed_dim))  # remember to take pos_embed[1:] for tiling over time
        self.temporal_embed = nn.Parameter(torch.zeros(1, num_frames, embed_dim))

        self.pos_drop = nn.Dropout(p=drop_rate)

        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
        self.blocks = nn.ModuleList([
            SpaceTimeBlock(
                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
                drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer, time_init=time_init,
                attention_style=attention_style, dim_text=None if i < 6 else DIM_TEXT)
            for i in range(depth)])
        self.norm = norm_layer(embed_dim)

        # Representation layer
        if representation_size:
            self.num_features = representation_size
            self.pre_logits = nn.Sequential(OrderedDict([
                ('fc', nn.Linear(embed_dim, representation_size)),
                ('act', nn.Tanh())
            ]))
        else:
            self.pre_logits = nn.Identity()

        # Classifier head
        self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()

        trunc_normal_(self.pos_embed, std=.02)
        trunc_normal_(self.cls_token, std=.02)

        # if num_frames > 1, then we perform ViT inflation and initialise time attention to zero so not necessary.
        if num_frames == 1:
            self.apply(self._init_weights)

        ## einops transformations
        self.einops_from_space = 'b (f n) d'
        self.einops_to_space = '(b f) n d'
        self.einops_from_time = 'b (f n) d'
        self.einops_to_time = '(b n) f d'


    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    @torch.jit.ignore
    def no_weight_decay(self):
        return {'pos_embed', 'cls_token'}

    def get_classifier(self):
        return self.head

    def reset_classifier(self, num_classes, global_pool=''):
        self.num_classes = num_classes
        self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()

    def forward_features(self, x):
        b, curr_frames, channels, _, _ = x.shape
        x = self.patch_embed(x)
        x = x.flatten(2).transpose(2, 1)
        x = x.reshape(b, -1, self.patch_embed.embed_dim)

        BF = x.shape[0]
        cls_tokens = self.cls_token.expand(BF, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        # positional embed needs to be tiled for each frame (this does [1,2,3] --> [1,2,3,1,2,3]...)
        cls_embed = self.pos_embed[:, 0, :].unsqueeze(1)
        tile_pos_embed = self.pos_embed[:, 1:, :].repeat(1, self.num_frames, 1)
        # temporal embed needs to be repeated within each frame (this does [1,2,3] --> [1,1,1,2,2,2,3,3,3]...)
        tile_temporal_embed = self.temporal_embed.repeat_interleave(self.patches_per_frame, 1)
        total_pos_embed = tile_pos_embed + tile_temporal_embed
        total_pos_embed = torch.cat([cls_embed, total_pos_embed], dim=1)

        curr_patches = x.shape[1]
        x = x + total_pos_embed[:, :curr_patches]
        x = self.pos_drop(x)
        n = self.patches_per_frame
        f = curr_frames

        for blk in self.blocks:
            if config_yaml['use_checkpoint']:

                def create_custom_forward(module):
                    def custom_forward(*inputs):
                        return module(*inputs, time_n=n, space_f=f)

                    return custom_forward

                x = torch.utils.checkpoint.checkpoint(create_custom_forward(blk), x, self.einops_from_space, self.einops_to_space, self.einops_from_time,
                    self.einops_to_time)
            else:
                x = blk(x, self.einops_from_space, self.einops_to_space, self.einops_from_time,
                    self.einops_to_time, time_n=n, space_f=f)

        x = self.norm(x)[:, 0]
        x = self.pre_logits(x)

        return x

    def forward(self, x):
        x = self.forward_features(x)
        x = self.head(x)
        return x

if __name__ == "__main__":
    network = SpaceTimeTransformer(num_frames=4)
    data = torch.rand((3, 4, 3, 224, 224))
    network(data)

EgoTaskQA/model/video_transformer.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import os
import logging
import pdb
from pathlib import Path
from functools import reduce
from operator import getitem
from datetime import datetime
from logger import setup_logging
from utils import read_json, write_json
import time
import inspect


class ConfigParser:
    def __init__(self, args, options='', timestamp=True, test=False, eval_mode=None):
        # parse default and custom cli options
        for opt in options:
            args.add_argument(*opt.flags, default=None, type=opt.type)
        args = args.parse_args()
        self.args = args
        if args.device:
            os.environ["CUDA_VISIBLE_DEVICES"] = args.device
        if args.resume is None:
            msg_no_cfg = "Configuration file need to be specified. Add '-c config.json', for example."
            assert args.config is not None, msg_no_cfg
            self.cfg_fname = Path(args.config)
            config = read_json(self.cfg_fname)
            self.resume = None
        else:
            self.resume = Path(args.resume)
            resume_cfg_fname = self.resume.parent / 'config.json'
            if eval_mode == "epic":
                resume_cfg_fname = Path('configs/epic.json')
            if eval_mode == "charades":
                resume_cfg_fname = Path('configs/charades.json')
            if eval_mode == "nlq":
                resume_cfg_fname = Path('configs/nlq.json')
            if eval_mode == "mq":
                resume_cfg_fname = Path('configs/mq.json')

            config = read_json(resume_cfg_fname)
            if args.config is not None:
                config.update(read_json(Path(args.config)))

        # load config file and apply custom cli options
        self._config = _update_config(config, options, args)

        # set save_dir where trained model and log will be saved.
        #save_dir = Path(self.config['trainer']['save_dir'])
        save_dir = Path(args.save_dir)
        # timestamp = datetime.now().strftime(r'%m%d_%H%M%S') if timestamp else ''
        timestamp = datetime.now().strftime(r'%m%d_%H') if timestamp else ''

        exper_name = self.config['name']

        self._save_dir = save_dir / 'models' /  timestamp
        self._web_log_dir = save_dir / 'web' /  timestamp
        self._log_dir = save_dir / 'log' /  timestamp
        self._tf_dir = save_dir / 'tf' / timestamp

        if not test:
            self.save_dir.mkdir(parents=True, exist_ok=True)
            self.log_dir.mkdir(parents=True, exist_ok=True)
            self._tf_dir.mkdir(parents=True, exist_ok=True)

        # if set, remove all previous experiments with the current config
        if vars(args).get("purge_exp_dir", False):
            for dirpath in (self._save_dir, self._log_dir, self._web_log_dir):
                config_dir = dirpath.parent
                existing = list(config_dir.glob("*"))
                print(f"purging {len(existing)} directories from config_dir...")
                tic = time.time()
                os.system(f"rm -rf {config_dir}")
                print(f"Finished purge in {time.time() - tic:.3f}s")

        # save updated config file to the checkpoint dir
        if not test:
            write_json(self.config, self.save_dir / 'config.json')

            # configure logging module
            setup_logging(self.log_dir)
            self.log_levels = {
                0: logging.WARNING,
                1: logging.INFO,
                2: logging.DEBUG
            }

    def initialize(self, name, module, *args, index=None, **kwargs):
        """
        finds a function handle with the name given as 'type' in config, and returns the
        instance initialized with corresponding keyword args given as 'args'.
        """
        if index is None:
            module_name = self[name]['type']
            module_args = dict(self[name]['args'])
            assert all([k not in module_args for k in kwargs]), 'Overwriting kwargs given in config file is not allowed'
            module_args.update(kwargs)
        else:
            module_name = self[name][index]['type']
            module_args = dict(self[name][index]['args'])

        # if parameter not in config subdict, then check if it's in global config.
        signature = inspect.signature(getattr(module, module_name).__init__)
        print(module_name)
        for param in signature.parameters.keys():
            if param not in module_args and param in self.config:
                module_args[param] = self[param]
            if module_name == 'FrozenInTime' and param == 'args':
                module_args[param] = self.args
            if module_name == 'MultiDistTextVideoDataLoader' and param == 'args':
                module_args[param] = self.args

        return getattr(module, module_name)(*args, **module_args)

    def __getitem__(self, name):
        return self.config[name]

    def get_logger(self, name, verbosity=2):
        msg_verbosity = 'verbosity option {} is invalid. Valid options are {}.'.format(verbosity,
                                                                                       self.log_levels.keys())
        assert verbosity in self.log_levels, msg_verbosity
        logger = logging.getLogger(name)
        logger.setLevel(self.log_levels[verbosity])
        return logger

    # setting read-only attributes
    @property
    def config(self):
        return self._config

    @property
    def save_dir(self):
        return self._save_dir

    @property
    def log_dir(self):
        return self._log_dir

    @property
    def tf_dir(self):
        return self._tf_dir

# helper functions used to update config dict with custom cli options
def _update_config(config, options, args):
    for opt in options:
        value = getattr(args, _get_opt_name(opt.flags))
        if value is not None:
            _set_by_path(config, opt.target, value)
    return config


def _get_opt_name(flags):
    for flg in flags:
        if flg.startswith('--'):
            return flg.replace('--', '')
    return flags[0].replace('--', '')


def _set_by_path(tree, keys, value):
    """Set a value in a nested object in tree by sequence of keys."""
    _get_by_path(tree, keys[:-1])[keys[-1]] = value


def _get_by_path(tree, keys):
    """Access a nested object in tree by sequence of keys."""
    return reduce(getitem, keys, tree)

EgoTaskQA/parse_config.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import json
import numpy as np
import pandas as pd
import pdb
import torch

def unique_cat(json_file_path):

    f = open(json_file_path)
    data = json.load(f)
    metadata = pd.DataFrame(data)

    l = []

    for idx in range(len(metadata)):
        l.append(metadata.iloc[idx]['reasoning_type'].split('$'))

    print(max([len(lst) for lst in l]))

    flat_list = [item for sublist in l for item in sublist]
    unique_list = np.unique(flat_list)

    return unique_list.tolist()


indirect_train = unique_cat('Data/qa/indirect/formatted_train_qas_encode.json')
indirect_test = unique_cat('Data/qa/indirect/formatted_test_qas_encode.json')
indirect_val = unique_cat('Data/qa/indirect/formatted_val_qas_encode.json')

direct_train = unique_cat('Data/qa/direct/formatted_train_qas_encode.json')
direct_test = unique_cat('Data/qa/direct/formatted_test_qas_encode.json')
direct_val = unique_cat('Data/qa/direct/formatted_val_qas_encode.json')

concat_unique_list = np.unique(list(indirect_train + indirect_test + indirect_val + direct_train + direct_test + direct_val)) 

unique_dict = {}
index = 0

for k in concat_unique_list:
    unique_dict.update({k: int(index)})
    index += 1


EgoTaskQA/reasoning_type_unique_cat.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


from torchvision import transforms
from torchvision.transforms._transforms_video import RandomCropVideo, RandomResizedCropVideo,CenterCropVideo, NormalizeVideo,ToTensorVideo,RandomHorizontalFlipVideo
import warnings
warnings.filterwarnings("ignore")

def init_transform_dict(input_res = 224,
                        center_crop = 256,
                        randcrop_scale=(0.5, 1.0),
                        color_jitter=(0, 0, 0),
                        norm_mean=(0.485, 0.456, 0.406),
                        norm_std=(0.229, 0.224, 0.225)):
    normalize = transforms.Normalize(mean=norm_mean, std=norm_std)
    tsfm_dict = {
        'train': transforms.Compose([
            transforms.RandomResizedCrop(input_res, scale=randcrop_scale),
            transforms.RandomHorizontalFlip(),
            transforms.ColorJitter(brightness=color_jitter[0], saturation=color_jitter[1], hue=color_jitter[2]),
            normalize,
        ]),
        'val': transforms.Compose([
            transforms.Resize(center_crop),
            transforms.CenterCrop(center_crop),
            transforms.Resize(input_res),
            normalize,
        ]),
        'test': transforms.Compose([
            transforms.Resize(center_crop),
            transforms.CenterCrop(center_crop),
            transforms.Resize(input_res),
            normalize,
        ])
    }
    return tsfm_dict

def init_video_transform_dict(input_res = 224,
                        center_crop = 256,
                        randcrop_scale=(0.5, 1.0),
                        color_jitter=(0, 0, 0),
                        norm_mean=(0.485, 0.456, 0.406),
                        norm_std=(0.229, 0.224, 0.225)):
    print('Video Transform is used!')
    normalize = NormalizeVideo(mean=norm_mean, std=norm_std)
    tsfm_dict = {
        'train': transforms.Compose([
            RandomResizedCropVideo(input_res, scale=randcrop_scale),
            RandomHorizontalFlipVideo(),
            transforms.ColorJitter(brightness=color_jitter[0], saturation=color_jitter[1], hue=color_jitter[2]),
            normalize,
        ]),
        'val': transforms.Compose([
            transforms.Resize(center_crop),
            transforms.CenterCrop(center_crop),
            transforms.Resize(input_res),
            normalize,
        ]),
        'test': transforms.Compose([
            transforms.Resize(center_crop),
            transforms.CenterCrop(center_crop),
            transforms.Resize(input_res),
            normalize,
        ])
    }
    return tsfm_dict

EgoTaskQA/transforms.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

from .util import *

EgoTaskQA/utils/__init__.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import numbers
from typing import List, Tuple

import torch
from torch import Tensor
from torchvision.transforms import functional_pil as F_pil, functional_tensor as F_t
from torchvision.transforms.functional import center_crop, crop


def _get_image_size(img: Tensor) -> List[int]:
    """Returns image size as [w, h]
    """
    if isinstance(img, torch.Tensor):
        return F_t._get_image_size(img)

    return F_pil._get_image_size(img)

def center_plus_four_crops(img: Tensor, size: List[int],
                              margin_h: int, margin_w: int) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
    """Crop the given image into four tiled borders and the central crop.
    """

    if isinstance(size, numbers.Number):
        size = (int(size), int(size))
    elif isinstance(size, (tuple, list)) and len(size) == 1:
        size = (size[0], size[0])

    if len(size) != 2:
        raise ValueError("Please provide only two dimensions (h, w) for size.")

    image_width, image_height = _get_image_size(img)

    crop_height, crop_width = size

    if crop_width > image_width or crop_height > image_height:
        msg = "Requested crop size {} is bigger than input size {}"
        raise ValueError(msg.format(size, (image_height, image_width)))

    if crop_width + margin_w > image_width:
        msg = "Requested margin size {} + input {} is bigger than input size {}"
        raise ValueError(msg.format((margin_h, margin_w), size, (image_height, image_width)))

    #vertical_border_height = image_height - crop_height
    #horizontal_border_height = image_width - crop_width

    #x1 = horizontal_border_height // 2
    x11 = (image_width - crop_width - 2 * margin_w) // 2
    x12 = x11 + margin_w
    x21 = x12 + crop_width
    x22 = x21 + margin_w

    y11 = (image_height - crop_height - 2 * margin_h) // 2
    y12 = y11 + margin_h
    y21 = y12 + crop_height
    y22 = y21 + margin_h

    tl = crop(img, y11, x11, margin_h, margin_w + crop_width)
    tr = crop(img, y11, x21, margin_h + crop_height, margin_w)
    bl = crop(img, y12, x11, margin_h + crop_height, margin_w)
    br = crop(img, y21, x12, margin_h, margin_w + crop_width)
    center = center_crop(img, [crop_height, crop_width])

    return tl, tr, bl, br, center



def center_plus_twohori_crops(img: Tensor, size: List[int],
                              margin_w: int) -> Tuple[Tensor, Tensor, Tensor]:
    """Crop the given image into four tiled borders and the central crop.
    """

    if isinstance(size, numbers.Number):
        size = (int(size), int(size))
    elif isinstance(size, (tuple, list)) and len(size) == 1:
        size = (size[0], size[0])

    if len(size) != 2:
        raise ValueError("Please provide only two dimensions (h, w) for size.")

    image_width, image_height = _get_image_size(img)

    crop_height, crop_width = size

    if crop_width > image_width or crop_height > image_height:
        msg = "Requested crop size {} is bigger than input size {}"
        raise ValueError(msg.format(size, (image_height, image_width)))

    if crop_width + margin_w > image_width :
        msg = "Requested margin size {} + input {} is bigger than input size {}"
        raise ValueError(msg.format((0, margin_w), size, (image_height, image_width)))

    x11 = (image_width - crop_width - 2 * margin_w) // 2
    x12 = x11 + margin_w
    x21 = x12 + crop_width

    y11 = (image_height - crop_height) // 2

    left = crop(img, y11, x11, crop_height, margin_w)
    right = crop(img, y11, x21, crop_height, margin_w)
    center = center_crop(img, [crop_height, crop_width])

    return left, right, center

from torch import nn
class TwoHoriCrop(nn.Module):
    def __init__(self, size, margin_w):
        super().__init__()
        self.size = size
        self.margin_w = margin_w

    def forward(self, x):
        return center_plus_twohori_crops(x, self.size, self.margin_w)

if __name__ == "__main__":
    from PIL import Image

    img = Image.open('visualisations/guitar.png')
    crops = center_plus_four_crops(img, [336, 336], 112, 112)
    order = ['tl', 'tr', 'bl', 'br', 'center']

    for idx, subimg in zip(order, crops):
        subimg.save(f'visualisations/guitar_{idx}.png')

    crops = center_plus_twohori_crops(img, [448, 448], 112)
    order = ['left', 'right', 'center2']

    for idx, subimg in zip(order, crops):
        subimg.save(f'visualisations/guitar_{idx}.png')

EgoTaskQA/utils/custom_transforms.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import os

import dominate
from dominate.tags import a, attr, br, h3, img, meta, p, source, span, table, td, tr, video


class HTML:
    """This HTML class allows us to save images and write texts into a single HTML file.

     It consists of functions such as <add_header> (add a text header to the HTML file),
     <add_images> (add a row of images to the HTML file), and <save> (save the HTML to the disk).
     It is based on Python library 'dominate', a Python library for creating and
     manipulating HTML documents using a DOM API.
    """

    def __init__(self, web_dir, title, refresh=0):
        """Initialize the HTML classes

        Parameters:
            web_dir (str) -- a directory that stores the webpage. HTML file will be
            created at <web_dir>/index.html; images will be saved at <web_dir/images/
            title (str)   -- the webpage name
            reflect (int) -- how often the website refresh itself; if 0; no refreshing
        """
        self.title = title
        self.web_dir = web_dir
        self.img_dir = os.path.join(self.web_dir, "images")
        if not os.path.exists(self.web_dir):
            os.makedirs(self.web_dir)
        if not os.path.exists(self.img_dir):
            os.makedirs(self.img_dir)

        self.doc = dominate.document(title=title)
        if refresh > 0:
            with self.doc.head:
                meta(http_equiv="refresh", content=str(refresh))

    def get_image_dir(self):
        """Return the directory that stores images"""
        return self.img_dir

    def add_header(self, text):
        """Insert a header to the HTML file

        Parameters:
            text (str) -- the header text
        """
        with self.doc:
            h3(text)

    def add_videos(self, vids, txts, links, width=400, hidden_tag="hidden"):
        """add images to the HTML file

        Parameters:
            vids (str list)   -- a list of image paths
            txts (str list)  -- a list of image names shown on the website
            links (str list) --  a list of hyperref links; when you click an image,
            it will redirect you to a new page
        """
        self.t = table(border=1, style="table-layout: fixed;")  # Insert a table
        self.doc.add(self.t)
        colors = ["red", "blue", "gold", "salman"]
        with self.t:
            with tr():
                for vid, txt, link in zip(vids, txts, links):
                    td_style = "word-wrap: break-word; width:{}px".format(width)
                    with td(style=td_style, halign="center", valign="top"):
                        with p():
                            vid_path = str(vid)
                            if vid_path == hidden_tag:
                                p_style = "font-weight: bold; width:{}px;"
                                p_style = p_style.format(width * 3)
                                p("hidden video", style=p_style)
                            else:
                                with a(href=str(link)):
                                    with video():
                                        attr(controls="controls")
                                        source(src=vid_path, type="video/mp4")
                            br()
                            rows = txt.split("<br>")
                            for idx, row in enumerate(rows):
                                color = colors[idx % len(colors)]
                                bold_tag = "<b>"
                                if not row.startswith(bold_tag):
                                    s_style = "color:{};".format(color)
                                else:
                                    s_style = "color:black; font-weight: bold;"
                                    row = row[len(bold_tag):]
                                span(row, style=s_style)

    def add_images(self, ims, txts, links, width=400):
        """add images to the HTML file

        Parameters:
            ims (str list)   -- a list of image paths
            txts (str list)  -- a list of image names shown on the website
            links (str list) --  a list of hyperref links; when you click an image,
            it will redirect you to a new page
        """
        self.t = table(border=1, style="table-layout: fixed;")  # Insert a table
        self.doc.add(self.t)
        colors = ["red", "blue", "gold", "salman"]
        with self.t:
            with tr():
                for im, txt, link in zip(ims, txts, links):
                    td_style = "word-wrap: break-word;"
                    with td(style=td_style, halign="center", valign="top"):
                        with p():
                            with a(href=link):
                                img(
                                    style="width:%dpx" % width,
                                    src=im,
                                )
                            br()
                            rows = txt.split("<br>")
                            for idx, row in enumerate(rows):
                                color = colors[idx % len(colors)]
                                bold_tag = "<b>"
                                if not row.startswith(bold_tag):
                                    s_style = "color:{};".format(color)
                                else:
                                    s_style = "color:black; font-weight: bold;"
                                    row = row[len(bold_tag):]
                                span(row, style=s_style)

    def save(self):
        """save the current content to the HMTL file"""
        html_file = "%s/index.html" % self.web_dir
        f = open(html_file, "wt")
        f.write(self.doc.render())
        f.close()


if __name__ == "__main__":  # we show an example usage here.
    html = HTML("web/", "test_html")
    html.add_header("hello world")

    ims, txts, links = [], [], []
    for n in range(4):
        ims.append("image_%d.png" % n)
        txts.append("text_%d" % n)
        links.append("image_%d.png" % n)
    html.add_images(ims, txts, links)
    html.save()

EgoTaskQA/utils/html.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import numpy as np

def calculate_mAP(sim_mat, relevancy_matrix):
    """
    Computes the mean average precision according to the following formula of
    average precision:
    \frac{\sum_{k=1}^n p(k) x rel(k)}{num_rel_docs}

    where p(k) is the precision at k, rel(k) is an indicator function
    determining whether the kth returned item is relevant or not and
    num_rel_docs is the number of relevant items to find within the search.

    The mean average precision is the mean of the average precision for each
    query item (i.e row in the matrix)

    This function takes in two parameters:
        - sim_mat: a NxM matrix which represents the similarity between two
        modalities (with modality 1 being of size N and modality 2 of size M).
        - relevancy_matrix: an NxM matrix which represents the relevancy between two
        modalities of items (with modality 1 being of size N and modality 2 of
        size M).
    """
    #Find the order of the items in modality 2 according to modality 1
    ranked_order = (-sim_mat).argsort()
    ranked_sim_mat = sim_mat[np.arange(sim_mat.shape[0])[:, None], ranked_order]
    #re-order the relevancy matrix to accommodate the proposals
    ranked_rel_mat = relevancy_matrix[np.arange(relevancy_matrix.shape[0])[:, None], ranked_order]

    #find the number of relevant items found at each k
    cumulative_rel_mat = np.cumsum(ranked_rel_mat, axis=1)
    #Mask this ensuring that it is non zero if the kth term is 1 (rel(k) above)
    cumulative_rel_mat[ranked_rel_mat != 1] = 0
    #find the divisor for p(k)
    divisor = np.arange(ranked_rel_mat.shape[1]) + 1

    #find the number of relevant docs per query item
    number_rel_docs = np.sum(ranked_rel_mat==1, axis=1)

    #find the average precision per query, within np.sum finds p(k) * rel(k)
    avg_precision = np.sum(cumulative_rel_mat / divisor, axis=1) / number_rel_docs
    mAP = np.mean(avg_precision)
    return mAP

EgoTaskQA/utils/mAP.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import numpy as np

def calculate_DCG(similarity_matrix, relevancy_matrix, k_counts):
    """
    Calculates the Discounted Cumulative Gain (DCG) between two modalities for
    the first modality.
    DCG = \sum_{i=1}^k \frac{rel_i}{log_2(i + 1)}
    i.e. the sum of the k relevant retrievals which is calculated as the scaled
    relevancy for the ith item. The scale is designed such that early
    retrievals are more important than later retrievals.
    Params:
        - similarity_matrix: matrix of size n1 x n2 where n1 is the number of
          items in the first modality and n2 is the number of items in the
          second modality. The [ith,jth] element is the predicted similarity
          between the ith item from the first modality and the jth item from
          the second modality.
        - relevancy_matrix: matrix of size n1 x n2 (see similarity_matrix
          above). The [ith, jth] element is the semantic relevancy between the
          ith item from the first modality and the jth item from the second
          modality.
        - k_counts: matrix of size n1 x n2 (see similarity_matrix above) which
          includes information on which items to use to calculate the DCG for
          (see calculate_k_counts for more info on this matrix).
    Returns:
        - The DCG for each item in the first modality, a n1 length vector.
    """
    x_sz, y_sz = similarity_matrix.shape
    ranks = np.argsort(similarity_matrix)[:, ::-1]
    #Create vector of size (n,) where n is the length of the last dimension in
    #similarity matrix
    #This vector is of the form log(i+1)
    logs = np.log2(np.arange(y_sz) + 2)
    #Convert logs into the divisor for the DCG calculation, of size similarity
    #matrix
    divisors = np.repeat(np.expand_dims(logs, axis=0), x_sz, axis=0)

    #mask out the sorted relevancy matrix to only use the first k relevant
    #retrievals for each item.
    columns = np.repeat(np.expand_dims(np.arange(x_sz), axis=1), y_sz, axis=1)
    numerators = relevancy_matrix[columns, ranks] * k_counts
    #Calculate the final DCG score (note that this isn't expected to sum to 1)
    return np.sum(numerators / divisors, axis=1)

def calculate_k_counts(relevancy_matrix):
    """
    Works out the maximum number of allowed retrievals when working out the
    Discounted Cumulative Gain. For each query the DCG only uses the first k
    items retrieved which constitute the k relevant items for that query
    (otherwise the nDCG scores can be deceptively high for bad rankings).
    Params:
        - relevancy_matrix: matrix of size n1 x n2 where n1 is the number of
          items in the first modality and n2 is the number of items in the
          second modality.  The [ith, jth] element is the semantic relevancy
          between the ith item from the first modality and the jth item from
          the second modality.
    Returns:
        - Matrix of size n1 x n2 (see relevancy matrix for more info). This is
          created as a mask such that if the [ith, jth] element is 1 it
          represents a valid item to use for the calculation of DCG for the
          ith item after sorting. For example, if relevancy matrix of:
        [[1, 0.5, 0],
          [0, 0  , 1]]
          is given, then the k_counts matrix will be:
        [[1, 1, 0],
         [1, 0, 0]]
         i.e. the first row has 2 non-zero items, so the first two retrieved
         items should be used in the calculation. In the second row there is
         only 1 relevant item, therefore only the first retrieved item should
         be used for the DCG calculation.
    """
    return (np.sort(relevancy_matrix)[:, ::-1] > 0).astype(int)


def calculate_IDCG(relevancy_matrix, k_counts):
    """
    Calculates the Ideal Discounted Cumulative Gain (IDCG) which is the value
    of the Discounted Cumulative Gain (DCG) for a perfect retrieval, i.e. the
    items in the second modality were retrieved in order of their descending
    relevancy.
    Params:
        - relevancy_matrix: matrix of size n1 x n2 where n1 is the number of
          items in the first modality and n2 is the number of items in the
          second modality. The [ith, jth] element is the semantic relevancy
          between the ith item from the first modality and the jth item from
          the second modality.
        - k_counts: matrix of size n1 x n2 (see similarity_matrix above) which
          includes information on which items to use to calculate the DCG for
          (see calculate_k_counts for more info on this matrix).
    """
    return calculate_DCG(relevancy_matrix, relevancy_matrix, k_counts)

def calculate_nDCG(similarity_matrix, relevancy_matrix, k_counts=None, IDCG=None, reduction='mean'):
    """
    Calculates the normalised Discounted Cumulative Gain (nDCG) between two
    modalities for the first modality using the Discounted Cumulative Gain
    (DCG) and the Ideal Discounted Cumulative Gain (IDCG).

    nDCG = \frac{DCG}{IDCG}

    Params:
        - similarity_matrix: matrix of size n1 x n2 where n1 is the number of
          items in the first modality and n2 is the number of items in the second
          modality. The [ith,jth] element is the predicted similarity between
          the ith item from the first modality and the jth item from the second
          modality.
        - relevancy_matrix: matrix of size n1 x n2 (see similarity_matrix
          above). The [ith, jth] element is the semantic relevancy between the
          ith item from the first modality and the jth item from the second
          modality.
        - k_counts: optional parameter: matrix of size n1 x n2 (see
          similarity_matrix above) which includes information on which items to
          use to calculate the DCG for (see calculate_k_counts for more info on
          this matrix). This will be calculated using calculate_IDCG if not
          present, but should be pre-processed for efficiency.
        - IDCG: Optional parameter which includes the pre-processed Ideal
          Discounted Cumulative Gain (IDCG). This is a vector of size n1 (see
          similarity_matrix above) which contains the IDCG value for each item
          from the first modality. This will be calculated using calculate_IDCG
          if not present, but should be pre-processed for efficiency.
        - reduction: what to use to reduce the different nDCG scores. By
          default this applies np.mean across all different queries.
    Returns:
        - The nDCG values for the first modality.
    """
    if k_counts is None:
        k_counts = calculate_k_counts(relevancy_matrix)
    DCG = calculate_DCG(similarity_matrix, relevancy_matrix, k_counts)
    if IDCG is None:
        IDCG = calculate_IDCG(relevancy_matrix, k_counts)
    if reduction == 'mean':
        return np.mean(DCG / IDCG)
    elif reduction is None:
        return DCG / IDCG


if __name__ == '__main__':
    sim_matrix = np.array([
        [1.0, 0.7, 0.4, 0.0],
        [0.3, 0.9, 0.6, 0.1],
        [0.2, 0.5, 0.8, 0.4]
    ])
    rel_matrix = np.array([
        [1.0, 0.5, 0.25, 0.0],
        [0.0, 1.0, 0.4, 0.0],
        [0.5, 0.3, 1.0, 0.0]
    ])
    k_counts = np.array([
        [1, 1, 1, 0],
        [1, 1, 0, 0],
        [1, 1, 1, 0]
    ])
    assert (k_counts == calculate_k_counts(rel_matrix)).all()
    nDCG = calculate_nDCG(sim_matrix, rel_matrix, k_counts)
    assert nDCG == 0.9371789900735429
    DCG = calculate_DCG(sim_matrix, rel_matrix, k_counts)
    IDCG = calculate_IDCG(rel_matrix, k_counts)

    assert nDCG == np.mean(DCG / IDCG)

    pre_nDCG = calculate_nDCG(sim_matrix, rel_matrix, k_counts, IDCG=IDCG)
    assert pre_nDCG == nDCG

    post_mean_nDCG = calculate_nDCG(sim_matrix, rel_matrix, k_counts, IDCG=IDCG, reduction=None)
    assert np.mean(post_mean_nDCG) == pre_nDCG

EgoTaskQA/utils/nDCG.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import functools
import json
import os
import socket
import time
from collections import OrderedDict
from datetime import datetime
from itertools import repeat
from pathlib import Path

import humanize
import numpy as np
import psutil


def replace_nested_dict_item(obj, key, replace_value):
    for k, v in obj.items():
        if isinstance(v, dict):
            obj[k] = replace_nested_dict_item(v, key, replace_value)
    if key in obj:
        obj[key] = replace_value
    return obj


def state_dict_data_parallel_fix(load_state_dict, curr_state_dict):
    load_keys = list(load_state_dict.keys())
    curr_keys = list(curr_state_dict.keys())

    redo_dp = False
    undo_dp = False
    if not curr_keys[0].startswith('module.') and load_keys[0].startswith('module.'):   # this
        undo_dp = True
    elif curr_keys[0].startswith('module.') and not load_keys[0].startswith('module.'):
        redo_dp = True

    if undo_dp: # this
        from collections import OrderedDict
        new_state_dict = OrderedDict()
        for k, v in load_state_dict.items():
            name = k[7:]  # remove `module.`
            new_state_dict[name] = v
        # load params
    elif redo_dp:
        from collections import OrderedDict
        new_state_dict = OrderedDict()
        for k, v in load_state_dict.items():
            name = 'module.' + k  # remove `module.`
            new_state_dict[name] = v
    else:
        new_state_dict = load_state_dict
    return new_state_dict

def print_numpy(x, val=True, shp=False):
    """Print the mean, min, max, median, std, and size of a numpy array
    Parameters:
        val (bool) -- if print the values of the numpy array
        shp (bool) -- if print the shape of the numpy array
    """
    x = x.astype(np.float64)
    if shp:
        print('shape,', x.shape)
    if val:
        x = x.flatten()
        print('mean = %3.3f, min = %3.3f, max = %3.3f, median = %3.3f, std=%3.3f' % (
            np.mean(x), np.min(x), np.max(x), np.median(x), np.std(x)))


def mkdirs(paths):
    """create empty directories if they don't exist
    Parameters:
        paths (str list) -- a list of directory paths
    """
    if isinstance(paths, list) and not isinstance(paths, str):
        for path in paths:
            mkdir(path)
    else:
        mkdir(paths)


def mkdir(path):
    """create a single empty directory if it didn't exist
    Parameters:
        path (str) -- a single directory path
    """
    if not os.path.exists(path):
        os.makedirs(path)

def read_json(fname):
    with fname.open('rt') as handle:
        return json.load(handle, object_hook=OrderedDict)

def write_json(content, fname):
    with fname.open('wt') as handle:
        json.dump(content, handle, indent=4, sort_keys=False)

def inf_loop(data_loader):
    ''' wrapper function for endless data loader. '''
    for loader in repeat(data_loader):
        yield from loader

def memory_summary():
    vmem = psutil.virtual_memory()
    msg = (
        f">>> Currently using {vmem.percent}% of system memory "
        f"{humanize.naturalsize(vmem.used)}/{humanize.naturalsize(vmem.available)}"
    )
    print(msg)

@functools.lru_cache(maxsize=64, typed=False)
def memcache(path):
    suffix = Path(path).suffix
    print(f"loading features >>>", end=" ")
    tic = time.time()
    if suffix == ".npy":
        res = np_loader(path)
    else:
        raise ValueError(f"unknown suffix: {suffix} for path {path}")
    print(f"[Total: {time.time() - tic:.1f}s] ({socket.gethostname() + ':' + str(path)})")
    return res

def np_loader(np_path, l2norm=False):
    with open(np_path, "rb") as f:
        data = np.load(f, encoding="latin1", allow_pickle=True)
    if isinstance(data, np.ndarray) and data.size == 1:
        data = data[()]  # handle numpy dict storage convnetion
    if l2norm:
        print("L2 normalizing features")
        if isinstance(data, dict):
            for key in data:
                feats_ = data[key]
                feats_ = feats_ / max(np.linalg.norm(feats_), 1E-6)
                data[key] = feats_
        elif data.ndim == 2:
            data_norm = np.linalg.norm(data, axis=1)
            data = data / np.maximum(data_norm.reshape(-1, 1), 1E-6)
        else:
            raise ValueError("unexpected data format {}".format(type(data)))
    return data


class Timer:
    def __init__(self):
        self.cache = datetime.now()

    def check(self):
        now = datetime.now()
        duration = now - self.cache
        self.cache = now
        return duration.total_seconds()

    def reset(self):
        self.cache = datetime.now()

EgoTaskQA/utils/util.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import random
import cv2
import numpy as np
import torch


def load_frames_from_video_path(path, num_frames, sample='rand'):
    cap = cv2.VideoCapture(path)
    assert (cap.isOpened())
    vlen = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    acc_samples = min(num_frames, vlen)
    intervals = np.linspace(start=0, stop=vlen, num=acc_samples + 1).astype(int)
    ranges = []
    for idx, interv in enumerate(intervals[:-1]):
        ranges.append((interv, intervals[idx + 1] - 1))
    if sample == 'rand':
        frame_idxs = [random.choice(range(x[0], x[1])) for x in ranges]
    elif sample == 'uniform':
        frame_idxs = [(x[0] + x[1]) // 2 for x in ranges]
    else:
        raise NotImplementedError

    frames = []
    for index in frame_idxs:
        cap.set(cv2.CAP_PROP_POS_FRAMES, index)
        ret, frame = cap.read()
        if ret:
            cv2.imwrite(f'images/{index}.jpg', frame)
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frame = torch.from_numpy(frame)
            # (H x W x C) to (C x H x W)
            frame = frame.permute(2, 0, 1)
            frames.append(frame)
        else:
            raise ValueError

    frames = torch.stack(frames).float() / 255
    cap.release()
    return frames, frame_idxs

EgoTaskQA/utils/video.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import os
import cv2
import shutil
import json
import sys
import multiprocessing
from multiprocessing import Pool
import subprocess
import pandas as pd
from csv import reader, writer

video_dir = './dataset/ego4d_256_multiprocessing'
output_dir = './dataset/ego4d_chunked_multiprocessing'

save_dir = os.path.join(output_dir)
if not os.path.exists(save_dir):
    os.makedirs(save_dir)

dur_limit = 600

def video2segments(infos):
    global count
    index, uid, dur = infos[0], infos[1], infos[2]
    input_path = os.path.join(video_dir, uid + '.mp4')

    output_uid_dir = os.path.join(output_dir, uid)
    if not os.path.exists(output_uid_dir):
        os.makedirs(output_uid_dir)

    #if index % num_partition != partition:
    #    return

    assert os.path.exists(input_path)

    cap = cv2.VideoCapture(input_path)
    rate = cap.get(5)
    frame_num = cap.get(7)
    duration = frame_num / rate

    if duration <= dur_limit:
        shutil.copyfile(input_path, os.path.join(output_uid_dir, '0.mp4'))
        return

    num_seg = duration // dur_limit

    s1_time = 0;
    s2_time = dur_limit
    num_finished = 0
    while num_finished <= num_seg:
        output_mp4_path = os.path.join(output_uid_dir, str(num_finished) + '.mp4')

        cmd = 'ffmpeg -y -i {} -ss {} -to {} -async 1 {}'.format(input_path, s1_time, s2_time,
                                                                                 output_mp4_path)
        # print(cmd)
        subprocess.call(cmd, shell=True)

        # Update for next steps
        s1_time = s2_time
        s2_time = min(s1_time + dur_limit, duration)
        num_finished += 1
    return

if __name__ == "__main__":
    with open('./dataset/ego4d/manifest.csv', 'r') as csv:
        csv_reader = list(reader(csv))[1:]

    downloaded = os.listdir('./dataset/ego4d')

    uid_list = []
    infos_list = []
    num_valid = 0
    for i_item, item in enumerate(csv_reader):
        uid, dur = item[0], float(item[2])
        existed = uid + '.mp4' in downloaded

        if not existed:
            continue

        uid_list.append(uid)
        infos_list.append([num_valid, uid, dur])
        num_valid += 1

    # for infos in infos_list:
    #     video2segments(infos)

    pool = Pool(multiprocessing.cpu_count())
    pool.map(video2segments, tuple(infos_list))

EgoTaskQA/utils/video_chunk.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import os
import time
import sys
import subprocess
import multiprocessing
from multiprocessing import Pool, Value

folder_path = './dataset/ego4d'
output_path = './dataset/ego4d_256_multiprocessing'

def videos_resize(videoinfos):
    global count

    videoid, videoname = videoinfos

    if os.path.exists(os.path.join(output_path, videoname)):
        print(f'{videoname} is resized.')
        return

    inname = folder_path + '/' + videoname
    outname = output_path + '/' + videoname

    cmd = "ffmpeg -y -i {} -filter:v scale=\"trunc(oh*a/2)*2:256\" -c:a copy {}".format( inname, outname)
    subprocess.call(cmd, shell=True)

    return

if __name__ == "__main__":
    file_list = []
    mp4_list = [item for item in os.listdir(folder_path) if item.endswith('.mp4')]
    for id, video in enumerate(mp4_list):
        file_list.append([id, video])

    pool = Pool(multiprocessing.cpu_count()) 
    pool.map(videos_resize, tuple(file_list))

    # for file in file_list:
    #    videos_resize(file)

EgoTaskQA/utils/video_resize.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import matplotlib
import numpy as np
import torch

matplotlib.use('Agg')


def visualise_path(pred, target, window):
    """
    :param pred: (P, 2) Tensor where P is the number of predictions, and 2 is the (i,j) coordinate
    :param target: (T, 2) Tensor where T is the number of targets, and 2 is the (i,j) coordinate
    :param dims: (H, W) tup/list the desired height and width of matrix (should be >= to max(i), max(j))
    :param assignment_method: Method of assignment (dtw, minimum etc.)
    :return: image, visualisation of path prediction and target.
    """
    tp = torch.Tensor((64, 191, 64))
    fp = torch.Tensor((191, 64, 64))
    gt = torch.Tensor((102, 153, 255))

    grid = torch.ones_like(window).unsqueeze(0).repeat(3, 1, 1) * 255
    inf = 130 * torch.ones_like(grid)
    grid = torch.where(torch.isnan(window), inf, grid)

    clip_idxs = [t[0] for t in target]
    local_idxs = np.unique(np.array(clip_idxs)).tolist()

    for t in target:
        local_idx = local_idxs.index(t[0])
        grid[:, local_idx, t[1]] = gt

    for p in pred:
        local_idx = local_idxs.index(p[0])
        if (grid[:, local_idx, p[1]] == gt).all():
            grid[:, local_idx, p[1]] = tp
        else:
            grid[:, local_idx, p[1]] = fp

    return grid / 255


def batch_path_vis(pred_dict, target, window):
    grids = []

    window = window.cpu()
    for key, pred in pred_dict.items():
        tmp_window = window
        if key == 'min_dist':
            tmp_window = torch.zeros_like(window)
        grids.append(visualise_path(pred, target, tmp_window))

    return torch.stack(grids)


if __name__ == "__main__":
    pred = [[1, 1], [2, 4]]
    gt = [[1, 1], [3, 4]]
    window = torch.zeros((5, 6))
    visualise_path(pred, gt, window)

EgoTaskQA/utils/visualisation.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


from .base_data_loader import *
from .base_dataset import *
from .base_model import *
from .base_trainer import *

EgoVLPv2/base/__init__.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import numpy as np
from torch.utils.data import DataLoader
from torch.utils.data.dataloader import default_collate
from torch.utils.data.sampler import SubsetRandomSampler
from torch.utils.data.distributed import DistributedSampler

class BaseDataLoader(DataLoader):
    """
    Base class for all data loaders
    """
    def __init__(self, dataset, batch_size, shuffle, validation_split, num_workers, collate_fn=default_collate):
        self.validation_split = validation_split
        self.shuffle = shuffle

        self.batch_idx = 0
        self.n_samples = len(dataset)

        self.sampler, self.valid_sampler = self._split_sampler(self.validation_split)

        self.init_kwargs = {
            'dataset': dataset,
            'batch_size': batch_size,
            'shuffle': self.shuffle,
            'collate_fn': collate_fn,
            'num_workers': num_workers
        }
        super().__init__(sampler=self.sampler, **self.init_kwargs)

    def _split_sampler(self, split):
        if split == 0.0:
            return None, None

        idx_full = np.arange(self.n_samples)

        np.random.seed(0)
        np.random.shuffle(idx_full)

        if isinstance(split, int):
            assert split > 0
            assert split < self.n_samples, "validation set size is configured to be larger than entire dataset."
            len_valid = split
        else:
            len_valid = int(self.n_samples * split)

        valid_idx = idx_full[0:len_valid]
        train_idx = np.delete(idx_full, np.arange(0, len_valid))

        train_sampler = SubsetRandomSampler(train_idx)
        valid_sampler = SubsetRandomSampler(valid_idx)
        # turn off shuffle option which is mutually exclusive with sampler
        self.shuffle = False
        self.n_samples = len(train_idx)

        return train_sampler, valid_sampler

    def split_validation(self, diff_kwargs=None):
        init_kwargs = self.init_kwargs
        if diff_kwargs is not None:
            init_kwargs.update(diff_kwargs)
        if self.valid_sampler is None:
            return None
        else:
            return DataLoader(sampler=self.valid_sampler, **self.init_kwargs)

    def num_samples(self):
        return len(self.sampler)


class BaseDataLoaderExplicitSplit(DataLoader):
    """
    Base class for all data loaders
    """
    def __init__(self, dataset, batch_size, shuffle, num_workers, collate_fn=default_collate):
        self.shuffle = shuffle

        self.batch_idx = 0
        self.n_samples = len(dataset)

        self.init_kwargs = {
            'dataset': dataset,
            'batch_size': batch_size,
            'shuffle': self.shuffle,
            'collate_fn': collate_fn,
            'num_workers': num_workers,
            'pin_memory': True
        }
        super().__init__(**self.init_kwargs)


class DistBaseDataLoaderExplicitSplit(DataLoader):
    """
    Base class for all data loaders
    """
    def __init__(self, dataset, batch_size, shuffle, num_workers, collate_fn=default_collate):
        self.shuffle = shuffle

        self.batch_idx = 0
        self.n_samples = len(dataset)
        self.train_sampler = DistributedSampler(dataset)
        self.init_kwargs = {
            'dataset': dataset,
            'batch_size': batch_size,
            'shuffle': False,
            'collate_fn': collate_fn,
            'num_workers': num_workers,
            'pin_memory': True,
            'sampler': self.train_sampler
        }
        super().__init__(**self.init_kwargs)


class MultiDistBaseDataLoaderExplicitSplit(DataLoader):
    """
    Base class for all data loaders
    """
    def __init__(self, args, dataset, batch_size, shuffle, num_workers, collate_fn=default_collate):
        self.shuffle = shuffle

        self.batch_idx = 0
        self.n_samples = len(dataset)
        self.args = args
        #print("args.world_size: ", self.args)
        #self.train_sampler = DistributedSampler(dataset, num_replicas=self.args.world_size, rank=self.args.rank, drop_last=False)
        self.train_sampler = DistributedSampler(dataset, drop_last=False, shuffle=True) ## Shuffle False only for debugging; change it later
        self.init_kwargs = {
            'dataset': dataset,
            'batch_size': batch_size,
            'shuffle': False,
            'collate_fn': collate_fn,
            'num_workers': num_workers,
            'pin_memory': True,
            'sampler': self.train_sampler
        }
        super().__init__(**self.init_kwargs)

class BaseMultiDataLoader:
    """
    Currently implemented as undersample the bigger dataloaders...
    """
    def __init__(self, dataloaders):
        self.dataloaders = dataloaders
        self.batch_size = self.dataloaders[0].batch_size

    def __getitem__(self, item):
        dl_idx = item % len(self.dataloaders)
        return next(iter(self.dataloaders[dl_idx]))

    def __len__(self):
        return min(len(x) for x in self.dataloaders) * len(self.dataloaders)

    def num_samples(self):
        return sum(len(x.sampler) for x in self.dataloaders)

EgoVLPv2/base/base_data_loader.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import os
import pdb

import tqdm
import random
from abc import abstractmethod

import av
import cv2
import decord
import ffmpeg
import numpy as np
import torch
from PIL import Image
from torch.utils.data import Dataset, get_worker_info
from torchvision import transforms


class TextVideoDataset(Dataset):
    def __init__(self,
                 dataset_name,
                 text_params,
                 video_params,
                 data_dir,
                 meta_dir=None,
                 split='train',
                 tsfms=None,
                 cut=None,
                 subsample=1,
                 sliding_window_stride=-1,
                 reader='decord',
                 neg_param=None,
                 ):
        self.dataset_name = dataset_name
        self.text_params = text_params
        self.video_params = video_params
        # check for environment variables
        self.data_dir = os.path.expandvars(data_dir)
        if meta_dir is not None:
            self.meta_dir = os.path.expandvars(meta_dir)
        else:
            self.meta_dir = self.data_dir
        self.split = split
        self.transforms = tsfms
        self.cut = cut
        self.subsample = subsample
        self.sliding_window_stride = sliding_window_stride
        self.video_reader = video_reader[reader]
        self.label_type = 'caption'
        self.neg_param = neg_param
        self._load_metadata()
        if self.sliding_window_stride != -1:
            if self.split != 'test':
                raise ValueError('Fixing frame sampling is for test time only. can remove but...')
            self._fix_temporal_samples()

    @abstractmethod
    def _load_metadata(self):
        raise NotImplementedError("Metadata loading must be implemented by subclass")

    @abstractmethod
    def _get_video_path(self, sample):
        raise NotImplementedError("Get video path function must be implemented by subclass")

    def _get_caption(self, sample):
        raise NotImplementedError("Get caption function must be implemented by subclass")

    def _get_video_lens(self):
        vlen_li = []
        for idx, row in self.metadata.iterrows():
            video_path = self._get_video_path(row)[0]
            vlen_li.append(get_video_len(video_path))

        return vlen_li

    def _fix_temporal_samples(self):
        self.metadata['vlen'] = self._get_video_lens()
        self.metadata['frame_intervals'] = self.metadata['vlen'].apply(
            lambda x: np.linspace(start=0, stop=x, num=min(x, self.video_params['num_frames']) + 1).astype(int))
        self.metadata['fix_start'] = self.metadata['frame_intervals'].apply(
            lambda x: np.arange(0, int(x[-1] / len(x - 1)), self.sliding_window_stride)
        )
        self.metadata = self.metadata.explode('fix_start')

    def __len__(self):
        return len(self.metadata)

    def __getitem__(self, item):
        item = item % len(self.metadata)
        sample = self.metadata.iloc[item]
        video_fp, rel_fp = self._get_video_path(sample)
        caption = self._get_caption(sample)

        video_loading = self.video_params.get('loading', 'strict')
        frame_sample = 'rand'
        fix_start = None
        if self.split == 'test':
            frame_sample = 'uniform'
        if self.sliding_window_stride != -1:
            fix_start = sample['fix_start']

        try:
            if os.path.isfile(video_fp):
                imgs, idxs = self.video_reader(video_fp, self.video_params['num_frames'], frame_sample,
                                               fix_start=fix_start)
            else:
                print(f"Warning: missing video file {video_fp}.")
                assert False
        except Exception as e:
            if video_loading == 'strict':
                raise ValueError(
                    f'Video loading failed for {video_fp}, video loading for this dataset is strict.') from e
            else:
                imgs = Image.new('RGB', (self.video_params['input_res'], self.video_params['input_res']), (0, 0, 0))
                imgs = transforms.ToTensor()(imgs).unsqueeze(0)

        # if self.transforms is not None:
        #     imgs = self.transforms(imgs)

        if self.transforms is not None:
            if self.video_params['num_frames'] > 1:
                imgs = imgs.transpose(0, 1)  # [T, C, H, W] ---> [C, T, H, W]
                imgs = self.transforms(imgs)
                imgs = imgs.transpose(0, 1)  # recover
            else:
                imgs = self.transforms(imgs)

        # if self.transforms is not None:
        #     # for video aug if T>1 else just do image aug
        #     if imgs.size(0) > 1:
        #         # added by Mr. YAN
        #         imgs = imgs.transpose(0, 1) # [T, C, H, W] ---> [C, T, H, W]
        #         imgs = self.transforms(imgs)
        #         imgs = imgs.transpose(0, 1) # recover
        #     else:
        #         imgs = self.transforms(imgs)    # t

        final = torch.zeros([self.video_params['num_frames'], 3, self.video_params['input_res'],
                             self.video_params['input_res']])
        final[:imgs.shape[0]] = imgs

        meta_arr = {'raw_captions': caption, 'paths': rel_fp, 'dataset': self.dataset_name}
        data = {'video': final, 'text': caption, 'meta': meta_arr}
        return data


class TextImageDataset(TextVideoDataset):

    def __getitem__(self, item):
        item = item % len(self.metadata)
        sample = self.metadata.iloc[item]
        video_fp, rel_fp = self._get_video_path(sample)
        caption = self._get_caption(sample)

        video_loading = self.video_params.get('loading', 'strict')

        try:
            img = Image.open(video_fp).convert("RGB")
        except:
            if video_loading == 'strict':
                raise ValueError(f'Image loading failed for {video_fp}, image loading for this dataset is strict.')
            else:
                img = Image.new('RGB', (self.video_params['input_res'], self.video_params['input_res']), (0, 0, 0))

        # convert to tensor because video transforms don't, expand such that its a 1-frame video.
        img = transforms.ToTensor()(img).unsqueeze(0)
        if self.transforms is not None:
            img = self.transforms(img)
        meta_arr = {'raw_captions': caption, 'paths': rel_fp, 'dataset': self.dataset_name}
        data = {'video': img, 'text': caption, 'meta': meta_arr}
        return data


def sample_frames(num_frames, vlen, sample='rand', fix_start=None):
    acc_samples = min(num_frames, vlen)
    intervals = np.linspace(start=0, stop=vlen, num=acc_samples + 1).astype(int)
    ranges = []
    for idx, interv in enumerate(intervals[:-1]):
        ranges.append((interv, intervals[idx + 1] - 1))
    if sample == 'rand':
        frame_idxs = [random.choice(range(x[0], x[1])) for x in ranges]
    elif fix_start is not None:
        frame_idxs = [x[0] + fix_start for x in ranges]
    elif sample == 'uniform':
        frame_idxs = [(x[0] + x[1]) // 2 for x in ranges]
    else:
        raise NotImplementedError

    return frame_idxs

def sample_frames_clips(start, end, vlen, acc_samples):
    start = max(0, start)
    end = min(vlen, end)

    intervals = np.linspace(start=start, stop=end, num=int(acc_samples) + 1).astype(int)
    ranges = []
    for idx, interv in enumerate(intervals[:-1]):
        ranges.append((interv, intervals[idx + 1] - 1))
        frame_idxs = [(x[0] + x[1]) // 2 for x in ranges
                      ]
    return frame_idxs

def sample_frames_start_end(num_frames, start, end, sample='rand', fix_start=None):
    acc_samples = min(num_frames, end)
    intervals = np.linspace(start=start, stop=end, num=acc_samples + 1).astype(int)
    ranges = []
    for idx, interv in enumerate(intervals[:-1]):
        ranges.append((interv, intervals[idx + 1] - 1))
    if sample == 'rand':
        frame_idxs = [random.choice(range(x[0], x[1])) for x in ranges]
    elif fix_start is not None:
        frame_idxs = [x[0] + fix_start for x in ranges]
    elif sample == 'uniform':
        frame_idxs = [(x[0] + x[1]) // 2 for x in ranges]
    else:
        raise NotImplementedError

    return frame_idxs

def read_frames_cv2(video_path, num_frames, sample='rand', fix_start=None):
    cap = cv2.VideoCapture(video_path)
    assert (cap.isOpened())
    vlen = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    # get indexes of sampled frames
    frame_idxs = sample_frames(num_frames, vlen, sample=sample, fix_start=fix_start)
    frames = []
    success_idxs = []
    for index in frame_idxs:
        cap.set(cv2.CAP_PROP_POS_FRAMES, index - 1)
        ret, frame = cap.read()
        if ret:
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frame = torch.from_numpy(frame)
            # (H x W x C) to (C x H x W)
            frame = frame.permute(2, 0, 1)
            frames.append(frame)
            success_idxs.append(index)
        else:
            pass
            # print(frame_idxs, ' fail ', index, f'  (vlen {vlen})')

    frames = torch.stack(frames).float() / 255
    cap.release()
    return frames, success_idxs

def read_frames_cv2_egoclip(video_path_1, video_path_2, num_frames, sample,
                            start_sec, end_sec, bound_sec):
    if video_path_1 == video_path_2:
        cap1 = cv2.VideoCapture(video_path_1)
        cap2 = cap1
        vlen1 = int(cap1.get(cv2.CAP_PROP_FRAME_COUNT))
        vlen2 = vlen1
        assert (cap1.isOpened())
    else:   # some clips may span two segments.
        cap1 = cv2.VideoCapture(video_path_1)
        cap2 = cv2.VideoCapture(video_path_2)
        vlen1 = int(cap1.get(cv2.CAP_PROP_FRAME_COUNT))
        vlen2 = int(cap2.get(cv2.CAP_PROP_FRAME_COUNT))
        assert (cap1.isOpened())
        assert (cap2.isOpened())

    # get indexes of sampled frames
    start_f = max(0, int(start_sec * 30))
    end_f = max(0, int(end_sec * 30))
    bound_f = int(bound_sec * 30)
    frame_idxs = sample_frames_start_end(num_frames, start_f, end_f, sample=sample)

    frames = []
    success_idxs = []
    for index in frame_idxs:
        _index = index % (600 * 30)
        if index > bound_f: # frame from the last video
            _index = min(_index, vlen2)
            cap2.set(cv2.CAP_PROP_POS_FRAMES, _index - 1)
            ret, frame = cap2.read()
        else:   # frame from the first video
            _index = min(_index, vlen1)
            cap1.set(cv2.CAP_PROP_POS_FRAMES, _index - 1)
            ret, frame = cap1.read()

        if ret:
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frame = torch.from_numpy(frame)
            # (H x W x C) to (C x H x W)
            frame = frame.permute(2, 0, 1)
            frames.append(frame)
            success_idxs.append(index)
        else:
            pass

    while len(frames) < num_frames: # complete the frame
        frames.append(frames[-1])

    frames = torch.stack(frames).float() / 255
    cap1.release()
    cap2.release()
    return frames, success_idxs

def read_frames_cv2_epic(video_path, start_frame, stop_frame, num_frames, sample='rand', fix_start=None):
    # get indexes of sampled frames
    frame_idxs = sample_frames_start_end(num_frames, start_frame, stop_frame, sample=sample, fix_start=fix_start)
    frames = []
    success_idxs = []
    for index in frame_idxs:
        img_name = 'frame_' + str(index).zfill(10) + '.jpg'
        frame = cv2.imread(os.path.join(video_path, img_name),cv2.COLOR_BGR2RGB)

        frame = torch.from_numpy(frame)
        # (H x W x C) to (C x H x W)
        frame = frame.permute(2, 0, 1)
        frames.append(frame)
        success_idxs.append(index)

    frames = torch.stack(frames).float() / 255
    return frames, success_idxs

def read_frames_cv2_charades(video_path, num_frames, sample, start_sec=None, end_sec=None):
    cap = cv2.VideoCapture(video_path)
    assert (cap.isOpened())
    vlen = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = cap.get(5)

    # get indexes of sampled frames
    if not start_sec and not end_sec:
        frame_idxs = sample_frames(num_frames, vlen, sample=sample)
    else:
        start_f = max(0, int(start_sec * fps))
        end_f = min(int(end_sec * fps), vlen)
        frame_idxs = sample_frames_start_end(num_frames, start_f, end_f, sample=sample)

    frames = []
    success_idxs = []
    for index in frame_idxs:
        cap.set(cv2.CAP_PROP_POS_FRAMES, index - 1)
        ret, frame = cap.read()
        if ret:
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frame = torch.from_numpy(frame)
            # (H x W x C) to (C x H x W)
            frame = frame.permute(2, 0, 1)
            frames.append(frame)
            success_idxs.append(index)
        else:
            pass

    frames = torch.stack(frames).float() / 255
    cap.release()
    return frames, success_idxs

def read_frames_av(video_path, num_frames, sample='rand', fix_start=None):
    reader = av.open(video_path)
    try:
        frames = []
        frames = [torch.from_numpy(f.to_rgb().to_ndarray()) for f in reader.decode(video=0)]
    except (RuntimeError, ZeroDivisionError) as exception:
        print('{}: WEBM reader cannot open {}. Empty '
              'list returned.'.format(type(exception).__name__, video_path))
    vlen = len(frames)
    frame_idxs = sample_frames(num_frames, vlen, sample=sample, fix_start=fix_start)
    frames = torch.stack([frames[idx] for idx in frame_idxs]).float() / 255
    frames = frames.permute(0, 3, 1, 2)
    return frames, frame_idxs

decord.bridge.set_bridge("torch")

def read_frames_decord(video_path, num_frames, sample='rand', fix_start=None):
    video_reader = decord.VideoReader(video_path, num_threads=1)
    vlen = len(video_reader)
    frame_idxs = sample_frames(num_frames, vlen, sample=sample, fix_start=fix_start)
    video_reader.skip_frames(1)
    frames = video_reader.get_batch(frame_idxs)

    frames = frames.float() / 255
    frames = frames.permute(0, 3, 1, 2)
    return frames, frame_idxs

def read_frames_decord_start_end(video_path, start, end, num_frames):
    video_reader = decord.VideoReader(video_path, num_threads=1)
    vlen = len(video_reader)
    frame_idxs = sample_frames_clips(start, end, vlen, num_frames + 1)
    video_reader.skip_frames(1)
    frames = video_reader.get_batch(frame_idxs)

    frames = frames.float() / 255
    frames = frames.permute(0, 3, 1, 2)
    return frames, frame_idxs

def get_video_len(video_path):
    cap = cv2.VideoCapture(video_path)
    if not (cap.isOpened()):
        return False
    vlen = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    cap.release()
    return vlen

video_reader = {
    'av': read_frames_av,
    'cv2': read_frames_cv2,
    'cv2_epic': read_frames_cv2_epic,
    'cv2_charades': read_frames_cv2_charades,
    'cv2_egoclip': read_frames_cv2_egoclip,
    'decord': read_frames_decord,
    'decord_start_end': read_frames_decord_start_end,
}

EgoVLPv2/base/base_dataset.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import torch.nn as nn
import numpy as np
from abc import abstractmethod


class BaseModel(nn.Module):
    """
    Base class for all models
    """
    @abstractmethod
    def forward(self, *inputs):
        """
        Forward pass logic

        :return: Model output
        """
        raise NotImplementedError

    def __str__(self):
        """
        Model prints with number of trainable parameters
        """
        model_parameters = filter(lambda p: p.requires_grad, self.parameters())
        params = sum(np.prod(p.size()) for p in model_parameters)
        return super().__str__() + '\nTrainable parameters: {}'.format(params)

EgoVLPv2/base/base_model.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


from abc import abstractmethod
import torch
from numpy import inf
from pathlib import Path

class BaseTrainer:
    """
    Base class for all trainers
    """
    def __init__(self, model, loss, metrics, optimizer, config, writer=None, init_val=False):
        self.config = config
        self.logger = config.get_logger('trainer', config['trainer']['verbosity'])
        self.init_val = init_val
        # setup GPU device if available, move model into configured device
        self.device, device_ids = self._prepare_device(config['n_gpu'])
        self.model = model.to(self.device)
        self.model.device = self.device
        if len(device_ids) > 1:
            self.model = torch.nn.DataParallel(model, device_ids=device_ids)

        loss = loss.to(self.device)
        self.loss = loss
        self.metrics = metrics
        self.optimizer = optimizer

        cfg_trainer = config['trainer']
        self.epochs = cfg_trainer['epochs']
        self.save_period = cfg_trainer['save_period']
        self.monitor = cfg_trainer.get('monitor', 'off')
        self.init_val = cfg_trainer.get('init_val', True)

        # configuration to monitor model performance and save best
        if self.monitor == 'off':
            self.mnt_mode = 'off'
            self.mnt_best = 0
        else:
            self.mnt_mode, self.mnt_metric = self.monitor.split()
            assert self.mnt_mode in ['min', 'max']

            self.mnt_best = inf if self.mnt_mode == 'min' else -inf
            self.early_stop = cfg_trainer.get('early_stop', inf)

        self.start_epoch = 1

        self.checkpoint_dir = config.save_dir

        # setup visualization writer instance                
        #self.writer = TensorboardWriter(config.log_dir, self.logger, cfg_trainer['tensorboard'])
        self.writer = writer

        if config.resume is not None:
            self._resume_checkpoint(config.resume)

    @abstractmethod
    def _train_epoch(self, epoch):
        """
        Training logic for an epoch

        :param epoch: Current epoch number
        """
        raise NotImplementedError

    @abstractmethod
    def _valid_epoch(self, epoch):
        """
        Training logic for an epoch

        :param epoch: Current epoch number
        """
        raise NotImplementedError


    def train(self):
        """
        Full training logic
        """
        not_improved_count = 0

        if self.init_val:
            _ = self._valid_epoch(-1)

        for epoch in range(self.start_epoch, self.epochs + 1):
            result = self._train_epoch(epoch)

            # save logged informations into log dict

            # save logged informations into log dict
            log = {'epoch': epoch}
            for key, value in result.items():
                if key == 'metrics':
                    log.update({mtr.__name__: value[i]
                                for i, mtr in enumerate(self.metrics)})
                elif key == 'val_metrics':
                    log.update({'val_' + mtr.__name__: value[i]
                                for i, mtr in enumerate(self.metrics)})
                elif key == 'nested_val_metrics':
                    # NOTE: currently only supports two layers of nesting
                    for subkey, subval in value.items():
                        for subsubkey, subsubval in subval.items():
                            for subsubsubkey, subsubsubval in subsubval.items():
                                log[f"val_{subkey}_{subsubkey}_{subsubsubkey}"] = subsubsubval
                else:
                    log[key] = value

            # print logged informations to the screen
            for key, value in log.items():
                self.logger.info('    {:15s}: {}'.format(str(key), value))

            # evaluate model performance according to configured metric, save best checkpoint as model_best
            best = False
            if self.mnt_mode != 'off':
                try:
                    # check whether model performance improved or not, according to specified metric(mnt_metric)
                    improved = (self.mnt_mode == 'min' and log[self.mnt_metric] <= self.mnt_best) or \
                               (self.mnt_mode == 'max' and log[self.mnt_metric] >= self.mnt_best)
                except KeyError:
                    self.logger.warning("Warning: Metric '{}' is not found. "
                                        "Model performance monitoring is disabled.".format(self.mnt_metric))
                    self.mnt_mode = 'off'
                    improved = False

                if improved:
                    self.mnt_best = log[self.mnt_metric]
                    not_improved_count = 0
                    best = True
                else:
                    not_improved_count += 1

                if not_improved_count > self.early_stop:
                    self.logger.info("Validation performance didn\'t improve for {} epochs. "
                                     "Training stops.".format(self.early_stop))
                    break

            if epoch % self.save_period == 0 or best:
            #if best:
                self._save_checkpoint(epoch, save_best=best)

    def _prepare_device(self, n_gpu_use):
        """
        setup GPU device if available, move model into configured device
        """
        n_gpu = torch.cuda.device_count()
        if n_gpu_use > 0 and n_gpu == 0:
            self.logger.warning("Warning: There\'s no GPU available on this machine,"
                                "training will be performed on CPU.")
            n_gpu_use = 0
        if n_gpu_use > n_gpu:
            self.logger.warning("Warning: The number of GPU\'s configured to use is {}, but only {} are available "
                                "on this machine.".format(n_gpu_use, n_gpu))
            n_gpu_use = n_gpu
        device = torch.device('cuda:0' if n_gpu_use > 0 else 'cpu')
        list_ids = list(range(n_gpu_use))
        return device, list_ids

    def _save_checkpoint(self, epoch, save_best=False):
        """
        Saving checkpoints

        :param epoch: current epoch number
        :param log: logging information of the epoch
        :param save_best: if True, rename the saved checkpoint to 'model_best.pth'
        """
        arch = type(self.model).__name__
        state = {
            'arch': arch,
            'epoch': epoch,
            'state_dict': self.model.state_dict(),
            'optimizer': self.optimizer.state_dict(),
            'monitor_best': self.mnt_best,
            'config': self.config
        }
        filename = str(self.checkpoint_dir / 'checkpoint-epoch{}.pth'.format(epoch))
        torch.save(state, filename)
        self.logger.info("Saving checkpoint: {} ...".format(filename))
        if save_best:
            best_path = str(self.checkpoint_dir / 'model_best.pth')
            torch.save(state, best_path)
            self.logger.info("Saving current best: model_best.pth ...")

    def _resume_checkpoint(self, resume_path):
        """
        Resume from saved checkpoints

        :param resume_path: Checkpoint path to be resumed
        """
        resume_path = str(resume_path)
        self.logger.info("Loading checkpoint: {} ...".format(resume_path))
        checkpoint = torch.load(resume_path)
        self.start_epoch = checkpoint['epoch'] + 1
        self.mnt_best = checkpoint['monitor_best']

        # load architecture params from checkpoint.
        if checkpoint['config']['arch'] != self.config['arch']:
            self.logger.warning("Warning: Architecture configuration given in config file is different from that of "
                                "checkpoint. This may yield an exception while state_dict is being loaded.")

        state_dict = checkpoint['state_dict']

        load_state_dict_keys = list(state_dict.keys())
        curr_state_dict_keys = list(self.model.state_dict().keys())
        redo_dp = False
        if not curr_state_dict_keys[0].startswith('module.') and load_state_dict_keys[0].startswith('module.'):
            undo_dp = True
        elif curr_state_dict_keys[0].startswith('module.') and not load_state_dict_keys[0].startswith('module.'):
            redo_dp = True
            undo_dp = False
        else:
            undo_dp = False

        if undo_dp:
            from collections import OrderedDict
            new_state_dict = OrderedDict()
            for k, v in state_dict.items():
                name = k[7:]  # remove `module.`
                new_state_dict[name] = v
            # load params
        elif redo_dp:
            from collections import OrderedDict
            new_state_dict = OrderedDict()
            for k, v in state_dict.items():
                name = 'module.' + k  # remove `module.`
                new_state_dict[name] = v
        else:
            new_state_dict = state_dict

        self.model.load_state_dict(new_state_dict)

        # load optimizer state from checkpoint only when optimizer type is not changed.
        if checkpoint['config']['optimizer']['type'] != self.config['optimizer']['type']:
            self.logger.warning("Warning: Optimizer type given in config file is different from that of checkpoint. "
                                "Optimizer parameters not being resumed.")
        else:
            self.optimizer.load_state_dict(checkpoint['optimizer'])

        self.logger.info("Checkpoint loaded. Resume training from epoch {}".format(self.start_epoch))


class Multi_BaseTrainer_dist:
    """
    Base class for all trainers
    """
    def __init__(self, args, model, loss, metrics, optimizer, scheduler, gpu, config, writer=None, init_val=False):
        self.config = config
        self.logger = config.get_logger('trainer', config['trainer']['verbosity'])
        self.init_val = init_val
        self.args = args
        # setup GPU device if available, move model into configured device
        self.device, device_ids = self._prepare_device(config['n_gpu'])

        self.model = model.cuda(gpu)
        #self.model.device = self.device

        # if len(device_ids) > 1:
            #self.model = torch.nn.parallel.DistributedDataParallel(model)
            #self.model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.rank])
            # self.model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank])
            #self.model = torch.nnç.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], find_unused_parameters=True)
            #self.model = torch.nn.DataParallel(model, device_ids=device_ids)

        ddp_args={}
        ddp_args['static_graph'] = True
        self.model = torch.nn.parallel.DistributedDataParallel(self.model, device_ids=[gpu], **ddp_args)  

        loss = loss.cuda(gpu)
        self.loss = loss
        self.metrics = metrics
        self.optimizer = optimizer
        self.scheduler = scheduler

        cfg_trainer = config['trainer']
        self.epochs = cfg_trainer['epochs']
        self.save_period = cfg_trainer['save_period']
        self.monitor = cfg_trainer.get('monitor', 'off')
        self.init_val = cfg_trainer.get('init_val', True)

        # configuration to monitor model performance and save best
        if self.monitor == 'off':
            self.mnt_mode = 'off'
            self.mnt_best = 0
        else:
            self.mnt_mode, self.mnt_metric = self.monitor.split()
            assert self.mnt_mode in ['min', 'max']

            self.mnt_best = inf if self.mnt_mode == 'min' else -inf
            self.early_stop = cfg_trainer.get('early_stop', inf)

        self.start_epoch = 1

        self.checkpoint_dir = Path(args.save_dir)

        # setup visualization writer instance
        #self.writer = TensorboardWriter(config.log_dir, self.logger, cfg_trainer['tensorboard'])
        self.writer = writer

        if config.resume is not None:
            self._resume_checkpoint(config.resume)

    @abstractmethod
    def _train_epoch(self, epoch, gpu):
        """
        Training logic for an epoch

        :param epoch: Current epoch number
        """
        raise NotImplementedError

    @abstractmethod
    def _valid_epoch(self, epoch, gpu):
        """
        Training logic for an epoch

        :param epoch: Current epoch number
        """
        raise NotImplementedError


    def train(self, gpu):
        """
        Full training logic
        """
        not_improved_count = 0

        if self.init_val:
            _ = self._valid_epoch(-1, gpu)
            # pass

        scaler = torch.cuda.amp.GradScaler()

        for epoch in range(self.start_epoch, self.epochs + 1):
            result = self._train_epoch(epoch, scaler, gpu)

            # save logged informations into log dict

            # save logged informations into log dict
            log = {'epoch': epoch}
            for key, value in result.items():
              if self.args.rank == 0:
                if key == 'metrics':
                    log.update({mtr.__name__: value[i]
                                for i, mtr in enumerate(self.metrics)})
                elif key == 'val_metrics':
                    log.update({'val_' + mtr.__name__: value[i]
                                for i, mtr in enumerate(self.metrics)})
                elif key == 'nested_val_metrics':
                    # NOTE: currently only supports two layers of nesting
                    for subkey, subval in value.items():
                        for subsubkey, subsubval in subval.items():
                            for subsubsubkey, subsubsubval in subsubval.items():
                                log[f"val_{subkey}_{subsubkey}_{subsubsubkey}"] = subsubsubval
                else:
                    log[key] = value

            # print logged informations to the screen
            for key, value in log.items():
              if self.args.rank == 0:
                self.logger.info('    {:15s}: {}'.format(str(key), value))

            # evaluate model performance according to configured metric, save best checkpoint as model_best
            best = False
            if self.mnt_mode != 'off' and self.args.rank == 0:
                try:
                    # check whether model performance improved or not, according to specified metric(mnt_metric)
                    improved = (self.mnt_mode == 'min' and log[self.mnt_metric] <= self.mnt_best) or \
                               (self.mnt_mode == 'max' and log[self.mnt_metric] >= self.mnt_best)
                except KeyError:
                    self.logger.warning("Warning: Metric '{}' is not found. "
                                        "Model performance monitoring is disabled.".format(self.mnt_metric))
                    self.mnt_mode = 'off'
                    improved = False

                if improved:
                    self.mnt_best = log[self.mnt_metric]
                    not_improved_count = 0
                    best = True
                else:
                    not_improved_count += 1

                #if not_improved_count > self.early_stop:
                #    self.logger.info("Validation performance didn\'t improve for {} epochs. "
                #                     "Training stops.".format(self.early_stop))
                #    break

            if epoch % self.save_period == 0 or best:
                if self.args.rank == 0:
                #if best:
                    self._save_checkpoint(epoch, save_best=best)

    def _prepare_device(self, n_gpu_use):
        """
        setup GPU device if available, move model into configured device
        """
        n_gpu = torch.cuda.device_count()
        if n_gpu_use > 0 and n_gpu == 0:
            self.logger.warning("Warning: There\'s no GPU available on this machine,"
                                "training will be performed on CPU.")
            n_gpu_use = 0
        if n_gpu_use > n_gpu:
            self.logger.warning("Warning: The number of GPU\'s configured to use is {}, but only {} are available "
                                "on this machine.".format(n_gpu_use, n_gpu))
            n_gpu_use = n_gpu
        device = torch.device('cuda' if n_gpu_use > 0 else 'cpu')
        list_ids = list(range(n_gpu_use))
        return device, list_ids

    def _save_checkpoint(self, epoch, save_best=False):
        """
        Saving checkpoints

        :param epoch: current epoch number
        :param log: logging information of the epoch
        :param save_best: if True, rename the saved checkpoint to 'model_best.pth'
        """
        arch = type(self.model).__name__
        state = {
            'arch': arch,
            'epoch': epoch,
            'state_dict': self.model.state_dict(),
            'optimizer': self.optimizer.state_dict(),
            'monitor_best': self.mnt_best,
            'config': self.config
        }
        filename = str(self.checkpoint_dir / 'checkpoint-epoch{}.pth'.format(epoch))
        torch.save(state, filename)
        self.logger.info("Saving checkpoint: {} ...".format(filename))
        if save_best:
            best_path = str(self.checkpoint_dir / 'model_best.pth')
            torch.save(state, best_path)
            self.logger.info("Saving current best: model_best.pth ...")

    def _resume_checkpoint(self, resume_path):
        """
        Resume from saved checkpoints

        :param resume_path: Checkpoint path to be resumed
        """
        resume_path = str(resume_path)
        self.logger.info("Loading checkpoint: {} ...".format(resume_path))
        #checkpoint = torch.load(resume_path)
        checkpoint = torch.load(resume_path, map_location='cuda:{}'.format(self.args.rank))
        self.start_epoch = checkpoint['epoch'] + 1
        self.mnt_best = checkpoint['monitor_best']

        # load architecture params from checkpoint.
        if checkpoint['config']['arch'] != self.config['arch']:
            self.logger.warning("Warning: Architecture configuration given in config file is different from that of "
                                "checkpoint. This may yield an exception while state_dict is being loaded.")

        state_dict = checkpoint['state_dict']

        load_state_dict_keys = list(state_dict.keys())
        curr_state_dict_keys = list(self.model.state_dict().keys())
        redo_dp = False
        if not curr_state_dict_keys[0].startswith('module.') and load_state_dict_keys[0].startswith('module.'):
            undo_dp = True
        elif curr_state_dict_keys[0].startswith('module.') and not load_state_dict_keys[0].startswith('module.'):
            redo_dp = True
            undo_dp = False
        else:
            undo_dp = False

        if undo_dp:
            from collections import OrderedDict
            new_state_dict = OrderedDict()
            for k, v in state_dict.items():
                name = k[7:]  # remove `module.`
                new_state_dict[name] = v
            # load params
        elif redo_dp:
            from collections import OrderedDict
            new_state_dict = OrderedDict()
            for k, v in state_dict.items():
                name = 'module.' + k  # remove `module.`
                new_state_dict[name] = v
        else:
            new_state_dict = state_dict

        self.model.load_state_dict(new_state_dict)

        # load optimizer state from checkpoint only when optimizer type is not changed.
        if checkpoint['config']['optimizer']['type'] != self.config['optimizer']['type']:
            self.logger.warning("Warning: Optimizer type given in config file is different from that of checkpoint. "
                                "Optimizer parameters not being resumed.")
        else:
            self.optimizer.load_state_dict(checkpoint['optimizer'])

        self.logger.info("Checkpoint loaded. Resume training from epoch {}".format(self.start_epoch))

EgoVLPv2/base/base_trainer.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import os
import sys
import csv
import pandas as pd
sys.path.append('/apdcephfs/private_qinghonglin/video_codebase/EgoVLP/')

from base.base_dataset import TextVideoDataset
try:
    from transforms import init_transform_dict, init_video_transform_dict
except:
    pass

import torch
from PIL import Image
from torchvision import transforms

class CharadesEgo(TextVideoDataset):
    def _load_metadata(self):
        split_files = {
            'train': 'metadata_train.csv',
            'val': 'CharadesEgo_v1_test_only1st.csv',
            'test': 'CharadesEgo_v1_test_only1st.csv'
        }
        target_split_fp = split_files[self.split]
        if self.split == 'train':
            metadata = pd.read_csv(os.path.join(self.meta_dir, target_split_fp), delimiter='\t')
        else:
            metadata = pd.read_csv(os.path.join(self.meta_dir, target_split_fp))
        if self.subsample < 1:
            metadata = metadata.sample(frac=self.subsample)

        self.metadata = metadata
        if not self.split == 'train':
            self.label = self._parse_charades_csv(os.path.join(self.meta_dir, target_split_fp))

    def _parse_charades_csv(self, filename):
        labels = {}
        with open(filename) as f:
            reader = csv.DictReader(f)
            for row in reader:
                vid = row['id']
                actions = row['actions']
                if actions == '':
                    actions = []
                else:
                    actions = [a.split(' ') for a in actions.split(';')]
                    actions = [{'class': x, 'start': float(
                        y), 'end': float(z)} for x, y, z in actions]
                labels[vid] = actions
        return labels

    def _get_video_path(self, sample):
        rel_video_fp = sample['id'] + '.mp4'
        full_video_fp = os.path.join(self.data_dir, rel_video_fp)

        return full_video_fp, rel_video_fp

    def _get_caption(self, sample):
        if self.split in ['val', 'test']:
            return sample[6]
        else:
            return sample['narration']

    def _cls2int(self, x):
        return int(x[1:])

    def __getitem__(self, item):
        if self.split in ['val', 'test']:
            return self._get_val(item)
        else:
            return self._get_train(item)

    def _get_train(self, item):
        item = item % len(self.metadata)
        sample = self.metadata.iloc[item]
        video_fp, rel_fp = self._get_video_path(sample)
        caption = self._get_caption(sample)

        start_sec, end_sec = sample['t_start'],  sample['t_end']

        video_loading = self.video_params.get('loading', 'non_strict')
        frame_sample = 'rand'
        if self.split in ['test', 'val']:
            frame_sample = 'uniform'

        try:
            if os.path.isfile(video_fp):
                imgs, idxs = self.video_reader(video_path=video_fp, num_frames=self.video_params['num_frames'], sample=frame_sample,
                                               start_sec=start_sec, end_sec=end_sec)
            else:
                print(f"Warning: missing video file {video_fp}.")
                assert False
        except Exception as e:
            if video_loading == 'strict':
                raise ValueError(
                    f'Video loading failed for {video_fp}, video loading for this dataset is strict.') from e
            else:
                imgs = Image.new('RGB', (self.video_params['input_res'], self.video_params['input_res']), (0, 0, 0))
                imgs = transforms.ToTensor()(imgs).unsqueeze(0)

        if self.transforms is not None:
            if self.video_params['num_frames'] > 1:
                imgs = imgs.transpose(0, 1)  # [T, C, H, W] ---> [C, T, H, W]
                imgs = self.transforms(imgs)
                imgs = imgs.transpose(0, 1)  # recover
            else:
                imgs = self.transforms(imgs)

        final = torch.zeros([self.video_params['num_frames'], 3, self.video_params['input_res'],
                             self.video_params['input_res']])
        final[:imgs.shape[0]] = imgs

        meta_arr = {'raw_captions': caption, 'paths': rel_fp, 'dataset': self.dataset_name}
        data = {'video': final, 'text': caption, 'meta': meta_arr, 'target': sample['cls']}
        return data

    def _get_val(self, item):
        item = item % len(self.metadata)
        sample = self.metadata.iloc[item]
        video_fp, rel_fp = self._get_video_path(sample)
        caption = self._get_caption(sample)

        # construct label
        label = self.label[sample['id']]
        target = torch.IntTensor(157).zero_()
        for x in label:
            target[self._cls2int(x['class'])] = 1

        video_loading = self.video_params.get('loading', 'strict')
        frame_sample = 'rand'
        if self.split in ['val', 'test']:
            frame_sample = 'uniform'

        try:
            if os.path.isfile(video_fp):
                # only supported for read_frames_decord_online
                imgs, idxs = self.video_reader(video_path=video_fp, num_frames=self.video_params['num_frames'], sample=frame_sample)
            else:
                print(f"Warning: missing video file {video_fp}.")
                assert False
        except Exception as e:
            if video_loading == 'strict':
                raise ValueError(
                    f'Video loading failed for {video_fp}, video loading for this dataset is strict.') from e
            else:
                imgs = Image.new('RGB', (self.video_params['input_res'], self.video_params['input_res']), (0, 0, 0))
                imgs = transforms.ToTensor()(imgs).unsqueeze(0)

        if self.transforms is not None:
            if self.video_params['num_frames'] > 1:
                imgs = imgs.transpose(0, 1)  # [T, C, H, W] ---> [C, T, H, W]
                imgs = self.transforms(imgs)
                imgs = imgs.transpose(0, 1)  # recover
            else:
                imgs = self.transforms(imgs)

        final = torch.zeros([self.video_params['num_frames'], 3, self.video_params['input_res'],
                             self.video_params['input_res']])
        final[:imgs.shape[0]] = imgs

        meta_arr = {'raw_captions': caption, 'paths': rel_fp, 'dataset': self.dataset_name}
        data = {'video': final, 'text': caption, 'meta': meta_arr, 'target': target}
        return data

if __name__ == "__main__":
    kwargs = dict(
        dataset_name="CharadesEgo",
        text_params={
            "input": "text"
        },
        video_params={
        "input_res": 224,
        "num_frames": 4,
        "loading": "lax"
        },
        data_dir="/apdcephfs/private_qinghonglin/video_dataset/charades/CharadesEgo_v1_480",
        meta_dir="/apdcephfs/private_qinghonglin/video_dataset/charades/CharadesEgo",
        tsfms=init_video_transform_dict()['test'],
        reader='cv2_charades',
        split='val'
    )
    dataset = CharadesEgo(**kwargs)
    for i in range(100):
        item = dataset[i]
        print(item.keys())

EgoVLPv2/data_loader/CharadesEgo_dataset.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import os
import sys
import json
import pandas as pd
from base.base_dataset import TextVideoDataset
from data_loader.transforms import init_transform_dict, init_video_transform_dict

class MomentQueries(TextVideoDataset):
    def _load_metadata(self):
        split_files = {
            'train': 'moments_train.json',
            'val': 'moments_val.json',            # there is no test
            'test': 'moments_test_unannotated.json'
        }
        target_split_fp = split_files[self.split]

        ann_file = os.path.join(self.meta_dir, target_split_fp)
        with open(ann_file) as f:
            anno_json = json.load(f)

        self.metadata = pd.DataFrame(columns=['video_uid', 'clip_uid',
                                              'video_start_sec', 'video_end_sec',
                                              'query'])

        for anno_video in anno_json["videos"]:
            for anno_clip in anno_video["clips"]:
                clip_times = float(anno_clip["video_start_sec"]), float(anno_clip["video_end_sec"])
                clip_duration = clip_times[1] - clip_times[0]
                new = pd.DataFrame({
                    'video_uid': anno_video['video_uid'],
                    'clip_uid': anno_clip['clip_uid'],
                    'video_start_sec': clip_times[0],
                    'video_end_sec': clip_times[1]}, index=[1])
                self.metadata = self.metadata.append(new, ignore_index=True)

        self.transforms = init_video_transform_dict()['test']

    def _get_video_path(self, sample):
        rel_video_fp = sample[0]
        full_video_fp = os.path.join(self.data_dir, rel_video_fp + '.mp4')
        return full_video_fp, rel_video_fp

    def _get_caption(self, sample):
        caption = sample['query']
        return caption

    def __getitem__(self, item):
        sample = self.metadata.iloc [item]
        video_fp, rel_fp = self._get_video_path(sample)

        fps = 1.87
        try:
            imgs, idxs = self.video_reader(video_fp, sample[2]*30, sample[3]*30,
                                               (sample[3]-sample[2]) * fps * self.video_params['num_frames'])
        except:
            print(f"Warning: missing video file {video_fp}.")

        if self.transforms is not None:
            imgs = imgs.transpose(0, 1)  # [T, C, H, W] ---> [C, T, H, W]
            imgs = self.transforms(imgs)
            imgs = imgs.transpose(0, 1)  # recover

        meta_arr = {'video_uid': sample[0], 'clip_uid': sample[1], 'dataset': self.dataset_name}
        data = {'video': imgs, 'meta' : meta_arr}
        return data

if __name__ == "__main__":
    split = 'train'
    kwargs = dict(
        dataset_name="Ego4d_MQ",
        text_params={
            "input": "text"
        },
        video_params={
        "input_res": 224,
        "num_frames": 4,
        "loading": "lax"
        },
        data_dir="/apdcephfs/private_qinghonglin/video_dataset/ego4d_256/data",
        meta_dir="/apdcephfs/private_qinghonglin/video_dataset/ego4d/benchmark_splits/mq/",
        tsfms=init_video_transform_dict()['test'],
        reader='decord_start_end',
        split=split,
    )
    dataset = MomentQueries(**kwargs)
    print(len(dataset))
    # for i in range(1000):
    #     item = dataset[i]
    #     # print(item.keys())
    #     print(item)

EgoVLPv2/data_loader/Ego4D_MQ_dataset.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import os
import sys
import json
import pdb
import pandas as pd
from base.base_dataset import TextVideoDataset
from data_loader.transforms import init_transform_dict, init_video_transform_dict

import torch
from PIL import Image
from torchvision import transforms

class EgoClip_EgoMCQ(TextVideoDataset):
    def _load_metadata(self):
        split_files = {
            'train': 'egoclip.csv',
            'val': 'egomcq.json',
            'test': 'egomcq.json'
        }
        target_split_fp = split_files[self.split]

        self.chunk_sec = 600  # Each segment is up to 600s
        self.noun_dim = 582  # num of nouns of ego4d taxonomy dictionary
        self.verb_dim = 118  # num of verbs of ego4d taxonomy dictionary

        if self.split == 'train':
            self.metadata = pd.read_csv(os.path.join(self.meta_dir, target_split_fp), sep='\t',error_bad_lines=False)
            self.frame_sample = 'rand'

            if self.neg_param:
                self.metadata['chunk_id'] = self.metadata['narration_time'] // self.neg_param
                self.metadata['chunk_id'] = self.metadata['chunk_id'].astype(str)
                self.metadata['segment_id'] = self.metadata['video_uid'] + '_' + self.metadata['chunk_id']

        elif self.split in ['val', 'test']:
            self.frame_sample = 'uniform'
            with open(os.path.join(self.meta_dir, target_split_fp), 'r') as load_f:
                self.metadata = json.load(load_f)

    def _get_video_path(self, sample):
        video_uid = sample['video_uid']
        video_start_sec = max(float(sample['clip_start']), 0)
        video_end_sec   = max(float(sample['clip_end']), 0)

        chunk_start_id = int(video_start_sec // self.chunk_sec)
        chunk_end_id = int(video_end_sec // self.chunk_sec)

        full_video_start_fp = os.path.join(self.data_dir, video_uid, str(chunk_start_id) + ".mp4")
        full_video_end_fp = os.path.join(self.data_dir, video_uid, str(chunk_end_id) + ".mp4")

        video_fp = [full_video_start_fp, full_video_end_fp]
        video_sec = [video_start_sec, video_end_sec]
        bound_sec = (chunk_start_id + 1) * self.chunk_sec
        return video_fp, video_sec, bound_sec

    def _get_video_frames(self, video_fp, video_sec, bound_sec):
        video_loading = self.video_params.get('loading', 'strict')
        try:
            if os.path.isfile(video_fp[0]) and os.path.isfile(video_fp[1]):
                imgs, idxs = self.video_reader(video_fp[0], video_fp[1], self.video_params['num_frames'], self.frame_sample,
                                               start_sec=video_sec[0], end_sec=video_sec[1], bound_sec=bound_sec)
            else:
                print(f"Warning: missing video file {video_fp}.")
                assert False
        except Exception as e:
            if video_loading == 'strict':
                raise ValueError(
                    f'Video loading failed for {video_fp}, video loading for this dataset is strict.') from e
            else:
                imgs = Image.new('RGB', (self.video_params['input_res'], self.video_params['input_res']), (0, 0, 0))
                imgs = transforms.ToTensor()(imgs).unsqueeze(0)

        if self.transforms is not None:
            if self.video_params['num_frames'] > 1:
                imgs = imgs.transpose(0, 1)  # [T, C, H, W] ---> [C, T, H, W]
                imgs = self.transforms(imgs)
                imgs = imgs.transpose(0, 1)  # recover
            else:
                imgs = self.transforms(imgs)

        final = torch.zeros([self.video_params['num_frames'], 3, self.video_params['input_res'],
                             self.video_params['input_res']])
        final[:imgs.shape[0]] = imgs
        return final

    def _get_caption(self, sample):
        noun_vec = torch.zeros(self.noun_dim)
        verb_vec = torch.zeros(self.verb_dim)
        noun_idx = eval(sample['tag_noun'])
        verb_idx = eval(sample['tag_verb'])
        for i in noun_idx:
            noun_vec[i] = 1
        for i in verb_idx:
            verb_vec[i] = 1

        return sample['clip_text'], noun_vec, verb_vec

    def _get_train_item(self, item):
        item = item % len(self.metadata)
        sample = self.metadata.iloc[item]
        video_fp, video_sec, bound_sec = self._get_video_path(sample)
        caption, noun_vec, verb_vec = self._get_caption(sample)
        final = self._get_video_frames(video_fp, video_sec, bound_sec)

        # Scene-aware negative sampling
        if self.neg_param:
            # sample_neg = self.metadata[(self.metadata.video_uid==sample.video_uid)].sample(1).iloc[0] # variant of negative sample from same video
            sample_neg = self.metadata[self.metadata.segment_id==sample.segment_id].sample(1).iloc[0]
            video_fp_neg, video_sec_neg, bound_sec_neg = self._get_video_path(sample_neg)
            caption_neg, noun_vec_neg, verb_vec_neg = self._get_caption(sample_neg)
            final_neg = self._get_video_frames(video_fp_neg, video_sec_neg, bound_sec_neg)

        meta_arr = {'raw_captions': caption, 'paths': video_fp, 'dataset': self.dataset_name}
        if self.neg_param:
            return {'video': final, 'text': caption,
                    'video_neg': final_neg, 'text_neg': caption_neg,
                    'meta': meta_arr,
                    'noun_vec': noun_vec, 'verb_vec': verb_vec,
                    'noun_vec_neg': noun_vec_neg, 'verb_vec_neg': verb_vec_neg}
        else:
            return {'video': final, 'text': caption,
                'meta': meta_arr,
                'noun_vec': noun_vec, 'verb_vec': verb_vec}

    def _get_val_item(self, item):
        item = item % len(self.metadata)
        itemMCQ = self.metadata[str(item)]

        answerIndex = itemMCQ['answer']
        sampleQuery = itemMCQ['query']
        textQuery, _, _ = self._get_caption(sampleQuery)

        sampleOptions = itemMCQ['choices']
        num_options = len(sampleOptions)
        textOptions = []
        videoOptions = torch.zeros([num_options, self.video_params['num_frames'], 3, self.video_params['input_res'],
                             self.video_params['input_res']])

        for id, option in enumerate(sampleOptions):
            sampleOptioni = sampleOptions[option]
            video_fp, video_sec, bound_sec = self._get_video_path(sampleOptioni)
            caption, _, _ = self._get_caption(sampleOptioni)
            textOptions.append(caption)

            imgs = self._get_video_frames(video_fp, video_sec, bound_sec)
            videoOptions[id] = imgs

        type =  itemMCQ['types']    # 1 for inter; 2 for intra
        data = {'video': videoOptions, 'text': textQuery, 'text_ops':textOptions, 'correct': answerIndex, 'type': type}
        return data

    def __len__(self):
        return len(self.metadata)

    def __getitem__(self, item):
        if self.split == 'train':
            return self._get_train_item(item)
        elif self.split in ['val', 'test']:
            return self._get_val_item(item)

if __name__ == "__main__":
    kwargs = dict(
        dataset_name="EgoClip_dataset",
        text_params={
            "input": "text"
        },
        video_params={
        "input_res": 224,
        "num_frames": 4,
        "loading": "lax"
        },
        data_dir="/fsx/spraman3/Video_Language_Pretraining/Pre-training/EgoVLP/dataset/ego4d_chunked_multiprocessing",
        meta_dir="/fsx/spraman3/Video_Language_Pretraining/Pre-training",
        tsfms=init_video_transform_dict()['test'],
        reader='cv2_egoclip',
        split='train',
        neg_param=60
    )
    dataset = EgoClip_EgoMCQ(**kwargs)
    for i in range(100):
        item = dataset[i]
        pdb.set_trace() 
        print(item.keys())

EgoVLPv2/data_loader/EgoClip_EgoMCQ_dataset.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import os
import sys
import random
import pickle
import numpy as np
import pandas as pd
from base.base_dataset import TextVideoDataset
from data_loader.transforms import init_transform_dict, init_video_transform_dict

import torch
from PIL import Image
from torchvision import transforms
import decord
import torchvision.transforms as transforms
import torchvision.transforms._transforms_video as transforms_video


class MultiInstanceRetrieval(TextVideoDataset):
    def _load_metadata(self):
        split_files = {
            'train': 'EPIC_100_retrieval_train.csv',
            'val': 'EPIC_100_retrieval_test.csv',            # there is no test
            'test': 'EPIC_100_retrieval_test.csv'
        }
        split_files_sentence = {
            'train': 'EPIC_100_retrieval_train_sentence.csv',
            'val': 'EPIC_100_retrieval_test_sentence.csv',  # there is no test
            'test': 'EPIC_100_retrieval_test_sentence.csv'
        }
        target_split_fp = split_files[self.split]
        metadata = pd.read_csv(os.path.join(self.meta_dir, target_split_fp))

        target_split_sentence_fp = split_files_sentence[self.split]
        metadata_sentence = pd.read_csv(os.path.join(self.meta_dir, target_split_sentence_fp))

        if self.split == 'train':
            path_relevancy = os.path.join(self.meta_dir, 'relevancy/caption_relevancy_EPIC_100_retrieval_train.pkl')
        elif self.split in ['val', 'test']:
            path_relevancy = os.path.join(self.meta_dir, 'relevancy/caption_relevancy_EPIC_100_retrieval_test.pkl')

        pkl_file = open(path_relevancy, 'rb')
        self.relevancy = 0.1
        self.relevancy_mat = pickle.load(pkl_file)

        self.metadata = metadata
        self.metadata_sentence = metadata_sentence

    def _get_video_path(self, sample):

        pid, vid = sample[1:3]
        rel_video_fp = '{}/{}.MP4'.format(pid, vid)
        full_video_fp = os.path.join(self.data_dir, rel_video_fp)

        return full_video_fp, rel_video_fp

    def _get_caption(self, idx, sample):
        # return sentence, relevancy score, idx
        if self.split == 'train':
            positive_list = np.where(self.relevancy_mat[idx] > self.relevancy)[0].tolist()
            if positive_list != []:
                pos = random.sample(positive_list, min(len(positive_list), 1))[0]
                if pos < len(self.metadata_sentence) and pos < self.relevancy_mat.shape[1]:
                    return self.metadata_sentence.iloc[pos][1], self.relevancy_mat[idx][pos], pos
            return sample[8], 1, 0

        elif self.split in ['val', 'test']:
            return sample[8], 1, -1

    def get_frame_ids(self, start_frame, end_frame, num_segments=32, jitter=True):
        seg_size = float(end_frame - start_frame - 1) / num_segments
        seq = []
        for i in range(num_segments):
            start = int(np.round(seg_size * i) + start_frame)
            end = int(np.round(seg_size * (i + 1)) + start_frame)
            end = min(end, end_frame)
            if jitter:
                frame_id = np.random.randint(low=start, high=(end + 1))
            else:
                frame_id = (start + end) // 2
            seq.append(frame_id)
        return seq

    def video_loader_by_frames(self, vid, frame_ids):
        vr = decord.VideoReader(vid)
        try:
            frames = vr.get_batch(frame_ids).numpy()
            frames = [torch.tensor(frame, dtype=torch.float32) for frame in frames]
        except (IndexError, decord.DECORDError) as error:
            print(error)
            print("Erroneous video: ", vid)
            raise ValueError()
            frames = [torch.zeros((240, 320, 3)) for _ in range(len(frame_ids))]
        return torch.stack(frames, dim=0)


    def datetime2sec(self, st):
        hh, mm, ss = st.split(':')
        return int(hh) * 3600 + int(mm) * 60 + float(ss)

    def __getitem__(self, item):
        item = item % len(self.metadata)
        sample = self.metadata.iloc[item]
        video_fp, rel_fp = self._get_video_path(sample)
        caption, relation, idx = self._get_caption(item, sample)

        video_loading = self.video_params.get('loading', 'strict')
        start_timestamp, end_timestamp = self.datetime2sec(sample[4]), self.datetime2sec(sample[5])
        fps = decord.VideoReader(video_fp).get_avg_fps()
        start_frame = int(np.round(fps * start_timestamp))
        end_frame = int(np.ceil(fps * end_timestamp))

        #start_frame, stop_frame = int(sample[6]), int(sample[7])

        frame_sample = 'rand'
        if self.split in ['test', 'val']:
            frame_sample = 'uniform'
        fix_start = None

        if True:
            if os.path.exists(video_fp):
                # only supported for read_frames_decord_online
            #    imgs, idxs = self.video_reader(video_fp, start_frame, stop_frame, self.video_params['num_frames'], frame_sample,
            #                                   fix_start=fix_start)
                frame_ids = self.get_frame_ids(start_frame, end_frame, num_segments=self.video_params['num_frames'], jitter=(self.split == 'train'))
                imgs = self.video_loader_by_frames(video_fp, frame_ids)
                #print("video_fp: ", video_fp)
                #print("imgs.shape: ", imgs.shape)
                #print(abcd)
            else:
                print(f"Warning: missing video file {video_fp}.")
                assert False
        #except Exception as e:
        #    if video_loading == 'strict':
        #        raise ValueError(
        #            f'Video loading failed for {video_fp}, video loading for this dataset is strict.') from e
        #    else:
        #        imgs = Image.new('RGB', (self.video_params['input_res'], self.video_params['input_res']), (0, 0, 0))
        #        imgs = transforms.ToTensor()(imgs).unsqueeze(0)


        if self.split in ['test', 'val']:
            crop_size = self.video_params["input_res"] 
            self.transforms = transforms.Compose([
                transforms.Resize(crop_size),
                transforms.CenterCrop(crop_size),
                transforms_video.NormalizeVideo(mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375]),
                ])
        else:
            crop_size = self.video_params["input_res"]
            self.transforms = transforms.Compose([
                transforms.RandomResizedCrop(crop_size, scale=(0.5, 1.0)),
                transforms_video.NormalizeVideo(mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375]),
                ])


        if self.transforms is not None:
            if self.video_params['num_frames'] > 1:
                #imgs = imgs.transpose(0, 1)  # [T, C, H, W] ---> [C, T, H, W]
                imgs = imgs.permute(3,0,1,2) # T H W C -> C T H W
                imgs = self.transforms(imgs)
                imgs = imgs.permute(1,0,2,3)
                #imgs = imgs.transpose(0, 1)  # recover
            else:
                imgs = self.transforms(imgs)

        #final = torch.zeros([self.video_params['num_frames'], 3, self.video_params['input_res'], self.video_params['input_res']])
        #final[:imgs.shape[0]] = imgs

        meta_arr = {'raw_captions': caption, 'paths': item, 'dataset': self.dataset_name}
        data = {'video': imgs, 'text': caption, 'meta': meta_arr, 'relation': relation, 'item_v': item, 'item_t': idx}
        return data

if __name__ == "__main__":
    kwargs = dict(
        dataset_name="MultiInstanceRetrieval",
        text_params={
            "input": "text"
        },
        video_params={
        "input_res": 224,
        "num_frames": 4,
        "loading": "lax"
        },
        data_dir="/apdcephfs/private_qinghonglin/video_dataset/epic-kitchens/epic-kitchens-rgb-frames",
        meta_dir="/apdcephfs/private_qinghonglin/video_dataset/epic-kitchens/epic-kitchens-100-annotations-master/retrieval_annotations",
        tsfms=init_video_transform_dict()['test'],
        reader='cv2_epic',
        split='train'
    )
    dataset = MultiInstanceRetrieval(**kwargs)
    for i in range(100):
        item = dataset[i]
        print(item.keys())

EgoVLPv2/data_loader/EpicKitchens_MIR_dataset.py
facebookresearch/EgoVLPv2
from .data_loader import *
from .EgoClip_EgoMCQ_dataset import *
from .EpicKitchens_MIR_dataset import *
from .CharadesEgo_dataset import *
from .Ego4D_MQ_dataset import *
from .transforms import *

EgoVLPv2/data_loader/__init__.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import pdb
import sys

from base import BaseDataLoaderExplicitSplit, BaseMultiDataLoader, \
    DistBaseDataLoaderExplicitSplit, MultiDistBaseDataLoaderExplicitSplit
from data_loader.EgoClip_EgoMCQ_dataset import EgoClip_EgoMCQ
from data_loader.EpicKitchens_MIR_dataset import MultiInstanceRetrieval
from data_loader.CharadesEgo_dataset import CharadesEgo
from data_loader.Ego4D_MQ_dataset import MomentQueries
from data_loader.transforms import init_transform_dict, init_video_transform_dict

def dataset_loader(dataset_name,
                   text_params,
                   video_params,
                   data_dir,
                   meta_dir=None,
                   split='train',
                   tsfms=None,
                   cut=None,
                   subsample=1,
                   sliding_window_stride=-1,
                   reader='decord',
                   neg_param=None):
    kwargs = dict(
        dataset_name=dataset_name,
        text_params=text_params,
        video_params=video_params,
        data_dir=data_dir,
        meta_dir=meta_dir,
        split=split,
        tsfms=tsfms,
        cut=cut,
        subsample=subsample,
        sliding_window_stride=sliding_window_stride,
        reader=reader,
        neg_param=neg_param,
    )

    # TODO: change to...
    #  dataset = globals()[dataset_name]
    #  ...is this safe / or just lazy?
    if dataset_name == "EgoClip":
        dataset = EgoClip_EgoMCQ(**kwargs)
    elif dataset_name == "EpicKitchens_MIR":
        dataset = MultiInstanceRetrieval(**kwargs)
    elif dataset_name == "CharadesEgo":
        dataset = CharadesEgo(**kwargs)
    elif dataset_name == "Ego4D_MQ":
        dataset = MomentQueries(**kwargs)
    else:
        raise NotImplementedError(f"Dataset: {dataset_name} not found.")

    return dataset


class TextVideoDataLoader(BaseDataLoaderExplicitSplit):
    def __init__(self,
                 dataset_name,
                 text_params,
                 video_params,
                 data_dir,
                 meta_dir=None,
                 split='train',
                 tsfm_params=None,
                 tsfm_split=None,
                 cut=None,
                 subsample=1,
                 sliding_window_stride=-1,
                 reader='decord',
                 neg_param=None,
                 batch_size=1,
                 num_workers=1,
                 shuffle=True):
        if tsfm_params is None:
            tsfm_params = {}
        if video_params['num_frames'] > 1:
            # video data can not do flip, crop aug
            tsfm_dict = init_video_transform_dict(**tsfm_params)
        else:
            tsfm_dict = init_transform_dict(**tsfm_params)
        if tsfm_split is None:
            tsfm_split = split
        tsfm = tsfm_dict[tsfm_split]
        dataset = dataset_loader(dataset_name, text_params, video_params, data_dir, meta_dir, split, tsfm, cut,
                                 subsample, sliding_window_stride, reader, neg_param)

        super().__init__(dataset, batch_size, shuffle, num_workers)
        self.dataset_name = dataset_name

class DistTextVideoDataLoader(DistBaseDataLoaderExplicitSplit):
    def __init__(self,
                 dataset_name,
                 text_params,
                 video_params,
                 data_dir,
                 meta_dir=None,
                 split='train',
                 tsfm_params=None,
                 tsfm_split=None,
                 cut=None,
                 subsample=1,
                 sliding_window_stride=-1,
                 reader='cv2',
                 neg_param=None,
                 batch_size=1,
                 num_workers=1,
                 shuffle=True):
        if tsfm_params is None:
            tsfm_params = {}
        if video_params['num_frames'] > 1:
            # video data can not do flip, crop aug
            tsfm_dict = init_video_transform_dict(**tsfm_params)
        else:
            tsfm_dict = init_transform_dict(**tsfm_params)

        # Updated
        if tsfm_split is None:
            tsfm_split = split
        tsfm = tsfm_dict[tsfm_split]

        dataset = dataset_loader(dataset_name, text_params, video_params, data_dir, meta_dir, split, tsfm, cut,
                                 subsample, sliding_window_stride, reader, neg_param)
        super().__init__(dataset, batch_size, shuffle, num_workers)
        self.dataset_name = dataset_name

class MultiDistTextVideoDataLoader(MultiDistBaseDataLoaderExplicitSplit):
    def __init__(self,
                 args,
                 dataset_name,
                 text_params,
                 video_params,
                 data_dir,
                 meta_dir=None,
                 split='train',
                 tsfm_params=None,
                 tsfm_split=None,
                 cut=None,
                 subsample=1,
                 sliding_window_stride=-1,
                 reader='cv2',
                 neg_param=None,
                 batch_size=1,
                 num_workers=1,
                 shuffle=True):
        if tsfm_params is None:
            tsfm_params = {}
        if video_params['num_frames'] > 1:
            # video data can not do flip, crop aug
            tsfm_dict = init_video_transform_dict(**tsfm_params)
        else:
            tsfm_dict = init_transform_dict(**tsfm_params)

        if tsfm_split is None:
            tsfm_split = split
        tsfm = tsfm_dict[tsfm_split]

        dataset = dataset_loader(dataset_name, text_params, video_params, data_dir, meta_dir, split, tsfm, cut,
                                 subsample, sliding_window_stride, reader, neg_param)
        super().__init__(args, dataset, batch_size, shuffle, num_workers)
        self.dataset_name = dataset_name

class TextVideoMultiDataLoader(BaseMultiDataLoader):
    # TODO: figure out neat way to have N data_loaders
    # TODO: also add N weighted sampler
    def __init__(self, data_loader1, data_loader2):
        # get class from "type" in dict
        dls_cfg = [data_loader1, data_loader2]
        dls = []
        for dcfg in dls_cfg:
            dl = globals()[dcfg['type']](**dcfg['args'])
            dls.append(dl)
        super().__init__(dls)

EgoVLPv2/data_loader/data_loader.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


from torchvision import transforms
from torchvision.transforms._transforms_video import RandomCropVideo, RandomResizedCropVideo,CenterCropVideo, NormalizeVideo,ToTensorVideo,RandomHorizontalFlipVideo
import warnings
warnings.filterwarnings("ignore")

def init_transform_dict(input_res=224,
                        center_crop=256,
                        randcrop_scale=(0.5, 1.0),
                        color_jitter=(0, 0, 0),
                        norm_mean=(0.485, 0.456, 0.406),
                        norm_std=(0.229, 0.224, 0.225)):
    normalize = transforms.Normalize(mean=norm_mean, std=norm_std)
    tsfm_dict = {
        'train': transforms.Compose([
            transforms.RandomResizedCrop(input_res, scale=randcrop_scale),
            transforms.RandomHorizontalFlip(),
            transforms.ColorJitter(brightness=color_jitter[0], saturation=color_jitter[1], hue=color_jitter[2]),
            normalize,
        ]),
        'val': transforms.Compose([
            transforms.Resize(center_crop),
            transforms.CenterCrop(center_crop),
            transforms.Resize(input_res),
            normalize,
        ]),
        'test': transforms.Compose([
            transforms.Resize(center_crop),
            transforms.CenterCrop(center_crop),
            transforms.Resize(input_res),
            normalize,
        ])
    }
    return tsfm_dict

def init_video_transform_dict(input_res=224,
                        center_crop=256,
                        randcrop_scale=(0.5, 1.0),
                        color_jitter=(0, 0, 0),
                        norm_mean=(0.485, 0.456, 0.406),
                        norm_std=(0.229, 0.224, 0.225)):
    print('Video Transform is used!')
    normalize = NormalizeVideo(mean=norm_mean, std=norm_std)
    tsfm_dict = {
        'train': transforms.Compose([
            RandomResizedCropVideo(input_res, scale=randcrop_scale),
            RandomHorizontalFlipVideo(),
            transforms.ColorJitter(brightness=color_jitter[0], saturation=color_jitter[1], hue=color_jitter[2]),
            normalize,
        ]),
        'val': transforms.Compose([
            transforms.Resize(center_crop),
            transforms.CenterCrop(center_crop),
            transforms.Resize(input_res),
            normalize,
        ]),
        'test': transforms.Compose([
            transforms.Resize(center_crop),
            transforms.CenterCrop(center_crop),
            transforms.Resize(input_res),
            normalize,
        ])
    }
    return tsfm_dict

EgoVLPv2/data_loader/transforms.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

from .logger import *
from .visualization import *

EgoVLPv2/logger/__init__.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import logging
import logging.config
from pathlib import Path

from utils import read_json


def setup_logging(save_dir, log_config='logger/logger_config.json', default_level=logging.INFO):
    """
    Setup logging configuration
    """
    log_config = Path(log_config)
    if log_config.is_file():
        config = read_json(log_config)
        # modify logging paths based on run config
        for _, handler in config['handlers'].items():
            if 'filename' in handler:
                handler['filename'] = str(save_dir / handler['filename'])

        logging.config.dictConfig(config)
    else:
        print("Warning: logging configuration file is not found in {}.".format(log_config))
        logging.basicConfig(level=default_level)

EgoVLPv2/logger/logger.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import importlib
from utils import Timer


class TensorboardWriter:
    def __init__(self, log_dir, logger, enabled):
        self.writer = None
        self.selected_module = ""

        if enabled:
            log_dir = str(log_dir)

            # Retrieve visualization writer.
            for module in ["torch.utils.tensorboard", "tensorboardX"]:
                try:
                    self.writer = importlib.import_module(module).SummaryWriter(log_dir)
                    succeeded = True
                    break
                except ImportError:
                    succeeded = False
                self.selected_module = module

            if not succeeded:
                message = "Warning: visualization (Tensorboard) is configured to use, but currently not installed on " \
                    "this machine. Please install either TensorboardX with 'pip install tensorboardx', upgrade " \
                    "PyTorch to version >= 1.1 for using 'torch.utils.tensorboard' or turn off the option in " \
                    "the 'config.json' file."
                logger.warning(message)

        self.step = 0
        self.mode = ''

        self.tb_writer_ftns = {
            'add_scalar', 'add_scalars', 'add_image', 'add_images', 'add_audio',
            'add_text', 'add_histogram', 'add_pr_curve', 'add_embedding'
        }
        self.tag_mode_exceptions = {'add_histogram', 'add_embedding'}
            
        self.timer = Timer()

    def set_step(self, step, mode='train'):
        self.mode = mode
        self.step = step
        if step == 0:
            self.timer.reset()
        else:
            duration = self.timer.check()
            self.add_scalar('steps_per_sec', 1 / duration)

    def __getattr__(self, name):
        """
        If visualization is configured to use:
            return add_data() methods of tensorboard with additional information (step, tag) added.
        Otherwise:
            return a blank function handle that does nothing
        """
        if name in self.tb_writer_ftns:
            add_data = getattr(self.writer, name, None)

            def wrapper(tag, data, *args, **kwargs):
                if add_data is not None:
                    # add mode(train/valid) tag
                    if name not in self.tag_mode_exceptions:
                        tag = '{}/{}'.format(tag, self.mode)
                    add_data(tag, data, self.step, *args, **kwargs)
            return wrapper
        else:
            # default action for returning methods defined in this class, set_step() for instance.
            try:
                attr = object.__getattr__(name)
            except AttributeError:
                raise AttributeError("type object '{}' has no attribute '{}'".format(self.selected_module, name))
            return attr


class SacredNeptuneWriter:
    def __init__(self):
        raise NotImplementedError

EgoVLPv2/logger/visualization.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import torch
import torch.nn as nn
import torch.nn.functional as F

from transformers.models.bert.modeling_bert import BertPredictionHeadTransform


class Pooler(nn.Module):
    def __init__(self, hidden_size):
        super().__init__()
        self.dense = nn.Linear(hidden_size, hidden_size)
        self.activation = nn.Tanh()

    def forward(self, hidden_states):
        first_token_tensor = hidden_states #[:, 0]
        pooled_output = self.dense(first_token_tensor)
        pooled_output = self.activation(pooled_output)
        return pooled_output


class ITMHead(nn.Module):
    def __init__(self, hidden_size):
        super().__init__()
        self.fc = nn.Linear(hidden_size, 2)

    def forward(self, x):
        x = self.fc(x)
        return x


class MLMHead(nn.Module):
    def __init__(self, config, weight=None):
        super().__init__()
        self.transform = BertPredictionHeadTransform(config)
        self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.bias = nn.Parameter(torch.zeros(config.vocab_size))
        if weight is not None:
            self.decoder.weight = weight

    def forward(self, x):
        x = self.transform(x)
        x = self.decoder(x) + self.bias
        return x

EgoVLPv2/model/heads.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import pdb
import torch
import torch.nn.functional as F
from torch import nn
import pickle

class NormSoftmaxLoss(nn.Module):
    def __init__(self, temperature=0.05):
        super().__init__()

        self.temperature = temperature

    def forward(self, x):
        "Assumes input x is similarity matrix of N x M \in [-1, 1], computed using the cosine similarity between normalised vectors"
        i_logsm = F.log_softmax(x/self.temperature, dim=1)
        j_logsm = F.log_softmax(x.t()/self.temperature, dim=1)

        # sum over positives
        idiag = torch.diag(i_logsm)
        loss_i = idiag.sum() / len(idiag)

        jdiag = torch.diag(j_logsm)
        loss_j = jdiag.sum() / len(jdiag)

        return - loss_i - loss_j, self.temperature

class EgoNCE(nn.Module):
    def __init__(self, temperature=0.05, noun=True, verb=True):
        super().__init__()
        self.noun = noun
        self.verb = verb
        self.temperature = temperature

    def forward(self, x, mask_v, mask_n):
        mask_diag = torch.eye(x.shape[0]).cuda()
        if self.noun and self.verb:
            mask = mask_v * mask_n + mask_diag
        elif self.noun:
            mask = mask_n + mask_diag
        else:
            mask = mask_v + mask_diag

        "Assumes input x is similarity matrix of N x M \in [-1, 1], computed using the cosine similarity between normalised vectors"
        i_sm = F.softmax(x/self.temperature, dim=1)
        j_sm = F.softmax(x.t()/self.temperature, dim=1)

        mask_bool = mask > 0
        idiag = torch.log(torch.sum(i_sm * mask_bool, dim=1) )
        loss_i = idiag.sum() / len(idiag)

        jdiag = torch.log(torch.sum(j_sm * mask_bool, dim=1) )
        loss_j = jdiag.sum() / len(jdiag)
        return - loss_i - loss_j, mask_bool, self.temperature
        #return - loss_i - loss_j


class MaxMarginRankingLoss(nn.Module):

    def __init__(self, margin=0.2, fix_norm=True):
        super().__init__()
        self.fix_norm = fix_norm
        self.loss = nn.MarginRankingLoss(margin)
        self.margin = margin

    def forward(self, x, weight=None):
        n = x.size()[0]

        x1 = torch.diag(x)
        x1 = x1.unsqueeze(1)
        x1 = x1.expand(n, n)
        x1 = x1.contiguous().view(-1, 1)
        x1 = torch.cat((x1, x1), 0)

        x2 = x.view(-1, 1)
        x3 = x.transpose(0, 1).contiguous().view(-1, 1)

        x2 = torch.cat((x2, x3), 0)
        max_margin = F.relu(self.margin - (x1 - x2))

        if self.fix_norm:
            # remove the elements from the diagonal
            keep = torch.ones(x.shape) - torch.eye(x.shape[0])  # 128 x 128
            keep1 = keep.view(-1, 1)
            keep2 = keep.transpose(0, 1).contiguous().view(-1, 1)
            keep_idx = torch.nonzero(torch.cat((keep1, keep2), 0).flatten()).flatten()
            if x1.is_cuda:
                keep_idx = keep_idx.cuda()
            x1_ = torch.index_select(x1, dim=0, index=keep_idx)
            x2_ = torch.index_select(x2, dim=0, index=keep_idx)
            max_margin = F.relu(self.margin - (x1_ - x2_))

        return max_margin.mean()

class AdaptiveMaxMarginRankingLoss(nn.Module):

    def __init__(self, margin=0.4, fix_norm=True):
        super().__init__()
        self.fix_norm = fix_norm
        self.loss = nn.MarginRankingLoss(margin)
        self.margin = margin

    def forward(self, x, weight=None):
        n = x.size()[0]

        x1 = torch.diag(x)
        x1 = x1.unsqueeze(1)
        x1 = x1.expand(n, n)
        x1 = x1.contiguous().view(-1, 1)
        x1 = torch.cat((x1, x1), 0)

        w1 = weight.unsqueeze(1)
        w1 = w1.expand(n, n)
        w1 = w1.contiguous().view(-1, 1)
        w1 = torch.cat((w1, w1), 0)

        x2 = x.view(-1, 1)
        x3 = x.transpose(0, 1).contiguous().view(-1, 1)

        x2 = torch.cat((x2, x3), 0)
        max_margin = F.relu(  w1 * self.margin - (x1 - x2))

        if self.fix_norm:
            # remove the elements from the diagonal
            keep = torch.ones(x.shape) - torch.eye(x.shape[0])  # 128 x 128
            keep1 = keep.view(-1, 1)
            keep2 = keep.transpose(0, 1).contiguous().view(-1, 1)
            keep_idx = torch.nonzero(torch.cat((keep1, keep2), 0).flatten()).flatten()
            if x1.is_cuda:
                keep_idx = keep_idx.cuda()
            x1_ = torch.index_select(x1, dim=0, index=keep_idx)
            w1_ = torch.index_select(w1, dim=0, index=keep_idx)
            x2_ = torch.index_select(x2, dim=0, index=keep_idx)
            max_margin =  F.relu( w1_ * self.margin - (x1_ - x2_))

        return max_margin.mean()

class CrossEntropy(nn.Module):
    def __init__(self):
        super().__init__()
        self.loss = nn.CrossEntropyLoss()

    def forward(self, output, target):
        return self.loss(output, target)

EgoVLPv2/model/loss.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


"""Module for computing performance metrics

"""
import math
import numbers
from pathlib import Path
#import ipdb
import numpy as np
import torch
import scipy.stats
from sklearn.metrics import average_precision_score
#import ipdb
#import pdb
import os
import pandas as pd
import pickle
from utils import nDCG, mAP

def t2v_metrics(sims, query_masks=None):
    """Compute retrieval metrics from a similiarity matrix.

    Args:
        sims (th.Tensor): N x M matrix of similarities between embeddings, where
             x_{i,j} = <text_embd[i], vid_embed[j]>
        query_masks (th.Tensor): mask any missing queries from the dataset (two videos
             in MSRVTT only have 19, rather than 20 captions)

    Returns:
        (dict[str:float]): retrieval metrics
    """
    assert sims.ndim == 2, "expected a matrix"
    num_queries, num_vids = sims.shape
    dists = -sims
    sorted_dists = np.sort(dists, axis=1)

    # The indices are computed such that they slice out the ground truth distances
    # from the psuedo-rectangular dist matrix
    queries_per_video = num_queries // num_vids
    gt_idx = [[np.ravel_multi_index([ii, jj], (num_queries, num_vids))
               for ii in range(jj * queries_per_video, (jj + 1) * queries_per_video)]
              for jj in range(num_vids)]
    gt_idx = np.array(gt_idx)
    gt_dists = dists.reshape(-1)[gt_idx.reshape(-1)]
    gt_dists = gt_dists[:, np.newaxis]
    rows, cols = np.where((sorted_dists - gt_dists) == 0)  # find column position of GT

    # --------------------------------
    # NOTE: Breaking ties
    # --------------------------------
    # We sometimes need to break ties (in general, these should occur extremely rarely,
    # but there are pathological cases when they can distort the scores, such as when
    # the similarity matrix is all zeros). Previous implementations (e.g. the t2i
    # evaluation function used
    # here: https://github.com/niluthpol/multimodal_vtt/blob/master/evaluation.py and
    # here: https://github.com/linxd5/VSE_Pytorch/blob/master/evaluation.py#L87) generally
    # break ties "optimistically".  However, if the similarity matrix is constant this
    # can evaluate to a perfect ranking. A principled option is to average over all
    # possible partial orderings implied by the ties. See # this paper for a discussion:
    #    McSherry, Frank, and Marc Najork,
    #    "Computing information retrieval performance measures efficiently in the presence
    #    of tied scores." European conference on information retrieval. Springer, Berlin,
    #    Heidelberg, 2008.
    # http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.145.8892&rep=rep1&type=pdf

    break_ties = "optimistically"
    #break_ties = "averaging"

    if rows.size > num_queries:
        assert np.unique(rows).size == num_queries, "issue in metric evaluation"
        if break_ties == "optimistically":
            _, idx = np.unique(rows, return_index=True)
            cols = cols[idx]
        elif break_ties == "averaging":
            # fast implementation, based on this code:
            # https://stackoverflow.com/a/49239335
            locs = np.argwhere((sorted_dists - gt_dists) == 0)

            # Find the split indices
            steps = np.diff(locs[:, 0])
            splits = np.nonzero(steps)[0] + 1
            splits = np.insert(splits, 0, 0)

            # Compute the result columns
            summed_cols = np.add.reduceat(locs[:, 1], splits)
            counts = np.diff(np.append(splits, locs.shape[0]))
            avg_cols = summed_cols / counts
            if False:
                print("Running slower code to verify rank averaging across ties")
                # slow, but more interpretable version, used for testing
                avg_cols_slow = [np.mean(cols[rows == idx]) for idx in range(num_queries)]
                assert np.array_equal(avg_cols, avg_cols_slow), "slow vs fast difference"
                print("passed num check")
            cols = avg_cols

    msg = "expected ranks to match queries ({} vs {}) "
    if cols.size != num_queries:
        #import ipdb;
        #ipdb.set_trace()
        pass
    assert cols.size == num_queries, msg

    if False:
        # overload mask to check that we can recover the scores for single-query
        # retrieval
        print("DEBUGGING MODE")
        query_masks = np.zeros_like(query_masks)
        query_masks[:, 0] = 1  # recover single query score

    if query_masks is not None:
        # remove invalid queries
        assert query_masks.size == num_queries, "invalid query mask shape"
        cols = cols[query_masks.reshape(-1).astype(np.bool)]
        assert cols.size == query_masks.sum(), "masking was not applied correctly"
        # update number of queries to account for those that were missing
        num_queries = query_masks.sum()

    if False:
        # sanity check against old logic for square matrices
        gt_dists_old = np.diag(dists)
        gt_dists_old = gt_dists_old[:, np.newaxis]
        _, cols_old = np.where((sorted_dists - gt_dists_old) == 0)
        assert np.array_equal(cols_old, cols), "new metric doesn't match"

    return cols2metrics(cols, num_queries)


def v2t_metrics(sims, query_masks=None):
    """Compute retrieval metrics from a similiarity matrix.

    Args:
        sims (th.Tensor): N x M matrix of similarities between embeddings, where
             x_{i,j} = <text_embd[i], vid_embed[j]>
        query_masks (th.Tensor): mask any missing captions from the dataset

    Returns:
        (dict[str:float]): retrieval metrics

    NOTES: We find the closest "GT caption" in the style of VSE, which corresponds
    to finding the rank of the closest relevant caption in embedding space:
    github.com/ryankiros/visual-semantic-embedding/blob/master/evaluation.py#L52-L56
    """
    # switch axes of text and video
    sims = sims.T

    if False:
        # experiment with toy example
        sims = np.ones((3, 3))
        sims[0, 0] = 2
        sims[1, 1:2] = 2
        sims[2, :] = 2
        query_masks = None

    assert sims.ndim == 2, "expected a matrix"
    num_queries, num_caps = sims.shape
    dists = -sims
    caps_per_video = num_caps // num_queries
    break_ties = "averaging"

    MISSING_VAL = 1E8
    query_ranks = []
    for ii in range(num_queries):
        row_dists = dists[ii, :]
        if query_masks is not None:
            # Set missing queries to have a distance of infinity.  A missing query
            # refers to a query position `n` for a video that had less than `n`
            # captions (for example, a few MSRVTT videos only have 19 queries)
            row_dists[np.logical_not(query_masks.reshape(-1))] = MISSING_VAL

        # NOTE: Using distance subtraction to perform the ranking is easier to make
        # deterministic than using argsort, which suffers from the issue of defining
        # "stability" for equal distances.  Example of distance subtraction code:
        # github.com/antoine77340/Mixture-of-Embedding-Experts/blob/master/train.py
        sorted_dists = np.sort(row_dists)

        min_rank = np.inf
        for jj in range(ii * caps_per_video, (ii + 1) * caps_per_video):
            if row_dists[jj] == MISSING_VAL:
                # skip rankings of missing captions
                continue
            ranks = np.where((sorted_dists - row_dists[jj]) == 0)[0]
            if break_ties == "optimistically":
                rank = ranks[0]
            elif break_ties == "averaging":
                # NOTE: If there is more than one caption per video, its possible for the
                # method to do "worse than chance" in the degenerate case when all
                # similarities are tied.  TODO(Samuel): Address this case.
                rank = ranks.mean()
            if rank < min_rank:
                min_rank = rank
        query_ranks.append(min_rank)
    query_ranks = np.array(query_ranks)

    # sanity check against old version of code
    if False:
        sorted_dists = np.sort(dists, axis=1)
        gt_dists_old = np.diag(dists)
        gt_dists_old = gt_dists_old[:, np.newaxis]
        rows_old, cols_old = np.where((sorted_dists - gt_dists_old) == 0)
        if rows_old.size > num_queries:
            _, idx = np.unique(rows_old, return_index=True)
            cols_old = cols_old[idx]
        num_diffs = (1 - (cols_old == query_ranks)).sum()
        msg = f"new metric doesn't match in {num_diffs} places"
        assert np.array_equal(cols_old, query_ranks), msg

        # visualise the distance matrix
        import sys
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
        sys.path.insert(0, str(Path.home() / "coding/src/zsvision/python"))
        from zsvision.zs_iterm import zs_dispFig  # NOQA
        plt.matshow(dists)
        zs_dispFig()

    return cols2metrics(query_ranks, num_queries)

def egomcq_accuracy_metrics_ensemble(preds, labels, types):
    metrics = {}
    type_list = torch.unique(types)
    group_list = ["Intra-video", "Inter-video"]
    for type_i, group_i in zip(type_list, group_list):
        correct = 0
        total = 0
        for pred, label, type in zip(preds, labels, types):
            if type == type_i:
                pred_ = torch.argmax(pred)
                if pred_.item() == label.item():
                    correct += 1
                total += 1
        accuracy = correct/total
        # metrics['accuracy_p' + str(type_i.item())] = accuracy * 100
        metrics[group_i] = accuracy * 100
    return metrics

def egomcq_accuracy_metrics_vtm(preds, labels, types):
    metrics = {}
    type_list = torch.unique(types)
    group_list = ["Intra-video", "Inter-video"]
    for type_i, group_i in zip(type_list, group_list):
        correct = 0
        total = 0
        for pred, label, type in zip(preds, labels, types):
            if type == type_i:
                pred_ = torch.argmax(pred)
                if pred_.item() == label.item():
                    correct += 1
                total += 1
        accuracy = correct/total
        # metrics['accuracy_p' + str(type_i.item())] = accuracy * 100
        metrics[group_i] = accuracy * 100
    return metrics


def initialise_nDCG_values(relevancy_matrix):
    vis_k_counts = nDCG.calculate_k_counts(relevancy_matrix)
    txt_k_counts = nDCG.calculate_k_counts(relevancy_matrix.T)

    vis_IDCG = nDCG.calculate_IDCG(relevancy_matrix, vis_k_counts)
    txt_IDCG = nDCG.calculate_IDCG(relevancy_matrix.T, txt_k_counts)

    k_counts_dict = {'v': vis_k_counts, 't': txt_k_counts}
    IDCG_dict = {'v': vis_IDCG, 't': txt_IDCG}

    return IDCG_dict, k_counts_dict

def initialise_jpose_nDCG_values(relevancy_matrix):
    action_IDCG, action_k_values = initialise_nDCG_values(relevancy_matrix)

    dataset = {}
    dataset['action'] = {}
    dataset['action']['IDCG'] = action_IDCG
    dataset['action']['k_values'] = action_k_values
    return dataset

def mir_metrics_vtc(similarity_matrix, idx_arr):
    # considered unique narrations for evaluation of EPIC
    metrics = {}

    path_dataframes = '/cis/home/shraman/EPIC-KITCHENS_rgbframes/retrieval_annotations'
    video_id=pd.read_csv(os.path.join(path_dataframes , "EPIC_100_retrieval_test.csv")).values[:,0]
    text_id=pd.read_csv(os.path.join(path_dataframes , "EPIC_100_retrieval_test_sentence.csv")).values[:,0]
    similarity_matrix = (similarity_matrix + 1) / 2

    indexes=[]
    for elem in text_id:
        try:
            indexes.append(video_id.tolist().index(elem))
        except:
            print(f"error happened when index of {elem}.")

    order = []
    idx_arr_list = idx_arr.tolist()
    for i in range(len(video_id)):
        order.append(idx_arr_list.index(i))
    similarity_matrix = similarity_matrix[order,:][:, order]

    similarity_matrix = similarity_matrix.T[:,indexes]

    path_relevancy = "/cis/home/shraman/EPIC-KITCHENS_rgbframes/retrieval_annotations/relevancy/caption_relevancy_EPIC_100_retrieval_test.pkl"
    pkl_file = open(path_relevancy, 'rb')
    relevancy = pickle.load(pkl_file)

    dataset = initialise_jpose_nDCG_values(relevancy)
    vis_nDCG = nDCG.calculate_nDCG(similarity_matrix,relevancy, dataset['action']['k_values']['v'],IDCG=dataset['action']['IDCG']['v'])
    txt_nDCG = nDCG.calculate_nDCG(similarity_matrix.T,relevancy.T, dataset['action']['k_values']['t'],IDCG=dataset['action']['IDCG']['t'])
    metrics['nDCG_V2T'] = vis_nDCG * 100
    metrics['nDCG_T2V'] = txt_nDCG * 100
    metrics['nDCG_AVG'] = 100 * (vis_nDCG + txt_nDCG) / 2

    vis_mAP = mAP.calculate_mAP(similarity_matrix,
                                relevancy)
    txt_mAP = mAP.calculate_mAP(similarity_matrix.T,
                                relevancy.T)
    metrics['mAP_V2T'] = vis_mAP * 100
    metrics['mAP_T2V'] = txt_mAP * 100
    metrics['mAP_AVG'] =  100 * (vis_mAP + txt_mAP) / 2
    return metrics

def map(submission_array, gt_array):
    """ Returns mAP, weighted mAP, and AP array """
    m_aps = []
    n_classes = submission_array.shape[1]
    for oc_i in range(n_classes):
        sorted_idxs = np.argsort(-submission_array[:, oc_i])
        tp = gt_array[:, oc_i][sorted_idxs] == 1
        fp = np.invert(tp)
        n_pos = tp.sum()
        if n_pos < 0.1:
            m_aps.append(float('nan'))
            continue
        fp.sum()
        f_pcs = np.cumsum(fp)
        t_pcs = np.cumsum(tp)
        prec = t_pcs / (f_pcs+t_pcs).astype(float)
        avg_prec = 0
        for i in range(submission_array.shape[0]):
            if tp[i]:
                avg_prec += prec[i]
        m_aps.append(avg_prec / n_pos.astype(float))
    m_aps = np.array(m_aps)
    #m_ap = np.mean(m_aps)
    m_ap = np.nanmean(m_aps)
    w_ap = (m_aps * gt_array.sum(axis=0) / gt_array.sum().sum().astype(float))
    return m_ap, w_ap, m_aps

def charades_metrics_vtc(submission_array, gt_array):
    """
    Approximate version of the charades evaluation function
    For precise numbers, use the submission file with the official matlab script
    """
    metrics = {}
    fix = submission_array.copy()
    empty = np.sum(gt_array, axis=1) == 0
    fix[empty, :] = np.NINF
    m_ap, w_ap, m_aps = map(fix, gt_array)
    metrics['mAP'] = m_ap
    # metrics['wAP'] = w_ap
    # metrics['mAPs'] = m_aps
    return metrics


def charades_metrics_vtm(submission_array, gt_array):
    """
    Approximate version of the charades evaluation function
    For precise numbers, use the submission file with the official matlab script
    """
    metrics = {}
    fix = submission_array.copy()
    empty = np.sum(gt_array, axis=1) == 0
    fix[empty, :] = np.NINF
    m_ap, w_ap, m_aps = map(fix, gt_array)
    metrics['mAP'] = m_ap
    # metrics['wAP'] = w_ap
    # metrics['mAPs'] = m_aps
    return metrics


def oscc_metrics(preds, labels):
    metrics = {}
    correct = 0
    total = 0
    for pred, label in zip(preds, labels):
        pred_ = torch.argmax(pred)
        if pred_.item() == label.item():
            correct += 1
        total += 1
    accuracy = correct/total
    metrics['accuracy'] = accuracy * 100
    return metrics

def pnr_metrics(
    preds,
    labels,
    sc_labels,
    fps,
    parent_start_frames,
    parent_end_frames,
    parent_pnr_frames,
):
    metrics = {}
    distance_list = list()
    for pred, label, sc_label, \
        parent_start_frame, parent_end_frame, parent_pnr_frame, \
        ind_fps in zip(
        preds,
        labels,
        sc_labels,
        parent_start_frames,
        parent_end_frames,
        parent_pnr_frames,
        fps
    ):
        # import pdb; pdb.set_trace()
        if sc_label.item() == 1:
            keyframe_loc_pred = torch.argmax(pred).item()
            # import pdb; pdb.set_trace()
            keyframe_loc_pred_mapped = (parent_end_frame - parent_start_frame) / 16 * keyframe_loc_pred
            keyframe_loc_pred_mapped = keyframe_loc_pred_mapped.item()
            gt = parent_pnr_frame.item() - parent_start_frame.item()
            err_frame = abs(keyframe_loc_pred_mapped - gt)
            err_sec = err_frame/ind_fps.item()
            # print(keyframe_loc_pred_mapped, gt,  err_frame, err_sec)
            distance_list.append(err_sec)
    if len(distance_list) == 0:
        # If evaluating the trained model, use this
        # Otherwise, Lightning expects us to give a number.
        # Due to this, the Tensorboard graphs' results for keyframe distance
        # will be a little inaccurate.
        metrics['keyframe_distance'] = np.mean(0.0)
    metrics['keyframe_distance'] = np.mean(distance_list)
    # import pdb;
    # pdb.set_trace()
    return metrics

EgoVLPv2/model/metric.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import os
import pdb

import timm
import torch
import yaml
import torch.nn as nn
import numpy as np
import torch.nn.functional as F
from transformers import AutoModel
from einops import rearrange, repeat

from base import BaseModel
from model import video_transformer
from model.video_transformer import SpaceTimeTransformer
from utils.util import state_dict_data_parallel_fix

from model import roberta
from model.roberta import RobertaModel, _prepare_decoder_attention_mask
from model import heads
from transformers import RobertaConfig
from functools import partial
import copy
import torch.distributed as dist

with open('./EgoNCE_MLM_ITM_Config.yml') as f:
    config = yaml.load(f, Loader=yaml.FullLoader)

def init_weights(module):
    if isinstance(module, (nn.Linear, nn.Embedding)):
        module.weight.data.normal_(mean=0.0, std=0.02)
    elif isinstance(module, nn.LayerNorm):
        module.bias.data.zero_()
        module.weight.data.fill_(1.0)

    if isinstance(module, nn.Linear) and module.bias is not None:
        module.bias.data.zero_()


class FrozenInTime(BaseModel):
    def __init__(self,
                 video_params,
                 text_params,
                 projection_dim=4096,
                 load_checkpoint=None,
                 projection='minimal',
                 load_temporal_fix='bilinear',
                 config = config,
                 task_names = 'EgoNCE_ITM_MLM',
                 norm_layer = None,
                 embed_dim=768):
        super().__init__()

        self.video_params = video_params
        self.text_params = text_params
        self.load_temporal_fix = load_temporal_fix
        self.config = config
        self.task_names = task_names
        if not text_params['pretrained']:
            raise NotImplementedError("Huggingface text models require pretrained init.")

        if self.text_params['model'].startswith('roberta'):
            self.text_model = RobertaModel.from_pretrained("roberta-base")
        self.text_model.train()

        pretrained = video_params['pretrained']
        if video_params['model'] == "SpaceTimeTransformer":
            self.num_frames = video_params['num_frames']
            time_init = 'zeros'
            attention_style = 'frozen-in-time'
            arch_config = 'base_patch16_224'
            vit_init = 'imagenet-21k'
            if arch_config == 'base_patch16_224':
                vit_model = torch.load("/cis/home/shraman/works_meta_2022/pre-training/EgoVLP_Fused_HardNegITM_Checkpoint_multinode/frozen-in-time-main/pretrained/jx_vit_base_p16_224-80ecf9dd.pth", map_location="cpu")
                model = SpaceTimeTransformer(num_frames=self.num_frames,
                                            time_init=time_init,
                                            attention_style=attention_style)
            else:
                raise NotImplementedError

            model.head = nn.Identity()
            model.pre_logits = nn.Identity()
            ftr_dim = model.embed_dim
           
            if load_checkpoint in ["", None]:
                vit_checkpoint = vit_model
                new_vit_dict = state_dict_data_parallel_fix(vit_checkpoint, model.state_dict())
                model.load_state_dict(new_vit_dict, strict=False)
            self.video_model = model
        else:
            raise NotImplementedError(f"{video_params['model']} not implemented")

        # for backwards compatibility (old models)
        self.video_model.fc = nn.Identity()

        # Project to a common embedding
        if projection == 'minimal':

            txt_proj = nn.Sequential(
                nn.Linear(self.text_model.config.hidden_size, projection_dim, bias=False),
                nn.ReLU(inplace=True), nn.Linear(projection_dim, projection_dim, bias=True),
                nn.ReLU(inplace=True), nn.Linear(projection_dim, projection_dim, bias=True)
            )

            vid_proj = nn.Sequential(
                nn.Linear(ftr_dim, projection_dim, bias=False),
                nn.ReLU(inplace=True), nn.Linear(projection_dim, projection_dim, bias=True),
                nn.ReLU(inplace=True), nn.Linear(projection_dim, projection_dim, bias=True)
            )

        elif projection == '':
            txt_proj = nn.Identity()
            vid_proj = nn.Identity()
        else:
            raise NotImplementedError
        self.txt_proj = txt_proj
        self.vid_proj = vid_proj

        if ('MLM' in self.task_names or 'ITM' in self.task_names):
            # for FIBER-like cross-attention

            bert_config = RobertaConfig(
                vocab_size=self.config["vocab_size"],
                hidden_size=self.config["hidden_size"],
                num_hidden_layers=self.config["num_layers"],
                num_attention_heads=self.config["num_heads"],
                intermediate_size=self.config["hidden_size"] * self.config["mlp_ratio"],
                #max_position_embeddings=maxlen, [was used in BTGOT script]
                hidden_dropout_prob=self.config["drop_rate"],
                attention_probs_dropout_prob=self.config["drop_rate"],
            )

            self.num_fuse_block=self.config["num_fuse_block"]
            self.num_text_layer=self.config["num_layers"]
            roberta.NUM_FUSE_BLOCK = self.video_model.NUM_FUSE_BLOCK=self.num_fuse_block
            roberta.DIM_IMG=self.config["input_image_embed_size"]
            self.video_model.DIM_TXT=self.config["input_text_embed_size"]

            self.cross_modal_text_transform = nn.Linear(self.config["input_text_embed_size"], self.config["hidden_size"])
            self.cross_modal_text_transform.apply(init_weights)
            self.cross_modal_video_transform = nn.Linear(self.config["input_image_embed_size"], self.config["hidden_size"])
            self.cross_modal_video_transform.apply(init_weights)

            self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))

            self.num_patches = self.video_model.patch_embed.num_patches
            self.patches_per_frame = self.num_patches//self.num_frames
            norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
            self.norm = norm_layer(embed_dim)
            self.pre_logits = nn.Identity()


            self.avgpool = nn.AdaptiveAvgPool1d(1)
            self.cross_modal_video_pooler = heads.Pooler(self.config["hidden_size"])
            self.cross_modal_video_pooler.apply(init_weights)
            self.cross_modal_text_pooler = heads.Pooler(self.config["hidden_size"])
            self.cross_modal_text_pooler.apply(init_weights)

            ## einops transformations
            self.einops_from_space = 'b (f n) d'
            self.einops_to_space = '(b f) n d'
            self.einops_from_time = 'b (f n) d'
            self.einops_to_time = '(b n) f d'

        if 'MLM' in self.task_names:
            self.mlm_score = heads.MLMHead(bert_config)
            self.mlm_score.apply(init_weights)

        if 'ITM' in self.task_names:
            self.itm_score = heads.ITMHead(self.config["hidden_size"] * 2)
            self.itm_score.apply(init_weights)

        if load_checkpoint not in ["", None]:
            checkpoint = torch.load(load_checkpoint, map_location='cpu')
            state_dict = checkpoint['state_dict']
            new_state_dict = state_dict_data_parallel_fix(state_dict, self.state_dict())
            new_state_dict = self._inflate_positional_embeds(new_state_dict)
            self.load_state_dict(new_state_dict, strict=False)

    def set_device(self, device):
        self.device = device

    def infer(self, data, video_only=False, return_embeds=True, task_names=None, ret={}):
        
        text_data = data['text']
        video_data = data['video']

        if task_names is not None:
            self.task_names = task_names


        if 'EgoNCE' in self.task_names: 

            text_embeddings = self.compute_text(text_data)
            video_embeddings = self.compute_video(video_data)


            if return_embeds:
                ret.update({'text_embeds':text_embeddings,
                'video_embeds':video_embeddings
                })

        if 'ITM' in self.task_names:

            b, curr_frames, channels, _, _ = video_data.shape
            video_data_itm = self.video_model.patch_embed(video_data)
            video_data_itm = video_data_itm.flatten(2).transpose(2, 1)
            video_data_itm = video_data_itm.reshape(b, -1, self.video_model.patch_embed.embed_dim)

            BF = video_data_itm.shape[0]
            cls_tokens = self.cls_token.expand(BF, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
            video_data_itm = torch.cat((cls_tokens, video_data_itm), dim=1)
            # positional embed needs to be tiled for each frame (this does [1,2,3] --> [1,2,3,1,2,3]...)
            cls_embed = self.video_model.pos_embed[:, 0, :].unsqueeze(1)
            tile_pos_embed = self.video_model.pos_embed[:, 1:, :].repeat(1, self.num_frames, 1)
            # temporal embed needs to be repeated within each frame (this does [1,2,3] --> [1,1,1,2,2,2,3,3,3]...)
            tile_temporal_embed = self.video_model.temporal_embed.repeat_interleave(self.patches_per_frame, 1)
            total_pos_embed = tile_pos_embed + tile_temporal_embed
            total_pos_embed = torch.cat([cls_embed, total_pos_embed], dim=1)

            n = self.patches_per_frame
            f = curr_frames

            curr_patches = video_data_itm.shape[1]
            video_data_itm = video_data_itm + total_pos_embed[:, :curr_patches]
            video_data_itm = self.video_model.pos_drop(video_data_itm)

            unfused_blocks = self.num_text_layer - self.num_fuse_block

            

            for blk_i, blk in enumerate(self.video_model.blocks[:unfused_blocks]):
                if self.config['use_checkpoint']:
                    video_data_itm = torch.utils.checkpoint.checkpoint(blk, video_data_itm, self.einops_from_space, self.einops_to_space, self.einops_from_time, self.einops_to_time,
                                n, f)
                else:
                    video_data_itm = blk(video_data_itm, self.einops_from_space, self.einops_to_space, self.einops_from_time, self.einops_to_time,
                                time_n=n, space_f=f)
                            
            
            text_embeds = self.text_model.embeddings(input_ids=text_data['input_ids']) # before it was input_ids=text_ids
            device = text_embeds.device
            text_masks = text_data['attention_mask']
            input_shape = text_masks.size()
            extend_text_masks = self.text_model.get_extended_attention_mask(text_masks, input_shape, device)
            for layer_i, layer in enumerate(self.text_model.encoder.layer[:unfused_blocks]):
                if self.config['use_checkpoint']:
                    text_embeds = torch.utils.checkpoint.checkpoint(layer, text_embeds, extend_text_masks)[0]
                else:
                    text_embeds = layer(text_embeds, extend_text_masks)[0]


            for blk_i, blk in enumerate(self.video_model.blocks[unfused_blocks:self.num_text_layer]):
                if self.config['use_checkpoint']:
                    

                    fuse_video_data = torch.utils.checkpoint.checkpoint(blk, video_data_itm, self.einops_from_space, self.einops_to_space, self.einops_from_time, self.einops_to_time, 
                                          n, f, text_embeds, extend_text_masks)
                    text_embeds = torch.utils.checkpoint.checkpoint(self.text_model.encoder.layer[blk_i + unfused_blocks],
                                          text_embeds, extend_text_masks, None, (video_data_itm), None, None, False, True)[0]
                else:
                    fuse_video_data = blk(video_data_itm, self.einops_from_space, self.einops_to_space, self.einops_from_time, self.einops_to_time, 
                                          y=text_embeds, y_mask=extend_text_masks, time_n=n, space_f=f)
                    text_embeds = self.text_model.encoder.layer[blk_i + unfused_blocks](text_embeds, extend_text_masks, encoder_hidden_states=(video_data_itm), last_norm=True)[0]
                video_data_itm = fuse_video_data

            
            #print("Shape of model output", video_data.shape)
            video_data_itm = self.norm(video_data_itm)[:, 0]
            video_data_itm = self.pre_logits(video_data_itm)

            text_embeds = text_embeds[:, 0]
            text_embeds = self.cross_modal_text_transform(text_embeds)
            video_embeds = self.cross_modal_video_transform(video_data_itm)

            cls_feats_text = self.cross_modal_text_pooler(text_embeds)
            
            cls_feats_video = self.cross_modal_video_pooler(video_embeds)

            cls_feats = torch.cat([cls_feats_text, cls_feats_video], dim=-1)

            ret.update({
                "cross_attn_itm_logits": self.itm_score(cls_feats)
            })


        if 'MLM' in self.task_names:
            
            b, curr_frames, channels, _, _ = video_data.shape
            video_data_mlm = self.video_model.patch_embed(video_data)
            video_data_mlm = video_data_mlm.flatten(2).transpose(2, 1)
            video_data_mlm = video_data_mlm.reshape(b, -1, self.video_model.patch_embed.embed_dim)

            BF = video_data_mlm.shape[0]
            cls_tokens = self.cls_token.expand(BF, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
            video_data_mlm = torch.cat((cls_tokens, video_data_mlm), dim=1)
            # positional embed needs to be tiled for each frame (this does [1,2,3] --> [1,2,3,1,2,3]...)
            cls_embed = self.video_model.pos_embed[:, 0, :].unsqueeze(1)
            tile_pos_embed = self.video_model.pos_embed[:, 1:, :].repeat(1, self.num_frames, 1)
            # temporal embed needs to be repeated within each frame (this does [1,2,3] --> [1,1,1,2,2,2,3,3,3]...)
            tile_temporal_embed = self.video_model.temporal_embed.repeat_interleave(self.patches_per_frame, 1)
            total_pos_embed = tile_pos_embed + tile_temporal_embed
            total_pos_embed = torch.cat([cls_embed, total_pos_embed], dim=1)

            #print("total_pos_embed shape: ", total_pos_embed.shape)

            n = self.patches_per_frame
            f = curr_frames

            curr_patches = video_data_mlm.shape[1]
            video_data_mlm = video_data_mlm + total_pos_embed[:, :curr_patches]
            video_data_mlm = self.video_model.pos_drop(video_data_mlm)

            #print("video_data_mlm shape: ", video_data_mlm.shape)

            unfused_blocks = self.num_text_layer - self.num_fuse_block

            
            for blk_i, blk in enumerate(self.video_model.blocks[:unfused_blocks]):
                if self.config['use_checkpoint']:
                    video_data_mlm = torch.utils.checkpoint.checkpoint(blk, video_data_mlm, self.einops_from_space, self.einops_to_space, self.einops_from_time, self.einops_to_time,
                                n, f)
                else:
                    video_data_mlm = blk(video_data_mlm, self.einops_from_space, self.einops_to_space, self.einops_from_time, self.einops_to_time,
                                time_n=n, space_f=f)
                       
            
            text_embeds = self.text_model.embeddings(input_ids=data['text_mlm_ids']) # before it was input_ids=text_ids
            device = text_embeds.device
            text_masks = text_data['attention_mask']
            input_shape = text_masks.size()
            extend_text_masks = self.text_model.get_extended_attention_mask(text_masks, input_shape, device)

            for layer_i, layer in enumerate(self.text_model.encoder.layer[:unfused_blocks]):
                if self.config['use_checkpoint']:
                    text_embeds = torch.utils.checkpoint.checkpoint(layer, text_embeds, extend_text_masks)[0]
                else:
                    text_embeds = layer(text_embeds, extend_text_masks)[0]

            for blk_i, blk in enumerate(self.video_model.blocks[unfused_blocks:self.num_text_layer]):
                if self.config['use_checkpoint']:

                    fuse_video_data = torch.utils.checkpoint.checkpoint(blk, video_data_mlm, self.einops_from_space, self.einops_to_space, self.einops_from_time, self.einops_to_time,
                            n, f, text_embeds, extend_text_masks)
                    text_embeds = torch.utils.checkpoint.checkpoint(self.text_model.encoder.layer[blk_i + unfused_blocks],
                                          text_embeds, extend_text_masks, None, (video_data_mlm), None, None, False, True)[0]
                else:
                    fuse_video_data = blk(video_data_mlm, self.einops_from_space, self.einops_to_space, self.einops_from_time, self.einops_to_time,
                                          y=text_embeds, y_mask=extend_text_masks, time_n=n, space_f=f)
                    text_embeds = self.text_model.encoder.layer[blk_i + unfused_blocks](text_embeds, extend_text_masks, encoder_hidden_states=(video_data_mlm), last_norm=True)[0]
                video_data_mlm = fuse_video_data


            text_embeds = text_embeds #[:, 0]
            text_embeds = self.cross_modal_text_transform(text_embeds)

            ret.update({
                "cross_attn_mlm_logits": self.mlm_score(text_embeds)
            })

        return ret

    
    def forward(self, data, n_embeds, v_embeds, allgather, n_gpu, args, config, loss_egonce, gpu, return_embeds=True, task_names='EgoNCE_ITM_MLM'):

        ret = {}
        loss_dict = {}

        if 'Feature_Extraction' in task_names:
            video_embeddings = self.compute_video(data['video'])
            return video_embeddings


        if 'EgoNCE' in task_names:

            ret = self.infer(data, task_names='EgoNCE')
            video_embeds = ret['video_embeds']
            text_embeds = ret['text_embeds']
            video_embeds = allgather(video_embeds, n_gpu, args)
            text_embeds = allgather(text_embeds, n_gpu, args)
            n_embeds = allgather(n_embeds, n_gpu, args)
            v_embeds = allgather(v_embeds, n_gpu, args)
            output = sim_matrix(text_embeds, video_embeds)

            if config['loss']['type'] == 'EgoNCE':
                sim_v = sim_matrix(v_embeds, v_embeds)
                sim_n = sim_matrix(n_embeds, n_embeds)
                loss, mask_bool, temp = loss_egonce(output, sim_v, sim_n)
            else:
                loss, mask_bool, temp = loss_egonce(output)

            ret.update({"sim_v2t": output, "sim_t2v": output.t(),})

            loss_dict.update({'EgoNCE': loss})

        
        # MLM
        if 'MLM' in task_names:

            ret = self.infer(data, task_names='MLM', ret=ret)

            mlm_logits = ret["cross_attn_mlm_logits"].view(-1, 50265)
            mlm_labels = data["text_mlm_labels"].view(-1)

            mlm_logits = allgather(mlm_logits, n_gpu, args)
            mlm_labels = allgather(mlm_labels, n_gpu, args)

            loss_mlm = torch.nn.functional.cross_entropy(
                                mlm_logits,
                                mlm_labels,
                                ignore_index=-100,
                                ).mean()

            loss = loss + loss_mlm

            loss_dict.update({"loss_mlm": loss_mlm})


        # ITM
        if 'ITM' in task_names:

            rank = dist.get_rank()

            all_video = allgather(data['video'], n_gpu, args)
            all_text_ids = allgather(data['text']['input_ids'], n_gpu, args)
            all_text_masks = allgather(data['text']['attention_mask'], n_gpu, args)

            pos_len = data['video'].size(0) // 2
            neg_len = data['video'].size(0) - pos_len
            itm_labels = torch.cat([torch.ones(pos_len), torch.zeros(neg_len)]).cuda(gpu, non_blocking=True)

            itm_labels = itm_labels[torch.randperm(itm_labels.size(0))]

            batch_size = len(itm_labels)

            with torch.no_grad():
                weights_v2t = F.softmax(ret['sim_v2t'][batch_size*rank : batch_size * (rank + 1), :]/temp, dim=1)
                weights_t2v = F.softmax(ret['sim_t2v'][batch_size*rank : batch_size * (rank + 1), :]/temp, dim=1)

                weights_v2t.masked_fill_(mask_bool[batch_size*rank : batch_size * (rank + 1), :], 0)
                weights_t2v.masked_fill_(mask_bool[batch_size*rank : batch_size * (rank + 1), :], 0)

            data_itm = copy.deepcopy(data)

            for idx in range(len(itm_labels)):
                if itm_labels[idx] == 1:
                    data_itm['video'][idx, :] = all_video[rank*batch_size + idx, :]
                    data_itm['text']['input_ids'][idx, :] = all_text_ids[rank*batch_size + idx, :]
                    data_itm['text']['attention_mask'][idx, :] = all_text_masks[rank*batch_size + idx, :]


                else:
                    if np.random.rand() > 0.5:
                        neg_idx = torch.multinomial(weights_t2v[idx] + 1e-9, 1).item()
                        data_itm['video'][idx, :] = all_video[neg_idx, :]
                        data_itm['text']['input_ids'][idx, :] = all_text_ids[rank*batch_size + idx, :]
                        data_itm['text']['attention_mask'][idx, :] = all_text_masks[rank*batch_size + idx, :]
                    else:
                        neg_idx = torch.multinomial(weights_v2t[idx] + 1e-9, 1).item()
                        data_itm['video'][idx, :] = all_video[rank*batch_size + idx, :]
                        data_itm['text']['input_ids'][idx, :] = all_text_ids[neg_idx, :]
                        data_itm['text']['attention_mask'][idx, :] = all_text_masks[neg_idx, :]


            ret = self.infer(data_itm, task_names='ITM', ret=ret)

            itm_logits = ret["cross_attn_itm_logits"]

            itm_logits = allgather(itm_logits, n_gpu, args)
            itm_labels = allgather(itm_labels, n_gpu, args)

            loss_itm = torch.nn.functional.cross_entropy(itm_logits, itm_labels.long()).mean()

            loss = loss + 2*loss_itm

            #print("ITM loss: ", loss_itm)
            loss_dict.update({"loss_itm": loss_itm})

        loss_dict.update({"loss_total": loss})

        return loss, loss_dict, ret



    def compute_text(self, text_data):
        if self.text_params['model'].startswith('bert'):
            text_embeddings = self.text_model(text_data['input_ids'], attention_mask=text_data['attention_mask'])[
                'pooler_output']
        elif self.text_params['model'].startswith('distilbert'):
            text_embeddings = self.text_model(**text_data).last_hidden_state[:, 0, :]
        elif self.text_params['model'].startswith('roberta'):
            text_embeddings = self.text_model(**text_data).last_hidden_state[:, 0, :]
        else:
            raise NotImplementedError
        if self.config['use_checkpoint']:
            text_embeddings = torch.utils.checkpoint.checkpoint(self.txt_proj, text_embeddings)
        else:
            text_embeddings = self.txt_proj(text_embeddings)
        return text_embeddings

    def compute_text_tokens(self, text_data):
        if self.text_params['model'].startswith('bert'):
            text_embeddings = self.text_model(text_data['input_ids'], attention_mask=text_data['attention_mask'])[
                'pooler_output']    # not implement for bert
        elif self.text_params['model'].startswith('distilbert'):
            text_embeddings = self.text_model(**text_data).last_hidden_state
        elif self.text_params['model'].startswith('roberta'):
            text_embeddings = self.text_model(**text_data).last_hidden_state
        else:
            raise NotImplementedError

        if self.config['use_checkpoint']:
            text_embeddings = torch.utils.checkpoint.checkpoint(self.txt_proj, text_embeddings)
        else:
            text_embeddings = self.txt_proj(text_embeddings)
        return text_embeddings

    def compute_video(self, video_data):
        video_embeddings = self.video_model(video_data)
        if self.config['use_checkpoint']:
            video_embeddings = torch.utils.checkpoint.checkpoint(self.vid_proj, video_embeddings)
        else:
            video_embeddings = self.vid_proj(video_embeddings)
        return video_embeddings

    def _inflate_positional_embeds(self, new_state_dict):
        # allow loading of timesformer with fewer num_frames
        curr_keys = list(self.state_dict().keys())
        if 'video_model.temporal_embed' in new_state_dict and 'video_model.temporal_embed' in curr_keys:
            load_temporal_embed = new_state_dict['video_model.temporal_embed']
            load_num_frames = load_temporal_embed.shape[1]
            curr_num_frames = self.video_params['num_frames']
            embed_dim = load_temporal_embed.shape[2]

            if load_num_frames != curr_num_frames:
                if load_num_frames > curr_num_frames:
                    print(f'### loaded {self.video_params["model"]} model has MORE frames than current...'
                          f'### loading weights, filling in the extras via {self.load_temporal_fix}')
                    new_temporal_embed = load_temporal_embed[:, :curr_num_frames, :]
                else:
                    print(f'### loaded {self.video_params["model"]} model has FEWER frames than current...'
                          f'### loading weights, filling in the extras via {self.load_temporal_fix}')
                    if self.load_temporal_fix == 'zeros':
                        new_temporal_embed = torch.zeros([load_temporal_embed.shape[0], curr_num_frames, embed_dim])
                        new_temporal_embed[:, :load_num_frames] = load_temporal_embed
                    elif self.load_temporal_fix in ['interp', 'bilinear']:
                        # interpolate
                        # unsqueeze so pytorch thinks its an image
                        mode = 'nearest'
                        if self.load_temporal_fix == 'bilinear':
                            mode = 'bilinear'
                        load_temporal_embed = load_temporal_embed.unsqueeze(0)
                        new_temporal_embed = F.interpolate(load_temporal_embed,
                                                           (curr_num_frames, embed_dim), mode=mode, align_corners=True).squeeze(0)
                    else:
                        raise NotImplementedError
                new_state_dict['video_model.temporal_embed'] = new_temporal_embed
        # allow loading with smaller spatial patches. assumes custom border crop, to append the
        # border patches to the input sequence
        if 'video_model.pos_embed' in new_state_dict and 'video_model.pos_embed' in curr_keys:
            load_pos_embed = new_state_dict['video_model.pos_embed']
            load_num_patches = load_pos_embed.shape[1]
            curr_pos_embed = self.state_dict()['video_model.pos_embed']
            if load_num_patches != curr_pos_embed.shape[1]:
                raise NotImplementedError(
                    'Loading models with different spatial resolution / patch number not yet implemented, sorry.')

        return new_state_dict

def sim_matrix(a, b, eps=1e-8):
    """
    added eps for numerical stability
    """
    a_n, b_n = a.norm(dim=1)[:, None], b.norm(dim=1)[:, None]
    a_norm = a / torch.max(a_n, eps * torch.ones_like(a_n))
    b_norm = b / torch.max(b_n, eps * torch.ones_like(b_n))
    sim_mt = torch.mm(a_norm, b_norm.transpose(0, 1))
    return sim_mt


def sim_matrix_batch_val(a, b, eps=1e-8):
    """
    added eps for numerical stability
    """
    a_n, b_n = a.norm(dim=-1).unsqueeze(-1), b.norm(dim=-1).unsqueeze(-1)
    a_norm = a / torch.max(a_n, eps * torch.ones_like(a_n))
    b_norm = b / torch.max(b_n, eps * torch.ones_like(b_n))
    sim_mt = torch.bmm(a_norm, b_norm.transpose(1, 2))
    return sim_mt


if __name__ == "__main__":
    pass

EgoVLPv2/model/model.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import os
import pdb

import timm
import torch
import yaml
import torch.nn as nn
import numpy as np
import torch.nn.functional as F
from transformers import AutoModel
from einops import rearrange, repeat

from base import BaseModel
from model import video_transformer
from model.video_transformer import SpaceTimeTransformer
from utils.util import state_dict_data_parallel_fix

from model import roberta
from model.roberta import RobertaModel, _prepare_decoder_attention_mask
from model import heads
from transformers import RobertaConfig
from functools import partial
import copy
import torch.distributed as dist
import pickle

with open('./EgoNCE_MLM_ITM_Config.yml') as f:
    config = yaml.load(f, Loader=yaml.FullLoader)

def init_weights(module):
    if isinstance(module, (nn.Linear, nn.Embedding)):
        module.weight.data.normal_(mean=0.0, std=0.02)
    elif isinstance(module, nn.LayerNorm):
        module.bias.data.zero_()
        module.weight.data.fill_(1.0)

    if isinstance(module, nn.Linear) and module.bias is not None:
        module.bias.data.zero_()


class FrozenInTime(BaseModel):
    def __init__(self,
                 video_params,
                 text_params,
                 projection_dim=4096,
                 load_checkpoint=None,
                 projection='minimal',
                 load_temporal_fix='bilinear',
                 config = config,
                 task_names = 'EgoNCE_ITM_MLM',
                 norm_layer = None,
                 embed_dim=768):
        super().__init__()

        self.video_params = video_params
        self.text_params = text_params
        self.load_temporal_fix = load_temporal_fix
        self.config = config
        self.task_names = task_names
        try:
            self.relevancy_mat_train_path = os.path.join(self.config['data_loader']['args']['meta_dir'], 'relevancy', 'caption_relevancy_EPIC_100_retrieval_train.pkl')
            pkl_file = open(self.relevancy_mat_train_path, 'rb')
            self.relevancy_mat_train = pickle.load(pkl_file)
        except:
            pass
        
        if not text_params['pretrained']:
            raise NotImplementedError("Huggingface text models require pretrained init.")

        if self.text_params['model'].startswith('roberta'):
            self.text_model = RobertaModel.from_pretrained("roberta-base")
        self.text_model.train()

        pretrained = video_params['pretrained']
        if video_params['model'] == "SpaceTimeTransformer":
            self.num_frames = video_params["num_frames"]
            self.drop_path_rate = video_params["drop_path_rate"]
            time_init = 'zeros'
            attention_style = 'frozen-in-time'
            arch_config = 'base_patch16_224'
            vit_init = 'imagenet-21k'
            if arch_config == 'base_patch16_224':
                #vit_model = timm.models.vision_transformer.vit_base_patch16_224(pretrained=pretrained)
                vit_model = torch.load("/cis/home/shraman/works_meta_2022/pre-training/EgoVLP_Fused_HardNegITM_Checkpoint_multinode/frozen-in-time-main/pretrained/jx_vit_base_p16_224-80ecf9dd.pth", map_location="cpu")
                model = SpaceTimeTransformer(num_frames=self.num_frames,
                                            time_init=time_init,
                                            attention_style=attention_style,
                                            drop_path_rate=self.drop_path_rate)
            else:
                raise NotImplementedError

            model.head = nn.Identity()
            model.pre_logits = nn.Identity()
            ftr_dim = model.embed_dim
           
            if load_checkpoint in ["", None]:
                # vit_checkpoint = vit_model.state_dict()
                # model.load_state_dict(vit_checkpoint, strict=False)
                vit_checkpoint = vit_model
                new_vit_dict = state_dict_data_parallel_fix(vit_checkpoint, model.state_dict())
                model.load_state_dict(new_vit_dict, strict=False)
            self.video_model = model
        else:
            raise NotImplementedError(f"{video_params['model']} not implemented")

        # for backwards compatibility (old models)
        self.video_model.fc = nn.Identity()

        # Project to a common embedding
        if projection == 'minimal':

            txt_proj = nn.Sequential(nn.ReLU(), nn.Linear(self.text_model.config.hidden_size, 256),)
            vid_proj = nn.Sequential(nn.Linear(ftr_dim, 256))

            #txt_proj = nn.Sequential(
            #    nn.Linear(self.text_model.config.hidden_size, projection_dim, bias=False),
            #    nn.ReLU(inplace=True), nn.Linear(projection_dim, projection_dim, bias=True),
            #    nn.ReLU(inplace=True), nn.Linear(projection_dim, projection_dim, bias=True)
            #)

            #vid_proj = nn.Sequential(
            #    nn.Linear(ftr_dim, projection_dim, bias=False),
            #    nn.ReLU(inplace=True), nn.Linear(projection_dim, projection_dim, bias=True),
            #    nn.ReLU(inplace=True), nn.Linear(projection_dim, projection_dim, bias=True)
            #)

        
        elif projection == '':
            txt_proj = nn.Identity()
            vid_proj = nn.Identity()
        else:
            raise NotImplementedError
        self.txt_proj = txt_proj
        self.vid_proj = vid_proj

        if ('MLM' in self.task_names or 'ITM' in self.task_names):
            # for FIBER-like cross-attention

            bert_config = RobertaConfig(
                vocab_size=self.config["vocab_size"],
                hidden_size=self.config["hidden_size"],
                num_hidden_layers=self.config["num_layers"],
                num_attention_heads=self.config["num_heads"],
                intermediate_size=self.config["hidden_size"] * config["mlp_ratio"],
                #max_position_embeddings=maxlen, [was used in BTGOT script]
                hidden_dropout_prob=self.config["drop_rate"],
                attention_probs_dropout_prob=self.config["drop_rate"],
            )

            self.num_fuse_block=self.config["num_fuse_block"]
            self.num_text_layer=self.config["num_layers"]
            roberta.NUM_FUSE_BLOCK = self.video_model.NUM_FUSE_BLOCK=self.num_fuse_block
            roberta.DIM_IMG=self.config["input_image_embed_size"]
            self.video_model.DIM_TXT=self.config["input_text_embed_size"]

            self.cross_modal_text_transform = nn.Linear(self.config["input_text_embed_size"], self.config["hidden_size"])
            self.cross_modal_text_transform.apply(init_weights)
            self.cross_modal_video_transform = nn.Linear(self.config["input_image_embed_size"], self.config["hidden_size"])
            self.cross_modal_video_transform.apply(init_weights)

            self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))

            self.num_patches = self.video_model.patch_embed.num_patches
            self.patches_per_frame = self.num_patches//self.num_frames
            norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
            self.norm = norm_layer(embed_dim)
            self.pre_logits = nn.Identity()


            self.avgpool = nn.AdaptiveAvgPool1d(1)
            self.cross_modal_video_pooler = heads.Pooler(config["hidden_size"])
            self.cross_modal_video_pooler.apply(init_weights)
            self.cross_modal_text_pooler = heads.Pooler(config["hidden_size"])
            self.cross_modal_text_pooler.apply(init_weights)

            ## einops transformations
            self.einops_from_space = 'b (f n) d'
            self.einops_to_space = '(b f) n d'
            self.einops_from_time = 'b (f n) d'
            self.einops_to_time = '(b n) f d'

        if 'MLM' in self.task_names:
            self.mlm_score = heads.MLMHead(bert_config)
            self.mlm_score.apply(init_weights)

        if 'ITM' in self.task_names:
            self.itm_score = heads.ITMHead(config["hidden_size"] * 2)
            self.itm_score.apply(init_weights)

        if load_checkpoint not in ["", None]:
            # checkpoint = torch.load(load_checkpoint)
            # local_rank = int(os.environ['LOCAL_RANK'])  # fixed by qinghong.
            # checkpoint = torch.load(load_checkpoint, map_location='cuda:{}'.format(local_rank))
            checkpoint = torch.load(load_checkpoint, map_location='cpu')
            state_dict = checkpoint['state_dict']
            new_state_dict = state_dict_data_parallel_fix(state_dict, self.state_dict())
            new_state_dict = self._inflate_positional_embeds(new_state_dict)
            self.load_state_dict(new_state_dict, strict=False)

    def set_device(self, device):
        self.device = device

    def infer(self, data, video_only=False, return_embeds=True, task_names=None, ret={}):

        text_data = data['text']
        video_data = data['video']

        if task_names is not None:
            self.task_names = task_names


        if 'Dual' in self.task_names: 

            text_embeddings = self.compute_text(text_data)
            video_embeddings = self.compute_video(video_data)

            if return_embeds:
                ret.update({'text_embeds':text_embeddings,
                'video_embeds':video_embeddings
                })

        if 'ITM' in self.task_names:

            b, curr_frames, channels, _, _ = video_data.shape
            video_data_itm = self.video_model.patch_embed(video_data)
            video_data_itm = video_data_itm.flatten(2).transpose(2, 1)
            video_data_itm = video_data_itm.reshape(b, -1, self.video_model.patch_embed.embed_dim)

            BF = video_data_itm.shape[0]
            cls_tokens = self.cls_token.expand(BF, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
            video_data_itm = torch.cat((cls_tokens, video_data_itm), dim=1)
            # positional embed needs to be tiled for each frame (this does [1,2,3] --> [1,2,3,1,2,3]...)
            cls_embed = self.video_model.pos_embed[:, 0, :].unsqueeze(1)
            tile_pos_embed = self.video_model.pos_embed[:, 1:, :].repeat(1, self.num_frames, 1)
            # temporal embed needs to be repeated within each frame (this does [1,2,3] --> [1,1,1,2,2,2,3,3,3]...)
            tile_temporal_embed = self.video_model.temporal_embed.repeat_interleave(self.patches_per_frame, 1)
            total_pos_embed = tile_pos_embed + tile_temporal_embed
            total_pos_embed = torch.cat([cls_embed, total_pos_embed], dim=1)

            n = self.patches_per_frame
            f = curr_frames

            curr_patches = video_data_itm.shape[1]
            video_data_itm = video_data_itm + total_pos_embed[:, :curr_patches]
            video_data_itm = self.video_model.pos_drop(video_data_itm)

            unfused_blocks = self.num_text_layer - self.num_fuse_block

            

            for blk_i, blk in enumerate(self.video_model.blocks[:unfused_blocks]):
                if config['use_checkpoint']:
                    video_data_itm = torch.utils.checkpoint.checkpoint(blk, video_data_itm, self.einops_from_space, self.einops_to_space, self.einops_from_time, self.einops_to_time,
                                n, f)
                else:
                    video_data_itm = blk(video_data_itm, self.einops_from_space, self.einops_to_space, self.einops_from_time, self.einops_to_time,
                                time_n=n, space_f=f)
                            
            
            text_embeds = self.text_model.embeddings(input_ids=text_data['input_ids']) # before it was input_ids=text_ids
            device = text_embeds.device
            #print(text_data.keys())
            text_masks = text_data['attention_mask']
            input_shape = text_masks.size()
            extend_text_masks = self.text_model.get_extended_attention_mask(text_masks, input_shape, device)
            for layer_i, layer in enumerate(self.text_model.encoder.layer[:unfused_blocks]):
                if config['use_checkpoint']:
                    text_embeds = torch.utils.checkpoint.checkpoint(layer, text_embeds, extend_text_masks)[0]
                else:
                    text_embeds = layer(text_embeds, extend_text_masks)[0]


            for blk_i, blk in enumerate(self.video_model.blocks[unfused_blocks:self.num_text_layer]):
                if config['use_checkpoint']:
                    

                    fuse_video_data = torch.utils.checkpoint.checkpoint(blk, video_data_itm, self.einops_from_space, self.einops_to_space, self.einops_from_time, self.einops_to_time, 
                                          n, f, text_embeds, extend_text_masks)
                    text_embeds = torch.utils.checkpoint.checkpoint(self.text_model.encoder.layer[blk_i + unfused_blocks],
                                          text_embeds, extend_text_masks, None, (video_data_itm), None, None, False, True)[0]
                else:
                    fuse_video_data = blk(video_data_itm, self.einops_from_space, self.einops_to_space, self.einops_from_time, self.einops_to_time, 
                                          y=text_embeds, y_mask=extend_text_masks, time_n=n, space_f=f)
                    text_embeds = self.text_model.encoder.layer[blk_i + unfused_blocks](text_embeds, extend_text_masks, encoder_hidden_states=(video_data_itm), last_norm=True)[0]
                video_data_itm = fuse_video_data

            
            #print("Shape of model output", video_data.shape)
            video_data_itm = self.norm(video_data_itm)[:, 0]
            video_data_itm = self.pre_logits(video_data_itm)

            text_embeds = text_embeds[:, 0]
            text_embeds = self.cross_modal_text_transform(text_embeds)
            video_embeds = self.cross_modal_video_transform(video_data_itm)

            cls_feats_text = self.cross_modal_text_pooler(text_embeds)
            
            cls_feats_video = self.cross_modal_video_pooler(video_embeds)

            cls_feats = torch.cat([cls_feats_text, cls_feats_video], dim=-1)

            ret.update({
                "cross_attn_itm_logits": self.itm_score(cls_feats)
            })

            #print("ITM Forward Pass Done!!")


        if 'MLM' in self.task_names:
            
            #print("MLM started")

            b, curr_frames, channels, _, _ = video_data.shape
            video_data_mlm = self.video_model.patch_embed(video_data)
            video_data_mlm = video_data_mlm.flatten(2).transpose(2, 1)
            video_data_mlm = video_data_mlm.reshape(b, -1, self.video_model.patch_embed.embed_dim)

            #print("video_data_mlm shape: ", video_data_mlm.shape)

            BF = video_data_mlm.shape[0]
            cls_tokens = self.cls_token.expand(BF, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
            video_data_mlm = torch.cat((cls_tokens, video_data_mlm), dim=1)
            # positional embed needs to be tiled for each frame (this does [1,2,3] --> [1,2,3,1,2,3]...)
            cls_embed = self.video_model.pos_embed[:, 0, :].unsqueeze(1)
            tile_pos_embed = self.video_model.pos_embed[:, 1:, :].repeat(1, self.num_frames, 1)
            # temporal embed needs to be repeated within each frame (this does [1,2,3] --> [1,1,1,2,2,2,3,3,3]...)
            tile_temporal_embed = self.video_model.temporal_embed.repeat_interleave(self.patches_per_frame, 1)
            total_pos_embed = tile_pos_embed + tile_temporal_embed
            total_pos_embed = torch.cat([cls_embed, total_pos_embed], dim=1)

            #print("total_pos_embed shape: ", total_pos_embed.shape)

            n = self.patches_per_frame
            f = curr_frames

            curr_patches = video_data_mlm.shape[1]
            video_data_mlm = video_data_mlm + total_pos_embed[:, :curr_patches]
            video_data_mlm = self.video_model.pos_drop(video_data_mlm)

            #print("video_data_mlm shape: ", video_data_mlm.shape)

            unfused_blocks = self.num_text_layer - self.num_fuse_block

            
            for blk_i, blk in enumerate(self.video_model.blocks[:unfused_blocks]):
                if config['use_checkpoint']:
                    video_data_mlm = torch.utils.checkpoint.checkpoint(blk, video_data_mlm, self.einops_from_space, self.einops_to_space, self.einops_from_time, self.einops_to_time,
                                n, f)
                else:
                    video_data_mlm = blk(video_data_mlm, self.einops_from_space, self.einops_to_space, self.einops_from_time, self.einops_to_time,
                                time_n=n, space_f=f)
            #print("video_data_mlm shape: ", video_data_mlm.shape)
                       
            
            text_embeds = self.text_model.embeddings(input_ids=data['text_mlm_ids']) # before it was input_ids=text_ids
            device = text_embeds.device
            text_masks = text_data['attention_mask']
            input_shape = text_masks.size()
            extend_text_masks = self.text_model.get_extended_attention_mask(text_masks, input_shape, device)

            for layer_i, layer in enumerate(self.text_model.encoder.layer[:unfused_blocks]):
                if config['use_checkpoint']:
                    text_embeds = torch.utils.checkpoint.checkpoint(layer, text_embeds, extend_text_masks)[0]
                else:
                    text_embeds = layer(text_embeds, extend_text_masks)[0]

            for blk_i, blk in enumerate(self.video_model.blocks[unfused_blocks:self.num_text_layer]):
                if config['use_checkpoint']:

                    fuse_video_data = torch.utils.checkpoint.checkpoint(blk, video_data_mlm, self.einops_from_space, self.einops_to_space, self.einops_from_time, self.einops_to_time,
                            n, f, text_embeds, extend_text_masks)
                    text_embeds = torch.utils.checkpoint.checkpoint(self.text_model.encoder.layer[blk_i + unfused_blocks],
                                          text_embeds, extend_text_masks, None, (video_data_mlm), None, None, False, True)[0]
                else:
                    fuse_video_data = blk(video_data_mlm, self.einops_from_space, self.einops_to_space, self.einops_from_time, self.einops_to_time,
                                          y=text_embeds, y_mask=extend_text_masks, time_n=n, space_f=f)
                    text_embeds = self.text_model.encoder.layer[blk_i + unfused_blocks](text_embeds, extend_text_masks, encoder_hidden_states=(video_data_mlm), last_norm=True)[0]
                video_data_mlm = fuse_video_data


            text_embeds = text_embeds #[:, 0]
            text_embeds = self.cross_modal_text_transform(text_embeds)

            ret.update({
                "cross_attn_mlm_logits": self.mlm_score(text_embeds)
            })

        return ret
        #return sim_matrix(text_embeddings, video_embeddings)

    
    def create_mask_relevancy(self, relevancy_mat, item_v, item_t):

        mask_rel = torch.zeros(len(item_v), len(item_t))
        item_t_list = item_t.tolist()
        for _i in range(len(item_v)):
            current_item_v = int(item_v[_i].item())
            mask_rel[current_item_v] = relevancy_mat[current_item_v][item_t_list] > 0.1

        return mask_rel



    def forward(self, data, allgather, n_gpu, args, config, loss_dual, gpu, return_embeds=True, task_names='Dual', dataset_name='charades'):

        ret = {}
        loss_dict = {}

        if 'Dual' in task_names:
            
            ret = self.infer(data, task_names='Dual')
            video_embeds = ret['video_embeds']
            text_embeds = ret['text_embeds']
            video_embeds = allgather(video_embeds, n_gpu, args)
            text_embeds = allgather(text_embeds, n_gpu, args)

            output = sim_matrix(text_embeds, video_embeds)

            if dataset_name == 'epic':
                w_embeds = data['relation']
                w_embeds = allgather(w_embeds, n_gpu, args)
                loss = loss_dual(output, w_embeds)
            elif dataset_name == 'charades':
                loss, temp = loss_dual(output)
            else:
                raise NameError()

            if dataset_name == 'epic':
                ret.update({"sim_v2t": output, "sim_t2v": output.t(), 'epic_relation': w_embeds})
            elif dataset_name == 'charades':
                ret.update({"sim_v2t": output, "sim_t2v": output.t()})
            else:
                raise NameError()

            loss_dict.update({'Dual': loss})

        return loss, loss_dict, ret


    def compute_text(self, text_data):
        if self.text_params['model'].startswith('bert'):
            text_embeddings = self.text_model(text_data['input_ids'], attention_mask=text_data['attention_mask'])[
                'pooler_output']
        elif self.text_params['model'].startswith('distilbert'):
            text_embeddings = self.text_model(**text_data).last_hidden_state[:, 0, :]
        elif self.text_params['model'].startswith('roberta'):
            text_embeddings = self.text_model(**text_data).last_hidden_state[:, 0, :]
        else:
            raise NotImplementedError
        if config['use_checkpoint']:
            text_embeddings = torch.utils.checkpoint.checkpoint(self.txt_proj, text_embeddings)
        else:
            text_embeddings = self.txt_proj(text_embeddings)
        return text_embeddings

    def compute_text_tokens(self, text_data):
        if self.text_params['model'].startswith('bert'):
            text_embeddings = self.text_model(text_data['input_ids'], attention_mask=text_data['attention_mask'])[
                'pooler_output']    # not implement for bert
        elif self.text_params['model'].startswith('distilbert'):
            text_embeddings = self.text_model(**text_data).last_hidden_state
        elif self.text_params['model'].startswith('roberta'):
            text_embeddings = self.text_model(**text_data).last_hidden_state
        else:
            raise NotImplementedError

        if config['use_checkpoint']:
            text_embeddings = torch.utils.checkpoint.checkpoint(self.txt_proj, text_embeddings)
        else:
            text_embeddings = self.txt_proj(text_embeddings)
        return text_embeddings

    def compute_video(self, video_data):
        video_embeddings = self.video_model(video_data)
        if config['use_checkpoint']:
            video_embeddings = torch.utils.checkpoint.checkpoint(self.vid_proj, video_embeddings)
        else:
            video_embeddings = self.vid_proj(video_embeddings)
        return video_embeddings

    def _inflate_positional_embeds(self, new_state_dict):
        # allow loading of timesformer with fewer num_frames
        curr_keys = list(self.state_dict().keys())
        if 'video_model.temporal_embed' in new_state_dict and 'video_model.temporal_embed' in curr_keys:
            load_temporal_embed = new_state_dict['video_model.temporal_embed']
            load_num_frames = load_temporal_embed.shape[1]
            curr_num_frames = self.video_params['num_frames']
            embed_dim = load_temporal_embed.shape[2]

            if load_num_frames != curr_num_frames:
                if load_num_frames > curr_num_frames:
                    print(f'### loaded {self.video_params["model"]} model has MORE frames than current...'
                          f'### loading weights, filling in the extras via {self.load_temporal_fix}')
                    new_temporal_embed = load_temporal_embed[:, :curr_num_frames, :]
                else:
                    print(f'### loaded {self.video_params["model"]} model has FEWER frames than current...'
                          f'### loading weights, filling in the extras via {self.load_temporal_fix}')
                    if self.load_temporal_fix == 'zeros':
                        new_temporal_embed = torch.zeros([load_temporal_embed.shape[0], curr_num_frames, embed_dim])
                        new_temporal_embed[:, :load_num_frames] = load_temporal_embed
                    elif self.load_temporal_fix in ['interp', 'bilinear']:
                        # interpolate
                        # unsqueeze so pytorch thinks its an image
                        mode = 'nearest'
                        if self.load_temporal_fix == 'bilinear':
                            mode = 'bilinear'
                        load_temporal_embed = load_temporal_embed.unsqueeze(0)
                        new_temporal_embed = F.interpolate(load_temporal_embed,
                                                           (curr_num_frames, embed_dim), mode=mode, align_corners=True).squeeze(0)
                    else:
                        raise NotImplementedError
                new_state_dict['video_model.temporal_embed'] = new_temporal_embed
        # allow loading with smaller spatial patches. assumes custom border crop, to append the
        # border patches to the input sequence
        if 'video_model.pos_embed' in new_state_dict and 'video_model.pos_embed' in curr_keys:
            load_pos_embed = new_state_dict['video_model.pos_embed']
            load_num_patches = load_pos_embed.shape[1]
            curr_pos_embed = self.state_dict()['video_model.pos_embed']
            if load_num_patches != curr_pos_embed.shape[1]:
                raise NotImplementedError(
                    'Loading models with different spatial resolution / patch number not yet implemented, sorry.')

        return new_state_dict

def sim_matrix(a, b, eps=1e-8):
    """
    added eps for numerical stability
    """
    a_n, b_n = a.norm(dim=1)[:, None], b.norm(dim=1)[:, None]
    a_norm = a / torch.max(a_n, eps * torch.ones_like(a_n))
    b_norm = b / torch.max(b_n, eps * torch.ones_like(b_n))
    sim_mt = torch.mm(a_norm, b_norm.transpose(0, 1))
    return sim_mt


def sim_matrix_batch_val(a, b, eps=1e-8):
    """
    added eps for numerical stability
    """
    a_n, b_n = a.norm(dim=-1).unsqueeze(-1), b.norm(dim=-1).unsqueeze(-1)
    a_norm = a / torch.max(a_n, eps * torch.ones_like(a_n))
    b_norm = b / torch.max(b_n, eps * torch.ones_like(b_n))
    sim_mt = torch.bmm(a_norm, b_norm.transpose(1, 2))
    return sim_mt


if __name__ == "__main__":
    pass

EgoVLPv2/model/model_epic_charades.py
facebookresearch/EgoVLPv2
# coding=utf-8
# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""PyTorch RoBERTa model. """

import yaml

with open('./EgoNCE_MLM_ITM_Config.yml') as f:
    config_yaml = yaml.load(f, Loader=yaml.FullLoader)

NUM_FUSE_BLOCK = config_yaml['num_fuse_block']
DIM_IMG = 768

import math
import inspect
from typing import Callable, List, Optional, Set, Tuple, Union

import torch
import torch.nn as nn
import torch.utils.checkpoint
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from transformers.activations import ACT2FN, gelu
from transformers.file_utils import (
    add_code_sample_docstrings,
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    replace_return_docstrings,
)
from transformers.modeling_outputs import (
    BaseModelOutputWithPastAndCrossAttentions,
    BaseModelOutputWithPoolingAndCrossAttentions,
    CausalLMOutputWithCrossAttentions,
    MaskedLMOutput,
    MultipleChoiceModelOutput,
    QuestionAnsweringModelOutput,
    SequenceClassifierOutput,
    TokenClassifierOutput,
)
from transformers.modeling_utils import (
    PreTrainedModel,
    find_pruneable_heads_and_indices,
    prune_linear_layer,
)
from transformers.utils import logging
from transformers.models.roberta.configuration_roberta import RobertaConfig


logger = logging.get_logger(__name__)

_CHECKPOINT_FOR_DOC = "roberta-base"
_CONFIG_FOR_DOC = "RobertaConfig"
_TOKENIZER_FOR_DOC = "RobertaTokenizer"

ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST = [
    "roberta-base",
    "roberta-large",
    "roberta-large-mnli",
    "distilroberta-base",
    "roberta-base-openai-detector",
    "roberta-large-openai-detector",
    # See all RoBERTa models at https://huggingface.co/models?filter=roberta
]


def apply_chunking_to_forward(
    forward_fn: Callable[..., torch.Tensor], chunk_size: int, chunk_dim: int, *input_tensors, last_norm=True
) -> torch.Tensor:
    """
    This function chunks the `input_tensors` into smaller input tensor parts of size `chunk_size` over the dimension
    `chunk_dim`. It then applies a layer `forward_fn` to each chunk independently to save memory.
    If the `forward_fn` is independent across the `chunk_dim` this function will yield the same result as directly
    applying `forward_fn` to `input_tensors`.
    Args:
        forward_fn (`Callable[..., torch.Tensor]`):
            The forward function of the model.
        chunk_size (`int`):
            The chunk size of a chunked tensor: `num_chunks = len(input_tensors[0]) / chunk_size`.
        chunk_dim (`int`):
            The dimension over which the `input_tensors` should be chunked.
        input_tensors (`Tuple[torch.Tensor]`):
            The input tensors of `forward_fn` which will be chunked
    Returns:
        `torch.Tensor`: A tensor with the same shape as the `forward_fn` would have given if applied`.
    Examples:
    ```python
    # rename the usual forward() fn to forward_chunk()
    def forward_chunk(self, hidden_states):
        hidden_states = self.decoder(hidden_states)
        return hidden_states
    # implement a chunked forward function
    def forward(self, hidden_states):
        return apply_chunking_to_forward(self.forward_chunk, self.chunk_size_lm_head, self.seq_len_dim, hidden_states)
    ```"""

    assert len(input_tensors) > 0, f"{input_tensors} has to be a tuple/list of tensors"

    # inspect.signature exist since python 3.5 and is a python method -> no problem with backward compatibility
    num_args_in_forward_chunk_fn = len(inspect.signature(forward_fn).parameters)
    # if num_args_in_forward_chunk_fn != len(input_tensors):
    #    raise ValueError(
    #        f"forward_chunk_fn expects {num_args_in_forward_chunk_fn} arguments, but only {len(input_tensors)} input "
    #        "tensors are given"
    #    )

    if chunk_size > 0:
        tensor_shape = input_tensors[0].shape[chunk_dim]
        for input_tensor in input_tensors:
            if input_tensor.shape[chunk_dim] != tensor_shape:
                raise ValueError(
                    f"All input tenors have to be of the same shape: {tensor_shape}, "
                    f"found shape {input_tensor.shape[chunk_dim]}"
                )

        if input_tensors[0].shape[chunk_dim] % chunk_size != 0:
            raise ValueError(
                f"The dimension to be chunked {input_tensors[0].shape[chunk_dim]} has to be a multiple of the chunk "
                f"size {chunk_size}"
            )

        num_chunks = input_tensors[0].shape[chunk_dim] // chunk_size

        # chunk input tensor into tuples
        input_tensors_chunks = tuple(input_tensor.chunk(num_chunks, dim=chunk_dim) for input_tensor in input_tensors)
        # apply forward fn to every tuple
        output_chunks = tuple(
            forward_fn(*input_tensors_chunk, last_norm=last_norm) for input_tensors_chunk in zip(*input_tensors_chunks)
        )
        # concatenate output at same dimension
        return torch.cat(output_chunks, dim=chunk_dim)

    return forward_fn(*input_tensors, last_norm=last_norm)


class RobertaEmbeddings(nn.Module):
    """
    Same as BertEmbeddings with a tiny tweak for positional embeddings indexing.
    """

    # Copied from transformers.models.bert.modeling_bert.BertEmbeddings.__init__
    def __init__(self, config):
        super().__init__()
        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
        self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)

        # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
        # any TensorFlow checkpoint file
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

        # position_ids (1, len position emb) is contiguous in memory and exported when serialized
        self.register_buffer("position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)))
        self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")

        # End copy
        self.padding_idx = config.pad_token_id
        self.position_embeddings = nn.Embedding(
            config.max_position_embeddings, config.hidden_size, padding_idx=self.padding_idx
        )

    def forward(
        self, input_ids=None, token_type_ids=None, position_ids=None, inputs_embeds=None, past_key_values_length=0
    ):
        if position_ids is None:
            if input_ids is not None:
                # Create the position ids from the input token ids. Any padded tokens remain padded.
                position_ids = create_position_ids_from_input_ids(
                    input_ids, self.padding_idx, past_key_values_length
                ).to(input_ids.device)
            else:
                position_ids = self.create_position_ids_from_inputs_embeds(inputs_embeds)

        if input_ids is not None:
            input_shape = input_ids.size()
        else:
            input_shape = inputs_embeds.size()[:-1]

        if token_type_ids is None:
            token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)

        if inputs_embeds is None:
            inputs_embeds = self.word_embeddings(input_ids)
        token_type_embeddings = self.token_type_embeddings(token_type_ids)

        embeddings = inputs_embeds + token_type_embeddings
        if self.position_embedding_type == "absolute":
            position_embeddings = self.position_embeddings(position_ids)
            embeddings += position_embeddings
        embeddings = self.LayerNorm(embeddings)
        embeddings = self.dropout(embeddings)
        return embeddings

    def create_position_ids_from_inputs_embeds(self, inputs_embeds):
        """
        We are provided embeddings directly. We cannot infer which are padded so just generate sequential position ids.
        Args:
            inputs_embeds: torch.Tensor
        Returns: torch.Tensor
        """
        input_shape = inputs_embeds.size()[:-1]
        sequence_length = input_shape[1]

        position_ids = torch.arange(
            self.padding_idx + 1, sequence_length + self.padding_idx + 1, dtype=torch.long, device=inputs_embeds.device
        )
        return position_ids.unsqueeze(0).expand(input_shape)


# Copied from transformers.models.bert.modeling_bert.BertSelfAttention with Bert->Roberta
class RobertaSelfAttention(nn.Module):
    def __init__(self, config, layer_index=None):
        super().__init__()
        if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
            raise ValueError(
                f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
                f"heads ({config.num_attention_heads})"
            )

        self.num_attention_heads = config.num_attention_heads
        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
        self.all_head_size = self.num_attention_heads * self.attention_head_size

        self.query = nn.Linear(config.hidden_size, self.all_head_size)
        if layer_index is None:
            self.key = nn.Linear(config.hidden_size, self.all_head_size)
            self.value = nn.Linear(config.hidden_size, self.all_head_size)
        else:
            self.key = nn.Linear(DIM_IMG, self.all_head_size)
            self.value = nn.Linear(DIM_IMG, self.all_head_size)

        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
        self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")
        if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
            self.max_position_embeddings = config.max_position_embeddings
            self.distance_embedding = nn.Embedding(2 * config.max_position_embeddings - 1, self.attention_head_size)

        self.is_decoder = config.is_decoder

    def transpose_for_scores(self, x):
        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
        x = x.view(*new_x_shape)
        return x.permute(0, 2, 1, 3)

    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        past_key_value=None,
        output_attentions=False,
    ):
        mixed_query_layer = self.query(hidden_states)

        # If this is instantiated as a cross-attention module, the keys
        # and values come from an encoder; the attention mask needs to be
        # such that the encoder's padding tokens are not attended to.
        is_cross_attention = encoder_hidden_states is not None

        if encoder_hidden_states is not None:
            key_layer = self.transpose_for_scores(self.key(encoder_hidden_states))
            value_layer = self.transpose_for_scores(self.value(encoder_hidden_states))
            attention_mask = encoder_attention_mask
        else:
            key_layer = self.transpose_for_scores(self.key(hidden_states))
            value_layer = self.transpose_for_scores(self.value(hidden_states))

        query_layer = self.transpose_for_scores(mixed_query_layer)

        # Take the dot product between "query" and "key" to get the raw attention scores.
        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))

        if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
            seq_length = hidden_states.size()[1]
            position_ids_l = torch.arange(seq_length, dtype=torch.long, device=hidden_states.device).view(-1, 1)
            position_ids_r = torch.arange(seq_length, dtype=torch.long, device=hidden_states.device).view(1, -1)
            distance = position_ids_l - position_ids_r
            positional_embedding = self.distance_embedding(distance + self.max_position_embeddings - 1)
            positional_embedding = positional_embedding.to(dtype=query_layer.dtype)  # fp16 compatibility

            if self.position_embedding_type == "relative_key":
                relative_position_scores = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
                attention_scores = attention_scores + relative_position_scores
            elif self.position_embedding_type == "relative_key_query":
                relative_position_scores_query = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
                relative_position_scores_key = torch.einsum("bhrd,lrd->bhlr", key_layer, positional_embedding)
                attention_scores = attention_scores + relative_position_scores_query + relative_position_scores_key

        attention_scores = attention_scores / math.sqrt(self.attention_head_size)
        if attention_mask is not None:
            # Apply the attention mask is (precomputed for all layers in RobertaModel forward() function)
            attention_scores = attention_scores + attention_mask

        # Normalize the attention scores to probabilities.
        attention_probs = nn.Softmax(dim=-1)(attention_scores)

        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.
        attention_probs = self.dropout(attention_probs)

        # Mask heads if we want to
        if head_mask is not None:
            attention_probs = attention_probs * head_mask

        context_layer = torch.matmul(attention_probs, value_layer)

        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
        context_layer = context_layer.view(*new_context_layer_shape)

        outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)

        return outputs


# Copied from transformers.models.bert.modeling_bert.BertSelfOutput
class RobertaSelfOutput(nn.Module):
    def __init__(self, config, layer_index=None):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        if layer_index is None:
            self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, hidden_states):
        #print("Shape of hidden_states: ", hidden_states.shape)
        hidden_states = self.dense(hidden_states)
        hidden_states = self.dropout(hidden_states)
        return hidden_states


# Copied from transformers.models.bert.modeling_bert.BertAttention with Bert->Roberta
class RobertaAttention(nn.Module):
    def __init__(self, config, layer_index=None):
        super().__init__()
        self.self = RobertaSelfAttention(config, layer_index=layer_index)
        self.output = RobertaSelfOutput(config, layer_index=layer_index)
        self.pruned_heads = set()

    def prune_heads(self, heads):
        if len(heads) == 0:
            return
        heads, index = find_pruneable_heads_and_indices(
            heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads
        )

        # Prune linear layers
        self.self.query = prune_linear_layer(self.self.query, index)
        self.self.key = prune_linear_layer(self.self.key, index)
        self.self.value = prune_linear_layer(self.self.value, index)
        self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)

        # Update hyper params and store pruned heads
        self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
        self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads
        self.pruned_heads = self.pruned_heads.union(heads)

    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        past_key_value=None,
        output_attentions=False,
    ):
        self_outputs = self.self(
            hidden_states,
            attention_mask,
            head_mask,
            encoder_hidden_states,
            encoder_attention_mask,
            past_key_value,
            output_attentions,
        )
        attention_output = self.output(self_outputs[0])
        outputs = (attention_output,) + self_outputs[1:]  # add attentions if we output them
        return outputs


# Copied from transformers.models.bert.modeling_bert.BertIntermediate
class RobertaIntermediate(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
        if isinstance(config.hidden_act, str):
            self.intermediate_act_fn = ACT2FN[config.hidden_act]
        else:
            self.intermediate_act_fn = config.hidden_act

    def forward(self, hidden_states):
        hidden_states = self.dense(hidden_states)
        hidden_states = self.intermediate_act_fn(hidden_states)
        return hidden_states


# Copied from transformers.models.bert.modeling_bert.BertOutput
class RobertaOutput(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, hidden_states, input_tensor, last_norm=True):
        hidden_states = self.dense(hidden_states)
        hidden_states = self.dropout(hidden_states)
        hidden_states = hidden_states + input_tensor
        if last_norm:
            hidden_states = self.LayerNorm(hidden_states)
        return hidden_states


# Copied from transformers.models.bert.modeling_bert.BertLayer with Bert->Roberta
class RobertaLayer(nn.Module):
    def __init__(self, config, layer_index=None):
        super().__init__()
        self.chunk_size_feed_forward = config.chunk_size_feed_forward
        self.seq_len_dim = 1
        self.attention = RobertaAttention(config)
        self.is_decoder = config.is_decoder
        self.add_cross_attention = config.add_cross_attention
        if layer_index >= 12 - NUM_FUSE_BLOCK:
            self.crossattention_t2i = RobertaAttention(config, layer_index=layer_index)
            self.alpha_t2i = nn.Parameter(torch.Tensor([0]))
        self.intermediate = RobertaIntermediate(config)
        self.output = RobertaOutput(config)

    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        past_key_value=None,
        output_attentions=False,
        last_norm=True,
    ):
        # decoder uni-directional self-attention cached key/values tuple is at positions 1,2
        self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
        self_attention_outputs = self.attention(
            hidden_states,
            attention_mask,
            head_mask,
            output_attentions=output_attentions,
            past_key_value=self_attn_past_key_value,
        )
        attention_output = self_attention_outputs[0]

        # if decoder, the last output is tuple of self-attn cache
        outputs = self_attention_outputs[1:]  # add self attentions if we output attention weights

        cross_attn_present_key_value = None
        if encoder_hidden_states is not None:
            assert hasattr(
                self, "crossattention_t2i"
            ), f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers by setting `config.add_cross_attention=True`"

            # cross_attn cached key/values tuple is at positions 3,4 of past_key_value tuple
            cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None
            cross_attention_outputs = self.crossattention_t2i(
                attention_output,
                attention_mask,
                head_mask,
                encoder_hidden_states,
                encoder_attention_mask,
                cross_attn_past_key_value,
                output_attentions,
            )
            attention_output = self.alpha_t2i * cross_attention_outputs[0] + attention_output

        attention_output = self.attention.output.LayerNorm(attention_output + hidden_states)

        layer_output = apply_chunking_to_forward(
            self.feed_forward_chunk,
            self.chunk_size_feed_forward,
            self.seq_len_dim,
            attention_output,
            last_norm=last_norm,
        )
        outputs = (layer_output,) + outputs

        # if decoder, return the attn key/values as the last output
        return outputs

    def feed_forward_chunk(self, attention_output, last_norm=True):
        intermediate_output = self.intermediate(attention_output)
        layer_output = self.output(intermediate_output, attention_output, last_norm=last_norm)
        return layer_output


# Copied from transformers.models.bert.modeling_bert.BertEncoder with Bert->Roberta
class RobertaEncoder(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.config.gradient_checkpointing = config_yaml['use_checkpoint']
        self.layer = nn.ModuleList(
            [RobertaLayer(config, layer_index=layer_i) for layer_i in range(config.num_hidden_layers)]
        )

    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        past_key_values=None,
        use_cache=None,
        output_attentions=False,
        output_hidden_states=False,
        return_dict=True,
    ):
        all_hidden_states = () if output_hidden_states else None
        all_self_attentions = () if output_attentions else None
        all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None

        next_decoder_cache = () if use_cache else None
        for i, layer_module in enumerate(self.layer):
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

            layer_head_mask = head_mask[i] if head_mask is not None else None
            past_key_value = past_key_values[i] if past_key_values is not None else None

            if getattr(self.config, "gradient_checkpointing", False) and self.training:

                if use_cache:
                    logger.warning(
                        "`use_cache=True` is incompatible with `config.gradient_checkpointing=True`. Setting "
                        "`use_cache=False`..."
                    )
                    use_cache = False

                #def create_custom_forward(module):
                    #def custom_forward(*inputs):
                        #return module(*inputs, past_key_value, output_attentions)
                    #return custom_forward

                layer_outputs = torch.utils.checkpoint.checkpoint(
                    layer_module,
                    hidden_states,
                    attention_mask,
                    layer_head_mask,
                    encoder_hidden_states,
                    encoder_attention_mask,
                    past_key_value, 
                    output_attentions
                )
            else:
                layer_outputs = layer_module(
                    hidden_states,
                    attention_mask,
                    layer_head_mask,
                    encoder_hidden_states,
                    encoder_attention_mask,
                    past_key_value,
                    output_attentions,
                )


            hidden_states = layer_outputs[0]
            if use_cache:
                next_decoder_cache += (layer_outputs[-1],)
            if output_attentions:
                all_self_attentions = all_self_attentions + (layer_outputs[1],)
                if self.config.add_cross_attention:
                    all_cross_attentions = all_cross_attentions + (layer_outputs[2],)

        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        if not return_dict:
            return tuple(
                v
                for v in [
                    hidden_states,
                    next_decoder_cache,
                    all_hidden_states,
                    all_self_attentions,
                    all_cross_attentions,
                ]
                if v is not None
            )
        return BaseModelOutputWithPastAndCrossAttentions(
            last_hidden_state=hidden_states,
            past_key_values=next_decoder_cache,
            hidden_states=all_hidden_states,
            attentions=all_self_attentions,
            cross_attentions=all_cross_attentions,
        )

class RobertaPooler(nn.Module):
    def __init__(self, config):
        super().__init__()
        #self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        self.activation = nn.Tanh()

    def forward(self, hidden_states):
        # We "pool" the model by simply taking the hidden state corresponding
        # to the first token.
        first_token_tensor = hidden_states[:, 0]
        #pooled_output = self.dense(first_token_tensor)
        pooled_output = self.activation(first_token_tensor)
        return pooled_output


class RobertaPreTrainedModel(PreTrainedModel):
    """
    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
    models.
    """

    config_class = RobertaConfig
    base_model_prefix = "roberta"

    # Copied from transformers.models.bert.modeling_bert.BertPreTrainedModel._init_weights
    def _init_weights(self, module):
        """Initialize the weights"""
        if isinstance(module, nn.Linear):
            # Slightly different from the TF version which uses truncated_normal for initialization
            # cf https://github.com/pytorch/pytorch/pull/5617
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)


ROBERTA_START_DOCSTRING = r"""
    This model inherits from :class:`~transformers.PreTrainedModel`. Check the superclass documentation for the generic
    methods the library implements for all its model (such as downloading or saving, resizing the input embeddings,
    pruning heads etc.)
    This model is also a PyTorch `torch.nn.Module <https://pytorch.org/docs/stable/nn.html#torch.nn.Module>`__
    subclass. Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to
    general usage and behavior.
    Parameters:
        config (:class:`~transformers.RobertaConfig`): Model configuration class with all the parameters of the
            model. Initializing with a config file does not load the weights associated with the model, only the
            configuration. Check out the :meth:`~transformers.PreTrainedModel.from_pretrained` method to load the model
            weights.
"""

ROBERTA_INPUTS_DOCSTRING = r"""
    Args:
        input_ids (:obj:`torch.LongTensor` of shape :obj:`({0})`):
            Indices of input sequence tokens in the vocabulary.
            Indices can be obtained using :class:`~transformers.RobertaTokenizer`. See
            :meth:`transformers.PreTrainedTokenizer.encode` and :meth:`transformers.PreTrainedTokenizer.__call__` for
            details.
            `What are input IDs? <../glossary.html#input-ids>`__
        attention_mask (:obj:`torch.FloatTensor` of shape :obj:`({0})`, `optional`):
            Mask to avoid performing attention on padding token indices. Mask values selected in ``[0, 1]``:
            - 1 for tokens that are **not masked**,
            - 0 for tokens that are **masked**.
            `What are attention masks? <../glossary.html#attention-mask>`__
        token_type_ids (:obj:`torch.LongTensor` of shape :obj:`({0})`, `optional`):
            Segment token indices to indicate first and second portions of the inputs. Indices are selected in ``[0,
            1]``:
            - 0 corresponds to a `sentence A` token,
            - 1 corresponds to a `sentence B` token.
            `What are token type IDs? <../glossary.html#token-type-ids>`_
        position_ids (:obj:`torch.LongTensor` of shape :obj:`({0})`, `optional`):
            Indices of positions of each input sequence tokens in the position embeddings. Selected in the range ``[0,
            config.max_position_embeddings - 1]``.
            `What are position IDs? <../glossary.html#position-ids>`_
        head_mask (:obj:`torch.FloatTensor` of shape :obj:`(num_heads,)` or :obj:`(num_layers, num_heads)`, `optional`):
            Mask to nullify selected heads of the self-attention modules. Mask values selected in ``[0, 1]``:
            - 1 indicates the head is **not masked**,
            - 0 indicates the head is **masked**.
        inputs_embeds (:obj:`torch.FloatTensor` of shape :obj:`({0}, hidden_size)`, `optional`):
            Optionally, instead of passing :obj:`input_ids` you can choose to directly pass an embedded representation.
            This is useful if you want more control over how to convert :obj:`input_ids` indices into associated
            vectors than the model's internal embedding lookup matrix.
        output_attentions (:obj:`bool`, `optional`):
            Whether or not to return the attentions tensors of all attention layers. See ``attentions`` under returned
            tensors for more detail.
        output_hidden_states (:obj:`bool`, `optional`):
            Whether or not to return the hidden states of all layers. See ``hidden_states`` under returned tensors for
            more detail.
        return_dict (:obj:`bool`, `optional`):
            Whether or not to return a :class:`~transformers.file_utils.ModelOutput` instead of a plain tuple.
"""


@add_start_docstrings(
    "The bare RoBERTa Model transformer outputting raw hidden-states without any specific head on top.",
    ROBERTA_START_DOCSTRING,
)
class RobertaModel(RobertaPreTrainedModel):
    """
    The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of
    cross-attention is added between the self-attention layers, following the architecture described in `Attention is
    all you need`_ by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz
    Kaiser and Illia Polosukhin.
    To behave as an decoder the model needs to be initialized with the :obj:`is_decoder` argument of the configuration
    set to :obj:`True`. To be used in a Seq2Seq model, the model needs to initialized with both :obj:`is_decoder`
    argument and :obj:`add_cross_attention` set to :obj:`True`; an :obj:`encoder_hidden_states` is then expected as an
    input to the forward pass.
    .. _`Attention is all you need`: https://arxiv.org/abs/1706.03762
    """

    _keys_to_ignore_on_load_missing = [r"position_ids"]

    # Copied from transformers.models.bert.modeling_bert.BertModel.__init__ with Bert->Roberta
    def __init__(self, config, add_pooling_layer=True):
        super().__init__(config)
        self.config = config

        self.embeddings = RobertaEmbeddings(config)
        self.encoder = RobertaEncoder(config)

        self.pooler = RobertaPooler(config) if add_pooling_layer else None

        self.init_weights()

    def get_input_embeddings(self):
        return self.embeddings.word_embeddings

    def set_input_embeddings(self, value):
        self.embeddings.word_embeddings = value

    def _prune_heads(self, heads_to_prune):
        """
        Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
        class PreTrainedModel
        """
        for layer, heads in heads_to_prune.items():
            self.encoder.layer[layer].attention.prune_heads(heads)

    @add_start_docstrings_to_model_forward(ROBERTA_INPUTS_DOCSTRING.format("(batch_size, sequence_length)"))
    @add_code_sample_docstrings(
        processor_class=_TOKENIZER_FOR_DOC,
        checkpoint=_CHECKPOINT_FOR_DOC,
        output_type=BaseModelOutputWithPoolingAndCrossAttentions,
        config_class=_CONFIG_FOR_DOC,
    )
    # Copied from transformers.models.bert.modeling_bert.BertModel.forward
    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        past_key_values=None,
        use_cache=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
    ):
        r"""
        encoder_hidden_states  (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`):
            Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
            the model is configured as a decoder.
        encoder_attention_mask (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
            Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
            the cross-attention if the model is configured as a decoder. Mask values selected in ``[0, 1]``:
            - 1 for tokens that are **not masked**,
            - 0 for tokens that are **masked**.
        past_key_values (:obj:`tuple(tuple(torch.FloatTensor))` of length :obj:`config.n_layers` with each tuple having 4 tensors of shape :obj:`(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
            Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
            If :obj:`past_key_values` are used, the user can optionally input only the last :obj:`decoder_input_ids`
            (those that don't have their past key value states given to this model) of shape :obj:`(batch_size, 1)`
            instead of all :obj:`decoder_input_ids` of shape :obj:`(batch_size, sequence_length)`.
        use_cache (:obj:`bool`, `optional`):
            If set to :obj:`True`, :obj:`past_key_values` key value states are returned and can be used to speed up
            decoding (see :obj:`past_key_values`).
        """
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        use_cache = False

        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot secify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
            input_shape = input_ids.size()
            batch_size, seq_length = input_shape
        elif inputs_embeds is not None:
            input_shape = inputs_embeds.size()[:-1]
            batch_size, seq_length = input_shape
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        device = input_ids.device if input_ids is not None else inputs_embeds.device

        # past_key_values_length
        past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0

        if attention_mask is None:
            attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length)), device=device)
        if token_type_ids is None:
            token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)

        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
        # ourselves in which case we just need to make it broadcastable to all heads.
        extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape, device)

        # If a 2D or 3D attention mask is provided for the cross-attention
        # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
        if encoder_hidden_states is not None:
            encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
            encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
            if encoder_attention_mask is None:
                encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
            encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
        else:
            encoder_extended_attention_mask = None

        # Prepare head mask if needed
        # 1.0 in head_mask indicate we keep the head
        # attention_probs has shape bsz x n_heads x N x N
        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)

        embedding_output = self.embeddings(
            input_ids=input_ids,
            position_ids=position_ids,
            token_type_ids=token_type_ids,
            inputs_embeds=inputs_embeds,
            past_key_values_length=past_key_values_length,
        )
        encoder_outputs = self.encoder(
            embedding_output,
            attention_mask=extended_attention_mask,
            head_mask=head_mask,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_extended_attention_mask,
            past_key_values=past_key_values,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        sequence_output = encoder_outputs[0]
        pooled_output = self.pooler(sequence_output) if self.pooler is not None else None

        if not return_dict:
            return (sequence_output, pooled_output) + encoder_outputs[1:]

        return BaseModelOutputWithPoolingAndCrossAttentions(
            last_hidden_state=sequence_output,
            pooler_output=pooled_output,
            past_key_values=encoder_outputs.past_key_values,
            hidden_states=encoder_outputs.hidden_states,
            attentions=encoder_outputs.attentions,
            cross_attentions=encoder_outputs.cross_attentions,
        )


def create_position_ids_from_input_ids(input_ids, padding_idx, past_key_values_length=0):
    """
    Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding symbols
    are ignored. This is modified from fairseq's `utils.make_positions`.
    Args:
        x: torch.Tensor x:
    Returns: torch.Tensor
    """
    # The series of casts and type-conversions here are carefully balanced to both work with ONNX export and XLA.
    mask = input_ids.ne(padding_idx).int()
    incremental_indices = (torch.cumsum(mask, dim=1).type_as(mask) + past_key_values_length) * mask
    return incremental_indices.long() + padding_idx


def _make_causal_mask(input_ids_shape: torch.Size, dtype: torch.dtype, past_key_values_length: int = 0):
    """
    Make causal mask used for bi-directional self-attention.
    """
    bsz, tgt_len = input_ids_shape
    mask = torch.full((tgt_len, tgt_len), float("-inf"))
    mask_cond = torch.arange(mask.size(-1))
    mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
    mask = mask.to(dtype)

    if past_key_values_length > 0:
        mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype), mask], dim=-1)
    return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)


def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len=None):
    """
    Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
    """
    bsz, src_len = mask.size()
    tgt_len = tgt_len if tgt_len is not None else src_len

    expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)

    inverted_mask = 1.0 - expanded_mask

    return inverted_mask.masked_fill(inverted_mask.bool(), torch.finfo(dtype).min)


def _prepare_decoder_attention_mask(
    attention_mask, input_shape, inputs_embeds, device, past_key_values_length: int = 0
):
    # create causal mask
    # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
    combined_attention_mask = None
    if input_shape[-1] > 1:
        combined_attention_mask = _make_causal_mask(
            input_shape, inputs_embeds.dtype, past_key_values_length=past_key_values_length
        ).to(device)

    if attention_mask is not None:
        # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
        expanded_attn_mask = _expand_mask(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1])
        combined_attention_mask = (
            expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask
        )

    return combined_attention_mask

EgoVLPv2/model/roberta.py
facebookresearch/EgoVLPv2
"""
Implementations of Video Transformers in PyTorch

A PyTorch implementation of space-time transformer as described in
'Frozen in Time: A Joint Image and Video Encoder for End-to-End Retrieval' - https://arxiv.org/abs/2104.00650

A PyTorch implementation of timesformer as described in
'Is Space-Time Attention All You Need for Video Understanding?' - https://arxiv.org/abs/2102.05095

Acknowledgments:
- This code builds on Ross Wightman's vision_transformer code in pytorch-image-models:
https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py

- It is also inspired by lucidrains timesformer implementation:
https://github.com/lucidrains/TimeSformer-pytorch

Hacked together by Max Bain
"""

from collections import OrderedDict
from functools import partial
import yaml

import torch
from einops import rearrange, repeat
from timm.models.layers import DropPath, to_2tuple, trunc_normal_
from torch import einsum, nn

with open('./EgoNCE_MLM_ITM_Config.yml') as f:
    config_yaml = yaml.load(f, Loader=yaml.FullLoader)

NUM_FUSE_BLOCK = config_yaml['num_fuse_block']
DIM_TEXT = 768

def attn(q, k, v):
    sim = einsum('b i d, b j d -> b i j', q, k)
    attn = sim.softmax(dim=-1)
    out = einsum('b i j, b j d -> b i d', attn, v)
    return out


class Mlp(nn.Module):
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x


class VideoPatchEmbed(nn.Module):
    """ Video to Patch Embedding
    """

    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768,
                 num_frames=8):
        super().__init__()
        img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)
        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0]) * num_frames
        self.img_size = img_size
        self.patch_size = patch_size
        self.num_patches = num_patches
        self.num_frames = num_frames
        self.embed_dim = embed_dim
        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)

    def forward(self, x):
        B, F, C, H, W = x.shape
        assert F == self.num_frames, print(F, self.num_frames)
        x = x.view(-1, C, H, W)
        x = self.proj(x)
        return x


class VarAttention(nn.Module):
    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.,
                 initialize='random', dim_text=None, norm_layer=nn.LayerNorm, space_attn=True):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        # NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights
        self.scale = qk_scale or head_dim ** -0.5
        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.proj = nn.Linear(dim, dim)
        if initialize == 'zeros':
            self.qkv.weight.data.fill_(0)
            self.qkv.bias.data.fill_(0)
            # fill proj weight with 1 here to improve training dynamics. Otherwise temporal attention inputs
            # are multiplied by 0*0, which is hard for the model to move out of.
            self.proj.weight.data.fill_(1)
            self.proj.bias.data.fill_(0)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj_drop = nn.Dropout(proj_drop)

        self.softmax = nn.Softmax(dim=-1)

        if dim_text is not None and space_attn:
            self.qkv_text_i2t = nn.Linear(dim_text, dim * 2, bias=qkv_bias)
            self.qkv_i2t = nn.Linear(dim, dim, bias=qkv_bias)
            self.attn_drop_i2t = nn.Dropout(attn_drop)
            self.proj_i2t = nn.Linear(dim, dim)
            self.proj_drop_i2t = nn.Dropout(proj_drop)
            self.alpha_i2t = nn.Parameter(torch.Tensor([0]))
            self.norm_i2t_i = norm_layer(dim)

    def forward(self, x, einops_from, einops_to, y=None, y_mask=None, **einops_dims):
        h = self.num_heads
        # project x to q, k, v vaalues
        q, k, v = self.qkv(x).chunk(3, dim=-1)
        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))

        q = q*self.scale

        # splice out CLS token at index 1
        (cls_q, q_), (cls_k, k_), (cls_v, v_) = map(lambda t: (t[:, 0:1], t[:, 1:]), (q, k, v))

        # let CLS token attend to key / values of all patches across time and space
        cls_out = attn(cls_q, k, v)
        # rearrange across time or space
        q_, k_, v_ = map(lambda t: rearrange(t, f'{einops_from} -> {einops_to}', **einops_dims), (q_, k_, v_))

        # expand cls token keys and values across time or space and concat
        r = q_.shape[0] // cls_k.shape[0]
        cls_k, cls_v = map(lambda t: repeat(t, 'b () d -> (b r) () d', r=r), (cls_k, cls_v))

        k_ = torch.cat((cls_k, k_), dim=1)
        v_ = torch.cat((cls_v, v_), dim=1)

        # attention
        out = attn(q_, k_, v_)

        # merge back time or space
        out = rearrange(out, f'{einops_to} -> {einops_from}', **einops_dims)

        # concat back the cls token
        out = torch.cat((cls_out, out), dim=1)

        # merge back the heads
        out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
        ## to out
        x = self.proj(out)
        x = self.proj_drop(x)

        if y is not None:
            B_, N, C = x.shape
            B_text, N_text, C_text = y.shape

            kv_text = (
                self.qkv_text_i2t(y)
                .reshape(B_text, N_text, 2, self.num_heads, C // self.num_heads)
                .permute(2, 0, 3, 1, 4)
            )
            k_text, v_text = kv_text[0], kv_text[1]

            q_i2t = self.qkv_i2t(self.norm_i2t_i(x))
            q_i2t = q_i2t.reshape(B_, N, 1, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
            q_i2t = q_i2t[0]

            # image to text attention
            text_scale = k_text.size(-1) ** -0.5
            q_i2t = q_i2t * text_scale
            attn_i2t = q_i2t @ k_text.transpose(-2, -1)  # B_, nH, N, N_text

            # add image to text bias and text_mask
            if y_mask is not None:
                mask_and_i2t_bias = y_mask.view(B_text, 1, 1, N_text)
                attn_i2t = attn_i2t + mask_and_i2t_bias

            attn_i2t = self.softmax(attn_i2t)
            attn_i2t = self.attn_drop_i2t(attn_i2t)
            y = (attn_i2t @ v_text).transpose(1, 2).reshape(B_, N, C)
            y = self.proj_i2t(y)
            y = self.proj_drop_i2t(y)
            x = x + self.alpha_i2t * y

        return x


class SpaceTimeBlock(nn.Module):

    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, time_init='zeros',
                 attention_style='frozen-in-time', dim_text=None):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.attn = VarAttention(
            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop, dim_text=dim_text, 
            norm_layer=norm_layer, space_attn=True)

        self.timeattn = VarAttention(
            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop,
            initialize=time_init, dim_text=dim_text, norm_layer=norm_layer, space_attn=False)

        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
        self.norm3 = norm_layer(dim)

        self.attention_style = attention_style

    def forward(self, x, einops_from_space, einops_to_space, einops_from_time, einops_to_time,
                time_n, space_f, y=None, y_mask=None):

        time_output = self.timeattn(self.norm3(x), einops_from_time, einops_to_time, n=time_n, y=None, y_mask=None)
        time_residual = x + time_output
        space_output = self.attn(self.norm1(time_residual), einops_from_space,
                                 einops_to_space, f=space_f, y=y, y_mask=y_mask)
        if self.attention_style == 'frozen-in-time':
            space_residual = x + self.drop_path(space_output)
        else:
            raise NotImplementedError

        x = space_residual + self.drop_path(self.mlp(self.norm2(space_residual)))

        return x


class SpaceTimeTransformer(nn.Module):
    """ Vision Transformer

    A PyTorch impl of : `Space-Time Transformer` from Frozen-in-time  - by Max Bain.
        https://arxiv.org/abs/2104.00650

    Based off:
     - ViT implementation from the timm library [https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py]
    lucidrains timesformer implementation [https://github.com/lucidrains/TimeSformer-pytorch].

    Notable differences:
     - allows for variable length input frames (<= num_frames)
     - allows for variable length input resolution  (<= (img_size, img_size)) [UNTESTED]
     - different attention block mechanism
    """

    def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12,
                 num_heads=12, mlp_ratio=4., qkv_bias=True, qk_scale=None, representation_size=None,
                 drop_rate=0., attn_drop_rate=0., drop_path_rate=0., hybrid_backbone=None,
                 num_frames=8, time_init='rand', attention_style='frozen-in-time', norm_layer=nn.LayerNorm, dim_text=None):
        """
        Args:
            img_size (int, tuple): input image size
            patch_size (int, tuple): patch size
            in_chans (int): number of input channels
            num_classes (int): number of classes for classification head
            embed_dim (int): embedding dimension
            depth (int): depth of transformer
            num_heads (int): number of attention heads
            mlp_ratio (int): ratio of mlp hidden dim to embedding dim
            qkv_bias (bool): enable bias for qkv if True
            qk_scale (float): override default qk scale of head_dim ** -0.5 if set
            representation_size (Optional[int]): enable and set representation layer (pre-logits) to this value if set
            drop_rate (float): dropout rate
            attn_drop_rate (float): attention dropout rate
            drop_path_rate (float): stochastic depth rate
            hybrid_backbone (nn.Module): CNN backbone to use in-place of PatchEmbed module
            norm_layer: (nn.Module): normalization layer
            num_frames: (int) maximum number of frames expected as input
            time_init: (str) how to initialise the time attention layer, 'zeros' allows for the timesformer to start off
                        as ViT.
            attention_style: (str) how to attend to space and time.
        """
        super().__init__()
        self.num_classes = num_classes
        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
        self.num_frames = num_frames
        self.embed_dim = embed_dim
        norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
        print("######USING ATTENTION STYLE: ", attention_style)
        if hybrid_backbone is not None:
            raise NotImplementedError('hybrid backbone not implemented')
        else:
            self.patch_embed = VideoPatchEmbed(
                img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim, num_frames=num_frames)
        num_patches = self.patch_embed.num_patches
        self.patches_per_frame = num_patches // num_frames

        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.pos_embed = nn.Parameter(
            torch.zeros(1, self.patches_per_frame + 1,
                        embed_dim))  # remember to take pos_embed[1:] for tiling over time
        self.temporal_embed = nn.Parameter(torch.zeros(1, num_frames, embed_dim))

        self.pos_drop = nn.Dropout(p=drop_rate)

        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
        self.blocks = nn.ModuleList([
            SpaceTimeBlock(
                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
                drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer, time_init=time_init,
                attention_style=attention_style, dim_text=None if i < 6 else DIM_TEXT)
            for i in range(depth)])
        self.norm = norm_layer(embed_dim)

        # Representation layer
        if representation_size:
            self.num_features = representation_size
            self.pre_logits = nn.Sequential(OrderedDict([
                ('fc', nn.Linear(embed_dim, representation_size)),
                ('act', nn.Tanh())
            ]))
        else:
            self.pre_logits = nn.Identity()

        # Classifier head
        self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()

        trunc_normal_(self.pos_embed, std=.02)
        trunc_normal_(self.cls_token, std=.02)

        # if num_frames > 1, then we perform ViT inflation and initialise time attention to zero so not necessary.
        if num_frames == 1:
            self.apply(self._init_weights)

        ## einops transformations
        self.einops_from_space = 'b (f n) d'
        self.einops_to_space = '(b f) n d'
        self.einops_from_time = 'b (f n) d'
        self.einops_to_time = '(b n) f d'


    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    @torch.jit.ignore
    def no_weight_decay(self):
        return {'pos_embed', 'cls_token'}

    def get_classifier(self):
        return self.head

    def reset_classifier(self, num_classes, global_pool=''):
        self.num_classes = num_classes
        self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()

    def forward_features(self, x):
        b, curr_frames, channels, _, _ = x.shape
        x = self.patch_embed(x)
        x = x.flatten(2).transpose(2, 1)
        x = x.reshape(b, -1, self.patch_embed.embed_dim)

        BF = x.shape[0]
        cls_tokens = self.cls_token.expand(BF, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        # positional embed needs to be tiled for each frame (this does [1,2,3] --> [1,2,3,1,2,3]...)
        cls_embed = self.pos_embed[:, 0, :].unsqueeze(1)
        tile_pos_embed = self.pos_embed[:, 1:, :].repeat(1, self.num_frames, 1)
        # temporal embed needs to be repeated within each frame (this does [1,2,3] --> [1,1,1,2,2,2,3,3,3]...)
        tile_temporal_embed = self.temporal_embed.repeat_interleave(self.patches_per_frame, 1)
        total_pos_embed = tile_pos_embed + tile_temporal_embed
        total_pos_embed = torch.cat([cls_embed, total_pos_embed], dim=1)

        curr_patches = x.shape[1]
        x = x + total_pos_embed[:, :curr_patches]
        x = self.pos_drop(x)
        n = self.patches_per_frame
        f = curr_frames

        for blk in self.blocks:
            if config_yaml['use_checkpoint']:

                def create_custom_forward(module):
                    def custom_forward(*inputs):
                        return module(*inputs, time_n=n, space_f=f)

                    return custom_forward

                x = torch.utils.checkpoint.checkpoint(create_custom_forward(blk), x, self.einops_from_space, self.einops_to_space, self.einops_from_time,
                    self.einops_to_time)
            else:
                x = blk(x, self.einops_from_space, self.einops_to_space, self.einops_from_time,
                    self.einops_to_time, time_n=n, space_f=f)

        x = self.norm(x)[:, 0]
        x = self.pre_logits(x)

        return x

    def forward(self, x):
        x = self.forward_features(x)
        x = self.head(x)
        return x

if __name__ == "__main__":
    network = SpaceTimeTransformer(num_frames=4)
    data = torch.rand((3, 4, 3, 224, 224))
    network(data)

EgoVLPv2/model/video_transformer.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import os
import sys
import argparse
import collections
import transformers
import signal
import subprocess

from set_optim_schedule import set_schedule
import yaml

with open('./EgoNCE_MLM_ITM_Config.yml') as f:
    config_yaml = yaml.load(f, Loader=yaml.FullLoader)


import torch
import data_loader.data_loader as module_data
import model.loss as module_loss
import model.metric as module_metric
import model.model_epic_charades as module_arch
import utils.visualizer as module_vis
from parse_config import ConfigParser
from trainer.trainer_charades import Multi_Trainer_dist_Charades
from utils.util import replace_nested_dict_item
from tensorboardX import SummaryWriter
import warnings
warnings.filterwarnings("ignore")


parser = argparse.ArgumentParser(description='PyTorch Template')
parser.add_argument('--task_names', default='EgoNCE_MLM_ITM', type=str, help='Task_Names')
parser.add_argument('-c', '--config', type=str,
                help='config file path (default: None)')
parser.add_argument('-r', '--resume', default=None, type=str,
                help='path to latest checkpoint (default: None)')
parser.add_argument('-d', '--device', default=None, type=str,
                help='indices of GPUs to enable (default: all)')
parser.add_argument('-o', '--observe', action='store_true',
                help='Whether to observe (neptune)')
parser.add_argument('-l', '--launcher', choices=['none', 'pytorch'], default='none',help='job launcher')
parser.add_argument('-lr1', '--learning_rate1', type=float, default=2e-4)
parser.add_argument('-sc', '--schedule', default=[60, 80])
parser.add_argument('--print_freq', type=int, default=100, help="print loss after this number of steps")
parser.add_argument('--save_dir', type=str, help="dirctory for model saving")

CustomArgs = collections.namedtuple('CustomArgs', 'flags type target')
config = ConfigParser(parser)


def handle_sigusr1(signum, frame):
    os.system(f'scontrol requeue {os.getenv("SLURM_JOB_ID")}')
    exit()

def handle_sigterm(signum, frame):
    pass

def main():
    
    args = parser.parse_args()
    args.ngpus_per_node = torch.cuda.device_count()
    if 'SLURM_JOB_ID' in os.environ:
        print("If is being executed")
        # single-node and multi-node distributed training on SLURM cluster
        # requeue job on SLURM preemption
        signal.signal(signal.SIGUSR1, handle_sigusr1)
        signal.signal(signal.SIGTERM, handle_sigterm)
        # find a common host name on all nodes
        # assume scontrol returns hosts in the same order on all nodes
        cmd = 'scontrol show hostnames ' + os.getenv('SLURM_JOB_NODELIST')
        stdout = subprocess.check_output(cmd.split())
        host_name = stdout.decode().splitlines()[0]
        args.rank = int(os.getenv('SLURM_NODEID')) * args.ngpus_per_node
        args.world_size = int(os.getenv('SLURM_NNODES')) * args.ngpus_per_node
        args.dist_url = f'tcp://{host_name}:58472'
    else:
        # single-node distributed training
        print("Else is being executed")
        args.rank = 0
        args.dist_url = 'tcp://localhost:58472'
        args.world_size = args.ngpus_per_node
    torch.multiprocessing.spawn(main_worker, (args,), args.ngpus_per_node)


def main_worker(gpu, args):

    print("main worker started")
    args.rank += gpu
    torch.distributed.init_process_group(
        backend='nccl', init_method=args.dist_url,
        world_size=args.world_size, rank=args.rank)
    print("init processed finished")

    torch.cuda.set_device(gpu)
    torch.backends.cudnn.benchmark = True

    
    logger = config.get_logger('train')
    os.environ['TOKENIZERS_PARALLELISM'] = "false"
    if config['visualizer']['type'] != "":
        visualizer = config.initialize(
            name='visualizer',
            module=module_vis,
            exp_name=config['name'],
            web_dir=config._web_log_dir
        )
    else:
        visualizer = None

    tokenizer = transformers.AutoTokenizer.from_pretrained(config['arch']['args']['text_params']['model'],
                                                               TOKENIZERS_PARALLELISM=False)
    data_loader, valid_data_loader = init_dataloaders(config, module_data)
    
    if args.rank == 0:
        print('Train dataset: ', [x.n_samples for x in data_loader], ' samples')
        print('Val dataset: ', [x.n_samples for x in valid_data_loader], ' samples')
    # build model architecture, then print to console
    
    model = config.initialize('arch', module_arch)
    
    # get function handles of loss and metrics
    loss = config.initialize(name="loss", module=module_loss)
    metrics = [getattr(module_metric, met) for met in config['metrics']]
    # build optimizer, learning rate scheduler. delete every lines containing lr_scheduler for disabling scheduler

    ## set_schedule will be modified; currently same as pretraining
    max_steps = int(len(data_loader[0]) * config['trainer']['epochs'])
    if max_steps==0:
        max_steps = int(len(data_loader[0]) * 10)
    warmup_steps = config_yaml["warmup_steps"]
    if isinstance(config_yaml["warmup_steps"], float):
        warmup_steps = int(max_steps * warmup_steps)

    optimizer, scheduler = set_schedule(model, config, config_yaml, max_steps, warmup_steps)
    
    if args.rank == 0:
        logger.info(model)

    lr_scheduler = None
    writer = None

    if args.rank == 0:
        writer = SummaryWriter(log_dir=str(config.tf_dir))

    print("trainer should being here")
    trainer = Multi_Trainer_dist_Charades(args, model, loss, metrics, optimizer, scheduler, gpu,
                      config=config,
                      data_loader=data_loader,
                      valid_data_loader=valid_data_loader,
                      lr_scheduler=lr_scheduler,
                      visualizer=visualizer,
                      writer=writer,
                      tokenizer=tokenizer,
                      max_samples_per_epoch=config['trainer']['max_samples_per_epoch']) 
    trainer.train(gpu)
    

def init_dataloaders(config, module_data):
    """
    We need a way to change split from 'train' to 'val'.
    """
    if "type" in config["data_loader"] and "args" in config["data_loader"]:
        # then its a single dataloader
        data_loader = [config.initialize("data_loader", module_data)]
        config['data_loader']['args'] = replace_nested_dict_item(config['data_loader']['args'], 'split', 'val')
        config['data_loader']['args'] = replace_nested_dict_item(config['data_loader']['args'], 'batch_size', 1) ## Code is wrong for batch_size != 1
        valid_data_loader = [config.initialize("data_loader", module_data)]
    elif isinstance(config["data_loader"], list):
        data_loader = [config.initialize('data_loader', module_data, index=idx) for idx in
                       range(len(config['data_loader']))]
        new_cfg_li = []
        for dl_cfg in config['data_loader']:
            dl_cfg['args'] = replace_nested_dict_item(dl_cfg['args'], 'split', 'val')
            dl_cfg['args'] = replace_nested_dict_item(dl_cfg['args'], 'batch_size', 1) ## Code is wrong for batch_size != 1
            new_cfg_li.append(dl_cfg)
        config._config['data_loader'] = new_cfg_li
        valid_data_loader = [config.initialize('data_loader', module_data, index=idx) for idx in
                             range(len(config['data_loader']))]
    else:
        raise ValueError("Check data_loader config, not correct format.")

    return data_loader, valid_data_loader


if __name__ == '__main__':
    main()

EgoVLPv2/multinode_train_charades.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import os
import sys
import argparse
import collections
import transformers
import signal
import subprocess

from set_optim_schedule import set_schedule
import yaml

with open('./EgoNCE_MLM_ITM_Config.yml') as f:
    config_yaml = yaml.load(f, Loader=yaml.FullLoader)


import torch
import data_loader.data_loader as module_data
import model.loss as module_loss
import model.metric as module_metric
import model.model as module_arch
import utils.visualizer as module_vis
from parse_config import ConfigParser
from trainer.trainer_egoclip import Multi_Trainer_dist
from utils.util import replace_nested_dict_item
from tensorboardX import SummaryWriter
import warnings
warnings.filterwarnings("ignore")


parser = argparse.ArgumentParser(description='PyTorch Template')
parser.add_argument('--task_names', default='EgoNCE_MLM_ITM', type=str, help='Task_Names')
parser.add_argument('-c', '--config', default='/fsx/spraman3/Video_Language_Pretraining/Pre-training/EgoVLP_multinode/configs/pt/egoclip.json', type=str,
                help='config file path (default: None)')
parser.add_argument('-r', '--resume', default=None, type=str,
                help='path to latest checkpoint (default: None)')
parser.add_argument('-d', '--device', default=None, type=str,
                help='indices of GPUs to enable (default: all)')
parser.add_argument('-o', '--observe', action='store_true',
                help='Whether to observe (neptune)')
parser.add_argument('-l', '--launcher', choices=['none', 'pytorch'], default='none',help='job launcher')
parser.add_argument('-lr1', '--learning_rate1', type=float, default=2e-4)
parser.add_argument('-sc', '--schedule', default=[60, 80])
parser.add_argument('--print_freq', type=int, default=100, help="print loss after this number of steps")
parser.add_argument('--save_dir', type=str, help="dirctory for model saving")

CustomArgs = collections.namedtuple('CustomArgs', 'flags type target')
config = ConfigParser(parser)


def handle_sigusr1(signum, frame):
    os.system(f'scontrol requeue {os.getenv("SLURM_JOB_ID")}')
    exit()

def handle_sigterm(signum, frame):
    pass

def main():
    
    args = parser.parse_args()
    args.ngpus_per_node = torch.cuda.device_count()
    if 'SLURM_JOB_ID' in os.environ:
        print("If is being executed")
        # single-node and multi-node distributed training on SLURM cluster
        # requeue job on SLURM preemption
        signal.signal(signal.SIGUSR1, handle_sigusr1)
        signal.signal(signal.SIGTERM, handle_sigterm)
        # find a common host name on all nodes
        # assume scontrol returns hosts in the same order on all nodes
        cmd = 'scontrol show hostnames ' + os.getenv('SLURM_JOB_NODELIST')
        stdout = subprocess.check_output(cmd.split())
        host_name = stdout.decode().splitlines()[0]
        args.rank = int(os.getenv('SLURM_NODEID')) * args.ngpus_per_node
        args.world_size = int(os.getenv('SLURM_NNODES')) * args.ngpus_per_node
        args.dist_url = f'tcp://{host_name}:58472'
    else:
        # single-node distributed training
        print("Else is being executed")
        args.rank = 0
        args.dist_url = 'tcp://localhost:58472'
        args.world_size = args.ngpus_per_node
    torch.multiprocessing.spawn(main_worker, (args,), args.ngpus_per_node)


def main_worker(gpu, args):

    print("main worker started")
    args.rank += gpu
    torch.distributed.init_process_group(
        backend='nccl', init_method=args.dist_url,
        world_size=args.world_size, rank=args.rank)
    print("init processed finished")

    torch.cuda.set_device(gpu)
    torch.backends.cudnn.benchmark = True

    
    logger = config.get_logger('train')
    os.environ['TOKENIZERS_PARALLELISM'] = "false"
    if config['visualizer']['type'] != "":
        visualizer = config.initialize(
            name='visualizer',
            module=module_vis,
            exp_name=config['name'],
            web_dir=config._web_log_dir
        )
    else:
        visualizer = None


    # build tokenizer
    tokenizer = transformers.AutoTokenizer.from_pretrained(config['arch']['args']['text_params']['model'],
                                                               TOKENIZERS_PARALLELISM=False)

    # setup data_loader instances
    data_loader, valid_data_loader = init_dataloaders(config, module_data)
    
    if args.rank == 0:
        print('Train dataset: ', [x.n_samples for x in data_loader], ' samples')
        print('Val dataset: ', [x.n_samples for x in valid_data_loader], ' samples')
    
    # build model architecture, then print to console
    model = config.initialize('arch', module_arch)

    if args.rank == 0:
        logger.info(model)

    
    # get function handles of loss and metrics
    loss = config.initialize(name="loss", module=module_loss)
    metrics = [getattr(module_metric, met) for met in config['metrics']]


    # build optimizer, learning rate scheduler. delete every lines containing lr_scheduler for disabling scheduler
    max_steps = int(len(data_loader[0]) * config['trainer']['epochs'])
    if max_steps==0:
        max_steps = int(len(data_loader[0]) * 10)
    warmup_steps = config_yaml["warmup_steps"]
    if isinstance(config_yaml["warmup_steps"], float):
        warmup_steps = int(max_steps * warmup_steps)

    optimizer, scheduler = set_schedule(model, config, config_yaml, max_steps, warmup_steps)

    lr_scheduler = None
    writer = None

    if args.rank == 0:
        writer = SummaryWriter(log_dir=str(config.tf_dir))

    print("trainer should being here")
    trainer = Multi_Trainer_dist(args, model, loss, metrics, optimizer, scheduler, gpu,
                      config=config,
                      data_loader=data_loader,
                      valid_data_loader=valid_data_loader,
                      lr_scheduler=lr_scheduler,
                      visualizer=visualizer,
                      writer=writer,
                      tokenizer=tokenizer,
                      max_samples_per_epoch=config['trainer']['max_samples_per_epoch'])

    
    trainer.train(gpu)
    

def init_dataloaders(config, module_data):
    """
    We need a way to change split from 'train' to 'val'.
    """
    if "type" in config["data_loader"] and "args" in config["data_loader"]:
        # then its a single dataloader
        data_loader = [config.initialize("data_loader", module_data)]
        config['data_loader']['args'] = replace_nested_dict_item(config['data_loader']['args'], 'split', 'val')
        config['data_loader']['args'] = replace_nested_dict_item(config['data_loader']['args'], 'batch_size', 1) ## Code is wrong for batch_size != 1
        valid_data_loader = [config.initialize("data_loader", module_data)]
    elif isinstance(config["data_loader"], list):
        data_loader = [config.initialize('data_loader', module_data, index=idx) for idx in
                       range(len(config['data_loader']))]
        new_cfg_li = []
        for dl_cfg in config['data_loader']:
            dl_cfg['args'] = replace_nested_dict_item(dl_cfg['args'], 'split', 'val')
            dl_cfg['args'] = replace_nested_dict_item(dl_cfg['args'], 'batch_size', 1) ## Code is wrong for batch_size != 1
            new_cfg_li.append(dl_cfg)
        config._config['data_loader'] = new_cfg_li
        valid_data_loader = [config.initialize('data_loader', module_data, index=idx) for idx in
                             range(len(config['data_loader']))]
    else:
        raise ValueError("Check data_loader config, not correct format.")

    return data_loader, valid_data_loader


if __name__ == '__main__':
    main()

EgoVLPv2/multinode_train_egoclip.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import os
import sys
import argparse
import collections
import transformers
import signal
import subprocess
from set_optim_schedule import set_schedule
import yaml

with open('./EgoNCE_MLM_ITM_Config.yml') as f:
    config_yaml = yaml.load(f, Loader=yaml.FullLoader)

import torch
import data_loader.data_loader as module_data
import model.loss as module_loss
import model.metric as module_metric
import model.model_epic_charades as module_arch
import utils.visualizer as module_vis
from parse_config import ConfigParser
from trainer.trainer_epic import Multi_Trainer_dist_MIR
from utils.util import replace_nested_dict_item
from tensorboardX import SummaryWriter
import warnings
warnings.filterwarnings("ignore")


parser = argparse.ArgumentParser(description='PyTorch Template')
parser.add_argument('--task_names', default='EgoNCE_ITM_MLM', type=str, help='Task_Names')
parser.add_argument('-c', '--config', default='/fsx/spraman3/Video_Language_Pretraining/Pre-training/EgoVLP_multinode/configs/pt/egoclip.json', type=str,
                help='config file path (default: None)')
parser.add_argument('-r', '--resume', default=None, type=str,
                help='path to latest checkpoint (default: None)')
parser.add_argument('-d', '--device', default=None, type=str,
                help='indices of GPUs to enable (default: all)')
parser.add_argument('-o', '--observe', action='store_true',
                help='Whether to observe (neptune)')
parser.add_argument('-l', '--launcher', choices=['none', 'pytorch'], default='none',help='job launcher')
parser.add_argument('-lr1', '--learning_rate1', type=float, default=2e-4)
parser.add_argument('-sc', '--schedule', default=[60, 80])
parser.add_argument('--print_freq', type=int, default=100, help="print loss after this number of steps")
parser.add_argument('--save_dir', type=str, help="dirctory for model saving")

CustomArgs = collections.namedtuple('CustomArgs', 'flags type target')
config = ConfigParser(parser)


def handle_sigusr1(signum, frame):
    os.system(f'scontrol requeue {os.getenv("SLURM_JOB_ID")}')
    exit()

def handle_sigterm(signum, frame):
    pass

def main():
    
    args = parser.parse_args()
    args.ngpus_per_node = torch.cuda.device_count()
    if 'SLURM_JOB_ID' in os.environ:
        print("If is being executed")
        signal.signal(signal.SIGUSR1, handle_sigusr1)
        signal.signal(signal.SIGTERM, handle_sigterm)
        cmd = 'scontrol show hostnames ' + os.getenv('SLURM_JOB_NODELIST')
        stdout = subprocess.check_output(cmd.split())
        host_name = stdout.decode().splitlines()[0]
        args.rank = int(os.getenv('SLURM_NODEID')) * args.ngpus_per_node
        args.world_size = int(os.getenv('SLURM_NNODES')) * args.ngpus_per_node
        args.dist_url = f'tcp://{host_name}:58472'
    else:
        print("Else is being executed")
        args.rank = 0
        args.dist_url = 'tcp://localhost:58472'
        args.world_size = args.ngpus_per_node
    torch.multiprocessing.spawn(main_worker, (args,), args.ngpus_per_node)


def main_worker(gpu, args):

    print("main worker started")
    args.rank += gpu
    torch.distributed.init_process_group(
        backend='nccl', init_method=args.dist_url,
        world_size=args.world_size, rank=args.rank)
    print("init processed finished")

    torch.cuda.set_device(gpu)
    torch.backends.cudnn.benchmark = True
    print("device set")

    
    logger = config.get_logger('train')
    os.environ['TOKENIZERS_PARALLELISM'] = "false"
    if config['visualizer']['type'] != "":
        visualizer = config.initialize(
            name='visualizer',
            module=module_vis,
            exp_name=config['name'],
            web_dir=config._web_log_dir
        )
    else:
        visualizer = None
        
    # build tokenizer
    tokenizer = transformers.AutoTokenizer.from_pretrained(config['arch']['args']['text_params']['model'],
                                                               TOKENIZERS_PARALLELISM=False)

    # setup data_loader instances
    data_loader, valid_data_loader = init_dataloaders(config, module_data)
    
    if args.rank == 0:
        print('Train dataset: ', [x.n_samples for x in data_loader], ' samples')
        print('Val dataset: ', [x.n_samples for x in valid_data_loader], ' samples')
    # build model architecture, then print to console
    
    print("dataloader initialized")

    
    model = config.initialize('arch', module_arch)

    if args.rank == 0:
        logger.info(model)

    
    # get function handles of loss and metrics
    loss = config.initialize(name="loss", module=module_loss)
    metrics = [getattr(module_metric, met) for met in config['metrics']]
    # build optimizer, learning rate scheduler. delete every lines containing lr_scheduler for disabling scheduler
    
    ## set_schedule will be modified; currently same as pretraining
    max_steps = int(len(data_loader[0]) * config['trainer']['epochs'])
    if max_steps==0:
        max_steps = int(len(data_loader[0]) * 10)
    warmup_steps = config_yaml["warmup_steps"]
    if isinstance(config_yaml["warmup_steps"], float):
        warmup_steps = int(max_steps * warmup_steps)

    optimizer, scheduler = set_schedule(model, config, config_yaml, max_steps, warmup_steps)
    
    lr_scheduler = None
    writer = None

    if args.rank == 0:
        writer = SummaryWriter(log_dir=str(config.tf_dir))

    print("trainer should being here")
    trainer = Multi_Trainer_dist_MIR(args, model, loss, metrics, optimizer, scheduler, gpu,
                      config=config,
                      data_loader=data_loader,
                      valid_data_loader=valid_data_loader,
                      lr_scheduler=lr_scheduler,
                      visualizer=visualizer,
                      writer=writer,
                      tokenizer=tokenizer,
                      max_samples_per_epoch=config['trainer']['max_samples_per_epoch'])

    
    trainer.train(gpu)
    

def init_dataloaders(config, module_data):
    """
    We need a way to change split from 'train' to 'val'.
    """
    if "type" in config["data_loader"] and "args" in config["data_loader"]:
        # then its a single dataloader
        data_loader = [config.initialize("data_loader", module_data)]
        config['data_loader']['args'] = replace_nested_dict_item(config['data_loader']['args'], 'split', 'val')
        valid_data_loader = [config.initialize("data_loader", module_data)]
    elif isinstance(config["data_loader"], list):
        data_loader = [config.initialize('data_loader', module_data, index=idx) for idx in
                       range(len(config['data_loader']))]
        new_cfg_li = []
        for dl_cfg in config['data_loader']:
            dl_cfg['args'] = replace_nested_dict_item(dl_cfg['args'], 'split', 'val')
            new_cfg_li.append(dl_cfg)
        config._config['data_loader'] = new_cfg_li
        valid_data_loader = [config.initialize('data_loader', module_data, index=idx) for idx in
                             range(len(config['data_loader']))]
    else:
        raise ValueError("Check data_loader config, not correct format.")

    return data_loader, valid_data_loader


if __name__ == '__main__':
    main()

EgoVLPv2/multinode_train_epic.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import os
import logging
import pdb
from pathlib import Path
from functools import reduce
from operator import getitem
from datetime import datetime
# import pdb; pdb.set_trace()
from logger import setup_logging
from utils import read_json, write_json
import time
import inspect


class ConfigParser:
    def __init__(self, args, options='', timestamp=True, test=False, eval_mode=None):
        # parse default and custom cli options
        for opt in options:
            args.add_argument(*opt.flags, default=None, type=opt.type)
        args = args.parse_args()
        self.args = args
        if args.device:
            os.environ["CUDA_VISIBLE_DEVICES"] = args.device
        if args.resume is None:
            msg_no_cfg = "Configuration file need to be specified. Add '-c config.json', for example."
            assert args.config is not None, msg_no_cfg
            self.cfg_fname = Path(args.config)
            config = read_json(self.cfg_fname)
            self.resume = None
        else:
            self.resume = Path(args.resume)
            resume_cfg_fname = self.resume.parent / 'config.json'
            if eval_mode == "epic":
                resume_cfg_fname = Path('configs/eval/epic.json')
            if eval_mode == "charades":
                resume_cfg_fname = Path('configs/eval/charades.json')
            if eval_mode == "nlq":
                resume_cfg_fname = Path('configs/eval/nlq.json')
            if eval_mode == "mq":
                resume_cfg_fname = Path('configs/eval/mq.json')

            config = read_json(resume_cfg_fname)
            if args.config is not None:
                config.update(read_json(Path(args.config)))

        # load config file and apply custom cli options
        self._config = _update_config(config, options, args)

        # set save_dir where trained model and log will be saved.
        #save_dir = Path(self.config['trainer']['save_dir'])
        save_dir = Path(args.save_dir)
        # timestamp = datetime.now().strftime(r'%m%d_%H%M%S') if timestamp else ''
        timestamp = datetime.now().strftime(r'%m%d_%H') if timestamp else ''

        exper_name = self.config['name']

        self._save_dir = save_dir / 'models' /  timestamp
        self._web_log_dir = save_dir / 'web' /  timestamp
        self._log_dir = save_dir / 'log' /  timestamp
        self._tf_dir = save_dir / 'tf' / timestamp

        if not test:
            self.save_dir.mkdir(parents=True, exist_ok=True)
            self.log_dir.mkdir(parents=True, exist_ok=True)
            self._tf_dir.mkdir(parents=True, exist_ok=True)

        # if set, remove all previous experiments with the current config
        if vars(args).get("purge_exp_dir", False):
            for dirpath in (self._save_dir, self._log_dir, self._web_log_dir):
                config_dir = dirpath.parent
                existing = list(config_dir.glob("*"))
                print(f"purging {len(existing)} directories from config_dir...")
                tic = time.time()
                os.system(f"rm -rf {config_dir}")
                print(f"Finished purge in {time.time() - tic:.3f}s")

        # save updated config file to the checkpoint dir
        if not test:
            write_json(self.config, self.save_dir / 'config.json')

            # configure logging module
            setup_logging(self.log_dir)
            self.log_levels = {
                0: logging.WARNING,
                1: logging.INFO,
                2: logging.DEBUG
            }

    def initialize(self, name, module, *args, index=None, **kwargs):
        """
        finds a function handle with the name given as 'type' in config, and returns the
        instance initialized with corresponding keyword args given as 'args'.
        """
        if index is None:
            module_name = self[name]['type']
            module_args = dict(self[name]['args'])
            assert all([k not in module_args for k in kwargs]), 'Overwriting kwargs given in config file is not allowed'
            module_args.update(kwargs)
        else:
            module_name = self[name][index]['type']
            module_args = dict(self[name][index]['args'])

        # if parameter not in config subdict, then check if it's in global config.
        signature = inspect.signature(getattr(module, module_name).__init__)
        print(module_name)
        for param in signature.parameters.keys():
            if param not in module_args and param in self.config:
                module_args[param] = self[param]
            if module_name == 'FrozenInTime' and param == 'args':
                module_args[param] = self.args
            if module_name == 'MultiDistTextVideoDataLoader' and param == 'args':
                module_args[param] = self.args

        return getattr(module, module_name)(*args, **module_args)

    def __getitem__(self, name):
        return self.config[name]

    def get_logger(self, name, verbosity=2):
        msg_verbosity = 'verbosity option {} is invalid. Valid options are {}.'.format(verbosity,
                                                                                       self.log_levels.keys())
        assert verbosity in self.log_levels, msg_verbosity
        logger = logging.getLogger(name)
        logger.setLevel(self.log_levels[verbosity])
        return logger

    # setting read-only attributes
    @property
    def config(self):
        return self._config

    @property
    def save_dir(self):
        return self._save_dir

    @property
    def log_dir(self):
        return self._log_dir

    @property
    def tf_dir(self):
        return self._tf_dir

# helper functions used to update config dict with custom cli options
def _update_config(config, options, args):
    for opt in options:
        value = getattr(args, _get_opt_name(opt.flags))
        if value is not None:
            _set_by_path(config, opt.target, value)
    return config


def _get_opt_name(flags):
    for flg in flags:
        if flg.startswith('--'):
            return flg.replace('--', '')
    return flags[0].replace('--', '')


def _set_by_path(tree, keys, value):
    """Set a value in a nested object in tree by sequence of keys."""
    _get_by_path(tree, keys[:-1])[keys[-1]] = value


def _get_by_path(tree, keys):
    """Access a nested object in tree by sequence of keys."""
    return reduce(getitem, keys, tree)

EgoVLPv2/parse_config.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


from transformers.optimization import AdamW
from transformers import (
    get_constant_schedule,
    get_polynomial_decay_schedule_with_warmup,
    get_cosine_schedule_with_warmup,
)


def set_schedule(model, config, config_yaml, max_steps, warmup_steps):
    lr = config["optimizer"]["args"]["lr"]
    wd = config["optimizer"]["args"]["weight_decay"]

    no_decay = [
        "bias",
        "LayerNorm.bias",
        "LayerNorm.weight",
        "norm.bias",
        "norm.weight",
        "norm1.bias",
        "norm1.weight",
        "norm2.bias",
        "norm2.weight",
    ]
    head_names = ["mlm_score", "itm_score", "txt_proj", "vid_proj"]
    cross_modal_names = ["cross_modal", "i2t", "t2i"]
    lr_mult_head = config["optimizer"]["args"]["lr_mult_head"]
    lr_mult_cross_modal = config["optimizer"]["args"]["lr_mult_cross_modal"]
    end_lr = config_yaml["end_lr"]
    decay_power = config_yaml["decay_power"]
    optim_type = config["optimizer"]["type"]
    optimizer_grouped_parameters = [
        {
            "params": [
                p
                for n, p in model.named_parameters()
                if not any(nd in n for nd in no_decay)
                and not any(bb in n for bb in head_names)
                and not any(ht in n for ht in cross_modal_names)
            ],
            "weight_decay": wd,
            "lr": lr,
        },
        {
            "params": [
                p
                for n, p in model.named_parameters()
                if any(nd in n for nd in no_decay)
                and not any(bb in n for bb in head_names)
                and not any(ht in n for ht in cross_modal_names)
            ],
            "weight_decay": 0.0,
            "lr": lr,
        },
        {
            "params": [
                p
                for n, p in model.named_parameters()
                if not any(nd in n for nd in no_decay)
                and any(bb in n for bb in head_names)
                and not any(ht in n for ht in cross_modal_names)
            ],
            "weight_decay": wd,
            "lr": lr * lr_mult_head,
        },
        {
            "params": [
                p
                for n, p in model.named_parameters()
                if any(nd in n for nd in no_decay)
                and any(bb in n for bb in head_names)
                and not any(ht in n for ht in cross_modal_names)
            ],
            "weight_decay": 0.0,
            "lr": lr * lr_mult_head,
        },
        {
            "params": [
                p
                for n, p in model.named_parameters()
                if not any(nd in n for nd in no_decay)
                and not any(bb in n for bb in head_names)
                and any(ht in n for ht in cross_modal_names)
            ],
            "weight_decay": wd,
            "lr": lr * lr_mult_cross_modal,
        },
        {
            "params": [
                p
                for n, p in model.named_parameters()
                if any(nd in n for nd in no_decay)
                and not any(bb in n for bb in head_names)
                and any(ht in n for ht in cross_modal_names)
            ],
            "weight_decay": 0.0,
            "lr": lr * lr_mult_cross_modal,
        },
    ]

    if optim_type == "AdamW":
        optimizer = AdamW(optimizer_grouped_parameters, lr=lr, eps=1e-8, betas=(0.9, 0.98))
    elif optim_type == "adam":
        optimizer = torch.optim.Adam(optimizer_grouped_parameters, lr=lr)
    elif optim_type == "sgd":
        optimizer = torch.optim.SGD(optimizer_grouped_parameters, lr=lr, momentum=0.9)

    if decay_power == "cosine":
        scheduler = get_cosine_schedule_with_warmup(
            optimizer,
            num_warmup_steps=warmup_steps,
            num_training_steps=max_steps,
        )
    else:
        scheduler = get_polynomial_decay_schedule_with_warmup(
            optimizer,
            num_warmup_steps=warmup_steps,
            num_training_steps=max_steps,
            lr_end=end_lr,
            power=decay_power,
        )

    return optimizer, scheduler


def set_schedule_constant(model, config, config_yaml, max_steps, warmup_steps):
    
    lr = config["optimizer"]["args"]["lr"]
    wd = config["optimizer"]["args"]["weight_decay"]

    no_decay = [
        "bias",
        "LayerNorm.bias",
        "LayerNorm.weight",
        "norm.bias",
        "norm.weight",
        "norm1.bias",
        "norm1.weight",
        "norm2.bias",
        "norm2.weight",
    ]
    head_names = ["mlm_score", "itm_score", "txt_proj", "vid_proj"]
    cross_modal_names = ["cross_modal", "i2t", "t2i"]
    lr_mult_head = config["optimizer"]["args"]["lr_mult_head"]
    lr_mult_cross_modal = config["optimizer"]["args"]["lr_mult_cross_modal"]
    end_lr = config_yaml["end_lr"]
    decay_power = config_yaml["decay_power"]
    optim_type = config["optimizer"]["type"]

    optimizer_grouped_parameters = [
        {
            "params": [
                p
                for n, p in model.named_parameters()
                if not any(nd in n for nd in no_decay)
                and not any(bb in n for bb in head_names)
                and not any(ht in n for ht in cross_modal_names)
            ],
            "weight_decay": wd,
            "lr": lr,
        },
        {
            "params": [
                p
                for n, p in model.named_parameters()
                if any(nd in n for nd in no_decay)
                and not any(bb in n for bb in head_names)
                and not any(ht in n for ht in cross_modal_names)
            ],
            "weight_decay": 0.0,
            "lr": lr,
        },
        {
            "params": [
                p
                for n, p in model.named_parameters()
                if not any(nd in n for nd in no_decay)
                and any(bb in n for bb in head_names)
                and not any(ht in n for ht in cross_modal_names)
            ],
            "weight_decay": wd,
            "lr": lr * lr_mult_head,
        },
        {
            "params": [
                p
                for n, p in model.named_parameters()
                if any(nd in n for nd in no_decay)
                and any(bb in n for bb in head_names)
                and not any(ht in n for ht in cross_modal_names)
            ],
            "weight_decay": 0.0,
            "lr": lr * lr_mult_head,
        },
        {
            "params": [
                p
                for n, p in model.named_parameters()
                if not any(nd in n for nd in no_decay)
                and not any(bb in n for bb in head_names)
                and any(ht in n for ht in cross_modal_names)
            ],
            "weight_decay": wd,
            "lr": lr * lr_mult_cross_modal,
        },
        {
            "params": [
                p
                for n, p in model.named_parameters()
                if any(nd in n for nd in no_decay)
                and not any(bb in n for bb in head_names)
                and any(ht in n for ht in cross_modal_names)
            ],
            "weight_decay": 0.0,
            "lr": lr * lr_mult_cross_modal,
        },
    ]

    if optim_type == "AdamW":
        optimizer = AdamW(model.parameters(), lr=lr)
        #optimizer = AdamW(list(model.txt_proj.parameters()) + list(model.vid_proj.parameters()), lr=lr, eps=1e-8, betas=(0.9, 0.98))
    elif optim_type == "adam":
        optimizer = torch.optim.Adam(optimizer_grouped_parameters, lr=lr)
    elif optim_type == "sgd":
        optimizer = torch.optim.SGD(optimizer_grouped_parameters, lr=lr, momentum=0.9)

    scheduler = get_constant_schedule(optimizer)

    return optimizer, scheduler

EgoVLPv2/set_optim_schedule.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import os
import sys
import tqdm
import argparse
import numpy as np
import transformers
import torch
import torch.nn as nn
import model.metric as module_metric
import data_loader.data_loader as module_data
from utils import state_dict_data_parallel_fix
from parse_config import ConfigParser
import pdb
from model.model import FrozenInTime
from transformers import RobertaTokenizer
import model.model as module_arch


def run():
    # setup data_loader instances
    config._config['data_loader']['type'] = 'TextVideoDataLoader'
    config._config['data_loader']['args']['split'] = args.split
    config._config['data_loader']['args']['tsfm_split'] = 'test'  # set transform to test split to remove augmentations
    config._config['data_loader']['args']['shuffle'] = False
    config._config['data_loader']['args']['batch_size'] = args.batch_size
    config._config['data_loader']['args']['sliding_window_stride'] = args.sliding_window_stride

    data_loader = config.initialize('data_loader', module_data)

    model = config.initialize('arch', module_arch)

    # get function handles of loss and metrics
    metric_fns = [getattr(module_metric, met) for met in config['metrics']]

    device = torch.device(args.cuda_base if torch.cuda.is_available() else 'cpu')
    model = model.to(device)
    model = nn.DataParallel(model, device_ids = args.device_ids)
    model.eval()


    print(len(data_loader))

    if not os.path.exists(args.save_dir):
        os.mkdir(args.save_dir)

    # extract clip features
    num_frame = config.config['data_loader']['args']['video_params']['num_frames']
    dim = config.config['arch']['args']['projection_dim']
    with torch.no_grad():
        for i, data in enumerate(tqdm.tqdm(data_loader)):
            if os.path.exists(os.path.join(args.save_dir, data['meta']['clip_uid'][0]+'.pt')):
                print(f"{data['meta']['clip_uid']} is already.")
                continue
            # this implementation is cautious, we use 4f video-encoder to extract featurs of whole clip.
            f, c, h, w = data['video'].shape[1], data['video'].shape[2], data['video'].shape[3], data['video'].shape[4]
            data['video'] = data['video'][0][:(f // num_frame * num_frame), ]
            data['video'] = data['video'].reshape(-1, num_frame, c, h, w)

            data['video'] = data['video'].to(device)
            outs = torch.zeros(data['video'].shape[0], dim)
            
            b_s = 64 ## Batch size inside the next loop, as (897, 16, 3, 224, 224) goes out of memory

            if data['video'].shape[0] % b_s == 0:
                times = data['video'].shape[0] // b_s
            else:
                times = 1 + data['video'].shape[0] // b_s  ##  '1 +' is added by us

            for j in range(times):
                start = j*b_s
                if (j+1) * b_s > data['video'].shape[0]:
                    end = data['video'].shape[0]
                else:
                    end = (j+1)*b_s
                    
                data_batch = {'video': data['video'][start:end,]} 
                video_embeds = model(data=data_batch, n_embeds=None, v_embeds=None, allgather=None, n_gpu=None, args=None, config=None, loss_egonce=None, gpu=None, task_names='Feature_Extraction')
                outs[start:end,] = video_embeds

            torch.save(outs, os.path.join(args.save_dir, data['meta']['clip_uid'][0]+'.pt'))

if __name__ == '__main__':
    args = argparse.ArgumentParser(description='PyTorch Template')

    args.add_argument('-r', '--resume',
                      help='path to latest checkpoint (default: None)')
    args.add_argument('-d', '--device', default=None, type=str,
                      help='indices of GPUs to enable (default: all)')
    args.add_argument('-c', '--config', default=None, type=str,
                      help='config file path (default: None)')
    args.add_argument('-s', '--sliding_window_stride', default=-1, type=int,
                      help='test time temporal augmentation, repeat samples with different start times.')
    args.add_argument('--split', default='test', choices=['train', 'val', 'test'],
                      help='split to evaluate on.')
    args.add_argument('--batch_size', default=1, type=int,
                      help='size of batch')
    args.add_argument('--save_dir',
                      help='path to store text & video feats, this is for saving embeddings if you want to do offline retrieval.')
    args.add_argument('--cuda_base', help="in form cuda:x")
    args.add_argument('--device_ids', help='delimited list input', type=lambda s: [int(item) for item in s.split(',')])

    config = ConfigParser(args, test=True, eval_mode='mq')
    # hack to get sliding into config
    args = args.parse_args()
    config._config['sliding_window_stride'] = args.sliding_window_stride

    run()

EgoVLPv2/test_mq.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

from .trainer_egoclip import *
from .trainer_epic import *
from .trainer_charades import *

EgoVLPv2/trainer/__init__.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import numpy as np
import torch
from torch import nn
from tqdm.auto import tqdm
import torch.distributed as dist

from base import BaseTrainer, Multi_BaseTrainer_dist
from model.model import sim_matrix
from utils import inf_loop
import os
from csv import reader
from pathlib import Path
import sys
import time
import json 
import copy
import torch.nn.functional as F

class AllGather_multi(torch.autograd.Function):
    """An autograd function that performs allgather on a tensor."""

    @staticmethod
    def forward(ctx, tensor, n_gpu, args):
        output = [torch.empty_like(tensor) for _ in range(args.world_size)]
        dist.all_gather(output, tensor)
        ctx.rank = args.rank
        ctx.batch_size = tensor.shape[0]
        return torch.cat(output, 0)

    @staticmethod
    def backward(ctx, grad_output):
        return (
            grad_output[ctx.batch_size * ctx.rank : ctx.batch_size * (ctx.rank + 1)],
            None, None,
        )

class Multi_Trainer_dist_Charades(Multi_BaseTrainer_dist):
    """
    Trainer class

    Note:
        Inherited from BaseTrainer.
    """

    def __init__(self, args, model, loss, metrics, optimizer, scheduler, gpu, config, data_loader,
                 valid_data_loader=None, lr_scheduler=None, len_epoch=None, writer=None,
                 visualizer=None, tokenizer=None, max_samples_per_epoch=50000):
        super().__init__(args, model, loss, metrics, optimizer, scheduler, gpu, config, writer)
        self.config = config
        self.args = args
        self.data_loader = data_loader
        if len_epoch is None:
            self.len_epoch = min([len(x) for x in data_loader])
        else:
            self.data_loader = inf_loop(data_loader)
            self.len_epoch = len_epoch

        self.valid_data_loader = valid_data_loader
        self.do_validation = self.valid_data_loader is not None
        self.lr_scheduler = lr_scheduler
        self.visualizer = visualizer
        self.val_chunking = True
        self.batch_size = self.data_loader[0].batch_size
        self.log_step = int(np.sqrt(self.batch_size))
        self.total_batch_sum = sum([x.batch_size for x in self.data_loader])
        self.tokenizer = tokenizer
        self.max_samples_per_epoch = max_samples_per_epoch
        self.n_gpu = self.args.world_size
        self.allgather = AllGather_multi.apply
        self.args.model_name = self.config["arch"]["args"]["load_checkpoint"].split("/")[-1]

    def _eval_metrics(self, output):
        acc_metrics = np.zeros(len(self.metrics))
        for i, metric in enumerate(self.metrics):
            acc_metrics[i] += metric(output)
        return acc_metrics

    def _adjust_learning_rate(self, optimizer, epoch, args):
        lr = args.learning_rate1
        for milestone in args.schedule:
            lr *= 0.1 if epoch >= milestone else 1.
        for param_group in optimizer.param_groups:
            param_group['lr'] = lr

    def _train_epoch(self, epoch, scaler, gpu):
        """
        Training logic for an epoch

        :param epoch: Current training epoch.
        :return: A log that contains all information you want to save.

        Note:
            If you have additional information to record, for example:
                > additional_log = {"x": x, "y": y}
            merge it with log before return. i.e.
                > log = {**log, **additional_log}
                > return log

            The metrics in log must have the key 'metrics'.
        """


        if dist.get_rank() == 0:
            Path(self.args.save_dir).mkdir(parents=True, exist_ok=True)
            stats_file = open(Path(self.args.save_dir) / str(self.args.model_name + "_" + 'stats_vtc.txt'), 'a', buffering=1)
            print(' '.join(sys.argv))
            print(' '.join(sys.argv), file=stats_file)


        log = {}
        
        if self.do_validation and epoch == 1:
            val_log = self._valid_epoch(0, gpu)
            if self.args.rank == 0:
                log.update(val_log)

        self.model.train()
        total_loss = [0] * len(self.data_loader)
        total_metrics = np.zeros(len(self.metrics))
        for loader in self.data_loader:
            loader.train_sampler.set_epoch(epoch)
        for batch_idx, data_li in enumerate(zip(*self.data_loader)):
            if (batch_idx + 1) * self.total_batch_sum > self.max_samples_per_epoch:
                break
            for dl_idx, data in enumerate(data_li):
                # then assume we must tokenize the input, e.g. its a string
                if self.tokenizer is not None:
                    data['text'] = self.tokenizer(data['text'], return_tensors='pt', padding='max_length', max_length=15,
                                                  truncation=True) 
                data['text'] = {key: val.cuda(gpu, non_blocking=True) for key, val in data['text'].items()}
                data['video'] = data['video'].cuda(gpu, non_blocking=True)

                self.optimizer.zero_grad()
                
                with torch.set_grad_enabled(True):
                    with torch.cuda.amp.autocast():
                        loss, loss_dict, ret = self.model(data, self.allgather, self.n_gpu, self.args, self.config, self.loss, gpu, task_names='Dual', dataset_name='charades')

                scaler.scale(loss).backward()
                scaler.step(self.optimizer)
                scaler.update()
                self.scheduler.step()

                if self.args.rank == 0:
                    if batch_idx % self.args.print_freq == 0:
                        #print('Train Step: [{}/{}] Loss: {:.4f}'.format(batch_idx+1, len(train_loader), loss.item()))
                        stats = dict(epoch=epoch, step=batch_idx,
                                 lr_weights=self.optimizer.param_groups[0]['lr'],
                                 loss=loss.item())
                        print("stats: ", stats)
                        print(json.dumps(stats), file=stats_file)


                if self.writer is not None and self.args.rank == 0:
                    # self.writer.log_scalar(f'loss_train_{dl_idx}', loss.detach().item())
                    total = int(self.data_loader[dl_idx].n_samples/self.n_gpu)
                    current = batch_idx * self.data_loader[dl_idx].batch_size
                    final_total = (epoch-1) * total + current
                    self.writer.add_scalar(f'Loss_training/loss_{dl_idx}', loss.detach().item(), final_total)

                total_loss[dl_idx] += loss.detach().item()

                # if batch_idx % self.log_step == 0 and self.args.local_rank == 0:
                if batch_idx % self.args.print_freq == 0 and self.args.rank == 0:
                    self.logger.info('Train Epoch: {} dl{} {} Loss: {:.6f}'.format(
                        epoch,
                        dl_idx,
                        self._progress(batch_idx, dl_idx),
                        loss.detach().item()))
                
                self.optimizer.zero_grad()
            if batch_idx == self.len_epoch:
                break

        log.update({f'loss_{dl_idx}': total_loss[dl_idx] / self.len_epoch for dl_idx in range(len(self.data_loader))})

        if self.writer is not None and self.args.rank == 0:
            for dl_idx in range(len(self.data_loader)):
                tl = total_loss[dl_idx] / self.len_epoch
                self.writer.add_scalar(f'Loss_training/loss_total_{dl_idx}', tl, epoch-1)

        if self.do_validation:
            val_log = self._valid_epoch(epoch, gpu)
            if self.args.rank == 0:
                log.update(val_log)

        self._adjust_learning_rate(self.optimizer, epoch, self.args)

        return log

    def _valid_epoch(self, epoch, gpu):
        """
        Validate after training an epoch

        :return: A log that contains information about validation

        Note:
            The validation metrics in log must have the key 'val_metrics'.
        """
        self.model.eval()
        total_val_loss = [0] * len(self.valid_data_loader)
        total_val_metrics = [np.zeros(len(self.metrics))] * len(self.valid_data_loader)
        meta_arr = {x: [] for x in range(len(self.valid_data_loader))}
        text_embed_arr = {x: [] for x in range(len(self.valid_data_loader))}
        vid_embed_arr = {x: [] for x in range(len(self.valid_data_loader))}
        target_arr = {x: [] for x in range(len(self.valid_data_loader))}

        # construct set of sentences.
        cls_arr = []
        with open('/cis/home/shraman/works_meta_2022/Datasets/CharadesEgo_v1_480/CharadesEgo/Charades_v1_classes.txt',
                  'r') as charades:
            csv_reader = list(reader(charades))
        for line in csv_reader:
            cls_arr.append(str(line[0][5:]))

        with torch.no_grad():
            os.environ["TOKENIZERS_PARALLELISM"] = "false"
            data_cls = self.tokenizer(cls_arr, return_tensors='pt', padding='max_length', max_length=30, truncation=True)
            data_cls = {key: val.cuda(gpu, non_blocking=True) for key, val in data_cls.items()}

            dict_cls = {'text': data_cls, 'video': torch.Tensor(1, self.config["data_loader"][0]["args"]["video_params"]["num_frames"], 3, 224, 224).cuda(gpu, non_blocking=True)} ## 'video' can be random here with proper shape
            ret = self.model.module.infer(dict_cls, return_embeds=True, task_names="Dual", ret={})
            text_embed = ret['text_embeds']
            text_embeds = text_embed.cpu().detach()

            for dl_idx, dl in enumerate(self.valid_data_loader):
                for batch_idx, data in enumerate(tqdm(dl)):
                    meta_arr[dl_idx].append(data['meta'])
                    
                    if self.tokenizer is not None:
                        data['text'] = self.tokenizer(data['text'], return_tensors='pt', padding='max_length', max_length=30, truncation=True)
                    data['text'] = {key: val.cuda(gpu, non_blocking=True) for key, val in data['text'].items()}
                    data['video'] = data['video'].cuda(gpu, non_blocking=True)
                    data_target = data['target'].cuda(gpu, non_blocking=True)

                    ret = self.model.module.infer(data, return_embeds=True, task_names="Dual", ret={})
                    vid_embed = ret['video_embeds'].contiguous()

                    vid_embed_all = [torch.zeros_like(vid_embed) for _ in range(self.n_gpu)]
                    torch.distributed.all_gather(vid_embed_all, vid_embed)
                    vid_embed_all = torch.cat(vid_embed_all, dim=0)
                    vid_embed_arr[dl_idx].append(vid_embed_all.cpu())

                    data_target_all = [torch.zeros_like(data_target) for _ in range(self.n_gpu)]
                    torch.distributed.all_gather(data_target_all, data_target)
                    data_target_all = torch.cat(data_target_all, dim=0)
                    target_arr[dl_idx].append(data_target_all.cpu())

            if self.writer is not None and self.args.rank == 0:
                for dl_idx in range(len(self.valid_data_loader)):
                    tl = total_val_loss[dl_idx] / len(self.valid_data_loader[dl_idx])
                    self.writer.add_scalar(f'Loss_val/loss_total_{dl_idx}', tl, epoch-1)

        for dl_idx in range(len(self.valid_data_loader)):
            nested_metrics = {x: {} for x in range(len(self.valid_data_loader))}

            vid_embeds = torch.cat(vid_embed_arr[dl_idx])
            target_embeds = torch.cat(target_arr[dl_idx])

            sims = sim_matrix(text_embeds, vid_embeds).numpy().T
            pred_arr_cat = {"charades_metrics_vtc": sims}

            targets = target_embeds.numpy()

            for metric in self.metrics:
                metric_name = metric.__name__
                res = metric(pred_arr_cat[metric_name], targets)
                if self.args.rank == 0:
                    self.logger.info( verbose(epoch=epoch, metrics=res, name=self.valid_data_loader[dl_idx].dataset_name,
                            mode=metric_name, args=self.args) )
                nested_metrics[dl_idx][metric_name] = res

                if self.writer is not None and self.args.rank == 0:
                    to_write = format_nested_metrics_for_writer(res, mode=metric_name,
                                                                name=self.valid_data_loader[dl_idx].dataset_name)
                    for key, val in to_write.items():
                        key = key.replace('[', '_').replace(']', '_')
                        self.writer.add_scalar(f'Val_metrics_{dl_idx}/{key}', val, epoch-1)

                if self.visualizer is not None and self.args.rank == 0:
                    meta_arr_cat = {key: [] for key in meta_arr[0]}
                    for meta in meta_arr:
                        for key, val in meta.items():
                            meta_arr_cat[key] += val
                    self.visualizer.visualize_ranking(sims, epoch, meta_arr_cat, nested_metrics)

        res_dict = {}
        if self.args.rank == 0:
            res_dict = {f'val_loss_{dl_idx}': total_val_loss[dl_idx] / len(self.valid_data_loader[dl_idx])
                        for dl_idx in range(len(self.valid_data_loader))}
            res_dict['nested_val_metrics'] = nested_metrics

        return res_dict

    def _progress(self, batch_idx, dl_idx):
        base = '[{}/{} ({:.0f}%)]'
        if hasattr(self.data_loader[dl_idx], 'n_samples'):
            current = batch_idx * self.data_loader[dl_idx].batch_size
            total = int(self.data_loader[dl_idx].n_samples / self.n_gpu)
        else:
            current = batch_idx
            total = self.len_epoch
        return base.format(current, total, 100.0 * current / total)

def verbose(epoch, metrics, mode, args, name="TEST"):
    
    if dist.get_rank() == 0:
        Path(args.save_dir).mkdir(parents=True, exist_ok=True)
        stats_file = open(Path(args.save_dir) / str(args.model_name + "_" + 'stats_vtc.txt'), 'a', buffering=1)
        print(' '.join(sys.argv))
        print(' '.join(sys.argv), file=stats_file)
    
    mAP = metrics["mAP"]
    msg = f"[{mode}]{name:s} epoch {epoch}, mAP: {mAP:.3f}"
    print(msg)

    if dist.get_rank()==0:
        #print('Train Step: [{}/{}] Loss: {:.4f} Time: {}'.format(step+1, len(train_loader), loss.item(), int(time.time() - start_time)))
        stats = dict(epoch=epoch, msg=msg)
        print(json.dumps(stats), file=stats_file)
    
    return msg

def format_nested_metrics_for_writer(metrics, mode, name="TEST"):
    res = {}
    for key, val in metrics.items():
        log_name = f"[{mode}]{name}_{key}"
        res[log_name] = val
    return res

EgoVLPv2/trainer/trainer_charades.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import numpy as np
import torch
from torch import nn
from tqdm.auto import tqdm
import torch.distributed as dist
import copy
import torch.nn.functional as F
from pathlib import Path
import sys
import json 

from base import Multi_BaseTrainer_dist
from model.model import sim_matrix, sim_matrix_batch_val
from utils import inf_loop
import torch.distributed as dist
from transformers import DataCollatorForLanguageModeling


class AllGather_multi(torch.autograd.Function):
    """An autograd function that performs allgather on a tensor."""

    @staticmethod
    def forward(ctx, tensor, n_gpu, args):
        output = [torch.empty_like(tensor) for _ in range(args.world_size)]
        dist.all_gather(output, tensor)
        ctx.rank = args.rank
        ctx.batch_size = tensor.shape[0]
        return torch.cat(output, 0)

    @staticmethod
    def backward(ctx, grad_output):
        return (
            grad_output[ctx.batch_size * ctx.rank : ctx.batch_size * (ctx.rank + 1)],
            None, None,
        )

class Multi_Trainer_dist(Multi_BaseTrainer_dist):
    """
    Trainer class

    Note:
        Inherited from BaseTrainer.
    """

    def __init__(self, args, model, loss, metrics, optimizer, scheduler, gpu, config, data_loader,
                 valid_data_loader=None, lr_scheduler=None, len_epoch=None, writer=None,
                 visualizer=None, tokenizer=None, max_samples_per_epoch=50000):
        super().__init__(args, model, loss, metrics, optimizer, scheduler, gpu, config, writer)
        self.config = config
        self.args = args
        self.data_loader = data_loader
        if len_epoch is None:
            # epoch-based training
            # take the min
            self.len_epoch = min([len(x) for x in data_loader])
        else:
            # iteration-based training
            self.data_loader = inf_loop(data_loader)
            self.len_epoch = len_epoch

        self.valid_data_loader = valid_data_loader
        self.do_validation = self.valid_data_loader is not None
        self.lr_scheduler = lr_scheduler
        self.visualizer = visualizer
        self.val_chunking = True
        self.batch_size = self.data_loader[0].batch_size
        self.log_step = self.args.print_freq #int(np.sqrt(self.batch_size))
        self.total_batch_sum = sum([x.batch_size for x in self.data_loader])
        self.tokenizer = tokenizer
        self.max_samples_per_epoch = max_samples_per_epoch
        self.n_gpu = self.args.world_size
        self.allgather = AllGather_multi.apply
        self.data_collator = DataCollatorForLanguageModeling(self.tokenizer, mlm=True, mlm_probability=0.15)
        # self.writer = writer

    def _eval_metrics(self, output):
        acc_metrics = np.zeros(len(self.metrics))
        for i, metric in enumerate(self.metrics):
            acc_metrics[i] += metric(output)
            # if self.writer is not None:
            #     self.writer.log_scalar('{}'.format(metric.__name__), acc_metrics[i])
        return acc_metrics

    def _adjust_learning_rate(self, optimizer, epoch, args):
        lr = args.learning_rate1
        for milestone in args.schedule:
            lr *= 0.1 if epoch >= milestone else 1.
        for param_group in optimizer.param_groups:
            param_group['lr'] = lr

    def _train_epoch(self, epoch, scaler, gpu):

        self.model.train()
        total_loss = [0] * len(self.data_loader)
        total_metrics = np.zeros(len(self.metrics))

        if dist.get_rank() == 0:
            Path(self.args.save_dir).mkdir(parents=True, exist_ok=True)
            stats_file = open(Path(self.args.save_dir) / 'stats.txt', 'a', buffering=1)
            print(' '.join(sys.argv))
            print(' '.join(sys.argv), file=stats_file)


        for loader in self.data_loader:
            loader.train_sampler.set_epoch(epoch)
        for batch_idx, data_li in enumerate(zip(*self.data_loader)):
            #print("Length of data_li: ", len(data_li))
            if (batch_idx + 1) * self.total_batch_sum > self.max_samples_per_epoch:
                break
            for dl_idx, data in enumerate(data_li):
                
                if 'video_neg' in data.keys():  # w/ negative sampling
                    data['text'] = data['text'] + data['text_neg']
                    data['video'] = torch.cat( (data['video'], data['video_neg']), axis = 0)
                    data['noun_vec'] = torch.cat((data['noun_vec'], data['noun_vec_neg']), axis=0)
                    data['verb_vec'] = torch.cat((data['verb_vec'], data['verb_vec_neg']), axis=0)


                if self.tokenizer is not None:
                    data['text'] = self.tokenizer(data['text'], return_tensors='pt', padding='max_length', max_length=15,
                                                  truncation=True)

                if 'MLM' in self.args.task_names:
                    
                    data_text_input_list = []
                    for idx in range(data['text']['input_ids'].size(0)):
                        data_text_input_list.append(data['text']['input_ids'][idx])
                    mlm_dict = self.data_collator(data_text_input_list)


                    data.update({'text_mlm_ids': mlm_dict['input_ids'], 'text_mlm_labels': mlm_dict['labels']})
                    data['text_mlm_ids'] = data['text_mlm_ids'].cuda(gpu, non_blocking=True)
                    data['text_mlm_labels'] = data['text_mlm_labels'].cuda(gpu, non_blocking=True)

                data['text'] = {key: val.cuda(gpu, non_blocking=True) for key, val in data['text'].items()}
                data['video'] = data['video'].cuda(gpu, non_blocking=True)
                n_embeds = data['noun_vec'].cuda(gpu, non_blocking=True)
                v_embeds = data['verb_vec'].cuda(gpu, non_blocking=True)

                self.optimizer.zero_grad()

                with torch.set_grad_enabled(True):
                    with torch.cuda.amp.autocast():

                        loss, loss_dict, ret = self.model(data, n_embeds, v_embeds, self.allgather, self.n_gpu, self.args, self.config, self.loss, gpu, task_names=self.args.task_names)

                scaler.scale(loss).backward()
                scaler.step(self.optimizer)
                scaler.update()
                self.scheduler.step()


                if dist.get_rank()==0:
                    if (batch_idx) % self.args.print_freq == 0:
                        #print('Train Step: [{}/{}] Loss: {:.4f} Time: {}'.format(step+1, len(train_loader), loss.item(), int(time.time() - start_time)))
                        stats = dict(epoch=epoch, step=batch_idx,
                                    lr_weights=self.optimizer.param_groups[0]['lr'],
                                    loss=loss.item())
                        print(json.dumps(stats), file=stats_file)



                if self.writer is not None and self.args.rank == 0:
                    # self.writer.log_scalar(f'loss_train_{dl_idx}', loss.detach().item())
                    total = int(self.data_loader[dl_idx].n_samples/self.n_gpu)
                    current = batch_idx * self.data_loader[dl_idx].batch_size
                    final_total = (epoch-1) * total + current
                    self.writer.add_scalar(f'Loss_training/loss_{dl_idx}', loss.detach().item(), final_total)

                total_loss[dl_idx] += loss.detach().item()

                # if batch_idx % self.log_step == 0 and self.args.local_rank == 0:
                if batch_idx % self.log_step == 0 and self.args.rank == 0:
                    self.logger.info('Train Epoch: {} dl{} {} Loss: {:.6f}'.format(
                        epoch,
                        dl_idx,
                        self._progress(batch_idx, dl_idx),
                        loss.detach().item()))

                self.optimizer.zero_grad()
            
            if batch_idx == self.len_epoch:
                break

        log = {
            f'loss_{dl_idx}': total_loss[dl_idx] / self.len_epoch for dl_idx in range(len(self.data_loader))
        }

        if self.writer is not None and self.args.rank == 0:
            for dl_idx in range(len(self.data_loader)):
                tl = total_loss[dl_idx] / self.len_epoch
                self.writer.add_scalar(f'Loss_training/loss_total_{dl_idx}', tl, epoch-1)

        if self.do_validation:
            val_log = self._valid_epoch(epoch, gpu)
            if self.args.rank == 0:
                log.update(val_log)

        self._adjust_learning_rate(self.optimizer, epoch, self.args)

        return log

    def _valid_epoch(self, epoch, gpu):
        
        self.model.eval()
        total_val_loss = [0] * len(self.valid_data_loader)
        total_val_metrics = [np.zeros(len(self.metrics))] * len(self.valid_data_loader)

        gt_arr = {x: [] for x in range(len(self.valid_data_loader))}
        pred_arr_ensemble = {x: [] for x in range(len(self.valid_data_loader))}
        pred_arr_vtm = {x: [] for x in range(len(self.valid_data_loader))}
        type_arr = {x: [] for x in range(len(self.valid_data_loader))}

        

        with torch.no_grad():
            for dl_idx, dl in enumerate(self.valid_data_loader):

                for batch_idx, data in enumerate(tqdm(dl)):
                    b1,b2,b3,b4,b5,b6= data['video'].shape #b1 -> batch_size, b2 -> 5, b3 -> 4
                    data['video'] = data['video'].reshape(b1*b2,b3,b4,b5,b6)
                    data['text'] = data['text']

                    
                    data_text_all = []
                    for _i in range(self.n_gpu):
                        if self.args.rank == _i:
                            data_text_all.append(data['text'])


                    if self.tokenizer is not None:
                        data['text'] = self.tokenizer(data['text'], return_tensors='pt', padding=True, truncation=True)
                    data['text'] = {key: val.cuda(gpu, non_blocking=True) for key, val in data['text'].items()}
                    data['video'] = data['video'].cuda(gpu, non_blocking=True)


                    ret = self.model.module.infer(data, return_embeds=True, task_names="EgoNCE", ret={})

                    data['text']['input_ids'] = torch.repeat_interleave(data['text']['input_ids'], b2, dim=0)
                    data['text']['attention_mask'] = torch.repeat_interleave(data['text']['attention_mask'], b2, dim=0)

                    ret = self.model.module.infer(data, return_embeds=True, task_names="ITM", ret=ret)

                    data_gt = data['correct'][0].cuda(gpu, non_blocking=True).unsqueeze(0)


                    ret['text_embeds'] = ret['text_embeds'].reshape(b1,1,ret['text_embeds'].shape[-1])
                    ret['video_embeds'] = ret['video_embeds'].reshape(b1,b2,ret['video_embeds'].shape[-1])
                    data_pred_vtc = sim_matrix_batch_val(ret['text_embeds'], ret['video_embeds']).squeeze(1)
                    data_pred_vtm = F.softmax(ret["cross_attn_itm_logits"],dim=1)[:,1:].t().reshape(1,b1,b2)[0].contiguous()
                    
                    data_type = data['type'][0].cuda(gpu, non_blocking=True).unsqueeze(0)
                    data_gt_all = [torch.zeros_like(data_gt) for _ in range(self.n_gpu)]
                    torch.distributed.all_gather(data_gt_all, data_gt)
                    data_gt_all = torch.cat(data_gt_all, dim=0)

                    data_pred_all_vtc = [torch.zeros_like(data_pred_vtc) for _ in range(self.n_gpu)]
                    torch.distributed.all_gather(data_pred_all_vtc, data_pred_vtc)
                    data_pred_all_vtc = torch.cat(data_pred_all_vtc, dim=0)

                    data_pred_all_vtm = [torch.zeros_like(data_pred_vtm) for _ in range(self.n_gpu)]
                    torch.distributed.all_gather(data_pred_all_vtm, data_pred_vtm)
                    data_pred_all_vtm = torch.cat(data_pred_all_vtm, dim=0)

                    #print("Shape of data_pred_all_vtc: ", data_pred_all_vtc.shape)
                    #print("Shape of data_pred_all_vtm: ", data_pred_all_vtm.shape)
                    data_pred_all_ensemble = data_pred_all_vtc + data_pred_all_vtm

                    data_type_all = [torch.zeros_like(data_type) for _ in range(self.n_gpu)]
                    torch.distributed.all_gather(data_type_all, data_type)
                    data_type_all = torch.cat(data_type_all, dim=0)

                    gt_arr[dl_idx].append(data_gt_all.cpu())
                    pred_arr_ensemble[dl_idx].append(data_pred_all_ensemble.cpu())
                    pred_arr_vtm[dl_idx].append(data_pred_all_vtm.cpu())
                    type_arr[dl_idx].append(data_type_all.cpu())


            if self.writer is not None and self.args.rank == 0:
                for dl_idx in range(len(self.valid_data_loader)):
                    tl = total_val_loss[dl_idx] / len(self.valid_data_loader[dl_idx])
                    self.writer.add_scalar(f'Loss_val/loss_total_{dl_idx}', tl, epoch-1)

        for dl_idx in range(len(self.valid_data_loader)):
            nested_metrics = {x: {} for x in range(len(self.valid_data_loader))}

            gt_arr_cat = torch.cat(gt_arr[dl_idx])
            pred_arr_cat = {"egomcq_accuracy_metrics_ensemble": torch.cat(pred_arr_ensemble[dl_idx]), "egomcq_accuracy_metrics_vtm": torch.cat(pred_arr_vtm[dl_idx])}
            type_cat = torch.cat(type_arr[dl_idx])

            for metric in self.metrics:
                metric_name = metric.__name__
                res = metric(pred_arr_cat[metric_name], gt_arr_cat, type_cat)
                if self.args.rank == 0:
                    self.logger.info(
                        verbose(epoch=epoch, metrics=res, args=self.args, name=self.valid_data_loader[dl_idx].dataset_name))
                nested_metrics[dl_idx][metric_name] = res

                if self.writer is not None and self.args.rank == 0:
                    to_write = format_nested_metrics_for_writer(res, mode=metric_name,
                                                                name=self.valid_data_loader[dl_idx].dataset_name)
                    # for key, val in to_write.items():
                    #     self.writer.log_scalar(key, val)
                    for key, val in to_write.items():
                        key = key.replace('[', '_').replace(']', '_')
                        self.writer.add_scalar(f'Val_metrics_{dl_idx}/{key}', val, epoch - 1)

        res_dict = {}
        if self.args.rank == 0:
            res_dict = {f'val_loss_{dl_idx}': total_val_loss[dl_idx] / len(self.valid_data_loader[dl_idx])
                        for dl_idx in range(len(self.valid_data_loader))}
            res_dict['nested_val_metrics'] = nested_metrics

        return res_dict

    def _progress(self, batch_idx, dl_idx):
        base = '[{}/{} ({:.0f}%)]'
        if hasattr(self.data_loader[dl_idx], 'n_samples'):
            current = batch_idx * self.data_loader[dl_idx].batch_size
            total = int(self.data_loader[dl_idx].n_samples / self.n_gpu)
        else:
            current = batch_idx
            total = self.len_epoch
        return base.format(current, total, 100.0 * current / total)

def verbose(epoch, metrics, args, name="TEST"):
    msg = ""
    if dist.get_rank() == 0:
        Path(args.save_dir).mkdir(parents=True, exist_ok=True)
        stats_file = open(Path(args.save_dir) / 'stats.txt', 'a', buffering=1)
        print(' '.join(sys.argv))
        print(' '.join(sys.argv), file=stats_file)
    
    for key in metrics.keys():
        acc = metrics[key]
        msg += f"{name:s} epoch {epoch}, {key:s}, Acc: {acc:.1f};    "
    print(msg)
    
    if dist.get_rank()==0:
        #print('Train Step: [{}/{}] Loss: {:.4f} Time: {}'.format(step+1, len(train_loader), loss.item(), int(time.time() - start_time)))
        stats = dict(epoch=epoch, msg=msg)
        print(json.dumps(stats), file=stats_file)
    return msg

def format_nested_metrics_for_writer(metrics, mode, name="TEST"):
    res = {}
    for key, val in metrics.items():
        log_name = f"[{mode}]{name}_{key}"
        res[log_name] = val
    return res

EgoVLPv2/trainer/trainer_egoclip.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import pdb
import numpy as np
import torch
from torch import nn
from tqdm.auto import tqdm
import torch.distributed as dist

from base import BaseTrainer, Multi_BaseTrainer_dist
from model.model import sim_matrix
from utils import inf_loop
from pathlib import Path
import sys
import json

class AllGather_multi(torch.autograd.Function):
    """An autograd function that performs allgather on a tensor."""

    @staticmethod
    def forward(ctx, tensor, n_gpu, args):
        output = [torch.empty_like(tensor) for _ in range(args.world_size)]
        dist.all_gather(output, tensor)
        ctx.rank = args.rank
        ctx.batch_size = tensor.shape[0]
        return torch.cat(output, 0)

    @staticmethod
    def backward(ctx, grad_output):
        return (
            grad_output[ctx.batch_size * ctx.rank : ctx.batch_size * (ctx.rank + 1)],
            None, None,
        )

class Multi_Trainer_dist_MIR(Multi_BaseTrainer_dist):
    """
    Trainer class

    Note:
        Inherited from BaseTrainer.
    """

    def __init__(self, args, model, loss, metrics, optimizer, scheduler, gpu, config, data_loader,
                 valid_data_loader=None, lr_scheduler=None, len_epoch=None, writer=None,
                 visualizer=None, tokenizer=None, max_samples_per_epoch=50000):
        super().__init__(args, model, loss, metrics, optimizer, scheduler, gpu, config, writer)
        self.config = config
        self.args = args
        self.data_loader = data_loader
        if len_epoch is None:
            # epoch-based training
            # take the min
            self.len_epoch = min([len(x) for x in data_loader])
        else:
            # iteration-based training
            self.data_loader = inf_loop(data_loader)
            self.len_epoch = len_epoch

        self.valid_data_loader = valid_data_loader
        self.do_validation = self.valid_data_loader is not None
        self.lr_scheduler = lr_scheduler
        self.visualizer = visualizer
        self.val_chunking = True
        self.batch_size = self.data_loader[0].batch_size
        self.log_step = int(np.sqrt(self.batch_size))
        self.total_batch_sum = sum([x.batch_size for x in self.data_loader])
        self.tokenizer = tokenizer
        self.max_samples_per_epoch = max_samples_per_epoch
        self.n_gpu = self.args.world_size
        self.allgather = AllGather_multi.apply
        # self.writer = writer

    def _eval_metrics(self, output):
        acc_metrics = np.zeros(len(self.metrics))
        for i, metric in enumerate(self.metrics):
            acc_metrics[i] += metric(output)
            # if self.writer is not None:
            #     self.writer.log_scalar('{}'.format(metric.__name__), acc_metrics[i])
        return acc_metrics

    def _adjust_learning_rate(self, optimizer, epoch, args):
        lr = args.learning_rate1
        for milestone in args.schedule:
            lr *= 0.1 if epoch >= milestone else 1.
        for param_group in optimizer.param_groups:
            param_group['lr'] = lr

    def _train_epoch(self, epoch, scaler, gpu):
        """
        Training logic for an epoch

        :param epoch: Current training epoch.
        :return: A log that contains all information you want to save.

        Note:
            If you have additional information to record, for example:
                > additional_log = {"x": x, "y": y}
            merge it with log before return. i.e.
                > log = {**log, **additional_log}
                > return log

            The metrics in log must have the key 'metrics'.
        """

        if dist.get_rank() == 0:
            Path(self.args.save_dir).mkdir(parents=True, exist_ok=True)
            stats_file = open(Path(self.args.save_dir) / str('stats_vtc.txt'), 'a', buffering=1)
            print(' '.join(sys.argv))
            print(' '.join(sys.argv), file=stats_file)

        log = {}

        if self.do_validation and epoch == 1:
            val_log = self._valid_epoch(0, gpu)
            if self.args.rank == 0:
                log.update(val_log)

        self.model.train()
        total_loss = [0] * len(self.data_loader)
        total_metrics = np.zeros(len(self.metrics))
        
        for loader in self.data_loader:
            loader.train_sampler.set_epoch(epoch)
        for batch_idx, data_li in enumerate(zip(*self.data_loader)):
            if (batch_idx + 1) * self.total_batch_sum > self.max_samples_per_epoch:
                break
            for dl_idx, data in enumerate(data_li):
                # then assume we must tokenize the input, e.g. its a string
                if self.tokenizer is not None:
                    data['text'] = self.tokenizer(data['text'], return_tensors='pt', padding='max_length', max_length=30,
                                                  truncation=True)
                data['text'] = {key: val.cuda(gpu, non_blocking=True) for key, val in data['text'].items()}
                data['video'] = data['video'].cuda(gpu, non_blocking=True)
                data['relation'] = data['relation'].cuda(gpu, non_blocking=True)

                self.optimizer.zero_grad()
                with torch.set_grad_enabled(True):
                    with torch.cuda.amp.autocast():
                        loss, loss_dict, _ = self.model(data, self.allgather, self.n_gpu, self.args, self.config, self.loss, gpu, task_names='Dual', dataset_name='epic')
                        assert loss == loss_dict['Dual']

                scaler.scale(loss).backward()
                scaler.step(self.optimizer)
                scaler.update()
                self.scheduler.step()

                if dist.get_rank()==0:
                    if (batch_idx) % self.args.print_freq == 0:
                        stats = dict(epoch=epoch, step=batch_idx,
                                lr_weights=self.optimizer.param_groups[0]['lr'],
                                loss=loss.item())
                        print(json.dumps(stats), file=stats_file)


                if self.writer is not None and self.args.rank == 0:
                    # self.writer.log_scalar(f'loss_train_{dl_idx}', loss.detach().item())
                    total = int(self.data_loader[dl_idx].n_samples/self.n_gpu)
                    current = batch_idx * self.data_loader[dl_idx].batch_size
                    final_total = (epoch-1) * total + current
                    self.writer.add_scalar(f'Loss_training/loss_{dl_idx}', loss.detach().item(), final_total)

                total_loss[dl_idx] += loss.detach().item()

                # if batch_idx % self.log_step == 0 and self.args.local_rank == 0:
                if batch_idx % self.args.print_freq == 0 and self.args.rank == 0:
                    self.logger.info('Train Epoch: {} dl{} {} Loss: {:.6f}'.format(
                        epoch,
                        dl_idx,
                        self._progress(batch_idx, dl_idx),
                        loss.detach().item()))

                self.optimizer.zero_grad()

            #if batch_idx == 100:
            #    break

            if batch_idx == self.len_epoch:
                break

        log.update({f'loss_{dl_idx}': total_loss[dl_idx] / self.len_epoch for dl_idx in range(len(self.data_loader))})

        if self.writer is not None and self.args.rank == 0:
            for dl_idx in range(len(self.data_loader)):
                tl = total_loss[dl_idx] / self.len_epoch
                self.writer.add_scalar(f'Loss_training/loss_total_{dl_idx}', tl, epoch-1)

        if self.do_validation:
            val_log = self._valid_epoch(epoch, gpu)
            if self.args.rank == 0:
                log.update(val_log)

        self._adjust_learning_rate(self.optimizer, epoch, self.args)

        return log

    def _valid_epoch(self, epoch, gpu):
        """
        Validate after training an epoch

        :return: A log that contains information about validation

        Note:
            The validation metrics in log must have the key 'val_metrics'.
        """
        self.model.eval()
        total_val_loss = [0] * len(self.valid_data_loader)
        total_val_metrics = [np.zeros(len(self.metrics))] * len(self.valid_data_loader)
        meta_arr = {x: [] for x in range(len(self.valid_data_loader))}
        text_embed_arr = {x: [] for x in range(len(self.valid_data_loader))}
        vid_embed_arr = {x: [] for x in range(len(self.valid_data_loader))}
        idx_embed_arr = {x: [] for x in range(len(self.valid_data_loader))}
        with torch.no_grad():
            # for validation we switch the nested loop order, because alternate batches not needed...
            # ... and dataloaders can be of different length
            for dl_idx, dl in enumerate(self.valid_data_loader):
                for batch_idx, data in enumerate(tqdm(dl)):
                    meta_arr[dl_idx].append(data['meta'])
                    # import pdb; pdb.set_trace()
                    idx_embed = data['meta']['paths'].cuda(gpu, non_blocking=True)
                    if self.tokenizer is not None:
                        data['text'] = self.tokenizer(data['text'], return_tensors='pt', padding=True, truncation=True)
                    data['text'] = {key: val.cuda(gpu, non_blocking=True) for key, val in data['text'].items()}
                    data['video'] = data['video'].cuda(gpu, non_blocking=True)

                    # text_embed, vid_embed = self.model(data, return_embeds=True)
                    #text_embed, vid_embed = self.model.module(data, return_embeds=True)
                    ret = self.model.module.infer(data, return_embeds=True, task_names="Dual", ret={})
                    vid_embed = ret['video_embeds']
                    text_embed = ret['text_embeds']

                    # if isinstance(self.model, nn.DataParallel) and data["video"].shape[0] < len(self.model.device_ids):
                    # Note that if some batch has size smaller than the GPU size, `DataParallel` will fail.
                    # It can happen with the last batch of the dataset, depending on its size.
                    # This avoids using `DataParallel` in this case, and supposes the entire batch fits in one GPU.
                    #    text_embed, vid_embed = self.model.module(data, return_embeds=True)
                    # else:
                    #    text_embed, vid_embed = self.model(data, return_embeds=True)
                    vid_embed_all = [torch.zeros_like(vid_embed) for _ in range(self.n_gpu)]
                    torch.distributed.all_gather(vid_embed_all, vid_embed)
                    vid_embed_all = torch.cat(vid_embed_all, dim=0)

                    text_embed_all = [torch.zeros_like(text_embed) for _ in range(self.n_gpu)]
                    torch.distributed.all_gather(text_embed_all, text_embed)
                    text_embed_all = torch.cat(text_embed_all, dim=0)

                    text_embed_arr[dl_idx].append(text_embed_all.cpu())
                    vid_embed_arr[dl_idx].append(vid_embed_all.cpu())

                    idx_embed_all = [torch.zeros_like(idx_embed) for _ in range(self.n_gpu)]
                    torch.distributed.all_gather(idx_embed_all, idx_embed)
                    idx_embed_all = torch.cat(idx_embed_all, dim=0)
                    idx_embed_arr[dl_idx].append(idx_embed_all.cpu())

            if self.writer is not None and self.args.rank == 0:
                for dl_idx in range(len(self.valid_data_loader)):
                    tl = total_val_loss[dl_idx] / len(self.valid_data_loader[dl_idx])
                    self.writer.add_scalar(f'Loss_val/loss_total_{dl_idx}', tl, epoch-1)

        for dl_idx in range(len(self.valid_data_loader)):
            # TODO: this needs a clean
            # if self.writer is not None and self.args.rank == 0:
            #     self.writer.log_scalar(f'loss_val_{dl_idx}',
            #                            total_val_loss[dl_idx] / len(self.valid_data_loader[dl_idx]))
            nested_metrics = {x: {} for x in range(len(self.valid_data_loader))}

            text_embeds = torch.cat(text_embed_arr[dl_idx])
            vid_embeds = torch.cat(vid_embed_arr[dl_idx])
            arr_embeds = torch.cat(idx_embed_arr[dl_idx])

            sims = sim_matrix(text_embeds, vid_embeds).detach().cpu().numpy()

            for metric in self.metrics:
                metric_name = metric.__name__
                res = metric(sims, arr_embeds)
                if self.args.rank == 0:
                    self.logger.info( verbose(epoch=epoch, metrics=res, name=self.valid_data_loader[dl_idx].dataset_name,
                            mode=metric_name, args=self.args) )
                nested_metrics[dl_idx][metric_name] = res

                if self.writer is not None and self.args.rank == 0:
                    to_write = format_nested_metrics_for_writer(res, mode=metric_name,
                                                                name=self.valid_data_loader[dl_idx].dataset_name)
                    # for key, val in to_write.items():
                    #     self.writer.log_scalar(key, val)
                    for key, val in to_write.items():
                        key = key.replace('[', '_').replace(']', '_')
                        self.writer.add_scalar(f'Val_metrics_{dl_idx}/{key}', val, epoch-1)

                if self.visualizer is not None and self.args.rank == 0:
                    meta_arr_cat = {key: [] for key in meta_arr[0]}
                    for meta in meta_arr:
                        for key, val in meta.items():
                            meta_arr_cat[key] += val
                    self.visualizer.visualize_ranking(sims, epoch, meta_arr_cat, nested_metrics)

        res_dict = {}
        if self.args.rank == 0:
            res_dict = {f'val_loss_{dl_idx}': total_val_loss[dl_idx] / len(self.valid_data_loader[dl_idx])
                        for dl_idx in range(len(self.valid_data_loader))}
            res_dict['nested_val_metrics'] = nested_metrics

        return res_dict

    def _progress(self, batch_idx, dl_idx):
        base = '[{}/{} ({:.0f}%)]'
        if hasattr(self.data_loader[dl_idx], 'n_samples'):
            current = batch_idx * self.data_loader[dl_idx].batch_size
            total = int(self.data_loader[dl_idx].n_samples / self.n_gpu)
        else:
            current = batch_idx
            total = self.len_epoch
        return base.format(current, total, 100.0 * current / total)

def verbose(epoch, metrics, mode, args, name="TEST"):

    if dist.get_rank() == 0:
        Path(args.save_dir).mkdir(parents=True, exist_ok=True)
        stats_file = open(Path(args.save_dir) / 'stats_vtc.txt', 'a', buffering=1)
        print(' '.join(sys.argv))
        print(' '.join(sys.argv), file=stats_file)

    ndcg_v2t, ndcg_tv2, ndcg_avg = metrics["nDCG_V2T"], metrics["nDCG_T2V"], metrics["nDCG_AVG"]
    map_v2t, map_t2v, map_avg = metrics["mAP_V2T"], metrics["mAP_T2V"], metrics["mAP_AVG"]

    msg = f"[{mode}]{name:s} epoch {epoch}, nDCG_V2T: {ndcg_v2t:.3f}, nDCG_T2V: {ndcg_tv2:.3f}, nDCG_AVG: {ndcg_avg:.3f},"
    msg += f", mAP_V2T: {map_v2t:.3f}, mAP_T2V: {map_t2v:.3f}, mAP_AVG: {map_avg:.3f}"

    print(msg)

    if dist.get_rank()==0:
        stats = dict(epoch=epoch, msg=msg)
        print(json.dumps(stats), file=stats_file)

    return msg

def format_nested_metrics_for_writer(metrics, mode, name="TEST"):
    res = {}
    for key, val in metrics.items():
        log_name = f"[{mode}]{name}_{key}"
        res[log_name] = val
    return res

EgoVLPv2/trainer/trainer_epic.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


from .util import *

EgoVLPv2/utils/__init__.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import os
from csv import reader
import csv
import pandas as pd

metadata_dir = '/cis/home/shraman/works_meta_2022/Datasets/CharadesEgo_v1_480/CharadesEgo/'
split_files = {
    'train': 'CharadesEgo_v1_train_only1st.csv',
    'val': 'CharadesEgo_v1_test_only1st.csv',  # there is no test
    'test': 'CharadesEgo_v1_test_only1st.csv'
}

split = 'train'
target_split_fp = split_files[split]

with open(os.path.join(metadata_dir, target_split_fp), 'r') as charades:
    csv_reader = list(reader(charades))[1:]

path_metadata = "/cis/home/shraman/works_meta_2022/Datasets/CharadesEgo_v1_480/CharadesEgo/metadata_" + split + ".csv"
csv_metadata = open(path_metadata, 'w')

metadata_reader = csv.writer(csv_metadata, delimiter='¥', lineterminator='\n',
                             quoting=csv.QUOTE_NONE, escapechar='¥')

header = 'id' + '\t' + 'cls' + '\t' + 't_start' + '\t' + 't_end' + '\t' + 'narration'

metadata_reader.writerow([header])

mapping = pd.read_csv('/cis/home/shraman/works_meta_2022/Datasets/CharadesEgo_v1_480/CharadesEgo/Charades_v1_classes.txt', header=None)

count = 0
for row in csv_reader:
    id = row[0]
    actions = row[9].split(';')
    if actions == ['']:
        continue
    for action in actions:
        action_seg = action.split(' ')
        cls, t_start, t_end = int(action_seg[0][1:]), float(action_seg[1]), float(action_seg[2])

        narration = mapping.iloc[cls][0][5:]

        clip_info = id + '\t' + str(cls) + '\t' + \
                    str(t_start) + '\t' + str(t_end) + '\t' + str(narration)
        count += 1
        print(count)
        metadata_reader.writerow([clip_info])

EgoVLPv2/utils/charades_meta.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import numbers
from typing import List, Tuple

import torch
from torch import Tensor
from torchvision.transforms import functional_pil as F_pil, functional_tensor as F_t
from torchvision.transforms.functional import center_crop, crop


def _get_image_size(img: Tensor) -> List[int]:
    """Returns image size as [w, h]
    """
    if isinstance(img, torch.Tensor):
        return F_t._get_image_size(img)

    return F_pil._get_image_size(img)

def center_plus_four_crops(img: Tensor, size: List[int],
                              margin_h: int, margin_w: int) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
    """Crop the given image into four tiled borders and the central crop.
    """

    if isinstance(size, numbers.Number):
        size = (int(size), int(size))
    elif isinstance(size, (tuple, list)) and len(size) == 1:
        size = (size[0], size[0])

    if len(size) != 2:
        raise ValueError("Please provide only two dimensions (h, w) for size.")

    image_width, image_height = _get_image_size(img)

    crop_height, crop_width = size

    if crop_width > image_width or crop_height > image_height:
        msg = "Requested crop size {} is bigger than input size {}"
        raise ValueError(msg.format(size, (image_height, image_width)))

    if crop_width + margin_w > image_width:
        msg = "Requested margin size {} + input {} is bigger than input size {}"
        raise ValueError(msg.format((margin_h, margin_w), size, (image_height, image_width)))

    #vertical_border_height = image_height - crop_height
    #horizontal_border_height = image_width - crop_width

    #x1 = horizontal_border_height // 2
    x11 = (image_width - crop_width - 2 * margin_w) // 2
    x12 = x11 + margin_w
    x21 = x12 + crop_width
    x22 = x21 + margin_w

    y11 = (image_height - crop_height - 2 * margin_h) // 2
    y12 = y11 + margin_h
    y21 = y12 + crop_height
    y22 = y21 + margin_h

    tl = crop(img, y11, x11, margin_h, margin_w + crop_width)
    tr = crop(img, y11, x21, margin_h + crop_height, margin_w)
    bl = crop(img, y12, x11, margin_h + crop_height, margin_w)
    br = crop(img, y21, x12, margin_h, margin_w + crop_width)
    center = center_crop(img, [crop_height, crop_width])

    return tl, tr, bl, br, center



def center_plus_twohori_crops(img: Tensor, size: List[int],
                              margin_w: int) -> Tuple[Tensor, Tensor, Tensor]:
    """Crop the given image into four tiled borders and the central crop.
    """

    if isinstance(size, numbers.Number):
        size = (int(size), int(size))
    elif isinstance(size, (tuple, list)) and len(size) == 1:
        size = (size[0], size[0])

    if len(size) != 2:
        raise ValueError("Please provide only two dimensions (h, w) for size.")

    image_width, image_height = _get_image_size(img)

    crop_height, crop_width = size

    if crop_width > image_width or crop_height > image_height:
        msg = "Requested crop size {} is bigger than input size {}"
        raise ValueError(msg.format(size, (image_height, image_width)))

    if crop_width + margin_w > image_width :
        msg = "Requested margin size {} + input {} is bigger than input size {}"
        raise ValueError(msg.format((0, margin_w), size, (image_height, image_width)))

    # vertical_border_height = image_height - crop_height
    # horizontal_border_height = image_width - crop_width

    # x1 = horizontal_border_height // 2
    x11 = (image_width - crop_width - 2 * margin_w) // 2
    x12 = x11 + margin_w
    x21 = x12 + crop_width

    y11 = (image_height - crop_height) // 2

    left = crop(img, y11, x11, crop_height, margin_w)
    right = crop(img, y11, x21, crop_height, margin_w)
    center = center_crop(img, [crop_height, crop_width])

    return left, right, center

from torch import nn
class TwoHoriCrop(nn.Module):
    def __init__(self, size, margin_w):
        super().__init__()
        self.size = size
        self.margin_w = margin_w

    def forward(self, x):
        return center_plus_twohori_crops(x, self.size, self.margin_w)

if __name__ == "__main__":
    from PIL import Image

    img = Image.open('visualisations/guitar.png')
    crops = center_plus_four_crops(img, [336, 336], 112, 112)
    order = ['tl', 'tr', 'bl', 'br', 'center']

    for idx, subimg in zip(order, crops):
        subimg.save(f'visualisations/guitar_{idx}.png')

    crops = center_plus_twohori_crops(img, [448, 448], 112)
    order = ['left', 'right', 'center2']

    for idx, subimg in zip(order, crops):
        subimg.save(f'visualisations/guitar_{idx}.png')

EgoVLPv2/utils/custom_transforms.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import os
import dominate
from dominate.tags import a, attr, br, h3, img, meta, p, source, span, table, td, tr, video


class HTML:
    """This HTML class allows us to save images and write texts into a single HTML file.

     It consists of functions such as <add_header> (add a text header to the HTML file),
     <add_images> (add a row of images to the HTML file), and <save> (save the HTML to the disk).
     It is based on Python library 'dominate', a Python library for creating and
     manipulating HTML documents using a DOM API.
    """

    def __init__(self, web_dir, title, refresh=0):
        """Initialize the HTML classes

        Parameters:
            web_dir (str) -- a directory that stores the webpage. HTML file will be
            created at <web_dir>/index.html; images will be saved at <web_dir/images/
            title (str)   -- the webpage name
            reflect (int) -- how often the website refresh itself; if 0; no refreshing
        """
        self.title = title
        self.web_dir = web_dir
        self.img_dir = os.path.join(self.web_dir, "images")
        if not os.path.exists(self.web_dir):
            os.makedirs(self.web_dir)
        if not os.path.exists(self.img_dir):
            os.makedirs(self.img_dir)

        self.doc = dominate.document(title=title)
        if refresh > 0:
            with self.doc.head:
                meta(http_equiv="refresh", content=str(refresh))

    def get_image_dir(self):
        """Return the directory that stores images"""
        return self.img_dir

    def add_header(self, text):
        """Insert a header to the HTML file

        Parameters:
            text (str) -- the header text
        """
        with self.doc:
            h3(text)

    def add_videos(self, vids, txts, links, width=400, hidden_tag="hidden"):
        """add images to the HTML file

        Parameters:
            vids (str list)   -- a list of image paths
            txts (str list)  -- a list of image names shown on the website
            links (str list) --  a list of hyperref links; when you click an image,
            it will redirect you to a new page
        """
        self.t = table(border=1, style="table-layout: fixed;")  # Insert a table
        self.doc.add(self.t)
        colors = ["red", "blue", "gold", "salman"]
        with self.t:
            with tr():
                for vid, txt, link in zip(vids, txts, links):
                    td_style = "word-wrap: break-word; width:{}px".format(width)
                    with td(style=td_style, halign="center", valign="top"):
                        with p():
                            vid_path = str(vid)
                            if vid_path == hidden_tag:
                                p_style = "font-weight: bold; width:{}px;"
                                p_style = p_style.format(width * 3)
                                p("hidden video", style=p_style)
                            else:
                                with a(href=str(link)):
                                    with video():
                                        attr(controls="controls")
                                        source(src=vid_path, type="video/mp4")
                            br()
                            rows = txt.split("<br>")
                            for idx, row in enumerate(rows):
                                color = colors[idx % len(colors)]
                                bold_tag = "<b>"
                                if not row.startswith(bold_tag):
                                    s_style = "color:{};".format(color)
                                else:
                                    s_style = "color:black; font-weight: bold;"
                                    row = row[len(bold_tag):]
                                span(row, style=s_style)

    def add_images(self, ims, txts, links, width=400):
        """add images to the HTML file

        Parameters:
            ims (str list)   -- a list of image paths
            txts (str list)  -- a list of image names shown on the website
            links (str list) --  a list of hyperref links; when you click an image,
            it will redirect you to a new page
        """
        self.t = table(border=1, style="table-layout: fixed;")  # Insert a table
        self.doc.add(self.t)
        colors = ["red", "blue", "gold", "salman"]
        with self.t:
            with tr():
                for im, txt, link in zip(ims, txts, links):
                    td_style = "word-wrap: break-word;"
                    with td(style=td_style, halign="center", valign="top"):
                        with p():
                            with a(href=link):
                                img(
                                    style="width:%dpx" % width,
                                    src=im,
                                )
                            br()
                            rows = txt.split("<br>")
                            for idx, row in enumerate(rows):
                                color = colors[idx % len(colors)]
                                bold_tag = "<b>"
                                if not row.startswith(bold_tag):
                                    s_style = "color:{};".format(color)
                                else:
                                    s_style = "color:black; font-weight: bold;"
                                    row = row[len(bold_tag):]
                                span(row, style=s_style)

    def save(self):
        """save the current content to the HMTL file"""
        html_file = "%s/index.html" % self.web_dir
        f = open(html_file, "wt")
        f.write(self.doc.render())
        f.close()


if __name__ == "__main__":  # we show an example usage here.
    html = HTML("web/", "test_html")
    html.add_header("hello world")

    ims, txts, links = [], [], []
    for n in range(4):
        ims.append("image_%d.png" % n)
        txts.append("text_%d" % n)
        links.append("image_%d.png" % n)
    html.add_images(ims, txts, links)
    html.save()

EgoVLPv2/utils/html.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import numpy as np

def calculate_mAP(sim_mat, relevancy_matrix):
    """
    Computes the mean average precision according to the following formula of
    average precision:
    \frac{\sum_{k=1}^n p(k) x rel(k)}{num_rel_docs}

    where p(k) is the precision at k, rel(k) is an indicator function
    determining whether the kth returned item is relevant or not and
    num_rel_docs is the number of relevant items to find within the search.

    The mean average precision is the mean of the average precision for each
    query item (i.e row in the matrix)

    This function takes in two parameters:
        - sim_mat: a NxM matrix which represents the similarity between two
        modalities (with modality 1 being of size N and modality 2 of size M).
        - relevancy_matrix: an NxM matrix which represents the relevancy between two
        modalities of items (with modality 1 being of size N and modality 2 of
        size M).
    """
    #Find the order of the items in modality 2 according to modality 1
    ranked_order = (-sim_mat).argsort()
    ranked_sim_mat = sim_mat[np.arange(sim_mat.shape[0])[:, None], ranked_order]
    #re-order the relevancy matrix to accommodate the proposals
    ranked_rel_mat = relevancy_matrix[np.arange(relevancy_matrix.shape[0])[:, None], ranked_order]

    #find the number of relevant items found at each k
    cumulative_rel_mat = np.cumsum(ranked_rel_mat, axis=1)
    #Mask this ensuring that it is non zero if the kth term is 1 (rel(k) above)
    cumulative_rel_mat[ranked_rel_mat != 1] = 0
    #find the divisor for p(k)
    divisor = np.arange(ranked_rel_mat.shape[1]) + 1

    #find the number of relevant docs per query item
    number_rel_docs = np.sum(ranked_rel_mat==1, axis=1)

    #find the average precision per query, within np.sum finds p(k) * rel(k)
    avg_precision = np.sum(cumulative_rel_mat / divisor, axis=1) / number_rel_docs
    mAP = np.mean(avg_precision)
    return mAP

EgoVLPv2/utils/mAP.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import numpy as np

def calculate_DCG(similarity_matrix, relevancy_matrix, k_counts):
    """
    Calculates the Discounted Cumulative Gain (DCG) between two modalities for
    the first modality.
    DCG = \sum_{i=1}^k \frac{rel_i}{log_2(i + 1)}
    i.e. the sum of the k relevant retrievals which is calculated as the scaled
    relevancy for the ith item. The scale is designed such that early
    retrievals are more important than later retrievals.
    Params:
        - similarity_matrix: matrix of size n1 x n2 where n1 is the number of
          items in the first modality and n2 is the number of items in the
          second modality. The [ith,jth] element is the predicted similarity
          between the ith item from the first modality and the jth item from
          the second modality.
        - relevancy_matrix: matrix of size n1 x n2 (see similarity_matrix
          above). The [ith, jth] element is the semantic relevancy between the
          ith item from the first modality and the jth item from the second
          modality.
        - k_counts: matrix of size n1 x n2 (see similarity_matrix above) which
          includes information on which items to use to calculate the DCG for
          (see calculate_k_counts for more info on this matrix).
    Returns:
        - The DCG for each item in the first modality, a n1 length vector.
    """
    x_sz, y_sz = similarity_matrix.shape
    ranks = np.argsort(similarity_matrix)[:, ::-1]
    #Create vector of size (n,) where n is the length of the last dimension in
    #similarity matrix
    #This vector is of the form log(i+1)
    logs = np.log2(np.arange(y_sz) + 2)
    #Convert logs into the divisor for the DCG calculation, of size similarity
    #matrix
    divisors = np.repeat(np.expand_dims(logs, axis=0), x_sz, axis=0)

    #mask out the sorted relevancy matrix to only use the first k relevant
    #retrievals for each item.
    columns = np.repeat(np.expand_dims(np.arange(x_sz), axis=1), y_sz, axis=1)
    numerators = relevancy_matrix[columns, ranks] * k_counts
    #Calculate the final DCG score (note that this isn't expected to sum to 1)
    return np.sum(numerators / divisors, axis=1)

def calculate_k_counts(relevancy_matrix):
    """
    Works out the maximum number of allowed retrievals when working out the
    Discounted Cumulative Gain. For each query the DCG only uses the first k
    items retrieved which constitute the k relevant items for that query
    (otherwise the nDCG scores can be deceptively high for bad rankings).
    Params:
        - relevancy_matrix: matrix of size n1 x n2 where n1 is the number of
          items in the first modality and n2 is the number of items in the
          second modality.  The [ith, jth] element is the semantic relevancy
          between the ith item from the first modality and the jth item from
          the second modality.
    Returns:
        - Matrix of size n1 x n2 (see relevancy matrix for more info). This is
          created as a mask such that if the [ith, jth] element is 1 it
          represents a valid item to use for the calculation of DCG for the
          ith item after sorting. For example, if relevancy matrix of:
        [[1, 0.5, 0],
          [0, 0  , 1]]
          is given, then the k_counts matrix will be:
        [[1, 1, 0],
         [1, 0, 0]]
         i.e. the first row has 2 non-zero items, so the first two retrieved
         items should be used in the calculation. In the second row there is
         only 1 relevant item, therefore only the first retrieved item should
         be used for the DCG calculation.
    """
    return (np.sort(relevancy_matrix)[:, ::-1] > 0).astype(int)


def calculate_IDCG(relevancy_matrix, k_counts):
    """
    Calculates the Ideal Discounted Cumulative Gain (IDCG) which is the value
    of the Discounted Cumulative Gain (DCG) for a perfect retrieval, i.e. the
    items in the second modality were retrieved in order of their descending
    relevancy.
    Params:
        - relevancy_matrix: matrix of size n1 x n2 where n1 is the number of
          items in the first modality and n2 is the number of items in the
          second modality. The [ith, jth] element is the semantic relevancy
          between the ith item from the first modality and the jth item from
          the second modality.
        - k_counts: matrix of size n1 x n2 (see similarity_matrix above) which
          includes information on which items to use to calculate the DCG for
          (see calculate_k_counts for more info on this matrix).
    """
    return calculate_DCG(relevancy_matrix, relevancy_matrix, k_counts)

def calculate_nDCG(similarity_matrix, relevancy_matrix, k_counts=None, IDCG=None, reduction='mean'):
    """
    Calculates the normalised Discounted Cumulative Gain (nDCG) between two
    modalities for the first modality using the Discounted Cumulative Gain
    (DCG) and the Ideal Discounted Cumulative Gain (IDCG).

    nDCG = \frac{DCG}{IDCG}

    Params:
        - similarity_matrix: matrix of size n1 x n2 where n1 is the number of
          items in the first modality and n2 is the number of items in the second
          modality. The [ith,jth] element is the predicted similarity between
          the ith item from the first modality and the jth item from the second
          modality.
        - relevancy_matrix: matrix of size n1 x n2 (see similarity_matrix
          above). The [ith, jth] element is the semantic relevancy between the
          ith item from the first modality and the jth item from the second
          modality.
        - k_counts: optional parameter: matrix of size n1 x n2 (see
          similarity_matrix above) which includes information on which items to
          use to calculate the DCG for (see calculate_k_counts for more info on
          this matrix). This will be calculated using calculate_IDCG if not
          present, but should be pre-processed for efficiency.
        - IDCG: Optional parameter which includes the pre-processed Ideal
          Discounted Cumulative Gain (IDCG). This is a vector of size n1 (see
          similarity_matrix above) which contains the IDCG value for each item
          from the first modality. This will be calculated using calculate_IDCG
          if not present, but should be pre-processed for efficiency.
        - reduction: what to use to reduce the different nDCG scores. By
          default this applies np.mean across all different queries.
    Returns:
        - The nDCG values for the first modality.
    """
    if k_counts is None:
        k_counts = calculate_k_counts(relevancy_matrix)
    DCG = calculate_DCG(similarity_matrix, relevancy_matrix, k_counts)
    if IDCG is None:
        IDCG = calculate_IDCG(relevancy_matrix, k_counts)
    if reduction == 'mean':
        return np.mean(DCG / IDCG)
    elif reduction is None:
        return DCG / IDCG


if __name__ == '__main__':
    sim_matrix = np.array([
        [1.0, 0.7, 0.4, 0.0],
        [0.3, 0.9, 0.6, 0.1],
        [0.2, 0.5, 0.8, 0.4]
    ])
    rel_matrix = np.array([
        [1.0, 0.5, 0.25, 0.0],
        [0.0, 1.0, 0.4, 0.0],
        [0.5, 0.3, 1.0, 0.0]
    ])
    k_counts = np.array([
        [1, 1, 1, 0],
        [1, 1, 0, 0],
        [1, 1, 1, 0]
    ])
    assert (k_counts == calculate_k_counts(rel_matrix)).all()
    nDCG = calculate_nDCG(sim_matrix, rel_matrix, k_counts)
    assert nDCG == 0.9371789900735429
    DCG = calculate_DCG(sim_matrix, rel_matrix, k_counts)
    IDCG = calculate_IDCG(rel_matrix, k_counts)

    assert nDCG == np.mean(DCG / IDCG)

    pre_nDCG = calculate_nDCG(sim_matrix, rel_matrix, k_counts, IDCG=IDCG)
    assert pre_nDCG == nDCG

    post_mean_nDCG = calculate_nDCG(sim_matrix, rel_matrix, k_counts, IDCG=IDCG, reduction=None)
    assert np.mean(post_mean_nDCG) == pre_nDCG

EgoVLPv2/utils/nDCG.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import functools
import json
import os
import socket
import time
from collections import OrderedDict
from datetime import datetime
from itertools import repeat
from pathlib import Path

import humanize
import numpy as np
import psutil


def replace_nested_dict_item(obj, key, replace_value):
    for k, v in obj.items():
        if isinstance(v, dict):
            obj[k] = replace_nested_dict_item(v, key, replace_value)
    if key in obj:
        obj[key] = replace_value
    return obj


def state_dict_data_parallel_fix(load_state_dict, curr_state_dict):
    load_keys = list(load_state_dict.keys())
    curr_keys = list(curr_state_dict.keys())

    redo_dp = False
    undo_dp = False
    if not curr_keys[0].startswith('module.') and load_keys[0].startswith('module.'):   # this
        undo_dp = True
    elif curr_keys[0].startswith('module.') and not load_keys[0].startswith('module.'):
        redo_dp = True

    if undo_dp: # this
        from collections import OrderedDict
        new_state_dict = OrderedDict()
        for k, v in load_state_dict.items():
            name = k[7:]  # remove `module.`
            new_state_dict[name] = v
        # load params
    elif redo_dp:
        from collections import OrderedDict
        new_state_dict = OrderedDict()
        for k, v in load_state_dict.items():
            name = 'module.' + k  # remove `module.`
            new_state_dict[name] = v
    else:
        new_state_dict = load_state_dict
    return new_state_dict

def print_numpy(x, val=True, shp=False):
    """Print the mean, min, max, median, std, and size of a numpy array
    Parameters:
        val (bool) -- if print the values of the numpy array
        shp (bool) -- if print the shape of the numpy array
    """
    x = x.astype(np.float64)
    if shp:
        print('shape,', x.shape)
    if val:
        x = x.flatten()
        print('mean = %3.3f, min = %3.3f, max = %3.3f, median = %3.3f, std=%3.3f' % (
            np.mean(x), np.min(x), np.max(x), np.median(x), np.std(x)))


def mkdirs(paths):
    """create empty directories if they don't exist
    Parameters:
        paths (str list) -- a list of directory paths
    """
    if isinstance(paths, list) and not isinstance(paths, str):
        for path in paths:
            mkdir(path)
    else:
        mkdir(paths)


def mkdir(path):
    """create a single empty directory if it didn't exist
    Parameters:
        path (str) -- a single directory path
    """
    if not os.path.exists(path):
        os.makedirs(path)

def read_json(fname):
    with fname.open('rt') as handle:
        return json.load(handle, object_hook=OrderedDict)

def write_json(content, fname):
    with fname.open('wt') as handle:
        json.dump(content, handle, indent=4, sort_keys=False)

def inf_loop(data_loader):
    ''' wrapper function for endless data loader. '''
    for loader in repeat(data_loader):
        yield from loader

def memory_summary():
    vmem = psutil.virtual_memory()
    msg = (
        f">>> Currently using {vmem.percent}% of system memory "
        f"{humanize.naturalsize(vmem.used)}/{humanize.naturalsize(vmem.available)}"
    )
    print(msg)

@functools.lru_cache(maxsize=64, typed=False)
def memcache(path):
    suffix = Path(path).suffix
    print(f"loading features >>>", end=" ")
    tic = time.time()
    if suffix == ".npy":
        res = np_loader(path)
    else:
        raise ValueError(f"unknown suffix: {suffix} for path {path}")
    print(f"[Total: {time.time() - tic:.1f}s] ({socket.gethostname() + ':' + str(path)})")
    return res

def np_loader(np_path, l2norm=False):
    with open(np_path, "rb") as f:
        data = np.load(f, encoding="latin1", allow_pickle=True)
    if isinstance(data, np.ndarray) and data.size == 1:
        data = data[()]  # handle numpy dict storage convnetion
    if l2norm:
        print("L2 normalizing features")
        if isinstance(data, dict):
            for key in data:
                feats_ = data[key]
                feats_ = feats_ / max(np.linalg.norm(feats_), 1E-6)
                data[key] = feats_
        elif data.ndim == 2:
            data_norm = np.linalg.norm(data, axis=1)
            data = data / np.maximum(data_norm.reshape(-1, 1), 1E-6)
        else:
            raise ValueError("unexpected data format {}".format(type(data)))
    return data


class Timer:
    def __init__(self):
        self.cache = datetime.now()

    def check(self):
        now = datetime.now()
        duration = now - self.cache
        self.cache = now
        return duration.total_seconds()

    def reset(self):
        self.cache = datetime.now()

EgoVLPv2/utils/util.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import random
import cv2
import numpy as np
import torch


def load_frames_from_video_path(path, num_frames, sample='rand'):
    cap = cv2.VideoCapture(path)
    assert (cap.isOpened())
    vlen = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    acc_samples = min(num_frames, vlen)
    intervals = np.linspace(start=0, stop=vlen, num=acc_samples + 1).astype(int)
    ranges = []
    for idx, interv in enumerate(intervals[:-1]):
        ranges.append((interv, intervals[idx + 1] - 1))
    if sample == 'rand':
        frame_idxs = [random.choice(range(x[0], x[1])) for x in ranges]
    elif sample == 'uniform':
        frame_idxs = [(x[0] + x[1]) // 2 for x in ranges]
    else:
        raise NotImplementedError

    frames = []
    for index in frame_idxs:
        cap.set(cv2.CAP_PROP_POS_FRAMES, index)
        ret, frame = cap.read()
        if ret:
            cv2.imwrite(f'images/{index}.jpg', frame)
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frame = torch.from_numpy(frame)
            # (H x W x C) to (C x H x W)
            frame = frame.permute(2, 0, 1)
            frames.append(frame)
        else:
            raise ValueError

    frames = torch.stack(frames).float() / 255
    cap.release()
    return frames, frame_idxs

EgoVLPv2/utils/video.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import os
import cv2
import shutil
import json
import sys
import multiprocessing
from multiprocessing import Pool
import subprocess
import pandas as pd
from csv import reader, writer

video_dir = './dataset/ego4d_256_multiprocessing'
output_dir = './dataset/ego4d_chunked_multiprocessing'

save_dir = os.path.join(output_dir)
if not os.path.exists(save_dir):
    os.makedirs(save_dir)

dur_limit = 600

def video2segments(infos):
    global count
    index, uid, dur = infos[0], infos[1], infos[2]
    input_path = os.path.join(video_dir, uid + '.mp4')

    output_uid_dir = os.path.join(output_dir, uid)
    if not os.path.exists(output_uid_dir):
        os.makedirs(output_uid_dir)

    #if index % num_partition != partition:
    #    return

    assert os.path.exists(input_path)

    cap = cv2.VideoCapture(input_path)
    rate = cap.get(5)
    frame_num = cap.get(7)
    duration = frame_num / rate

    if duration <= dur_limit:
        shutil.copyfile(input_path, os.path.join(output_uid_dir, '0.mp4'))
        return

    num_seg = duration // dur_limit

    s1_time = 0;
    s2_time = dur_limit
    num_finished = 0
    while num_finished <= num_seg:
        output_mp4_path = os.path.join(output_uid_dir, str(num_finished) + '.mp4')

        cmd = 'ffmpeg -y -i {} -ss {} -to {} -async 1 {}'.format(input_path, s1_time, s2_time,
                                                                                 output_mp4_path)
        # print(cmd)
        subprocess.call(cmd, shell=True)

        # Update for next steps
        s1_time = s2_time
        s2_time = min(s1_time + dur_limit, duration)
        num_finished += 1
    return

if __name__ == "__main__":
    with open('./dataset/ego4d/manifest.csv', 'r') as csv:
        csv_reader = list(reader(csv))[1:]

    downloaded = os.listdir('./dataset/ego4d')

    uid_list = []
    infos_list = []
    num_valid = 0
    for i_item, item in enumerate(csv_reader):
        uid, dur = item[0], float(item[2])
        existed = uid + '.mp4' in downloaded

        if not existed:
            continue

        uid_list.append(uid)
        infos_list.append([num_valid, uid, dur])
        num_valid += 1

    # for infos in infos_list:
    #     video2segments(infos)

    pool = Pool(multiprocessing.cpu_count())
    pool.map(video2segments, tuple(infos_list))

EgoVLPv2/utils/video_chunk.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import os
import time
import sys
import subprocess
import multiprocessing
from multiprocessing import Pool, Value

folder_path = './dataset/ego4d'
output_path = './dataset/ego4d_256_multiprocessing'

def videos_resize(videoinfos):
    global count

    videoid, videoname = videoinfos

    if os.path.exists(os.path.join(output_path, videoname)):
        print(f'{videoname} is resized.')
        return

    inname = folder_path + '/' + videoname
    outname = output_path + '/' + videoname

    cmd = "ffmpeg -y -i {} -filter:v scale=\"trunc(oh*a/2)*2:256\" -c:a copy {}".format( inname, outname)
    subprocess.call(cmd, shell=True)

    return

if __name__ == "__main__":
    file_list = []
    mp4_list = [item for item in os.listdir(folder_path) if item.endswith('.mp4')]
    for id, video in enumerate(mp4_list):
        file_list.append([id, video])

    pool = Pool(multiprocessing.cpu_count()) 
    pool.map(videos_resize, tuple(file_list))

    # for file in file_list:
    #    videos_resize(file)

EgoVLPv2/utils/video_resize.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import matplotlib
import numpy as np
import torch

matplotlib.use('Agg')


def visualise_path(pred, target, window):
    """
    :param pred: (P, 2) Tensor where P is the number of predictions, and 2 is the (i,j) coordinate
    :param target: (T, 2) Tensor where T is the number of targets, and 2 is the (i,j) coordinate
    :param dims: (H, W) tup/list the desired height and width of matrix (should be >= to max(i), max(j))
    :param assignment_method: Method of assignment (dtw, minimum etc.)
    :return: image, visualisation of path prediction and target.
    """
    tp = torch.Tensor((64, 191, 64))
    fp = torch.Tensor((191, 64, 64))
    gt = torch.Tensor((102, 153, 255))

    grid = torch.ones_like(window).unsqueeze(0).repeat(3, 1, 1) * 255
    inf = 130 * torch.ones_like(grid)
    grid = torch.where(torch.isnan(window), inf, grid)

    clip_idxs = [t[0] for t in target]
    local_idxs = np.unique(np.array(clip_idxs)).tolist()

    for t in target:
        local_idx = local_idxs.index(t[0])
        grid[:, local_idx, t[1]] = gt

    for p in pred:
        local_idx = local_idxs.index(p[0])
        if (grid[:, local_idx, p[1]] == gt).all():
            grid[:, local_idx, p[1]] = tp
        else:
            grid[:, local_idx, p[1]] = fp

    return grid / 255


def batch_path_vis(pred_dict, target, window):
    grids = []

    window = window.cpu()
    for key, pred in pred_dict.items():
        tmp_window = window
        if key == 'min_dist':
            tmp_window = torch.zeros_like(window)
        grids.append(visualise_path(pred, target, tmp_window))

    return torch.stack(grids)


if __name__ == "__main__":
    pred = [[1, 1], [2, 4]]
    gt = [[1, 1], [3, 4]]
    window = torch.zeros((5, 6))
    visualise_path(pred, gt, window)

EgoVLPv2/utils/visualisation.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

from .base_data_loader import *
from .base_dataset import *
from .base_model import *
from .base_trainer import *

QFVS/base/__init__.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import numpy as np
from torch.utils.data import DataLoader
from torch.utils.data.dataloader import default_collate
from torch.utils.data.sampler import SubsetRandomSampler
from torch.utils.data.distributed import DistributedSampler

class BaseDataLoader(DataLoader):
    """
    Base class for all data loaders
    """
    def __init__(self, dataset, batch_size, shuffle, validation_split, num_workers, collate_fn=default_collate):
        self.validation_split = validation_split
        self.shuffle = shuffle

        self.batch_idx = 0
        self.n_samples = len(dataset)

        self.sampler, self.valid_sampler = self._split_sampler(self.validation_split)

        self.init_kwargs = {
            'dataset': dataset,
            'batch_size': batch_size,
            'shuffle': self.shuffle,
            'collate_fn': collate_fn,
            'num_workers': num_workers
        }
        super().__init__(sampler=self.sampler, **self.init_kwargs)

    def _split_sampler(self, split):
        if split == 0.0:
            return None, None

        idx_full = np.arange(self.n_samples)

        np.random.seed(0)
        np.random.shuffle(idx_full)

        if isinstance(split, int):
            assert split > 0
            assert split < self.n_samples, "validation set size is configured to be larger than entire dataset."
            len_valid = split
        else:
            len_valid = int(self.n_samples * split)

        valid_idx = idx_full[0:len_valid]
        train_idx = np.delete(idx_full, np.arange(0, len_valid))

        train_sampler = SubsetRandomSampler(train_idx)
        valid_sampler = SubsetRandomSampler(valid_idx)
        # turn off shuffle option which is mutually exclusive with sampler
        self.shuffle = False
        self.n_samples = len(train_idx)

        return train_sampler, valid_sampler

    def split_validation(self, diff_kwargs=None):
        init_kwargs = self.init_kwargs
        if diff_kwargs is not None:
            init_kwargs.update(diff_kwargs)
        if self.valid_sampler is None:
            return None
        else:
            return DataLoader(sampler=self.valid_sampler, **self.init_kwargs)

    def num_samples(self):
        return len(self.sampler)


class BaseDataLoaderExplicitSplit(DataLoader):
    """
    Base class for all data loaders
    """
    def __init__(self, dataset, batch_size, shuffle, num_workers, collate_fn=default_collate):
        self.shuffle = shuffle

        self.batch_idx = 0
        self.n_samples = len(dataset)

        self.init_kwargs = {
            'dataset': dataset,
            'batch_size': batch_size,
            'shuffle': self.shuffle,
            'collate_fn': collate_fn,
            'num_workers': num_workers,
            'pin_memory': True
        }
        super().__init__(**self.init_kwargs)


class DistBaseDataLoaderExplicitSplit(DataLoader):
    """
    Base class for all data loaders
    """
    def __init__(self, dataset, batch_size, shuffle, num_workers, collate_fn=default_collate):
        self.shuffle = shuffle

        self.batch_idx = 0
        self.n_samples = len(dataset)
        self.train_sampler = DistributedSampler(dataset)
        self.init_kwargs = {
            'dataset': dataset,
            'batch_size': batch_size,
            'shuffle': False,
            'collate_fn': collate_fn,
            'num_workers': num_workers,
            'pin_memory': True,
            'sampler': self.train_sampler
        }
        super().__init__(**self.init_kwargs)


class MultiDistBaseDataLoaderExplicitSplit(DataLoader):
    """
    Base class for all data loaders
    """
    def __init__(self, args, dataset, batch_size, shuffle, num_workers, collate_fn=default_collate):
        self.shuffle = shuffle

        self.batch_idx = 0
        self.n_samples = len(dataset)
        self.args = args
        #print("args.world_size: ", self.args)
        #self.train_sampler = DistributedSampler(dataset, num_replicas=self.args.world_size, rank=self.args.rank, drop_last=False)
        self.train_sampler = DistributedSampler(dataset, drop_last=False)
        self.init_kwargs = {
            'dataset': dataset,
            'batch_size': batch_size,
            'shuffle': False,
            'collate_fn': collate_fn,
            'num_workers': num_workers,
            'pin_memory': True,
            'sampler': self.train_sampler
        }
        super().__init__(**self.init_kwargs)

class BaseMultiDataLoader:
    """
    Currently implemented as undersample the bigger dataloaders...
    """
    def __init__(self, dataloaders):
        self.dataloaders = dataloaders
        self.batch_size = self.dataloaders[0].batch_size

    def __getitem__(self, item):
        dl_idx = item % len(self.dataloaders)
        return next(iter(self.dataloaders[dl_idx]))

    def __len__(self):
        return min(len(x) for x in self.dataloaders) * len(self.dataloaders)

    def num_samples(self):
        return sum(len(x.sampler) for x in self.dataloaders)

QFVS/base/base_data_loader.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import os
import pdb

import tqdm
import random
from abc import abstractmethod

import av
import cv2
import decord
import ffmpeg
import numpy as np
import torch
from PIL import Image
from torch.utils.data import Dataset, get_worker_info
from torchvision import transforms


class TextVideoDataset(Dataset):
    def __init__(self,
                 dataset_name,
                 text_params,
                 video_params,
                 data_dir,
                 meta_dir=None,
                 split='train',
                 tsfms=None,
                 cut=None,
                 subsample=1,
                 sliding_window_stride=-1,
                 reader='decord',
                 neg_param=None,
                 ):
        self.dataset_name = dataset_name
        self.text_params = text_params
        self.video_params = video_params
        # check for environment variables
        self.data_dir = os.path.expandvars(data_dir)
        if meta_dir is not None:
            self.meta_dir = os.path.expandvars(meta_dir)
        else:
            self.meta_dir = self.data_dir
        self.split = split
        self.transforms = tsfms
        self.cut = cut
        self.subsample = subsample
        self.sliding_window_stride = sliding_window_stride
        self.video_reader = video_reader[reader]
        self.label_type = 'caption'
        self.neg_param = neg_param
        self._load_metadata()
        if self.sliding_window_stride != -1:
            if self.split != 'test':
                raise ValueError('Fixing frame sampling is for test time only. can remove but...')
            self._fix_temporal_samples()

    @abstractmethod
    def _load_metadata(self):
        raise NotImplementedError("Metadata loading must be implemented by subclass")

    @abstractmethod
    def _get_video_path(self, sample):
        raise NotImplementedError("Get video path function must be implemented by subclass")

    def _get_caption(self, sample):
        raise NotImplementedError("Get caption function must be implemented by subclass")

    def _get_video_lens(self):
        vlen_li = []
        for idx, row in self.metadata.iterrows():
            video_path = self._get_video_path(row)[0]
            vlen_li.append(get_video_len(video_path))

        return vlen_li

    def _fix_temporal_samples(self):
        self.metadata['vlen'] = self._get_video_lens()
        self.metadata['frame_intervals'] = self.metadata['vlen'].apply(
            lambda x: np.linspace(start=0, stop=x, num=min(x, self.video_params['num_frames']) + 1).astype(int))
        self.metadata['fix_start'] = self.metadata['frame_intervals'].apply(
            lambda x: np.arange(0, int(x[-1] / len(x - 1)), self.sliding_window_stride)
        )
        self.metadata = self.metadata.explode('fix_start')

    def __len__(self):
        return len(self.metadata)

    def __getitem__(self, item):
        item = item % len(self.metadata)
        sample = self.metadata.iloc[item]
        video_fp, rel_fp = self._get_video_path(sample)
        caption = self._get_caption(sample)

        video_loading = self.video_params.get('loading', 'strict')
        frame_sample = 'rand'
        fix_start = None
        if self.split == 'test':
            frame_sample = 'uniform'
        if self.sliding_window_stride != -1:
            fix_start = sample['fix_start']

        try:
            if os.path.isfile(video_fp):
                imgs, idxs = self.video_reader(video_fp, self.video_params['num_frames'], frame_sample,
                                               fix_start=fix_start)
            else:
                print(f"Warning: missing video file {video_fp}.")
                assert False
        except Exception as e:
            if video_loading == 'strict':
                raise ValueError(
                    f'Video loading failed for {video_fp}, video loading for this dataset is strict.') from e
            else:
                imgs = Image.new('RGB', (self.video_params['input_res'], self.video_params['input_res']), (0, 0, 0))
                imgs = transforms.ToTensor()(imgs).unsqueeze(0)

        # if self.transforms is not None:
        #     imgs = self.transforms(imgs)

        if self.transforms is not None:
            if self.video_params['num_frames'] > 1:
                imgs = imgs.transpose(0, 1)  # [T, C, H, W] ---> [C, T, H, W]
                imgs = self.transforms(imgs)
                imgs = imgs.transpose(0, 1)  # recover
            else:
                imgs = self.transforms(imgs)

        # if self.transforms is not None:
        #     # for video aug if T>1 else just do image aug
        #     if imgs.size(0) > 1:
        #         # added by Mr. YAN
        #         imgs = imgs.transpose(0, 1) # [T, C, H, W] ---> [C, T, H, W]
        #         imgs = self.transforms(imgs)
        #         imgs = imgs.transpose(0, 1) # recover
        #     else:
        #         imgs = self.transforms(imgs)    # t

        final = torch.zeros([self.video_params['num_frames'], 3, self.video_params['input_res'],
                             self.video_params['input_res']])
        final[:imgs.shape[0]] = imgs

        meta_arr = {'raw_captions': caption, 'paths': rel_fp, 'dataset': self.dataset_name}
        data = {'video': final, 'text': caption, 'meta': meta_arr}
        return data


class TextImageDataset(TextVideoDataset):

    def __getitem__(self, item):
        item = item % len(self.metadata)
        sample = self.metadata.iloc[item]
        video_fp, rel_fp = self._get_video_path(sample)
        caption = self._get_caption(sample)

        video_loading = self.video_params.get('loading', 'strict')

        try:
            img = Image.open(video_fp).convert("RGB")
        except:
            if video_loading == 'strict':
                raise ValueError(f'Image loading failed for {video_fp}, image loading for this dataset is strict.')
            else:
                img = Image.new('RGB', (self.video_params['input_res'], self.video_params['input_res']), (0, 0, 0))

        # convert to tensor because video transforms don't, expand such that its a 1-frame video.
        img = transforms.ToTensor()(img).unsqueeze(0)
        if self.transforms is not None:
            img = self.transforms(img)
        meta_arr = {'raw_captions': caption, 'paths': rel_fp, 'dataset': self.dataset_name}
        data = {'video': img, 'text': caption, 'meta': meta_arr}
        return data


def sample_frames(num_frames, vlen, sample='rand', fix_start=None):
    acc_samples = min(num_frames, vlen)
    intervals = np.linspace(start=0, stop=vlen, num=acc_samples + 1).astype(int)
    ranges = []
    for idx, interv in enumerate(intervals[:-1]):
        ranges.append((interv, intervals[idx + 1] - 1))
    if sample == 'rand':
        frame_idxs = [random.choice(range(x[0], x[1])) for x in ranges]
    elif fix_start is not None:
        frame_idxs = [x[0] + fix_start for x in ranges]
    elif sample == 'uniform':
        frame_idxs = [(x[0] + x[1]) // 2 for x in ranges]
    else:
        raise NotImplementedError

    return frame_idxs

def sample_frames_clips(start, end, vlen, acc_samples):
    start = max(0, start)
    end = min(vlen, end)

    intervals = np.linspace(start=start, stop=end, num=int(acc_samples) + 1).astype(int)
    ranges = []
    for idx, interv in enumerate(intervals[:-1]):
        ranges.append((interv, intervals[idx + 1] - 1))
        frame_idxs = [(x[0] + x[1]) // 2 for x in ranges
                      ]
    return frame_idxs

def sample_frames_start_end(num_frames, start, end, sample='rand', fix_start=None):
    acc_samples = min(num_frames, end)
    intervals = np.linspace(start=start, stop=end, num=acc_samples + 1).astype(int)
    ranges = []
    for idx, interv in enumerate(intervals[:-1]):
        ranges.append((interv, intervals[idx + 1] - 1))
    if sample == 'rand':
        frame_idxs = [random.choice(range(x[0], x[1])) for x in ranges]
    elif fix_start is not None:
        frame_idxs = [x[0] + fix_start for x in ranges]
    elif sample == 'uniform':
        frame_idxs = [(x[0] + x[1]) // 2 for x in ranges]
    else:
        raise NotImplementedError

    return frame_idxs

def read_frames_cv2(video_path, num_frames, sample='rand', fix_start=None):
    cap = cv2.VideoCapture(video_path)
    assert (cap.isOpened())
    vlen = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    # get indexes of sampled frames
    frame_idxs = sample_frames(num_frames, vlen, sample=sample, fix_start=fix_start)
    frames = []
    success_idxs = []
    for index in frame_idxs:
        cap.set(cv2.CAP_PROP_POS_FRAMES, index - 1)
        ret, frame = cap.read()
        if ret:
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frame = torch.from_numpy(frame)
            # (H x W x C) to (C x H x W)
            frame = frame.permute(2, 0, 1)
            frames.append(frame)
            success_idxs.append(index)
        else:
            pass
            # print(frame_idxs, ' fail ', index, f'  (vlen {vlen})')

    frames = torch.stack(frames).float() / 255
    cap.release()
    return frames, success_idxs

def read_frames_cv2_egoclip(video_path_1, video_path_2, num_frames, sample,
                            start_sec, end_sec, bound_sec):
    if video_path_1 == video_path_2:
        cap1 = cv2.VideoCapture(video_path_1)
        cap2 = cap1
        vlen1 = int(cap1.get(cv2.CAP_PROP_FRAME_COUNT))
        vlen2 = vlen1
        assert (cap1.isOpened())
    else:   # some clips may span two segments.
        cap1 = cv2.VideoCapture(video_path_1)
        cap2 = cv2.VideoCapture(video_path_2)
        vlen1 = int(cap1.get(cv2.CAP_PROP_FRAME_COUNT))
        vlen2 = int(cap2.get(cv2.CAP_PROP_FRAME_COUNT))
        assert (cap1.isOpened())
        assert (cap2.isOpened())

    # get indexes of sampled frames
    start_f = max(0, int(start_sec * 30))
    end_f = max(0, int(end_sec * 30))
    bound_f = int(bound_sec * 30)
    frame_idxs = sample_frames_start_end(num_frames, start_f, end_f, sample=sample)

    frames = []
    success_idxs = []
    for index in frame_idxs:
        _index = index % (600 * 30)
        if index > bound_f: # frame from the last video
            _index = min(_index, vlen2)
            cap2.set(cv2.CAP_PROP_POS_FRAMES, _index - 1)
            ret, frame = cap2.read()
        else:   # frame from the first video
            _index = min(_index, vlen1)
            cap1.set(cv2.CAP_PROP_POS_FRAMES, _index - 1)
            ret, frame = cap1.read()

        if ret:
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frame = torch.from_numpy(frame)
            # (H x W x C) to (C x H x W)
            frame = frame.permute(2, 0, 1)
            frames.append(frame)
            success_idxs.append(index)
        else:
            pass

    while len(frames) < num_frames: # complete the frame
        frames.append(frames[-1])

    frames = torch.stack(frames).float() / 255
    cap1.release()
    cap2.release()
    return frames, success_idxs

def read_frames_cv2_epic(video_path, start_frame, stop_frame, num_frames, sample='rand', fix_start=None):
    # get indexes of sampled frames
    frame_idxs = sample_frames_start_end(num_frames, start_frame, stop_frame, sample=sample, fix_start=fix_start)
    frames = []
    success_idxs = []
    for index in frame_idxs:
        img_name = 'frame_' + str(index).zfill(10) + '.jpg'
        frame = cv2.imread(os.path.join(video_path, img_name),cv2.COLOR_BGR2RGB)

        frame = torch.from_numpy(frame)
        # (H x W x C) to (C x H x W)
        frame = frame.permute(2, 0, 1)
        frames.append(frame)
        success_idxs.append(index)

    frames = torch.stack(frames).float() / 255
    return frames, success_idxs

def read_frames_cv2_charades(video_path, num_frames, sample, start_sec=None, end_sec=None):
    cap = cv2.VideoCapture(video_path)
    assert (cap.isOpened())
    vlen = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = cap.get(5)

    # get indexes of sampled frames
    if not start_sec and not end_sec:
        frame_idxs = sample_frames(num_frames, vlen, sample=sample)
    else:
        start_f = max(0, int(start_sec * fps))
        end_f = min(int(end_sec * fps), vlen)
        frame_idxs = sample_frames_start_end(num_frames, start_f, end_f, sample=sample)

    frames = []
    success_idxs = []
    for index in frame_idxs:
        cap.set(cv2.CAP_PROP_POS_FRAMES, index - 1)
        ret, frame = cap.read()
        if ret:
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frame = torch.from_numpy(frame)
            # (H x W x C) to (C x H x W)
            frame = frame.permute(2, 0, 1)
            frames.append(frame)
            success_idxs.append(index)
        else:
            pass

    frames = torch.stack(frames).float() / 255
    cap.release()
    return frames, success_idxs

def read_frames_av(video_path, num_frames, sample='rand', fix_start=None):
    reader = av.open(video_path)
    try:
        frames = []
        frames = [torch.from_numpy(f.to_rgb().to_ndarray()) for f in reader.decode(video=0)]
    except (RuntimeError, ZeroDivisionError) as exception:
        print('{}: WEBM reader cannot open {}. Empty '
              'list returned.'.format(type(exception).__name__, video_path))
    vlen = len(frames)
    frame_idxs = sample_frames(num_frames, vlen, sample=sample, fix_start=fix_start)
    frames = torch.stack([frames[idx] for idx in frame_idxs]).float() / 255
    frames = frames.permute(0, 3, 1, 2)
    return frames, frame_idxs

decord.bridge.set_bridge("torch")

def read_frames_decord(video_path, num_frames, sample='rand', fix_start=None):
    video_reader = decord.VideoReader(video_path, num_threads=1)
    vlen = len(video_reader)
    frame_idxs = sample_frames(num_frames, vlen, sample=sample, fix_start=fix_start)
    video_reader.skip_frames(1)
    frames = video_reader.get_batch(frame_idxs)

    frames = frames.float() / 255
    frames = frames.permute(0, 3, 1, 2)
    return frames, frame_idxs

def read_frames_decord_start_end(video_path, start, end, num_frames):
    video_reader = decord.VideoReader(video_path, num_threads=1)
    vlen = len(video_reader)
    frame_idxs = sample_frames_clips(start, end, vlen, num_frames + 1)
    video_reader.skip_frames(1)
    frames = video_reader.get_batch(frame_idxs)

    frames = frames.float() / 255
    frames = frames.permute(0, 3, 1, 2)
    return frames, frame_idxs

def get_video_len(video_path):
    cap = cv2.VideoCapture(video_path)
    if not (cap.isOpened()):
        return False
    vlen = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    cap.release()
    return vlen

video_reader = {
    'av': read_frames_av,
    'cv2': read_frames_cv2,
    'cv2_epic': read_frames_cv2_epic,
    'cv2_charades': read_frames_cv2_charades,
    'cv2_egoclip': read_frames_cv2_egoclip,
    'decord': read_frames_decord,
    'decord_start_end': read_frames_decord_start_end,
}

QFVS/base/base_dataset.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import torch.nn as nn
import numpy as np
from abc import abstractmethod


class BaseModel(nn.Module):
    """
    Base class for all models
    """
    @abstractmethod
    def forward(self, *inputs):
        """
        Forward pass logic

        :return: Model output
        """
        raise NotImplementedError

    def __str__(self):
        """
        Model prints with number of trainable parameters
        """
        model_parameters = filter(lambda p: p.requires_grad, self.parameters())
        params = sum(np.prod(p.size()) for p in model_parameters)
        return super().__str__() + '\nTrainable parameters: {}'.format(params)

QFVS/base/base_model.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


from abc import abstractmethod

import torch
from numpy import inf
from pathlib import Path

class BaseTrainer:
    """
    Base class for all trainers
    """
    def __init__(self, model, loss, metrics, optimizer, config, writer=None, init_val=False):
        self.config = config
        self.logger = config.get_logger('trainer', config['trainer']['verbosity'])
        self.init_val = init_val
        # setup GPU device if available, move model into configured device
        self.device, device_ids = self._prepare_device(config['n_gpu'])
        self.model = model.to(self.device)
        self.model.device = self.device
        if len(device_ids) > 1:
            self.model = torch.nn.DataParallel(model, device_ids=device_ids)

        loss = loss.to(self.device)
        self.loss = loss
        self.metrics = metrics
        self.optimizer = optimizer

        cfg_trainer = config['trainer']
        self.epochs = cfg_trainer['epochs']
        self.save_period = cfg_trainer['save_period']
        self.monitor = cfg_trainer.get('monitor', 'off')
        self.init_val = cfg_trainer.get('init_val', True)

        # configuration to monitor model performance and save best
        if self.monitor == 'off':
            self.mnt_mode = 'off'
            self.mnt_best = 0
        else:
            self.mnt_mode, self.mnt_metric = self.monitor.split()
            assert self.mnt_mode in ['min', 'max']

            self.mnt_best = inf if self.mnt_mode == 'min' else -inf
            self.early_stop = cfg_trainer.get('early_stop', inf)

        self.start_epoch = 1

        self.checkpoint_dir = config.save_dir

        # setup visualization writer instance                
        #self.writer = TensorboardWriter(config.log_dir, self.logger, cfg_trainer['tensorboard'])
        self.writer = writer

        if config.resume is not None:
            self._resume_checkpoint(config.resume)

    @abstractmethod
    def _train_epoch(self, epoch):
        """
        Training logic for an epoch

        :param epoch: Current epoch number
        """
        raise NotImplementedError

    @abstractmethod
    def _valid_epoch(self, epoch):
        """
        Training logic for an epoch

        :param epoch: Current epoch number
        """
        raise NotImplementedError


    def train(self):
        """
        Full training logic
        """
        not_improved_count = 0

        if self.init_val:
            _ = self._valid_epoch(-1)

        for epoch in range(self.start_epoch, self.epochs + 1):
            result = self._train_epoch(epoch)

            # save logged informations into log dict

            # save logged informations into log dict
            log = {'epoch': epoch}
            for key, value in result.items():
                if key == 'metrics':
                    log.update({mtr.__name__: value[i]
                                for i, mtr in enumerate(self.metrics)})
                elif key == 'val_metrics':
                    log.update({'val_' + mtr.__name__: value[i]
                                for i, mtr in enumerate(self.metrics)})
                elif key == 'nested_val_metrics':
                    # NOTE: currently only supports two layers of nesting
                    for subkey, subval in value.items():
                        for subsubkey, subsubval in subval.items():
                            for subsubsubkey, subsubsubval in subsubval.items():
                                log[f"val_{subkey}_{subsubkey}_{subsubsubkey}"] = subsubsubval
                else:
                    log[key] = value

            # print logged informations to the screen
            for key, value in log.items():
                self.logger.info('    {:15s}: {}'.format(str(key), value))

            # evaluate model performance according to configured metric, save best checkpoint as model_best
            best = False
            if self.mnt_mode != 'off':
                try:
                    # check whether model performance improved or not, according to specified metric(mnt_metric)
                    improved = (self.mnt_mode == 'min' and log[self.mnt_metric] <= self.mnt_best) or \
                               (self.mnt_mode == 'max' and log[self.mnt_metric] >= self.mnt_best)
                except KeyError:
                    self.logger.warning("Warning: Metric '{}' is not found. "
                                        "Model performance monitoring is disabled.".format(self.mnt_metric))
                    self.mnt_mode = 'off'
                    improved = False

                if improved:
                    self.mnt_best = log[self.mnt_metric]
                    not_improved_count = 0
                    best = True
                else:
                    not_improved_count += 1

                if not_improved_count > self.early_stop:
                    self.logger.info("Validation performance didn\'t improve for {} epochs. "
                                     "Training stops.".format(self.early_stop))
                    break

            if epoch % self.save_period == 0 or best:
            #if best:
                self._save_checkpoint(epoch, save_best=best)

    def _prepare_device(self, n_gpu_use):
        """
        setup GPU device if available, move model into configured device
        """
        n_gpu = torch.cuda.device_count()
        if n_gpu_use > 0 and n_gpu == 0:
            self.logger.warning("Warning: There\'s no GPU available on this machine,"
                                "training will be performed on CPU.")
            n_gpu_use = 0
        if n_gpu_use > n_gpu:
            self.logger.warning("Warning: The number of GPU\'s configured to use is {}, but only {} are available "
                                "on this machine.".format(n_gpu_use, n_gpu))
            n_gpu_use = n_gpu
        device = torch.device('cuda:0' if n_gpu_use > 0 else 'cpu')
        list_ids = list(range(n_gpu_use))
        return device, list_ids

    def _save_checkpoint(self, epoch, save_best=False):
        """
        Saving checkpoints

        :param epoch: current epoch number
        :param log: logging information of the epoch
        :param save_best: if True, rename the saved checkpoint to 'model_best.pth'
        """
        arch = type(self.model).__name__
        state = {
            'arch': arch,
            'epoch': epoch,
            'state_dict': self.model.state_dict(),
            'optimizer': self.optimizer.state_dict(),
            'monitor_best': self.mnt_best,
            'config': self.config
        }
        filename = str(self.checkpoint_dir / 'checkpoint-epoch{}.pth'.format(epoch))
        torch.save(state, filename)
        self.logger.info("Saving checkpoint: {} ...".format(filename))
        if save_best:
            best_path = str(self.checkpoint_dir / 'model_best.pth')
            torch.save(state, best_path)
            self.logger.info("Saving current best: model_best.pth ...")

    def _resume_checkpoint(self, resume_path):
        """
        Resume from saved checkpoints

        :param resume_path: Checkpoint path to be resumed
        """
        resume_path = str(resume_path)
        self.logger.info("Loading checkpoint: {} ...".format(resume_path))
        checkpoint = torch.load(resume_path)
        self.start_epoch = checkpoint['epoch'] + 1
        self.mnt_best = checkpoint['monitor_best']

        # load architecture params from checkpoint.
        if checkpoint['config']['arch'] != self.config['arch']:
            self.logger.warning("Warning: Architecture configuration given in config file is different from that of "
                                "checkpoint. This may yield an exception while state_dict is being loaded.")

        state_dict = checkpoint['state_dict']

        load_state_dict_keys = list(state_dict.keys())
        curr_state_dict_keys = list(self.model.state_dict().keys())
        redo_dp = False
        if not curr_state_dict_keys[0].startswith('module.') and load_state_dict_keys[0].startswith('module.'):
            undo_dp = True
        elif curr_state_dict_keys[0].startswith('module.') and not load_state_dict_keys[0].startswith('module.'):
            redo_dp = True
            undo_dp = False
        else:
            undo_dp = False

        if undo_dp:
            from collections import OrderedDict
            new_state_dict = OrderedDict()
            for k, v in state_dict.items():
                name = k[7:]  # remove `module.`
                new_state_dict[name] = v
            # load params
        elif redo_dp:
            from collections import OrderedDict
            new_state_dict = OrderedDict()
            for k, v in state_dict.items():
                name = 'module.' + k  # remove `module.`
                new_state_dict[name] = v
        else:
            new_state_dict = state_dict

        self.model.load_state_dict(new_state_dict)

        # load optimizer state from checkpoint only when optimizer type is not changed.
        if checkpoint['config']['optimizer']['type'] != self.config['optimizer']['type']:
            self.logger.warning("Warning: Optimizer type given in config file is different from that of checkpoint. "
                                "Optimizer parameters not being resumed.")
        else:
            self.optimizer.load_state_dict(checkpoint['optimizer'])

        self.logger.info("Checkpoint loaded. Resume training from epoch {}".format(self.start_epoch))


class Multi_BaseTrainer_dist:
    """
    Base class for all trainers
    """
    def __init__(self, args, model, loss, metrics, optimizer, scheduler, gpu, config, writer=None, init_val=False):
        self.config = config
        self.logger = config.get_logger('trainer', config['trainer']['verbosity'])
        self.init_val = init_val
        self.args = args
        # setup GPU device if available, move model into configured device
        self.device, device_ids = self._prepare_device(config['n_gpu'])

        self.model = model.cuda(gpu)
        #self.model.device = self.device

        # if len(device_ids) > 1:
            #self.model = torch.nn.parallel.DistributedDataParallel(model)
            #self.model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.rank])
            # self.model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank])
            #self.model = torch.nnç.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], find_unused_parameters=True)
            #self.model = torch.nn.DataParallel(model, device_ids=device_ids)

        ddp_args={}
        ddp_args['static_graph'] = True
        self.model = torch.nn.parallel.DistributedDataParallel(self.model, device_ids=[gpu], **ddp_args)  

        loss = loss.cuda(gpu)
        self.loss = loss
        self.metrics = metrics
        self.optimizer = optimizer
        self.scheduler = scheduler

        cfg_trainer = config['trainer']
        self.epochs = cfg_trainer['epochs']
        self.save_period = cfg_trainer['save_period']
        self.monitor = cfg_trainer.get('monitor', 'off')
        self.init_val = cfg_trainer.get('init_val', True)

        # configuration to monitor model performance and save best
        if self.monitor == 'off':
            self.mnt_mode = 'off'
            self.mnt_best = 0
        else:
            self.mnt_mode, self.mnt_metric = self.monitor.split()
            assert self.mnt_mode in ['min', 'max']

            self.mnt_best = inf if self.mnt_mode == 'min' else -inf
            self.early_stop = cfg_trainer.get('early_stop', inf)

        self.start_epoch = 1

        self.checkpoint_dir = Path(args.save_dir)

        # setup visualization writer instance
        #self.writer = TensorboardWriter(config.log_dir, self.logger, cfg_trainer['tensorboard'])
        self.writer = writer

        if config.resume is not None:
            self._resume_checkpoint(config.resume)

    @abstractmethod
    def _train_epoch(self, epoch, gpu):
        """
        Training logic for an epoch

        :param epoch: Current epoch number
        """
        raise NotImplementedError

    @abstractmethod
    def _valid_epoch(self, epoch, gpu):
        """
        Training logic for an epoch

        :param epoch: Current epoch number
        """
        raise NotImplementedError


    def train(self, gpu):
        """
        Full training logic
        """
        not_improved_count = 0

        if self.init_val:
            _ = self._valid_epoch(-1, gpu)
            # pass

        scaler = torch.cuda.amp.GradScaler()

        for epoch in range(self.start_epoch, self.epochs + 1):
            result = self._train_epoch(epoch, scaler, gpu)

            # save logged informations into log dict

            # save logged informations into log dict
            log = {'epoch': epoch}
            for key, value in result.items():
              if self.args.rank == 0:
                if key == 'metrics':
                    log.update({mtr.__name__: value[i]
                                for i, mtr in enumerate(self.metrics)})
                elif key == 'val_metrics':
                    log.update({'val_' + mtr.__name__: value[i]
                                for i, mtr in enumerate(self.metrics)})
                elif key == 'nested_val_metrics':
                    # NOTE: currently only supports two layers of nesting
                    for subkey, subval in value.items():
                        for subsubkey, subsubval in subval.items():
                            for subsubsubkey, subsubsubval in subsubval.items():
                                log[f"val_{subkey}_{subsubkey}_{subsubsubkey}"] = subsubsubval
                else:
                    log[key] = value

            # print logged informations to the screen
            for key, value in log.items():
              if self.args.rank == 0:
                self.logger.info('    {:15s}: {}'.format(str(key), value))

            # evaluate model performance according to configured metric, save best checkpoint as model_best
            best = False
            if self.mnt_mode != 'off' and self.args.rank == 0:
                try:
                    # check whether model performance improved or not, according to specified metric(mnt_metric)
                    improved = (self.mnt_mode == 'min' and log[self.mnt_metric] <= self.mnt_best) or \
                               (self.mnt_mode == 'max' and log[self.mnt_metric] >= self.mnt_best)
                except KeyError:
                    self.logger.warning("Warning: Metric '{}' is not found. "
                                        "Model performance monitoring is disabled.".format(self.mnt_metric))
                    self.mnt_mode = 'off'
                    improved = False

                if improved:
                    self.mnt_best = log[self.mnt_metric]
                    not_improved_count = 0
                    best = True
                else:
                    not_improved_count += 1

                #if not_improved_count > self.early_stop:
                #    self.logger.info("Validation performance didn\'t improve for {} epochs. "
                #                     "Training stops.".format(self.early_stop))
                #    break

            if epoch % self.save_period == 0 or best:
                if self.args.rank == 0:
                #if best:
                    self._save_checkpoint(epoch, save_best=best)

    def _prepare_device(self, n_gpu_use):
        """
        setup GPU device if available, move model into configured device
        """
        n_gpu = torch.cuda.device_count()
        if n_gpu_use > 0 and n_gpu == 0:
            self.logger.warning("Warning: There\'s no GPU available on this machine,"
                                "training will be performed on CPU.")
            n_gpu_use = 0
        if n_gpu_use > n_gpu:
            self.logger.warning("Warning: The number of GPU\'s configured to use is {}, but only {} are available "
                                "on this machine.".format(n_gpu_use, n_gpu))
            n_gpu_use = n_gpu
        device = torch.device('cuda' if n_gpu_use > 0 else 'cpu')
        list_ids = list(range(n_gpu_use))
        return device, list_ids

    def _save_checkpoint(self, epoch, save_best=False):
        """
        Saving checkpoints

        :param epoch: current epoch number
        :param log: logging information of the epoch
        :param save_best: if True, rename the saved checkpoint to 'model_best.pth'
        """
        arch = type(self.model).__name__
        state = {
            'arch': arch,
            'epoch': epoch,
            'state_dict': self.model.state_dict(),
            'optimizer': self.optimizer.state_dict(),
            'monitor_best': self.mnt_best,
            'config': self.config
        }
        filename = str(self.checkpoint_dir / 'checkpoint-epoch{}.pth'.format(epoch))
        torch.save(state, filename)
        self.logger.info("Saving checkpoint: {} ...".format(filename))
        if save_best:
            best_path = str(self.checkpoint_dir / 'model_best.pth')
            torch.save(state, best_path)
            self.logger.info("Saving current best: model_best.pth ...")

    def _resume_checkpoint(self, resume_path):
        """
        Resume from saved checkpoints

        :param resume_path: Checkpoint path to be resumed
        """
        resume_path = str(resume_path)
        self.logger.info("Loading checkpoint: {} ...".format(resume_path))
        #checkpoint = torch.load(resume_path)
        checkpoint = torch.load(resume_path, map_location='cuda:{}'.format(self.args.rank))
        self.start_epoch = checkpoint['epoch'] + 1
        self.mnt_best = checkpoint['monitor_best']

        # load architecture params from checkpoint.
        if checkpoint['config']['arch'] != self.config['arch']:
            self.logger.warning("Warning: Architecture configuration given in config file is different from that of "
                                "checkpoint. This may yield an exception while state_dict is being loaded.")

        state_dict = checkpoint['state_dict']

        load_state_dict_keys = list(state_dict.keys())
        curr_state_dict_keys = list(self.model.state_dict().keys())
        redo_dp = False
        if not curr_state_dict_keys[0].startswith('module.') and load_state_dict_keys[0].startswith('module.'):
            undo_dp = True
        elif curr_state_dict_keys[0].startswith('module.') and not load_state_dict_keys[0].startswith('module.'):
            redo_dp = True
            undo_dp = False
        else:
            undo_dp = False

        if undo_dp:
            from collections import OrderedDict
            new_state_dict = OrderedDict()
            for k, v in state_dict.items():
                name = k[7:]  # remove `module.`
                new_state_dict[name] = v
            # load params
        elif redo_dp:
            from collections import OrderedDict
            new_state_dict = OrderedDict()
            for k, v in state_dict.items():
                name = 'module.' + k  # remove `module.`
                new_state_dict[name] = v
        else:
            new_state_dict = state_dict

        self.model.load_state_dict(new_state_dict)

        # load optimizer state from checkpoint only when optimizer type is not changed.
        if checkpoint['config']['optimizer']['type'] != self.config['optimizer']['type']:
            self.logger.warning("Warning: Optimizer type given in config file is different from that of checkpoint. "
                                "Optimizer parameters not being resumed.")
        else:
            self.optimizer.load_state_dict(checkpoint['optimizer'])

        self.logger.info("Checkpoint loaded. Resume training from epoch {}".format(self.start_epoch))

QFVS/base/base_trainer.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import os
import h5py
import torch
from torch.utils.data.dataset import Dataset
from utils_QFVS import load_pickle, load_json
from torch.utils.data import DataLoader
import transformers
import pdb

class UCTDataset(Dataset):
    def __init__(self, config, train_videos, all_features):
        self.config=config
        self.train_videos = train_videos
        self.all_features = all_features
        self.dataset=[]
        self.tokenizer = transformers.RobertaTokenizer.from_pretrained('roberta-base')
        for video_id in self.train_videos:
            for _ , _, files in os.walk("./Datasets/QFVS/data/origin_data/Query-Focused_Summaries/Oracle_Summaries/P0"+str(video_id)):
                for file in files:
                    self.dataset.append(file[:file.find("_oracle.txt")]+"_"+str(video_id))

    def __getitem__(self,index):
        video_id=self.dataset[index].split('_')[2]
        features = self.all_features[video_id]['feature']
        seg_len = self.all_features[video_id]['seg_len']

        transfer={"Cupglass":"Glass",
                  "Musicalinstrument":"Instrument",
                  "Petsanimal":"Animal"}

        concept1,concept2=self.dataset[index].split('_')[0:2]

        concept1_GT=torch.zeros(self.config["max_segment_num"]*self.config["max_frame_num"])
        concept2_GT=torch.zeros(self.config["max_segment_num"]*self.config["max_frame_num"])
        with open("./Datasets/QFVS/data/origin_data/Dense_per_shot_tags/P0"+video_id+"/P0"+video_id+".txt","r") as f:
            lines=f.readlines()
            for index,line in enumerate(lines):
                concepts=line.strip().split(',')
                if concept1 in concepts:
                    concept1_GT[index]=1
                if concept2 in concepts:
                    concept2_GT[index]=1

        shot_num=seg_len.sum()
        mask_GT=torch.zeros(self.config["max_segment_num"]*self.config["max_frame_num"],dtype=torch.bool)
        for i in range(shot_num):
            mask_GT[i]=1


        oracle_summary = torch.zeros(self.config["max_segment_num"]*self.config["max_frame_num"])
        GT_summary_shots = []
        with open("./Datasets/QFVS/data/origin_data/Query-Focused_Summaries/Oracle_Summaries/P0"+str(video_id)+"/"+str(concept1)+"_"+str(concept2)+"_"+"oracle.txt","r") as f:
            for line in f.readlines():
                GT_summary_shots.append(int(line.strip()))
        GT_summary_shots_modified = [x - 1 for x in GT_summary_shots]
        for element in GT_summary_shots_modified:
            oracle_summary[element] = 1

        if concept1 in transfer:
            concept1=transfer[concept1]
        if concept2 in transfer:
            concept2=transfer[concept2]


        concept1_prompt = "There is a " + str(concept1)
        concept2_prompt = "There is a " + str(concept2)
        query_prompt = "There is a " + str(concept1) + " and a " + str(concept2)
        
        concept1_embed = self.tokenizer(concept1_prompt, return_tensors='pt', padding='max_length', max_length=15, truncation=True)
        concept2_embed = self.tokenizer(concept2_prompt, return_tensors='pt', padding='max_length', max_length=15, truncation=True)
        query_embed = self.tokenizer(query_prompt, return_tensors='pt', padding='max_length', max_length=15, truncation=True)

        concept1_token, concept1_mask = concept1_embed['input_ids'], concept1_embed['attention_mask']
        concept2_token, concept2_mask = concept2_embed['input_ids'], concept2_embed['attention_mask']
        query_token, query_mask = query_embed['input_ids'], query_embed['attention_mask']

        return features, seg_len, concept1_token[0], concept1_mask[0], concept2_token[0], concept2_mask[0], query_token[0], query_mask[0], concept1_GT, concept2_GT, mask_GT, oracle_summary, video_id

    def __len__(self):
        return len(self.dataset)



if __name__=='__main__':

    config=load_json("./config/config.json")



QFVS/dataset_prompt.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
import torchvision
import os
from torchvision import transforms,models
import torch.nn as nn
from PIL import Image
import clip
import pdb
import argparse
import glob
import h5py
import json
from tqdm import tqdm
import model.model_video_unfused as module_arch
from pathlib import Path
from collections import OrderedDict
from segment import cpd_auto

max_segment_num = 20
max_frame_num = 200


import warnings
warnings.filterwarnings("ignore")

def read_json(fname):
    with fname.open('rt') as handle:
        return json.load(handle, object_hook=OrderedDict)

config = read_json(Path('qfvs.json'))

class Image_Dataset(Dataset):
    def __init__(self, image_list, transform=None):

        self.image_list = image_list
        self.transform = transform

    def __len__(self):
        return len(self.image_list)

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()
            
        img_name = self.image_list[idx]
        
        image = Image.open(img_name).convert('RGB')

        if self.transform:
            sample = self.transform(image)
          
        return sample



def extract_video_features(args, video_num):
    

    args.root_dir = './Datasets/QFVS/data/video_frames/P0' + str(video_num) + '-moviepy/'

    image_list = sorted(filter(os.path.isfile,
                        glob.glob(os.path.join(args.root_dir,'*.jpg'))))


    transform = transforms.Compose([
                transforms.RandomResizedCrop(224, interpolation=transforms.InterpolationMode.BICUBIC),
                transforms.ToTensor(),
                transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])


    model = getattr(module_arch, config['arch']['type'])(config['arch']['args']['video_params'], config['arch']['args']['text_params'], 
                                 projection=config['arch']['args']["projection"], load_checkpoint=config['arch']['args']["load_checkpoint"])

    data = Image_Dataset(image_list, transform=transform)
    dataloader = DataLoader(data, batch_size=5, drop_last=True, shuffle=False, num_workers=4)

    device = torch.device(args.cuda_base if torch.cuda.is_available() else 'cpu')
    model = model.to(device)
    model = nn.DataParallel(model, device_ids = args.device_ids)

    feature = []

    with torch.no_grad():
        for i, image in enumerate(tqdm(dataloader)):
            image = image.to(device)
            b, c, h, w = image.shape
            image_reshape = image.reshape(int(b/5), 5, c, h, w)
            output = model(image_reshape)
            output = output.detach().to('cpu').numpy()
            feature.append(output)
            
        feature = np.concatenate(feature, axis=0)
        frame_num=feature.shape[0]
        print("Number of Frames: ", frame_num)

        K=feature[:,0,:]
        K=np.dot(K,K.T)

        cps,_=cpd_auto(K,max_segment_num-1,1,desc_rate=1,verbose=False,lmax=max_frame_num-1) #int(K.shape[0]/25)
        seg_num=len(cps)+1

        assert seg_num<=max_segment_num

        seg_points=cps
        seg_points=np.insert(seg_points,0,0)
        seg_points=np.append(seg_points,frame_num)

        segments=[]
        for i in range(seg_num):
            segments.append(np.arange(seg_points[i],seg_points[i+1],1,dtype=np.int32))

        assert len(segments)<=max_segment_num

        for seg in segments:
            assert len(seg)<=max_frame_num

        seg_len=np.zeros((max_segment_num),dtype=np.int32)
        for index,seg in enumerate(segments):
            seg_len[index]=len(seg)

        # features

        for seg in segments:
            for frame in seg:
                assert frame<frame_num

        feature_dim=768

        features=torch.zeros((max_segment_num, max_frame_num, feature.shape[1], feature_dim))
        for seg_index,seg in enumerate(segments):
            for frame_index,frame in enumerate(seg):
                features[seg_index,frame_index]=torch.tensor(feature[frame])
        
        
        print("Feature Extraction Done!")

        return features, seg_len


QFVS/extract_features.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import os
import h5py
import numpy as np
import torch
import transformers
from torch.utils.data import DataLoader
import pdb
import json
from collections import OrderedDict
from tqdm import tqdm

from dataset_prompt import UCTDataset
from utils_QFVS import load_pickle
import math
import torch.nn as nn
import torch.nn.init as init
from pathlib import Path
import model.model_text_unfused as module_arch_text
import model.model_fused as module_arch_fused


class Extract_Multimodal_Features():
    def __init__(self,config, last_bias, train_videos, test_video, cuda_base, device_ids, features):
        self.config=config
        self.last_bias=last_bias
        self.train_videos = train_videos
        self.test_video = test_video
        self.cuda_base = cuda_base
        self.device_ids = device_ids
        self.features = features
        self._build_dataloader()
        self._bulid_model()

    def read_json(self, fname):
        with fname.open('rt') as handle:
            return json.load(handle, object_hook=OrderedDict)


    def _bulid_model(self):

        config_egoclip = self.read_json(Path('qfvs.json'))

        self.model_text_unfused = getattr(module_arch_text, config_egoclip['arch']['type'])(config_egoclip['arch']['args']['video_params'], config_egoclip['arch']['args']['text_params'],
                                projection=config_egoclip['arch']['args']["projection"], load_checkpoint=config_egoclip['arch']['args']["load_checkpoint"])

        self.model_fused = getattr(module_arch_fused, config_egoclip['arch']['type'])(config_egoclip['arch']['args']['video_params'], config_egoclip['arch']['args']['text_params'],
                                projection=config_egoclip['arch']['args']["projection"], load_checkpoint=config_egoclip['arch']['args']["load_checkpoint"])
        
        self.config_egoclip = config_egoclip

    def _build_dataset(self):
        
        return UCTDataset(self.config, self.train_videos + [self.test_video], self.features)

    def _build_dataloader(self):
        dataset=self._build_dataset()
        self.dataloader=DataLoader(dataset, batch_size=1, shuffle=False, num_workers=0)
        print(len(self.dataloader))
    

    def extract(self):

        device = torch.device(self.cuda_base if torch.cuda.is_available() else 'cpu')
        criterion=torch.nn.BCEWithLogitsLoss()
        
        self.model_fused = self.model_fused.to(device)
        self.model_text_unfused = self.model_text_unfused.to(device)
        self.model_fused = nn.DataParallel(self.model_fused, device_ids = self.device_ids)
        self.model_text_unfused = nn.DataParallel(self.model_text_unfused, device_ids = self.device_ids)
        
        self.model_fused.eval()
        self.model_text_unfused.eval()

        features_dict = {'1':[], '2':[], '3':[], '4':[]}

        for epoch in range(1):

            step = epoch*len(self.dataloader)
            batch_count=0
            
            for _idx, (features, seg_len, concept1_token, concept1_mask, concept2_token, concept2_mask, query_token, query_mask, concept1_GT, concept2_GT, mask_GT, oracle_summary, video_id) in enumerate(tqdm(self.dataloader)):
                
                
                train_num=seg_len.shape[0]
                batch_count+=1

                mask=torch.zeros(train_num,self.config["max_segment_num"],self.config["max_frame_num"],dtype=torch.bool).cuda()
                for i in range(train_num):
                    for j in range(len(seg_len[i])):
                        for k in range(seg_len[i][j]):
                            mask[i][j][k]=1

                
                fusedfeat_concept1_batched_list = []
                fusedfeat_concept2_batched_list = []
                fusedfeat_oracle_batched_list = []


                with torch.no_grad():
                    
                    concept1_embed, extend_text_mask_concept1 = self.model_text_unfused(concept1_token, concept1_mask)
                    concept2_embed, extend_text_mask_concept2 = self.model_text_unfused(concept2_token, concept2_mask)
                    query_embed, extend_text_mask_query = self.model_text_unfused(query_token, query_mask)

                    concept1_embed = concept1_embed.unsqueeze(1).repeat(1,200,1,1)
                    concept2_embed = concept2_embed.unsqueeze(1).repeat(1,200,1,1)
                    query_embed = query_embed.unsqueeze(1).repeat(1,200,1,1)
                    extend_text_mask_concept1 = extend_text_mask_concept1.unsqueeze(1).repeat(1,200,1,1,1)
                    extend_text_mask_concept2 = extend_text_mask_concept2.unsqueeze(1).repeat(1,200,1,1,1)
                    extend_text_mask_query = extend_text_mask_query.unsqueeze(1).repeat(1,200,1,1,1)

                    for j in range(concept1_embed.size(0)):

                        fusedfeat_concept1_list = []
                        fusedfeat_concept2_list = []
                        fusedfeat_oracle_list = []
                        
                        for idx in range(features.size(1)):

                            concept1_fusedfeat = self.model_fused(features[j, idx], concept1_embed[j], extend_text_mask_concept1[j])
                            concept2_fusedfeat = self.model_fused(features[j, idx], concept2_embed[j], extend_text_mask_concept2[j])
                            oracle_fusedfeat = self.model_fused(features[j, idx], query_embed[j], extend_text_mask_query[j])

                            fusedfeat_concept1_list.append(concept1_fusedfeat)
                            fusedfeat_concept2_list.append(concept2_fusedfeat)
                            fusedfeat_oracle_list.append(oracle_fusedfeat)

                        fusedfeat_concept1_list = torch.stack(fusedfeat_concept1_list, 0)
                        fusedfeat_concept1_batched_list.append(fusedfeat_concept1_list)
                        fusedfeat_concept2_list = torch.stack(fusedfeat_concept2_list, 0)
                        fusedfeat_concept2_batched_list.append(fusedfeat_concept2_list)
                        fusedfeat_oracle_list = torch.stack(fusedfeat_oracle_list, 0)
                        fusedfeat_oracle_batched_list.append(fusedfeat_oracle_list)

                    fusedfeat_concept1_batched_list = torch.stack(fusedfeat_concept1_batched_list, 0)
                    fusedfeat_concept2_batched_list = torch.stack(fusedfeat_concept2_batched_list, 0)
                    fusedfeat_oracle_batched_list = torch.stack(fusedfeat_oracle_batched_list, 0)

                features_dict[str(video_id[0])].append({'seg_len':seg_len.squeeze(0), 'fusedfeat_concept1_batched_list':fusedfeat_concept1_batched_list.squeeze(0), 'fusedfeat_concept2_batched_list':fusedfeat_concept2_batched_list.squeeze(0), 'fusedfeat_oracle_batched_list':fusedfeat_oracle_batched_list.squeeze(0), 'concept1_GT':concept1_GT.squeeze(0), 'concept2_GT':concept2_GT.squeeze(0), 'mask_GT':mask_GT.squeeze(0), 'oracle_summary':oracle_summary.squeeze(0)})
                

        path1 = self.config_egoclip['arch']['args']["load_checkpoint"].split('.')[0]

        torch.save(features_dict, './multimodal_features/' + str(path1) + '_features.pt')
    

QFVS/extract_multimodal_features.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

from .logger import *
from .visualization import *

QFVS/logger/__init__.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import logging
import logging.config
from pathlib import Path

from utils import read_json


def setup_logging(save_dir, log_config='logger/logger_config.json', default_level=logging.INFO):
    """
    Setup logging configuration
    """
    log_config = Path(log_config)
    if log_config.is_file():
        config = read_json(log_config)
        # modify logging paths based on run config
        for _, handler in config['handlers'].items():
            if 'filename' in handler:
                handler['filename'] = str(save_dir / handler['filename'])

        logging.config.dictConfig(config)
    else:
        print("Warning: logging configuration file is not found in {}.".format(log_config))
        logging.basicConfig(level=default_level)

QFVS/logger/logger.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import importlib

from utils import Timer


class TensorboardWriter:
    def __init__(self, log_dir, logger, enabled):
        self.writer = None
        self.selected_module = ""

        if enabled:
            log_dir = str(log_dir)

            # Retrieve visualization writer.
            for module in ["torch.utils.tensorboard", "tensorboardX"]:
                try:
                    self.writer = importlib.import_module(module).SummaryWriter(log_dir)
                    succeeded = True
                    break
                except ImportError:
                    succeeded = False
                self.selected_module = module

            if not succeeded:
                message = "Warning: visualization (Tensorboard) is configured to use, but currently not installed on " \
                    "this machine. Please install either TensorboardX with 'pip install tensorboardx', upgrade " \
                    "PyTorch to version >= 1.1 for using 'torch.utils.tensorboard' or turn off the option in " \
                    "the 'config.json' file."
                logger.warning(message)

        self.step = 0
        self.mode = ''

        self.tb_writer_ftns = {
            'add_scalar', 'add_scalars', 'add_image', 'add_images', 'add_audio',
            'add_text', 'add_histogram', 'add_pr_curve', 'add_embedding'
        }
        self.tag_mode_exceptions = {'add_histogram', 'add_embedding'}
            
        self.timer = Timer()

    def set_step(self, step, mode='train'):
        self.mode = mode
        self.step = step
        if step == 0:
            self.timer.reset()
        else:
            duration = self.timer.check()
            self.add_scalar('steps_per_sec', 1 / duration)

    def __getattr__(self, name):
        """
        If visualization is configured to use:
            return add_data() methods of tensorboard with additional information (step, tag) added.
        Otherwise:
            return a blank function handle that does nothing
        """
        if name in self.tb_writer_ftns:
            add_data = getattr(self.writer, name, None)

            def wrapper(tag, data, *args, **kwargs):
                if add_data is not None:
                    # add mode(train/valid) tag
                    if name not in self.tag_mode_exceptions:
                        tag = '{}/{}'.format(tag, self.mode)
                    add_data(tag, data, self.step, *args, **kwargs)
            return wrapper
        else:
            # default action for returning methods defined in this class, set_step() for instance.
            try:
                attr = object.__getattr__(name)
            except AttributeError:
                raise AttributeError("type object '{}' has no attribute '{}'".format(self.selected_module, name))
            return attr


class SacredNeptuneWriter:
    def __init__(self):
        raise NotImplementedError

QFVS/logger/visualization.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import yaml
import pdb
import json
import argparse 
import torch
import numpy as np
from runner_train import Runner_summary_transformer
from extract_features import extract_video_features
from extract_multimodal_features import Extract_Multimodal_Features
from pathlib import Path
from collections import OrderedDict

parser = argparse.ArgumentParser(description='Process some integers.')
parser.add_argument('--train_videos', help='delimited list input', type=lambda s: [int(item) for item in s.split(',')])
parser.add_argument('--test_video', help='test video id', type=int)
parser.add_argument('--cuda_base', help="in form cuda:x")
parser.add_argument('--device_ids', help='delimited list input', type=lambda s: [int(item) for item in s.split(',')])

def read_json(fname):
    with fname.open('rt') as handle:
        return json.load(handle, object_hook=OrderedDict)

config_qfvs = read_json(Path('qfvs.json'))

args = parser.parse_args()
with open('./QFVS.yml') as f:
    config = yaml.load(f, Loader=yaml.FullLoader)

path1 = config_qfvs['arch']['args']["load_checkpoint"].split('.')[0]

try:
    feature_dict = torch.load('./multimodal_features/' + str(path1) + '_features.pt')
except:
    features={}
    all_videos = args.train_videos + [args.test_video]
    for video_id in all_videos:
        feature, seg_len = extract_video_features(args, video_id)
        features.update({str(video_id) : {'feature': feature, 'seg_len': seg_len}})
    extract_multimodal_features = Extract_Multimodal_Features(config=config, last_bias=-1.91, train_videos=args.train_videos, test_video=args.test_video, cuda_base=args.cuda_base, device_ids=args.device_ids, features=features)
    extract_multimodal_features.extract()
    feature_dict = torch.load('./multimodal_features/' + str(path1) + '_features.pt')
    print("In Except")
    print("=================")

runner = Runner_summary_transformer(config=config, last_bias=-1.91, train_videos=args.train_videos, test_video=args.test_video, cuda_base=args.cuda_base, device_ids=args.device_ids, feature_dict=feature_dict)

runner.train()
runner.output()


QFVS/main.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import torch
import torch.nn as nn
import torch.nn.functional as F

from transformers.models.bert.modeling_bert import BertPredictionHeadTransform


class Pooler(nn.Module):
    def __init__(self, hidden_size):
        super().__init__()
        self.dense = nn.Linear(hidden_size, hidden_size)
        self.activation = nn.Tanh()

    def forward(self, hidden_states):
        first_token_tensor = hidden_states #[:, 0]
        pooled_output = self.dense(first_token_tensor)
        pooled_output = self.activation(pooled_output)
        return pooled_output


class ITMHead(nn.Module):
    def __init__(self, hidden_size):
        super().__init__()
        self.fc = nn.Linear(hidden_size, 2)

    def forward(self, x):
        x = self.fc(x)
        return x


class MLMHead(nn.Module):
    def __init__(self, config, weight=None):
        super().__init__()
        self.transform = BertPredictionHeadTransform(config)
        self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.bias = nn.Parameter(torch.zeros(config.vocab_size))
        if weight is not None:
            self.decoder.weight = weight

    def forward(self, x):
        x = self.transform(x)
        x = self.decoder(x) + self.bias
        return x

QFVS/model/heads.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import os
import pdb

import timm
import torch
import yaml
import torch.nn as nn
import numpy as np
import torch.nn.functional as F
from transformers import AutoModel
from einops import rearrange, repeat

from base import BaseModel
from model import video_transformer
from model.video_transformer import SpaceTimeTransformer
from utils.util import state_dict_data_parallel_fix

from model import roberta
from model.roberta import RobertaModel, _prepare_decoder_attention_mask
from model import heads
from transformers import RobertaConfig
from functools import partial
import copy
import torch.distributed as dist

with open('./QFVS.yml') as f:
    config = yaml.load(f, Loader=yaml.FullLoader)

def init_weights(module):
    if isinstance(module, (nn.Linear, nn.Embedding)):
        module.weight.data.normal_(mean=0.0, std=0.02)
    elif isinstance(module, nn.LayerNorm):
        module.bias.data.zero_()
        module.weight.data.fill_(1.0)

    if isinstance(module, nn.Linear) and module.bias is not None:
        module.bias.data.zero_()


class FrozenInTime(BaseModel):
    def __init__(self,
                 video_params,
                 text_params,
                 projection_dim=256,
                 load_checkpoint=None,
                 projection='minimal',
                 load_temporal_fix='bilinear',
                 config = config,
                 task_names = 'EgoNCE_ITM_MLM',
                 norm_layer = None,
                 embed_dim=768):
        super().__init__()

        self.video_params = video_params
        self.text_params = text_params
        self.load_temporal_fix = load_temporal_fix
        self.config = config
        self.task_names = task_names
        if not text_params['pretrained']:
            raise NotImplementedError("Huggingface text models require pretrained init.")

        # pdb.set_trace()
        #if self.text_params['model'].startswith('distilbert'):
        #    self.text_model = AutoModel.from_pretrained('distilbert-base-uncased')
        #else:
        #    self.text_model = AutoModel.from_pretrained(text_params['model'])
        if self.text_params['model'].startswith('roberta'):
            self.text_model = RobertaModel.from_pretrained("roberta-base")
        self.text_model.train()

        pretrained = video_params['pretrained']
        if video_params['model'] == "SpaceTimeTransformer":
            self.num_frames = 5 #self.config["data_loader"][0]["args"]["video_params"]["num_frames"]
            time_init = 'zeros'
            attention_style = 'frozen-in-time'
            arch_config = 'base_patch16_224'
            vit_init = 'imagenet-21k'
            if arch_config == 'base_patch16_224':
                #vit_model = timm.models.vision_transformer.vit_base_patch16_224(pretrained=pretrained)
                vit_model = torch.load("/cis/home/shraman/works_meta_2022/pre-training/EgoVLP_Fused_HardNegITM_Checkpoint_multinode/frozen-in-time-main/pretrained/jx_vit_base_p16_224-80ecf9dd.pth", map_location="cpu")
                model = SpaceTimeTransformer(num_frames=self.num_frames,
                                            time_init=time_init,
                                            attention_style=attention_style)
            else:
                raise NotImplementedError

            model.head = nn.Identity()
            model.pre_logits = nn.Identity()
            ftr_dim = model.embed_dim
           
            if load_checkpoint in ["", None]:
                vit_checkpoint = vit_model
                new_vit_dict = state_dict_data_parallel_fix(vit_checkpoint, model.state_dict())
                model.load_state_dict(new_vit_dict, strict=False)
            self.video_model = model
        else:
            raise NotImplementedError(f"{video_params['model']} not implemented")

        # for backwards compatibility (old models)
        self.video_model.fc = nn.Identity()

        if ('MLM' in self.task_names or 'ITM' in self.task_names):
            # for FIBER-like cross-attention

            bert_config = RobertaConfig(
                vocab_size=self.config["vocab_size"],
                hidden_size=self.config["hidden_size"],
                num_hidden_layers=self.config["num_layers"],
                num_attention_heads=self.config["num_heads"],
                intermediate_size=self.config["hidden_size"] * config["mlp_ratio"],
                #max_position_embeddings=maxlen, [was used in BTGOT script]
                hidden_dropout_prob=self.config["drop_rate"],
                attention_probs_dropout_prob=self.config["drop_rate"],
            )

            self.num_fuse_block=self.config["num_fuse_block"]
            self.num_text_layer=self.config["num_layers"]
            roberta.NUM_FUSE_BLOCK = self.video_model.NUM_FUSE_BLOCK=self.num_fuse_block
            roberta.DIM_IMG=self.config["input_image_embed_size"]
            self.video_model.DIM_TXT=self.config["input_text_embed_size"]

            self.cross_modal_text_transform = nn.Linear(self.config["input_text_embed_size"], self.config["hidden_size"])
            self.cross_modal_text_transform.apply(init_weights)
            self.cross_modal_video_transform = nn.Linear(self.config["input_image_embed_size"], self.config["hidden_size"])
            self.cross_modal_video_transform.apply(init_weights)

            self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))

            self.num_patches = self.video_model.patch_embed.num_patches
            self.patches_per_frame = self.num_patches//self.num_frames
            norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
            self.norm = norm_layer(embed_dim)
            self.pre_logits = nn.Identity()


            self.avgpool = nn.AdaptiveAvgPool1d(1)
            self.cross_modal_video_pooler = heads.Pooler(config["hidden_size"])
            self.cross_modal_video_pooler.apply(init_weights)
            self.cross_modal_text_pooler = heads.Pooler(config["hidden_size"])
            self.cross_modal_text_pooler.apply(init_weights)

            ## einops transformations
            self.einops_from_space = 'b (f n) d'
            self.einops_to_space = '(b f) n d'
            self.einops_from_time = 'b (f n) d'
            self.einops_to_time = '(b n) f d'

        if 'MLM' in self.task_names:
            self.mlm_score = heads.MLMHead(bert_config)
            self.mlm_score.apply(init_weights)

        if 'ITM' in self.task_names:
            self.itm_score = heads.ITMHead(config["hidden_size"] * 2)
            self.itm_score.apply(init_weights)

        if load_checkpoint not in ["", None]:
            checkpoint = torch.load(load_checkpoint, map_location='cpu')
            state_dict = checkpoint['state_dict']
            new_state_dict = state_dict_data_parallel_fix(state_dict, self.state_dict())
            new_state_dict = self._inflate_positional_embeds(new_state_dict)
            self.load_state_dict(new_state_dict, strict=False)

    def set_device(self, device):
        self.device = device

    def forward(self, video_data_itm, text_embeds, extend_text_masks):

        n = self.patches_per_frame
        f = 5

        unfused_blocks = self.num_text_layer - self.num_fuse_block

        for blk_i, blk in enumerate(self.video_model.blocks[unfused_blocks:self.num_text_layer]):
            if config['use_checkpoint']:
                fuse_video_data = torch.utils.checkpoint.checkpoint(blk, video_data_itm, self.einops_from_space, self.einops_to_space, self.einops_from_time, self.einops_to_time, 
                                     n, f, text_embeds, extend_text_masks)
                text_embeds = torch.utils.checkpoint.checkpoint(self.text_model.encoder.layer[blk_i + unfused_blocks],
                                     text_embeds, extend_text_masks, None, (video_data_itm), None, None, False, True)[0]
            else:
                fuse_video_data = blk(video_data_itm, self.einops_from_space, self.einops_to_space, self.einops_from_time, self.einops_to_time, 
                                     y=text_embeds, y_mask=extend_text_masks, time_n=n, space_f=f)
                text_embeds = self.text_model.encoder.layer[blk_i + unfused_blocks](text_embeds, extend_text_masks, encoder_hidden_states=(video_data_itm), last_norm=True)[0]
            video_data_itm = fuse_video_data

            
        #print("Shape of model output", video_data.shape)
        video_data_itm = self.norm(video_data_itm)[:, 0]
        video_data_itm = self.pre_logits(video_data_itm)

        video_embeds = self.cross_modal_video_transform(video_data_itm)
        
        return video_embeds


    def compute_text(self, text_data):
        if self.text_params['model'].startswith('bert'):
            text_embeddings = self.text_model(text_data['input_ids'], attention_mask=text_data['attention_mask'])[
                'pooler_output']
        elif self.text_params['model'].startswith('distilbert'):
            text_embeddings = self.text_model(**text_data).last_hidden_state[:, 0, :]
        elif self.text_params['model'].startswith('roberta'):
            text_embeddings = self.text_model(**text_data).last_hidden_state[:, 0, :]
        else:
            raise NotImplementedError
        text_embeddings = self.txt_proj(text_embeddings)
        return text_embeddings

    def compute_text_tokens(self, text_data):
        if self.text_params['model'].startswith('bert'):
            text_embeddings = self.text_model(text_data['input_ids'], attention_mask=text_data['attention_mask'])[
                'pooler_output']    # not implement for bert
        elif self.text_params['model'].startswith('distilbert'):
            text_embeddings = self.text_model(**text_data).last_hidden_state
        elif self.text_params['model'].startswith('roberta'):
            text_embeddings = self.text_model(**text_data).last_hidden_state
        else:
            raise NotImplementedError

        text_embeddings = self.txt_proj(text_embeddings)
        return text_embeddings

    def compute_video(self, video_data):
        video_embeddings = self.video_model(video_data)
        video_embeddings = self.vid_proj(video_embeddings)
        return video_embeddings

    def _inflate_positional_embeds(self, new_state_dict):
        # allow loading of timesformer with fewer num_frames
        curr_keys = list(self.state_dict().keys())
        if 'video_model.temporal_embed' in new_state_dict and 'video_model.temporal_embed' in curr_keys:
            load_temporal_embed = new_state_dict['video_model.temporal_embed']
            load_num_frames = load_temporal_embed.shape[1]
            curr_num_frames = self.video_params['num_frames']
            embed_dim = load_temporal_embed.shape[2]

            if load_num_frames != curr_num_frames:
                if load_num_frames > curr_num_frames:
                    print(f'### loaded {self.video_params["model"]} model has MORE frames than current...'
                          f'### loading weights, filling in the extras via {self.load_temporal_fix}')
                    new_temporal_embed = load_temporal_embed[:, :curr_num_frames, :]
                else:
                    print(f'### loaded {self.video_params["model"]} model has FEWER frames than current...'
                          f'### loading weights, filling in the extras via {self.load_temporal_fix}')
                    if self.load_temporal_fix == 'zeros':
                        new_temporal_embed = torch.zeros([load_temporal_embed.shape[0], curr_num_frames, embed_dim])
                        new_temporal_embed[:, :load_num_frames] = load_temporal_embed
                    elif self.load_temporal_fix in ['interp', 'bilinear']:
                        # interpolate
                        # unsqueeze so pytorch thinks its an image
                        mode = 'nearest'
                        if self.load_temporal_fix == 'bilinear':
                            mode = 'bilinear'
                        load_temporal_embed = load_temporal_embed.unsqueeze(0)
                        new_temporal_embed = F.interpolate(load_temporal_embed,
                                                           (curr_num_frames, embed_dim), mode=mode, align_corners=True).squeeze(0)
                    else:
                        raise NotImplementedError
                new_state_dict['video_model.temporal_embed'] = new_temporal_embed
        # allow loading with smaller spatial patches. assumes custom border crop, to append the
        # border patches to the input sequence
        if 'video_model.pos_embed' in new_state_dict and 'video_model.pos_embed' in curr_keys:
            load_pos_embed = new_state_dict['video_model.pos_embed']
            load_num_patches = load_pos_embed.shape[1]
            curr_pos_embed = self.state_dict()['video_model.pos_embed']
            if load_num_patches != curr_pos_embed.shape[1]:
                raise NotImplementedError(
                    'Loading models with different spatial resolution / patch number not yet implemented, sorry.')

        return new_state_dict

def sim_matrix(a, b, eps=1e-8):
    """
    added eps for numerical stability
    """
    a_n, b_n = a.norm(dim=1)[:, None], b.norm(dim=1)[:, None]
    a_norm = a / torch.max(a_n, eps * torch.ones_like(a_n))
    b_norm = b / torch.max(b_n, eps * torch.ones_like(b_n))
    sim_mt = torch.mm(a_norm, b_norm.transpose(0, 1))
    return sim_mt


def sim_matrix_batch_val(a, b, eps=1e-8):
    """
    added eps for numerical stability
    """
    a_n, b_n = a.norm(dim=-1).unsqueeze(-1), b.norm(dim=-1).unsqueeze(-1)
    a_norm = a / torch.max(a_n, eps * torch.ones_like(a_n))
    b_norm = b / torch.max(b_n, eps * torch.ones_like(b_n))
    sim_mt = torch.bmm(a_norm, b_norm.transpose(1, 2))
    return sim_mt


if __name__ == "__main__":
    pass

QFVS/model/model_fused.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import torch.nn.functional as F
import torch.nn as nn
import numpy as np
import random
import torch
import math
import re
import pdb


class PositionalEmbedding(nn.Module):
    def __init__(self, d_model, max_seq_len = 200):
        super().__init__()
        self.d_model = d_model
        pe = torch.zeros(max_seq_len, d_model)
        pe.requires_grad = False
        for pos in range(max_seq_len):
            for i in range(0, d_model, 2):
                pe[pos, i] = math.sin(pos / (10000 ** ((2 * i)/d_model)))
                pe[pos, i + 1] = math.cos(pos / (10000 ** ((2 * (i + 1))/d_model)))
        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)

    def forward(self, x):
        return self.pe[:,:x.size(1)] #x.size(1) = seq_len



class Transformer_Encoder(nn.Module):
    def __init__(self, model_dim, nhead, num_layers):
        super(Transformer_Encoder, self).__init__()

        self.model_dim = model_dim
        self.pos = PositionalEmbedding(model_dim)

        self.encoder_layer = nn.TransformerEncoderLayer(d_model=model_dim, nhead=nhead)
        self.transformer_encoder = nn.TransformerEncoder(self.encoder_layer, num_layers=num_layers)

        self.dropout = nn.Dropout(p=0.2)
        self.relu = nn.ReLU()
        self.projector_1 = nn.Linear(model_dim, 8)
        self.summ_head = nn.Linear(8, 1)

    def create_mask(self, seg_len):
        
        masked_seg = torch.empty(seg_len.shape[0], seg_len.shape[1], 200)
        for batch in range(seg_len.shape[0]):
            for idx in range(len(seg_len[batch])):
                zeros = torch.zeros(seg_len[batch][idx].item())
                ones = torch.ones(200-seg_len[batch][idx].item())
                mask = torch.cat((zeros, ones))
                masked_seg[batch][idx] = mask
        return masked_seg #(batch_size, 20, 200)


    def forward(self, features, seg_len):

        # features -> (batch_size, 20, 200, 768)

        batch_size = features.shape[0]
        features = features.reshape(features.shape[0]*features.shape[1], features.shape[2], features.shape[3])

        features = features + self.pos(features) #(batch_size*20, 200, model_dim)

        masked_seg = self.create_mask(seg_len).reshape(self.create_mask(seg_len).shape[0]*self.create_mask(seg_len).shape[1], self.create_mask(seg_len).shape[2]).to(features.device)  #(batch_size*20, 200)

        features = features.permute(1,0,2)
        batch_out = self.transformer_encoder(features, src_key_padding_mask = masked_seg)
        batch_out = batch_out.permute(1,0,2) #(batch_size*20, 200, model_dim)

        batch_out = batch_out.reshape(batch_size, features.size(1), batch_out.shape[1], batch_out.shape[2])
        batch_out = self.summ_head(self.dropout(self.relu(self.projector_1(batch_out)))).squeeze(-1)

        return batch_out

QFVS/model/model_summary.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import os
import pdb

import timm
import torch
import yaml
import torch.nn as nn
import numpy as np
import torch.nn.functional as F
from transformers import AutoModel
from einops import rearrange, repeat

from base import BaseModel
from model import video_transformer
from model.video_transformer import SpaceTimeTransformer
from utils.util import state_dict_data_parallel_fix

from model import roberta
from model.roberta import RobertaModel, _prepare_decoder_attention_mask
from model import heads
from transformers import RobertaConfig
from functools import partial
import copy
import torch.distributed as dist

with open('./QFVS.yml') as f:
    config = yaml.load(f, Loader=yaml.FullLoader)

def init_weights(module):
    if isinstance(module, (nn.Linear, nn.Embedding)):
        module.weight.data.normal_(mean=0.0, std=0.02)
    elif isinstance(module, nn.LayerNorm):
        module.bias.data.zero_()
        module.weight.data.fill_(1.0)

    if isinstance(module, nn.Linear) and module.bias is not None:
        module.bias.data.zero_()


class FrozenInTime(BaseModel):
    def __init__(self,
                 video_params,
                 text_params,
                 projection_dim=256,
                 load_checkpoint=None,
                 projection='minimal',
                 load_temporal_fix='bilinear',
                 config = config,
                 task_names = 'EgoNCE_ITM_MLM',
                 norm_layer = None,
                 embed_dim=768):
        super().__init__()

        self.video_params = video_params
        self.text_params = text_params
        self.load_temporal_fix = load_temporal_fix
        self.config = config
        self.task_names = task_names
        if not text_params['pretrained']:
            raise NotImplementedError("Huggingface text models require pretrained init.")

        if self.text_params['model'].startswith('roberta'):
            self.text_model = RobertaModel.from_pretrained("roberta-base")
        self.text_model.train()

        pretrained = video_params['pretrained']
        if video_params['model'] == "SpaceTimeTransformer":
            self.num_frames = 5 #self.config["data_loader"][0]["args"]["video_params"]["num_frames"]
            time_init = 'zeros'
            attention_style = 'frozen-in-time'
            arch_config = 'base_patch16_224'
            vit_init = 'imagenet-21k'
            if arch_config == 'base_patch16_224':
                #vit_model = timm.models.vision_transformer.vit_base_patch16_224(pretrained=pretrained)
                vit_model = torch.load("/cis/home/shraman/works_meta_2022/pre-training/EgoVLP_Fused_HardNegITM_Checkpoint_multinode/frozen-in-time-main/pretrained/jx_vit_base_p16_224-80ecf9dd.pth", map_location="cpu")
                model = SpaceTimeTransformer(num_frames=self.num_frames,
                                            time_init=time_init,
                                            attention_style=attention_style)
            else:
                raise NotImplementedError

            model.head = nn.Identity()
            model.pre_logits = nn.Identity()
            ftr_dim = model.embed_dim
           
            if load_checkpoint in ["", None]:
                vit_checkpoint = vit_model
                new_vit_dict = state_dict_data_parallel_fix(vit_checkpoint, model.state_dict())
                model.load_state_dict(new_vit_dict, strict=False)
            self.video_model = model
        else:
            raise NotImplementedError(f"{video_params['model']} not implemented")

        # for backwards compatibility (old models)
        self.video_model.fc = nn.Identity()


        if ('MLM' in self.task_names or 'ITM' in self.task_names):

            bert_config = RobertaConfig(
                vocab_size=self.config["vocab_size"],
                hidden_size=self.config["hidden_size"],
                num_hidden_layers=self.config["num_layers"],
                num_attention_heads=self.config["num_heads"],
                intermediate_size=self.config["hidden_size"] * config["mlp_ratio"],
                #max_position_embeddings=maxlen, [was used in BTGOT script]
                hidden_dropout_prob=self.config["drop_rate"],
                attention_probs_dropout_prob=self.config["drop_rate"],
            )

            self.num_fuse_block=self.config["num_fuse_block"]
            self.num_text_layer=self.config["num_layers"]
            roberta.NUM_FUSE_BLOCK = self.video_model.NUM_FUSE_BLOCK=self.num_fuse_block
            roberta.DIM_IMG=self.config["input_image_embed_size"]
            self.video_model.DIM_TXT=self.config["input_text_embed_size"]

            self.cross_modal_text_transform = nn.Linear(self.config["input_text_embed_size"], self.config["hidden_size"])
            self.cross_modal_text_transform.apply(init_weights)
            self.cross_modal_video_transform = nn.Linear(self.config["input_image_embed_size"], self.config["hidden_size"])
            self.cross_modal_video_transform.apply(init_weights)

            self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))

            self.num_patches = self.video_model.patch_embed.num_patches
            self.patches_per_frame = self.num_patches//self.num_frames
            norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
            self.norm = norm_layer(embed_dim)
            self.pre_logits = nn.Identity()


            self.avgpool = nn.AdaptiveAvgPool1d(1)
            self.cross_modal_video_pooler = heads.Pooler(config["hidden_size"])
            self.cross_modal_video_pooler.apply(init_weights)
            self.cross_modal_text_pooler = heads.Pooler(config["hidden_size"])
            self.cross_modal_text_pooler.apply(init_weights)

            ## einops transformations
            self.einops_from_space = 'b (f n) d'
            self.einops_to_space = '(b f) n d'
            self.einops_from_time = 'b (f n) d'
            self.einops_to_time = '(b n) f d'

        if 'MLM' in self.task_names:
            self.mlm_score = heads.MLMHead(bert_config)
            self.mlm_score.apply(init_weights)

        if 'ITM' in self.task_names:
            self.itm_score = heads.ITMHead(config["hidden_size"] * 2)
            self.itm_score.apply(init_weights)

        if load_checkpoint not in ["", None]:
            checkpoint = torch.load(load_checkpoint, map_location='cpu')
            state_dict = checkpoint['state_dict']
            new_state_dict = state_dict_data_parallel_fix(state_dict, self.state_dict())
            new_state_dict = self._inflate_positional_embeds(new_state_dict)
            self.load_state_dict(new_state_dict, strict=False)

    def set_device(self, device):
        self.device = device

    def forward(self, text_data_token, text_masks):

        unfused_blocks = self.num_text_layer - self.num_fuse_block
            
        text_embeds = self.text_model.embeddings(input_ids=text_data_token) # before it was input_ids=text_ids
        device = text_embeds.device
        #print(text_data.keys())
        input_shape = text_masks.size()
        extend_text_masks = self.text_model.get_extended_attention_mask(text_masks, input_shape, device)


        for layer_i, layer in enumerate(self.text_model.encoder.layer[:unfused_blocks]):

            if config['use_checkpoint']:
                text_embeds = torch.utils.checkpoint.checkpoint(layer, text_embeds, extend_text_masks)[0]
            else:
                text_embeds = layer(text_embeds, extend_text_masks)[0]

        return text_embeds, extend_text_masks


    def _inflate_positional_embeds(self, new_state_dict):
        # allow loading of timesformer with fewer num_frames
        curr_keys = list(self.state_dict().keys())
        if 'video_model.temporal_embed' in new_state_dict and 'video_model.temporal_embed' in curr_keys:
            load_temporal_embed = new_state_dict['video_model.temporal_embed']
            load_num_frames = load_temporal_embed.shape[1]
            curr_num_frames = self.video_params['num_frames']
            embed_dim = load_temporal_embed.shape[2]

            if load_num_frames != curr_num_frames:
                if load_num_frames > curr_num_frames:
                    print(f'### loaded {self.video_params["model"]} model has MORE frames than current...'
                          f'### loading weights, filling in the extras via {self.load_temporal_fix}')
                    new_temporal_embed = load_temporal_embed[:, :curr_num_frames, :]
                else:
                    print(f'### loaded {self.video_params["model"]} model has FEWER frames than current...'
                          f'### loading weights, filling in the extras via {self.load_temporal_fix}')
                    if self.load_temporal_fix == 'zeros':
                        new_temporal_embed = torch.zeros([load_temporal_embed.shape[0], curr_num_frames, embed_dim])
                        new_temporal_embed[:, :load_num_frames] = load_temporal_embed
                    elif self.load_temporal_fix in ['interp', 'bilinear']:
                        # interpolate
                        # unsqueeze so pytorch thinks its an image
                        mode = 'nearest'
                        if self.load_temporal_fix == 'bilinear':
                            mode = 'bilinear'
                        load_temporal_embed = load_temporal_embed.unsqueeze(0)
                        new_temporal_embed = F.interpolate(load_temporal_embed,
                                                           (curr_num_frames, embed_dim), mode=mode, align_corners=True).squeeze(0)
                    else:
                        raise NotImplementedError
                new_state_dict['video_model.temporal_embed'] = new_temporal_embed
        # allow loading with smaller spatial patches. assumes custom border crop, to append the
        # border patches to the input sequence
        if 'video_model.pos_embed' in new_state_dict and 'video_model.pos_embed' in curr_keys:
            load_pos_embed = new_state_dict['video_model.pos_embed']
            load_num_patches = load_pos_embed.shape[1]
            curr_pos_embed = self.state_dict()['video_model.pos_embed']
            if load_num_patches != curr_pos_embed.shape[1]:
                raise NotImplementedError(
                    'Loading models with different spatial resolution / patch number not yet implemented, sorry.')

        return new_state_dict

def sim_matrix(a, b, eps=1e-8):
    """
    added eps for numerical stability
    """
    a_n, b_n = a.norm(dim=1)[:, None], b.norm(dim=1)[:, None]
    a_norm = a / torch.max(a_n, eps * torch.ones_like(a_n))
    b_norm = b / torch.max(b_n, eps * torch.ones_like(b_n))
    sim_mt = torch.mm(a_norm, b_norm.transpose(0, 1))
    return sim_mt


def sim_matrix_batch_val(a, b, eps=1e-8):
    """
    added eps for numerical stability
    """
    a_n, b_n = a.norm(dim=-1).unsqueeze(-1), b.norm(dim=-1).unsqueeze(-1)
    a_norm = a / torch.max(a_n, eps * torch.ones_like(a_n))
    b_norm = b / torch.max(b_n, eps * torch.ones_like(b_n))
    sim_mt = torch.bmm(a_norm, b_norm.transpose(1, 2))
    return sim_mt


if __name__ == "__main__":
    pass

QFVS/model/model_text_unfused.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import os
import pdb

import timm
import torch
import yaml
import torch.nn as nn
import numpy as np
import torch.nn.functional as F
from transformers import AutoModel
from einops import rearrange, repeat

from base import BaseModel
from model import video_transformer
from model.video_transformer import SpaceTimeTransformer
from utils.util import state_dict_data_parallel_fix

from model import roberta
from model.roberta import RobertaModel, _prepare_decoder_attention_mask
from model import heads
from transformers import RobertaConfig
from functools import partial
import copy
import torch.distributed as dist

with open('./QFVS.yml') as f:
    config = yaml.load(f, Loader=yaml.FullLoader)

def init_weights(module):
    if isinstance(module, (nn.Linear, nn.Embedding)):
        module.weight.data.normal_(mean=0.0, std=0.02)
    elif isinstance(module, nn.LayerNorm):
        module.bias.data.zero_()
        module.weight.data.fill_(1.0)

    if isinstance(module, nn.Linear) and module.bias is not None:
        module.bias.data.zero_()


class FrozenInTime(BaseModel):
    def __init__(self,
                 video_params,
                 text_params,
                 projection_dim=256,
                 load_checkpoint=None,
                 projection='minimal',
                 load_temporal_fix='bilinear',
                 config = config,
                 task_names = 'EgoNCE_ITM_MLM',
                 norm_layer = None,
                 embed_dim=768):
        super().__init__()

        self.video_params = video_params
        self.text_params = text_params
        self.load_temporal_fix = load_temporal_fix
        self.config = config
        self.task_names = task_names
        if not text_params['pretrained']:
            raise NotImplementedError("Huggingface text models require pretrained init.")

        if self.text_params['model'].startswith('roberta'):
            self.text_model = RobertaModel.from_pretrained("roberta-base")
        self.text_model.train()

        pretrained = video_params['pretrained']
        if video_params['model'] == "SpaceTimeTransformer":
            self.num_frames = 5 #self.config["data_loader"][0]["args"]["video_params"]["num_frames"]
            time_init = 'zeros'
            attention_style = 'frozen-in-time'
            arch_config = 'base_patch16_224'
            vit_init = 'imagenet-21k'
            if arch_config == 'base_patch16_224':
                #vit_model = timm.models.vision_transformer.vit_base_patch16_224(pretrained=pretrained)
                vit_model = torch.load("/cis/home/shraman/works_meta_2022/pre-training/EgoVLP_Fused_HardNegITM_Checkpoint_multinode/frozen-in-time-main/pretrained/jx_vit_base_p16_224-80ecf9dd.pth", map_location="cpu")
                model = SpaceTimeTransformer(num_frames=self.num_frames,
                                            time_init=time_init,
                                            attention_style=attention_style)
            else:
                raise NotImplementedError

            model.head = nn.Identity()
            model.pre_logits = nn.Identity()
            ftr_dim = model.embed_dim
           
            if load_checkpoint in ["", None]:
                vit_checkpoint = vit_model
                new_vit_dict = state_dict_data_parallel_fix(vit_checkpoint, model.state_dict())
                model.load_state_dict(new_vit_dict, strict=False)
            self.video_model = model
        else:
            raise NotImplementedError(f"{video_params['model']} not implemented")

        # for backwards compatibility (old models)
        self.video_model.fc = nn.Identity()


        if ('MLM' in self.task_names or 'ITM' in self.task_names):
            # for FIBER-like cross-attention

            bert_config = RobertaConfig(
                vocab_size=self.config["vocab_size"],
                hidden_size=self.config["hidden_size"],
                num_hidden_layers=self.config["num_layers"],
                num_attention_heads=self.config["num_heads"],
                intermediate_size=self.config["hidden_size"] * config["mlp_ratio"],
                #max_position_embeddings=maxlen, [was used in BTGOT script]
                hidden_dropout_prob=self.config["drop_rate"],
                attention_probs_dropout_prob=self.config["drop_rate"],
            )

            self.num_fuse_block=self.config["num_fuse_block"]
            self.num_text_layer=self.config["num_layers"]
            roberta.NUM_FUSE_BLOCK = self.video_model.NUM_FUSE_BLOCK=self.num_fuse_block
            roberta.DIM_IMG=self.config["input_image_embed_size"]
            self.video_model.DIM_TXT=self.config["input_text_embed_size"]

            self.cross_modal_text_transform = nn.Linear(self.config["input_text_embed_size"], self.config["hidden_size"])
            self.cross_modal_text_transform.apply(init_weights)
            self.cross_modal_video_transform = nn.Linear(self.config["input_image_embed_size"], self.config["hidden_size"])
            self.cross_modal_video_transform.apply(init_weights)

            self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))

            self.num_patches = self.video_model.patch_embed.num_patches
            self.patches_per_frame = self.num_patches//self.num_frames
            norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
            self.norm = norm_layer(embed_dim)
            self.pre_logits = nn.Identity()


            self.avgpool = nn.AdaptiveAvgPool1d(1)
            self.cross_modal_video_pooler = heads.Pooler(config["hidden_size"])
            self.cross_modal_video_pooler.apply(init_weights)
            self.cross_modal_text_pooler = heads.Pooler(config["hidden_size"])
            self.cross_modal_text_pooler.apply(init_weights)

            ## einops transformations
            self.einops_from_space = 'b (f n) d'
            self.einops_to_space = '(b f) n d'
            self.einops_from_time = 'b (f n) d'
            self.einops_to_time = '(b n) f d'

        if 'MLM' in self.task_names:
            self.mlm_score = heads.MLMHead(bert_config)
            self.mlm_score.apply(init_weights)

        if 'ITM' in self.task_names:
            self.itm_score = heads.ITMHead(config["hidden_size"] * 2)
            self.itm_score.apply(init_weights)

        if load_checkpoint not in ["", None]:
            checkpoint = torch.load(load_checkpoint, map_location='cpu')
            state_dict = checkpoint['state_dict']
            new_state_dict = state_dict_data_parallel_fix(state_dict, self.state_dict())
            new_state_dict = self._inflate_positional_embeds(new_state_dict)
            self.load_state_dict(new_state_dict, strict=False)

    def set_device(self, device):
        self.device = device

    def forward(self, video_data):

        b, curr_frames, channels, _, _ = video_data.shape
        video_data_itm = self.video_model.patch_embed(video_data)
        video_data_itm = video_data_itm.flatten(2).transpose(2, 1)
        video_data_itm = video_data_itm.reshape(b, -1, self.video_model.patch_embed.embed_dim)

        BF = video_data_itm.shape[0]
        cls_tokens = self.cls_token.expand(BF, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        video_data_itm = torch.cat((cls_tokens, video_data_itm), dim=1)
        # positional embed needs to be tiled for each frame (this does [1,2,3] --> [1,2,3,1,2,3]...)
        cls_embed = self.video_model.pos_embed[:, 0, :].unsqueeze(1)
        tile_pos_embed = self.video_model.pos_embed[:, 1:, :].repeat(1, self.num_frames, 1)
        # temporal embed needs to be repeated within each frame (this does [1,2,3] --> [1,1,1,2,2,2,3,3,3]...)
        tile_temporal_embed = self.video_model.temporal_embed.repeat_interleave(self.patches_per_frame, 1)
        total_pos_embed = tile_pos_embed + tile_temporal_embed
        total_pos_embed = torch.cat([cls_embed, total_pos_embed], dim=1)

        n = self.patches_per_frame
        f = curr_frames

        curr_patches = video_data_itm.shape[1]
        video_data_itm = video_data_itm + total_pos_embed[:, :curr_patches]
        video_data_itm = self.video_model.pos_drop(video_data_itm)

        unfused_blocks = self.num_text_layer - self.num_fuse_block



        for blk_i, blk in enumerate(self.video_model.blocks[:unfused_blocks]):
            if config['use_checkpoint']:
                video_data_itm = torch.utils.checkpoint.checkpoint(blk, video_data_itm, self.einops_from_space, self.einops_to_space, self.einops_from_time, self.einops_to_time,
                                n, f)
            else:
                video_data_itm = blk(video_data_itm, self.einops_from_space, self.einops_to_space, self.einops_from_time, self.einops_to_time,
                                time_n=n, space_f=f)

        return video_data_itm


    def compute_text(self, text_data):
        if self.text_params['model'].startswith('bert'):
            text_embeddings = self.text_model(text_data['input_ids'], attention_mask=text_data['attention_mask'])[
                'pooler_output']
        elif self.text_params['model'].startswith('distilbert'):
            text_embeddings = self.text_model(**text_data).last_hidden_state[:, 0, :]
        elif self.text_params['model'].startswith('roberta'):
            text_embeddings = self.text_model(**text_data).last_hidden_state[:, 0, :]
        else:
            raise NotImplementedError
        text_embeddings = self.txt_proj(text_embeddings)
        return text_embeddings

    def compute_text_tokens(self, text_data):
        if self.text_params['model'].startswith('bert'):
            text_embeddings = self.text_model(text_data['input_ids'], attention_mask=text_data['attention_mask'])[
                'pooler_output']    # not implement for bert
        elif self.text_params['model'].startswith('distilbert'):
            text_embeddings = self.text_model(**text_data).last_hidden_state
        elif self.text_params['model'].startswith('roberta'):
            text_embeddings = self.text_model(**text_data).last_hidden_state
        else:
            raise NotImplementedError

        text_embeddings = self.txt_proj(text_embeddings)
        return text_embeddings

    def compute_video(self, video_data):
        video_embeddings = self.video_model(video_data)
        video_embeddings = self.vid_proj(video_embeddings)
        return video_embeddings

    def _inflate_positional_embeds(self, new_state_dict):
        # allow loading of timesformer with fewer num_frames
        curr_keys = list(self.state_dict().keys())
        if 'video_model.temporal_embed' in new_state_dict and 'video_model.temporal_embed' in curr_keys:
            load_temporal_embed = new_state_dict['video_model.temporal_embed']
            load_num_frames = load_temporal_embed.shape[1]
            curr_num_frames = self.video_params['num_frames']
            embed_dim = load_temporal_embed.shape[2]

            if load_num_frames != curr_num_frames:
                if load_num_frames > curr_num_frames:
                    print(f'### loaded {self.video_params["model"]} model has MORE frames than current...'
                          f'### loading weights, filling in the extras via {self.load_temporal_fix}')
                    new_temporal_embed = load_temporal_embed[:, :curr_num_frames, :]
                else:
                    print(f'### loaded {self.video_params["model"]} model has FEWER frames than current...'
                          f'### loading weights, filling in the extras via {self.load_temporal_fix}')
                    if self.load_temporal_fix == 'zeros':
                        new_temporal_embed = torch.zeros([load_temporal_embed.shape[0], curr_num_frames, embed_dim])
                        new_temporal_embed[:, :load_num_frames] = load_temporal_embed
                    elif self.load_temporal_fix in ['interp', 'bilinear']:
                        # interpolate
                        # unsqueeze so pytorch thinks its an image
                        mode = 'nearest'
                        if self.load_temporal_fix == 'bilinear':
                            mode = 'bilinear'
                        load_temporal_embed = load_temporal_embed.unsqueeze(0)
                        new_temporal_embed = F.interpolate(load_temporal_embed,
                                                           (curr_num_frames, embed_dim), mode=mode, align_corners=True).squeeze(0)
                    else:
                        raise NotImplementedError
                new_state_dict['video_model.temporal_embed'] = new_temporal_embed
        # allow loading with smaller spatial patches. assumes custom border crop, to append the
        # border patches to the input sequence
        if 'video_model.pos_embed' in new_state_dict and 'video_model.pos_embed' in curr_keys:
            load_pos_embed = new_state_dict['video_model.pos_embed']
            load_num_patches = load_pos_embed.shape[1]
            curr_pos_embed = self.state_dict()['video_model.pos_embed']
            if load_num_patches != curr_pos_embed.shape[1]:
                raise NotImplementedError(
                    'Loading models with different spatial resolution / patch number not yet implemented, sorry.')

        return new_state_dict

def sim_matrix(a, b, eps=1e-8):
    """
    added eps for numerical stability
    """
    a_n, b_n = a.norm(dim=1)[:, None], b.norm(dim=1)[:, None]
    a_norm = a / torch.max(a_n, eps * torch.ones_like(a_n))
    b_norm = b / torch.max(b_n, eps * torch.ones_like(b_n))
    sim_mt = torch.mm(a_norm, b_norm.transpose(0, 1))
    return sim_mt


def sim_matrix_batch_val(a, b, eps=1e-8):
    """
    added eps for numerical stability
    """
    a_n, b_n = a.norm(dim=-1).unsqueeze(-1), b.norm(dim=-1).unsqueeze(-1)
    a_norm = a / torch.max(a_n, eps * torch.ones_like(a_n))
    b_norm = b / torch.max(b_n, eps * torch.ones_like(b_n))
    sim_mt = torch.bmm(a_norm, b_norm.transpose(1, 2))
    return sim_mt


if __name__ == "__main__":
    pass

QFVS/model/model_video_unfused.py
facebookresearch/EgoVLPv2
# coding=utf-8
# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""PyTorch RoBERTa model. """

import yaml

with open('./QFVS.yml') as f:
    config_yaml = yaml.load(f, Loader=yaml.FullLoader)

NUM_FUSE_BLOCK = config_yaml['num_fuse_block']
DIM_IMG = 768

import math
import inspect
from typing import Callable, List, Optional, Set, Tuple, Union

import torch
import torch.nn as nn
import torch.utils.checkpoint
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from transformers.activations import ACT2FN, gelu
from transformers.file_utils import (
    add_code_sample_docstrings,
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    replace_return_docstrings,
)
from transformers.modeling_outputs import (
    BaseModelOutputWithPastAndCrossAttentions,
    BaseModelOutputWithPoolingAndCrossAttentions,
    CausalLMOutputWithCrossAttentions,
    MaskedLMOutput,
    MultipleChoiceModelOutput,
    QuestionAnsweringModelOutput,
    SequenceClassifierOutput,
    TokenClassifierOutput,
)
from transformers.modeling_utils import (
    PreTrainedModel,
    find_pruneable_heads_and_indices,
    prune_linear_layer,
)
from transformers.utils import logging
from transformers.models.roberta.configuration_roberta import RobertaConfig


logger = logging.get_logger(__name__)

_CHECKPOINT_FOR_DOC = "roberta-base"
_CONFIG_FOR_DOC = "RobertaConfig"
_TOKENIZER_FOR_DOC = "RobertaTokenizer"

ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST = [
    "roberta-base",
    "roberta-large",
    "roberta-large-mnli",
    "distilroberta-base",
    "roberta-base-openai-detector",
    "roberta-large-openai-detector",
    # See all RoBERTa models at https://huggingface.co/models?filter=roberta
]


def apply_chunking_to_forward(
    forward_fn: Callable[..., torch.Tensor], chunk_size: int, chunk_dim: int, *input_tensors, last_norm=True
) -> torch.Tensor:
    """
    This function chunks the `input_tensors` into smaller input tensor parts of size `chunk_size` over the dimension
    `chunk_dim`. It then applies a layer `forward_fn` to each chunk independently to save memory.
    If the `forward_fn` is independent across the `chunk_dim` this function will yield the same result as directly
    applying `forward_fn` to `input_tensors`.
    Args:
        forward_fn (`Callable[..., torch.Tensor]`):
            The forward function of the model.
        chunk_size (`int`):
            The chunk size of a chunked tensor: `num_chunks = len(input_tensors[0]) / chunk_size`.
        chunk_dim (`int`):
            The dimension over which the `input_tensors` should be chunked.
        input_tensors (`Tuple[torch.Tensor]`):
            The input tensors of `forward_fn` which will be chunked
    Returns:
        `torch.Tensor`: A tensor with the same shape as the `forward_fn` would have given if applied`.
    Examples:
    ```python
    # rename the usual forward() fn to forward_chunk()
    def forward_chunk(self, hidden_states):
        hidden_states = self.decoder(hidden_states)
        return hidden_states
    # implement a chunked forward function
    def forward(self, hidden_states):
        return apply_chunking_to_forward(self.forward_chunk, self.chunk_size_lm_head, self.seq_len_dim, hidden_states)
    ```"""

    assert len(input_tensors) > 0, f"{input_tensors} has to be a tuple/list of tensors"

    # inspect.signature exist since python 3.5 and is a python method -> no problem with backward compatibility
    num_args_in_forward_chunk_fn = len(inspect.signature(forward_fn).parameters)
    # if num_args_in_forward_chunk_fn != len(input_tensors):
    #    raise ValueError(
    #        f"forward_chunk_fn expects {num_args_in_forward_chunk_fn} arguments, but only {len(input_tensors)} input "
    #        "tensors are given"
    #    )

    if chunk_size > 0:
        tensor_shape = input_tensors[0].shape[chunk_dim]
        for input_tensor in input_tensors:
            if input_tensor.shape[chunk_dim] != tensor_shape:
                raise ValueError(
                    f"All input tenors have to be of the same shape: {tensor_shape}, "
                    f"found shape {input_tensor.shape[chunk_dim]}"
                )

        if input_tensors[0].shape[chunk_dim] % chunk_size != 0:
            raise ValueError(
                f"The dimension to be chunked {input_tensors[0].shape[chunk_dim]} has to be a multiple of the chunk "
                f"size {chunk_size}"
            )

        num_chunks = input_tensors[0].shape[chunk_dim] // chunk_size

        # chunk input tensor into tuples
        input_tensors_chunks = tuple(input_tensor.chunk(num_chunks, dim=chunk_dim) for input_tensor in input_tensors)
        # apply forward fn to every tuple
        output_chunks = tuple(
            forward_fn(*input_tensors_chunk, last_norm=last_norm) for input_tensors_chunk in zip(*input_tensors_chunks)
        )
        # concatenate output at same dimension
        return torch.cat(output_chunks, dim=chunk_dim)

    return forward_fn(*input_tensors, last_norm=last_norm)


class RobertaEmbeddings(nn.Module):
    """
    Same as BertEmbeddings with a tiny tweak for positional embeddings indexing.
    """

    # Copied from transformers.models.bert.modeling_bert.BertEmbeddings.__init__
    def __init__(self, config):
        super().__init__()
        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
        self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)

        # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
        # any TensorFlow checkpoint file
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

        # position_ids (1, len position emb) is contiguous in memory and exported when serialized
        self.register_buffer("position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)))
        self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")

        # End copy
        self.padding_idx = config.pad_token_id
        self.position_embeddings = nn.Embedding(
            config.max_position_embeddings, config.hidden_size, padding_idx=self.padding_idx
        )

    def forward(
        self, input_ids=None, token_type_ids=None, position_ids=None, inputs_embeds=None, past_key_values_length=0
    ):
        if position_ids is None:
            if input_ids is not None:
                # Create the position ids from the input token ids. Any padded tokens remain padded.
                position_ids = create_position_ids_from_input_ids(
                    input_ids, self.padding_idx, past_key_values_length
                ).to(input_ids.device)
            else:
                position_ids = self.create_position_ids_from_inputs_embeds(inputs_embeds)

        if input_ids is not None:
            input_shape = input_ids.size()
        else:
            input_shape = inputs_embeds.size()[:-1]

        if token_type_ids is None:
            token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)

        if inputs_embeds is None:
            inputs_embeds = self.word_embeddings(input_ids)
        token_type_embeddings = self.token_type_embeddings(token_type_ids)

        embeddings = inputs_embeds + token_type_embeddings
        if self.position_embedding_type == "absolute":
            position_embeddings = self.position_embeddings(position_ids)
            embeddings += position_embeddings
        embeddings = self.LayerNorm(embeddings)
        embeddings = self.dropout(embeddings)
        return embeddings

    def create_position_ids_from_inputs_embeds(self, inputs_embeds):
        """
        We are provided embeddings directly. We cannot infer which are padded so just generate sequential position ids.
        Args:
            inputs_embeds: torch.Tensor
        Returns: torch.Tensor
        """
        input_shape = inputs_embeds.size()[:-1]
        sequence_length = input_shape[1]

        position_ids = torch.arange(
            self.padding_idx + 1, sequence_length + self.padding_idx + 1, dtype=torch.long, device=inputs_embeds.device
        )
        return position_ids.unsqueeze(0).expand(input_shape)


# Copied from transformers.models.bert.modeling_bert.BertSelfAttention with Bert->Roberta
class RobertaSelfAttention(nn.Module):
    def __init__(self, config, layer_index=None):
        super().__init__()
        if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
            raise ValueError(
                f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
                f"heads ({config.num_attention_heads})"
            )

        self.num_attention_heads = config.num_attention_heads
        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
        self.all_head_size = self.num_attention_heads * self.attention_head_size

        self.query = nn.Linear(config.hidden_size, self.all_head_size)
        if layer_index is None:
            self.key = nn.Linear(config.hidden_size, self.all_head_size)
            self.value = nn.Linear(config.hidden_size, self.all_head_size)
        else:
            self.key = nn.Linear(DIM_IMG, self.all_head_size)
            self.value = nn.Linear(DIM_IMG, self.all_head_size)

        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
        self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")
        if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
            self.max_position_embeddings = config.max_position_embeddings
            self.distance_embedding = nn.Embedding(2 * config.max_position_embeddings - 1, self.attention_head_size)

        self.is_decoder = config.is_decoder

    def transpose_for_scores(self, x):
        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
        x = x.view(*new_x_shape)
        return x.permute(0, 2, 1, 3)

    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        past_key_value=None,
        output_attentions=False,
    ):
        mixed_query_layer = self.query(hidden_states)

        # If this is instantiated as a cross-attention module, the keys
        # and values come from an encoder; the attention mask needs to be
        # such that the encoder's padding tokens are not attended to.
        is_cross_attention = encoder_hidden_states is not None

        if encoder_hidden_states is not None:
            key_layer = self.transpose_for_scores(self.key(encoder_hidden_states))
            value_layer = self.transpose_for_scores(self.value(encoder_hidden_states))
            attention_mask = encoder_attention_mask
        else:
            key_layer = self.transpose_for_scores(self.key(hidden_states))
            value_layer = self.transpose_for_scores(self.value(hidden_states))

        query_layer = self.transpose_for_scores(mixed_query_layer)

        # Take the dot product between "query" and "key" to get the raw attention scores.
        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))

        if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
            seq_length = hidden_states.size()[1]
            position_ids_l = torch.arange(seq_length, dtype=torch.long, device=hidden_states.device).view(-1, 1)
            position_ids_r = torch.arange(seq_length, dtype=torch.long, device=hidden_states.device).view(1, -1)
            distance = position_ids_l - position_ids_r
            positional_embedding = self.distance_embedding(distance + self.max_position_embeddings - 1)
            positional_embedding = positional_embedding.to(dtype=query_layer.dtype)  # fp16 compatibility

            if self.position_embedding_type == "relative_key":
                relative_position_scores = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
                attention_scores = attention_scores + relative_position_scores
            elif self.position_embedding_type == "relative_key_query":
                relative_position_scores_query = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
                relative_position_scores_key = torch.einsum("bhrd,lrd->bhlr", key_layer, positional_embedding)
                attention_scores = attention_scores + relative_position_scores_query + relative_position_scores_key

        attention_scores = attention_scores / math.sqrt(self.attention_head_size)
        if attention_mask is not None:
            # Apply the attention mask is (precomputed for all layers in RobertaModel forward() function)
            attention_scores = attention_scores + attention_mask

        # Normalize the attention scores to probabilities.
        attention_probs = nn.Softmax(dim=-1)(attention_scores)

        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.
        attention_probs = self.dropout(attention_probs)

        # Mask heads if we want to
        if head_mask is not None:
            attention_probs = attention_probs * head_mask

        context_layer = torch.matmul(attention_probs, value_layer)

        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
        context_layer = context_layer.view(*new_context_layer_shape)

        outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)

        return outputs


# Copied from transformers.models.bert.modeling_bert.BertSelfOutput
class RobertaSelfOutput(nn.Module):
    def __init__(self, config, layer_index=None):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        if layer_index is None:
            self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, hidden_states):
        #print("Shape of hidden_states: ", hidden_states.shape)
        hidden_states = self.dense(hidden_states)
        hidden_states = self.dropout(hidden_states)
        return hidden_states


# Copied from transformers.models.bert.modeling_bert.BertAttention with Bert->Roberta
class RobertaAttention(nn.Module):
    def __init__(self, config, layer_index=None):
        super().__init__()
        self.self = RobertaSelfAttention(config, layer_index=layer_index)
        self.output = RobertaSelfOutput(config, layer_index=layer_index)
        self.pruned_heads = set()

    def prune_heads(self, heads):
        if len(heads) == 0:
            return
        heads, index = find_pruneable_heads_and_indices(
            heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads
        )

        # Prune linear layers
        self.self.query = prune_linear_layer(self.self.query, index)
        self.self.key = prune_linear_layer(self.self.key, index)
        self.self.value = prune_linear_layer(self.self.value, index)
        self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)

        # Update hyper params and store pruned heads
        self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
        self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads
        self.pruned_heads = self.pruned_heads.union(heads)

    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        past_key_value=None,
        output_attentions=False,
    ):
        self_outputs = self.self(
            hidden_states,
            attention_mask,
            head_mask,
            encoder_hidden_states,
            encoder_attention_mask,
            past_key_value,
            output_attentions,
        )
        attention_output = self.output(self_outputs[0])
        outputs = (attention_output,) + self_outputs[1:]  # add attentions if we output them
        return outputs


# Copied from transformers.models.bert.modeling_bert.BertIntermediate
class RobertaIntermediate(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
        if isinstance(config.hidden_act, str):
            self.intermediate_act_fn = ACT2FN[config.hidden_act]
        else:
            self.intermediate_act_fn = config.hidden_act

    def forward(self, hidden_states):
        hidden_states = self.dense(hidden_states)
        hidden_states = self.intermediate_act_fn(hidden_states)
        return hidden_states


# Copied from transformers.models.bert.modeling_bert.BertOutput
class RobertaOutput(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, hidden_states, input_tensor, last_norm=True):
        hidden_states = self.dense(hidden_states)
        hidden_states = self.dropout(hidden_states)
        hidden_states = hidden_states + input_tensor
        if last_norm:
            hidden_states = self.LayerNorm(hidden_states)
        return hidden_states


# Copied from transformers.models.bert.modeling_bert.BertLayer with Bert->Roberta
class RobertaLayer(nn.Module):
    def __init__(self, config, layer_index=None):
        super().__init__()
        self.chunk_size_feed_forward = config.chunk_size_feed_forward
        self.seq_len_dim = 1
        self.attention = RobertaAttention(config)
        self.is_decoder = config.is_decoder
        self.add_cross_attention = config.add_cross_attention
        if layer_index >= 12 - NUM_FUSE_BLOCK:
            self.crossattention_t2i = RobertaAttention(config, layer_index=layer_index)
            self.alpha_t2i = nn.Parameter(torch.Tensor([0]))
        self.intermediate = RobertaIntermediate(config)
        self.output = RobertaOutput(config)

    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        past_key_value=None,
        output_attentions=False,
        last_norm=True,
    ):
        # decoder uni-directional self-attention cached key/values tuple is at positions 1,2
        self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
        self_attention_outputs = self.attention(
            hidden_states,
            attention_mask,
            head_mask,
            output_attentions=output_attentions,
            past_key_value=self_attn_past_key_value,
        )
        attention_output = self_attention_outputs[0]

        # if decoder, the last output is tuple of self-attn cache
        outputs = self_attention_outputs[1:]  # add self attentions if we output attention weights

        cross_attn_present_key_value = None
        if encoder_hidden_states is not None:
            assert hasattr(
                self, "crossattention_t2i"
            ), f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers by setting `config.add_cross_attention=True`"

            # cross_attn cached key/values tuple is at positions 3,4 of past_key_value tuple
            cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None
            cross_attention_outputs = self.crossattention_t2i(
                attention_output,
                attention_mask,
                head_mask,
                encoder_hidden_states,
                encoder_attention_mask,
                cross_attn_past_key_value,
                output_attentions,
            )
            attention_output = self.alpha_t2i * cross_attention_outputs[0] + attention_output

        attention_output = self.attention.output.LayerNorm(attention_output + hidden_states)

        layer_output = apply_chunking_to_forward(
            self.feed_forward_chunk,
            self.chunk_size_feed_forward,
            self.seq_len_dim,
            attention_output,
            last_norm=last_norm,
        )
        outputs = (layer_output,) + outputs

        # if decoder, return the attn key/values as the last output
        return outputs

    def feed_forward_chunk(self, attention_output, last_norm=True):
        intermediate_output = self.intermediate(attention_output)
        layer_output = self.output(intermediate_output, attention_output, last_norm=last_norm)
        return layer_output


# Copied from transformers.models.bert.modeling_bert.BertEncoder with Bert->Roberta
class RobertaEncoder(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.config.gradient_checkpointing = config_yaml['use_checkpoint']
        self.layer = nn.ModuleList(
            [RobertaLayer(config, layer_index=layer_i) for layer_i in range(config.num_hidden_layers)]
        )

    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        past_key_values=None,
        use_cache=None,
        output_attentions=False,
        output_hidden_states=False,
        return_dict=True,
    ):
        all_hidden_states = () if output_hidden_states else None
        all_self_attentions = () if output_attentions else None
        all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None

        next_decoder_cache = () if use_cache else None
        for i, layer_module in enumerate(self.layer):
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

            layer_head_mask = head_mask[i] if head_mask is not None else None
            past_key_value = past_key_values[i] if past_key_values is not None else None

            if getattr(self.config, "gradient_checkpointing", False) and self.training:

                if use_cache:
                    logger.warning(
                        "`use_cache=True` is incompatible with `config.gradient_checkpointing=True`. Setting "
                        "`use_cache=False`..."
                    )
                    use_cache = False

                #def create_custom_forward(module):
                    #def custom_forward(*inputs):
                        #return module(*inputs, past_key_value, output_attentions)
                    #return custom_forward

                layer_outputs = torch.utils.checkpoint.checkpoint(
                    layer_module,
                    hidden_states,
                    attention_mask,
                    layer_head_mask,
                    encoder_hidden_states,
                    encoder_attention_mask,
                    past_key_value, 
                    output_attentions
                )
            else:
                layer_outputs = layer_module(
                    hidden_states,
                    attention_mask,
                    layer_head_mask,
                    encoder_hidden_states,
                    encoder_attention_mask,
                    past_key_value,
                    output_attentions,
                )


            hidden_states = layer_outputs[0]
            if use_cache:
                next_decoder_cache += (layer_outputs[-1],)
            if output_attentions:
                all_self_attentions = all_self_attentions + (layer_outputs[1],)
                if self.config.add_cross_attention:
                    all_cross_attentions = all_cross_attentions + (layer_outputs[2],)

        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        if not return_dict:
            return tuple(
                v
                for v in [
                    hidden_states,
                    next_decoder_cache,
                    all_hidden_states,
                    all_self_attentions,
                    all_cross_attentions,
                ]
                if v is not None
            )
        return BaseModelOutputWithPastAndCrossAttentions(
            last_hidden_state=hidden_states,
            past_key_values=next_decoder_cache,
            hidden_states=all_hidden_states,
            attentions=all_self_attentions,
            cross_attentions=all_cross_attentions,
        )

class RobertaPooler(nn.Module):
    def __init__(self, config):
        super().__init__()
        #self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        self.activation = nn.Tanh()

    def forward(self, hidden_states):
        # We "pool" the model by simply taking the hidden state corresponding
        # to the first token.
        first_token_tensor = hidden_states[:, 0]
        #pooled_output = self.dense(first_token_tensor)
        pooled_output = self.activation(first_token_tensor)
        return pooled_output


class RobertaPreTrainedModel(PreTrainedModel):
    """
    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
    models.
    """

    config_class = RobertaConfig
    base_model_prefix = "roberta"

    # Copied from transformers.models.bert.modeling_bert.BertPreTrainedModel._init_weights
    def _init_weights(self, module):
        """Initialize the weights"""
        if isinstance(module, nn.Linear):
            # Slightly different from the TF version which uses truncated_normal for initialization
            # cf https://github.com/pytorch/pytorch/pull/5617
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)


ROBERTA_START_DOCSTRING = r"""
    This model inherits from :class:`~transformers.PreTrainedModel`. Check the superclass documentation for the generic
    methods the library implements for all its model (such as downloading or saving, resizing the input embeddings,
    pruning heads etc.)
    This model is also a PyTorch `torch.nn.Module <https://pytorch.org/docs/stable/nn.html#torch.nn.Module>`__
    subclass. Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to
    general usage and behavior.
    Parameters:
        config (:class:`~transformers.RobertaConfig`): Model configuration class with all the parameters of the
            model. Initializing with a config file does not load the weights associated with the model, only the
            configuration. Check out the :meth:`~transformers.PreTrainedModel.from_pretrained` method to load the model
            weights.
"""

ROBERTA_INPUTS_DOCSTRING = r"""
    Args:
        input_ids (:obj:`torch.LongTensor` of shape :obj:`({0})`):
            Indices of input sequence tokens in the vocabulary.
            Indices can be obtained using :class:`~transformers.RobertaTokenizer`. See
            :meth:`transformers.PreTrainedTokenizer.encode` and :meth:`transformers.PreTrainedTokenizer.__call__` for
            details.
            `What are input IDs? <../glossary.html#input-ids>`__
        attention_mask (:obj:`torch.FloatTensor` of shape :obj:`({0})`, `optional`):
            Mask to avoid performing attention on padding token indices. Mask values selected in ``[0, 1]``:
            - 1 for tokens that are **not masked**,
            - 0 for tokens that are **masked**.
            `What are attention masks? <../glossary.html#attention-mask>`__
        token_type_ids (:obj:`torch.LongTensor` of shape :obj:`({0})`, `optional`):
            Segment token indices to indicate first and second portions of the inputs. Indices are selected in ``[0,
            1]``:
            - 0 corresponds to a `sentence A` token,
            - 1 corresponds to a `sentence B` token.
            `What are token type IDs? <../glossary.html#token-type-ids>`_
        position_ids (:obj:`torch.LongTensor` of shape :obj:`({0})`, `optional`):
            Indices of positions of each input sequence tokens in the position embeddings. Selected in the range ``[0,
            config.max_position_embeddings - 1]``.
            `What are position IDs? <../glossary.html#position-ids>`_
        head_mask (:obj:`torch.FloatTensor` of shape :obj:`(num_heads,)` or :obj:`(num_layers, num_heads)`, `optional`):
            Mask to nullify selected heads of the self-attention modules. Mask values selected in ``[0, 1]``:
            - 1 indicates the head is **not masked**,
            - 0 indicates the head is **masked**.
        inputs_embeds (:obj:`torch.FloatTensor` of shape :obj:`({0}, hidden_size)`, `optional`):
            Optionally, instead of passing :obj:`input_ids` you can choose to directly pass an embedded representation.
            This is useful if you want more control over how to convert :obj:`input_ids` indices into associated
            vectors than the model's internal embedding lookup matrix.
        output_attentions (:obj:`bool`, `optional`):
            Whether or not to return the attentions tensors of all attention layers. See ``attentions`` under returned
            tensors for more detail.
        output_hidden_states (:obj:`bool`, `optional`):
            Whether or not to return the hidden states of all layers. See ``hidden_states`` under returned tensors for
            more detail.
        return_dict (:obj:`bool`, `optional`):
            Whether or not to return a :class:`~transformers.file_utils.ModelOutput` instead of a plain tuple.
"""


@add_start_docstrings(
    "The bare RoBERTa Model transformer outputting raw hidden-states without any specific head on top.",
    ROBERTA_START_DOCSTRING,
)
class RobertaModel(RobertaPreTrainedModel):
    """
    The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of
    cross-attention is added between the self-attention layers, following the architecture described in `Attention is
    all you need`_ by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz
    Kaiser and Illia Polosukhin.
    To behave as an decoder the model needs to be initialized with the :obj:`is_decoder` argument of the configuration
    set to :obj:`True`. To be used in a Seq2Seq model, the model needs to initialized with both :obj:`is_decoder`
    argument and :obj:`add_cross_attention` set to :obj:`True`; an :obj:`encoder_hidden_states` is then expected as an
    input to the forward pass.
    .. _`Attention is all you need`: https://arxiv.org/abs/1706.03762
    """

    _keys_to_ignore_on_load_missing = [r"position_ids"]

    # Copied from transformers.models.bert.modeling_bert.BertModel.__init__ with Bert->Roberta
    def __init__(self, config, add_pooling_layer=True):
        super().__init__(config)
        self.config = config

        self.embeddings = RobertaEmbeddings(config)
        self.encoder = RobertaEncoder(config)

        self.pooler = RobertaPooler(config) if add_pooling_layer else None

        self.init_weights()

    def get_input_embeddings(self):
        return self.embeddings.word_embeddings

    def set_input_embeddings(self, value):
        self.embeddings.word_embeddings = value

    def _prune_heads(self, heads_to_prune):
        """
        Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
        class PreTrainedModel
        """
        for layer, heads in heads_to_prune.items():
            self.encoder.layer[layer].attention.prune_heads(heads)

    @add_start_docstrings_to_model_forward(ROBERTA_INPUTS_DOCSTRING.format("(batch_size, sequence_length)"))
    @add_code_sample_docstrings(
        processor_class=_TOKENIZER_FOR_DOC,
        checkpoint=_CHECKPOINT_FOR_DOC,
        output_type=BaseModelOutputWithPoolingAndCrossAttentions,
        config_class=_CONFIG_FOR_DOC,
    )
    # Copied from transformers.models.bert.modeling_bert.BertModel.forward
    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        past_key_values=None,
        use_cache=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
    ):
        r"""
        encoder_hidden_states  (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`):
            Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
            the model is configured as a decoder.
        encoder_attention_mask (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
            Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
            the cross-attention if the model is configured as a decoder. Mask values selected in ``[0, 1]``:
            - 1 for tokens that are **not masked**,
            - 0 for tokens that are **masked**.
        past_key_values (:obj:`tuple(tuple(torch.FloatTensor))` of length :obj:`config.n_layers` with each tuple having 4 tensors of shape :obj:`(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
            Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
            If :obj:`past_key_values` are used, the user can optionally input only the last :obj:`decoder_input_ids`
            (those that don't have their past key value states given to this model) of shape :obj:`(batch_size, 1)`
            instead of all :obj:`decoder_input_ids` of shape :obj:`(batch_size, sequence_length)`.
        use_cache (:obj:`bool`, `optional`):
            If set to :obj:`True`, :obj:`past_key_values` key value states are returned and can be used to speed up
            decoding (see :obj:`past_key_values`).
        """
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        use_cache = False

        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot secify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
            input_shape = input_ids.size()
            batch_size, seq_length = input_shape
        elif inputs_embeds is not None:
            input_shape = inputs_embeds.size()[:-1]
            batch_size, seq_length = input_shape
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        device = input_ids.device if input_ids is not None else inputs_embeds.device

        # past_key_values_length
        past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0

        if attention_mask is None:
            attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length)), device=device)
        if token_type_ids is None:
            token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)

        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
        # ourselves in which case we just need to make it broadcastable to all heads.
        extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape, device)

        # If a 2D or 3D attention mask is provided for the cross-attention
        # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
        if encoder_hidden_states is not None:
            encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
            encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
            if encoder_attention_mask is None:
                encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
            encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
        else:
            encoder_extended_attention_mask = None

        # Prepare head mask if needed
        # 1.0 in head_mask indicate we keep the head
        # attention_probs has shape bsz x n_heads x N x N
        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)

        embedding_output = self.embeddings(
            input_ids=input_ids,
            position_ids=position_ids,
            token_type_ids=token_type_ids,
            inputs_embeds=inputs_embeds,
            past_key_values_length=past_key_values_length,
        )
        encoder_outputs = self.encoder(
            embedding_output,
            attention_mask=extended_attention_mask,
            head_mask=head_mask,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_extended_attention_mask,
            past_key_values=past_key_values,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        sequence_output = encoder_outputs[0]
        pooled_output = self.pooler(sequence_output) if self.pooler is not None else None

        if not return_dict:
            return (sequence_output, pooled_output) + encoder_outputs[1:]

        return BaseModelOutputWithPoolingAndCrossAttentions(
            last_hidden_state=sequence_output,
            pooler_output=pooled_output,
            past_key_values=encoder_outputs.past_key_values,
            hidden_states=encoder_outputs.hidden_states,
            attentions=encoder_outputs.attentions,
            cross_attentions=encoder_outputs.cross_attentions,
        )


def create_position_ids_from_input_ids(input_ids, padding_idx, past_key_values_length=0):
    """
    Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding symbols
    are ignored. This is modified from fairseq's `utils.make_positions`.
    Args:
        x: torch.Tensor x:
    Returns: torch.Tensor
    """
    # The series of casts and type-conversions here are carefully balanced to both work with ONNX export and XLA.
    mask = input_ids.ne(padding_idx).int()
    incremental_indices = (torch.cumsum(mask, dim=1).type_as(mask) + past_key_values_length) * mask
    return incremental_indices.long() + padding_idx


def _make_causal_mask(input_ids_shape: torch.Size, dtype: torch.dtype, past_key_values_length: int = 0):
    """
    Make causal mask used for bi-directional self-attention.
    """
    bsz, tgt_len = input_ids_shape
    mask = torch.full((tgt_len, tgt_len), float("-inf"))
    mask_cond = torch.arange(mask.size(-1))
    mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
    mask = mask.to(dtype)

    if past_key_values_length > 0:
        mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype), mask], dim=-1)
    return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)


def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len=None):
    """
    Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
    """
    bsz, src_len = mask.size()
    tgt_len = tgt_len if tgt_len is not None else src_len

    expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)

    inverted_mask = 1.0 - expanded_mask

    return inverted_mask.masked_fill(inverted_mask.bool(), torch.finfo(dtype).min)


def _prepare_decoder_attention_mask(
    attention_mask, input_shape, inputs_embeds, device, past_key_values_length: int = 0
):
    # create causal mask
    # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
    combined_attention_mask = None
    if input_shape[-1] > 1:
        combined_attention_mask = _make_causal_mask(
            input_shape, inputs_embeds.dtype, past_key_values_length=past_key_values_length
        ).to(device)

    if attention_mask is not None:
        # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
        expanded_attn_mask = _expand_mask(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1])
        combined_attention_mask = (
            expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask
        )

    return combined_attention_mask

QFVS/model/roberta.py
facebookresearch/EgoVLPv2
"""
Implementations of Video Transformers in PyTorch

A PyTorch implementation of space-time transformer as described in
'Frozen in Time: A Joint Image and Video Encoder for End-to-End Retrieval' - https://arxiv.org/abs/2104.00650

A PyTorch implementation of timesformer as described in
'Is Space-Time Attention All You Need for Video Understanding?' - https://arxiv.org/abs/2102.05095

Acknowledgments:
- This code builds on Ross Wightman's vision_transformer code in pytorch-image-models:
https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py

- It is also inspired by lucidrains timesformer implementation:
https://github.com/lucidrains/TimeSformer-pytorch

Hacked together by Max Bain
"""

from collections import OrderedDict
from functools import partial
import yaml

import torch
from einops import rearrange, repeat
from timm.models.layers import DropPath, to_2tuple, trunc_normal_
from torch import einsum, nn

with open('./QFVS.yml') as f:
    config_yaml = yaml.load(f, Loader=yaml.FullLoader)

NUM_FUSE_BLOCK = config_yaml['num_fuse_block']
DIM_TEXT = 768

def attn(q, k, v):
    sim = einsum('b i d, b j d -> b i j', q, k)
    attn = sim.softmax(dim=-1)
    out = einsum('b i j, b j d -> b i d', attn, v)
    return out


class Mlp(nn.Module):
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x


class VideoPatchEmbed(nn.Module):
    """ Video to Patch Embedding
    """

    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768,
                 num_frames=8):
        super().__init__()
        img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)
        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0]) * num_frames
        self.img_size = img_size
        self.patch_size = patch_size
        self.num_patches = num_patches
        self.num_frames = num_frames
        self.embed_dim = embed_dim
        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)

    def forward(self, x):
        B, F, C, H, W = x.shape
        assert F == self.num_frames, print(F, self.num_frames)
        x = x.view(-1, C, H, W)
        x = self.proj(x)
        return x


class VarAttention(nn.Module):
    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.,
                 initialize='random', dim_text=None, norm_layer=nn.LayerNorm, space_attn=True):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        # NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights
        self.scale = qk_scale or head_dim ** -0.5
        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.proj = nn.Linear(dim, dim)
        if initialize == 'zeros':
            self.qkv.weight.data.fill_(0)
            self.qkv.bias.data.fill_(0)
            # fill proj weight with 1 here to improve training dynamics. Otherwise temporal attention inputs
            # are multiplied by 0*0, which is hard for the model to move out of.
            self.proj.weight.data.fill_(1)
            self.proj.bias.data.fill_(0)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj_drop = nn.Dropout(proj_drop)

        self.softmax = nn.Softmax(dim=-1)

        if dim_text is not None and space_attn:
            self.qkv_text_i2t = nn.Linear(dim_text, dim * 2, bias=qkv_bias)
            self.qkv_i2t = nn.Linear(dim, dim, bias=qkv_bias)
            self.attn_drop_i2t = nn.Dropout(attn_drop)
            self.proj_i2t = nn.Linear(dim, dim)
            self.proj_drop_i2t = nn.Dropout(proj_drop)
            self.alpha_i2t = nn.Parameter(torch.Tensor([0]))
            self.norm_i2t_i = norm_layer(dim)

    def forward(self, x, einops_from, einops_to, y=None, y_mask=None, **einops_dims):
        h = self.num_heads
        # project x to q, k, v vaalues
        q, k, v = self.qkv(x).chunk(3, dim=-1)
        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))

        q = q*self.scale

        # splice out CLS token at index 1
        (cls_q, q_), (cls_k, k_), (cls_v, v_) = map(lambda t: (t[:, 0:1], t[:, 1:]), (q, k, v))

        # let CLS token attend to key / values of all patches across time and space
        cls_out = attn(cls_q, k, v)
        # rearrange across time or space
        q_, k_, v_ = map(lambda t: rearrange(t, f'{einops_from} -> {einops_to}', **einops_dims), (q_, k_, v_))

        # expand cls token keys and values across time or space and concat
        r = q_.shape[0] // cls_k.shape[0]
        cls_k, cls_v = map(lambda t: repeat(t, 'b () d -> (b r) () d', r=r), (cls_k, cls_v))

        k_ = torch.cat((cls_k, k_), dim=1)
        v_ = torch.cat((cls_v, v_), dim=1)

        # attention
        out = attn(q_, k_, v_)

        # merge back time or space
        out = rearrange(out, f'{einops_to} -> {einops_from}', **einops_dims)

        # concat back the cls token
        out = torch.cat((cls_out, out), dim=1)

        # merge back the heads
        out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
        ## to out
        x = self.proj(out)
        x = self.proj_drop(x)

        if y is not None:
            B_, N, C = x.shape
            B_text, N_text, C_text = y.shape

            kv_text = (
                self.qkv_text_i2t(y)
                .reshape(B_text, N_text, 2, self.num_heads, C // self.num_heads)
                .permute(2, 0, 3, 1, 4)
            )
            k_text, v_text = kv_text[0], kv_text[1]

            q_i2t = self.qkv_i2t(self.norm_i2t_i(x))
            q_i2t = q_i2t.reshape(B_, N, 1, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
            q_i2t = q_i2t[0]

            # image to text attention
            text_scale = k_text.size(-1) ** -0.5
            q_i2t = q_i2t * text_scale
            attn_i2t = q_i2t @ k_text.transpose(-2, -1)  # B_, nH, N, N_text

            # add image to text bias and text_mask
            if y_mask is not None:
                mask_and_i2t_bias = y_mask.view(B_text, 1, 1, N_text)
                attn_i2t = attn_i2t + mask_and_i2t_bias

            attn_i2t = self.softmax(attn_i2t)
            attn_i2t = self.attn_drop_i2t(attn_i2t)
            y = (attn_i2t @ v_text).transpose(1, 2).reshape(B_, N, C)
            y = self.proj_i2t(y)
            y = self.proj_drop_i2t(y)
            x = x + self.alpha_i2t * y

        return x


class SpaceTimeBlock(nn.Module):

    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, time_init='zeros',
                 attention_style='frozen-in-time', dim_text=None):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.attn = VarAttention(
            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop, dim_text=dim_text, 
            norm_layer=norm_layer, space_attn=True)

        self.timeattn = VarAttention(
            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop,
            initialize=time_init, dim_text=dim_text, norm_layer=norm_layer, space_attn=False)

        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
        self.norm3 = norm_layer(dim)

        self.attention_style = attention_style

    def forward(self, x, einops_from_space, einops_to_space, einops_from_time, einops_to_time,
                time_n, space_f, y=None, y_mask=None):

        time_output = self.timeattn(self.norm3(x), einops_from_time, einops_to_time, n=time_n, y=None, y_mask=None)
        time_residual = x + time_output
        space_output = self.attn(self.norm1(time_residual), einops_from_space,
                                 einops_to_space, f=space_f, y=y, y_mask=y_mask)
        if self.attention_style == 'frozen-in-time':
            space_residual = x + self.drop_path(space_output)
        else:
            raise NotImplementedError

        x = space_residual + self.drop_path(self.mlp(self.norm2(space_residual)))

        return x


class SpaceTimeTransformer(nn.Module):
    """ Vision Transformer

    A PyTorch impl of : `Space-Time Transformer` from Frozen-in-time  - by Max Bain.
        https://arxiv.org/abs/2104.00650

    Based off:
     - ViT implementation from the timm library [https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py]
    lucidrains timesformer implementation [https://github.com/lucidrains/TimeSformer-pytorch].

    Notable differences:
     - allows for variable length input frames (<= num_frames)
     - allows for variable length input resolution  (<= (img_size, img_size)) [UNTESTED]
     - different attention block mechanism
    """

    def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12,
                 num_heads=12, mlp_ratio=4., qkv_bias=True, qk_scale=None, representation_size=None,
                 drop_rate=0., attn_drop_rate=0., drop_path_rate=0., hybrid_backbone=None,
                 num_frames=8, time_init='rand', attention_style='frozen-in-time', norm_layer=nn.LayerNorm, dim_text=None):
        """
        Args:
            img_size (int, tuple): input image size
            patch_size (int, tuple): patch size
            in_chans (int): number of input channels
            num_classes (int): number of classes for classification head
            embed_dim (int): embedding dimension
            depth (int): depth of transformer
            num_heads (int): number of attention heads
            mlp_ratio (int): ratio of mlp hidden dim to embedding dim
            qkv_bias (bool): enable bias for qkv if True
            qk_scale (float): override default qk scale of head_dim ** -0.5 if set
            representation_size (Optional[int]): enable and set representation layer (pre-logits) to this value if set
            drop_rate (float): dropout rate
            attn_drop_rate (float): attention dropout rate
            drop_path_rate (float): stochastic depth rate
            hybrid_backbone (nn.Module): CNN backbone to use in-place of PatchEmbed module
            norm_layer: (nn.Module): normalization layer
            num_frames: (int) maximum number of frames expected as input
            time_init: (str) how to initialise the time attention layer, 'zeros' allows for the timesformer to start off
                        as ViT.
            attention_style: (str) how to attend to space and time.
        """
        super().__init__()
        self.num_classes = num_classes
        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
        self.num_frames = num_frames
        self.embed_dim = embed_dim
        norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
        print("######USING ATTENTION STYLE: ", attention_style)
        if hybrid_backbone is not None:
            raise NotImplementedError('hybrid backbone not implemented')
        else:
            self.patch_embed = VideoPatchEmbed(
                img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim, num_frames=num_frames)
        num_patches = self.patch_embed.num_patches
        self.patches_per_frame = num_patches // num_frames

        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.pos_embed = nn.Parameter(
            torch.zeros(1, self.patches_per_frame + 1,
                        embed_dim))  # remember to take pos_embed[1:] for tiling over time
        self.temporal_embed = nn.Parameter(torch.zeros(1, num_frames, embed_dim))

        self.pos_drop = nn.Dropout(p=drop_rate)

        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
        self.blocks = nn.ModuleList([
            SpaceTimeBlock(
                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
                drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer, time_init=time_init,
                attention_style=attention_style, dim_text=None if i < 6 else DIM_TEXT)
            for i in range(depth)])
        self.norm = norm_layer(embed_dim)

        # Representation layer
        if representation_size:
            self.num_features = representation_size
            self.pre_logits = nn.Sequential(OrderedDict([
                ('fc', nn.Linear(embed_dim, representation_size)),
                ('act', nn.Tanh())
            ]))
        else:
            self.pre_logits = nn.Identity()

        # Classifier head
        self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()

        trunc_normal_(self.pos_embed, std=.02)
        trunc_normal_(self.cls_token, std=.02)

        # if num_frames > 1, then we perform ViT inflation and initialise time attention to zero so not necessary.
        if num_frames == 1:
            self.apply(self._init_weights)

        ## einops transformations
        self.einops_from_space = 'b (f n) d'
        self.einops_to_space = '(b f) n d'
        self.einops_from_time = 'b (f n) d'
        self.einops_to_time = '(b n) f d'


    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    @torch.jit.ignore
    def no_weight_decay(self):
        return {'pos_embed', 'cls_token'}

    def get_classifier(self):
        return self.head

    def reset_classifier(self, num_classes, global_pool=''):
        self.num_classes = num_classes
        self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()

    def forward_features(self, x):
        b, curr_frames, channels, _, _ = x.shape
        x = self.patch_embed(x)
        x = x.flatten(2).transpose(2, 1)
        x = x.reshape(b, -1, self.patch_embed.embed_dim)

        BF = x.shape[0]
        cls_tokens = self.cls_token.expand(BF, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        # positional embed needs to be tiled for each frame (this does [1,2,3] --> [1,2,3,1,2,3]...)
        cls_embed = self.pos_embed[:, 0, :].unsqueeze(1)
        tile_pos_embed = self.pos_embed[:, 1:, :].repeat(1, self.num_frames, 1)
        # temporal embed needs to be repeated within each frame (this does [1,2,3] --> [1,1,1,2,2,2,3,3,3]...)
        tile_temporal_embed = self.temporal_embed.repeat_interleave(self.patches_per_frame, 1)
        total_pos_embed = tile_pos_embed + tile_temporal_embed
        total_pos_embed = torch.cat([cls_embed, total_pos_embed], dim=1)

        curr_patches = x.shape[1]
        x = x + total_pos_embed[:, :curr_patches]
        x = self.pos_drop(x)
        n = self.patches_per_frame
        f = curr_frames

        for blk in self.blocks:
            if config_yaml['use_checkpoint']:

                def create_custom_forward(module):
                    def custom_forward(*inputs):
                        return module(*inputs, time_n=n, space_f=f)

                    return custom_forward

                x = torch.utils.checkpoint.checkpoint(create_custom_forward(blk), x, self.einops_from_space, self.einops_to_space, self.einops_from_time,
                    self.einops_to_time)
            else:
                x = blk(x, self.einops_from_space, self.einops_to_space, self.einops_from_time,
                    self.einops_to_time, time_n=n, space_f=f)

        x = self.norm(x)[:, 0]
        x = self.pre_logits(x)

        return x

    def forward(self, x):
        x = self.forward_features(x)
        x = self.head(x)
        return x

if __name__ == "__main__":
    network = SpaceTimeTransformer(num_frames=4)
    data = torch.rand((3, 4, 3, 224, 224))
    network(data)

QFVS/model/video_transformer.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import os
import logging
import pdb
from pathlib import Path
from functools import reduce
from operator import getitem
from datetime import datetime
# import pdb; pdb.set_trace()
from logger import setup_logging
from utils import read_json, write_json
import time
import inspect


class ConfigParser:
    def __init__(self, args, options='', timestamp=True, test=False, eval_mode=None):
        # parse default and custom cli options
        for opt in options:
            args.add_argument(*opt.flags, default=None, type=opt.type)
        args = args.parse_args()
        self.args = args
        if args.cuda_base:
            os.environ["CUDA_VISIBLE_DEVICES"] = args.cuda_base
        if args.resume is None:
            msg_no_cfg = "Configuration file need to be specified. Add '-c config.json', for example."
            assert args.config is not None, msg_no_cfg
            self.cfg_fname = Path(args.config)
            config = read_json(self.cfg_fname)
            self.resume = None
        else:
            self.resume = Path(args.resume)
            resume_cfg_fname = self.resume.parent / 'config.json'
            if eval_mode == "epic":
                resume_cfg_fname = Path('')
            if eval_mode == "charades":
                resume_cfg_fname = Path('')
            if eval_mode == "nlq":
                resume_cfg_fname = Path('')
            if eval_mode == "mq":
                resume_cfg_fname = Path('')

            config = read_json(resume_cfg_fname)
            if args.config is not None:
                config.update(read_json(Path(args.config)))

        # load config file and apply custom cli options
        self._config = _update_config(config, options, args)

        # set save_dir where trained model and log will be saved.
        save_dir = Path(args.save_dir)
        # timestamp = datetime.now().strftime(r'%m%d_%H%M%S') if timestamp else ''
        timestamp = datetime.now().strftime(r'%m%d_%H') if timestamp else ''

        exper_name = self.config['name']

        self._save_dir = save_dir / 'models' /  timestamp
        self._web_log_dir = save_dir / 'web' /  timestamp
        self._log_dir = save_dir / 'log' /  timestamp
        self._tf_dir = save_dir / 'tf' / timestamp

        if not test:
            self.save_dir.mkdir(parents=True, exist_ok=True)
            self.log_dir.mkdir(parents=True, exist_ok=True)
            self._tf_dir.mkdir(parents=True, exist_ok=True)

        # if set, remove all previous experiments with the current config
        if vars(args).get("purge_exp_dir", False):
            for dirpath in (self._save_dir, self._log_dir, self._web_log_dir):
                config_dir = dirpath.parent
                existing = list(config_dir.glob("*"))
                print(f"purging {len(existing)} directories from config_dir...")
                tic = time.time()
                os.system(f"rm -rf {config_dir}")
                print(f"Finished purge in {time.time() - tic:.3f}s")

        # save updated config file to the checkpoint dir
        if not test:
            write_json(self.config, self.save_dir / 'config.json')

            # configure logging module
            setup_logging(self.log_dir)
            self.log_levels = {
                0: logging.WARNING,
                1: logging.INFO,
                2: logging.DEBUG
            }

    def initialize(self, name, module, *args, index=None, **kwargs):
        """
        finds a function handle with the name given as 'type' in config, and returns the
        instance initialized with corresponding keyword args given as 'args'.
        """
        if index is None:
            module_name = self[name]['type']
            module_args = dict(self[name]['args'])
            assert all([k not in module_args for k in kwargs]), 'Overwriting kwargs given in config file is not allowed'
            module_args.update(kwargs)
        else:
            module_name = self[name][index]['type']
            module_args = dict(self[name][index]['args'])

        # if parameter not in config subdict, then check if it's in global config.
        signature = inspect.signature(getattr(module, module_name).__init__)
        print(module_name)
        for param in signature.parameters.keys():
            if param not in module_args and param in self.config:
                module_args[param] = self[param]
            if module_name == 'FrozenInTime' and param == 'args':
                module_args[param] = self.args
            if module_name == 'MultiDistTextVideoDataLoader' and param == 'args':
                module_args[param] = self.args

        return getattr(module, module_name)(*args, **module_args)

    def __getitem__(self, name):
        return self.config[name]

    def get_logger(self, name, verbosity=2):
        msg_verbosity = 'verbosity option {} is invalid. Valid options are {}.'.format(verbosity,
                                                                                       self.log_levels.keys())
        assert verbosity in self.log_levels, msg_verbosity
        logger = logging.getLogger(name)
        logger.setLevel(self.log_levels[verbosity])
        return logger

    # setting read-only attributes
    @property
    def config(self):
        return self._config

    @property
    def save_dir(self):
        return self._save_dir

    @property
    def log_dir(self):
        return self._log_dir

    @property
    def tf_dir(self):
        return self._tf_dir

# helper functions used to update config dict with custom cli options
def _update_config(config, options, args):
    for opt in options:
        value = getattr(args, _get_opt_name(opt.flags))
        if value is not None:
            _set_by_path(config, opt.target, value)
    return config


def _get_opt_name(flags):
    for flg in flags:
        if flg.startswith('--'):
            return flg.replace('--', '')
    return flags[0].replace('--', '')


def _set_by_path(tree, keys, value):
    """Set a value in a nested object in tree by sequence of keys."""
    _get_by_path(tree, keys[:-1])[keys[-1]] = value


def _get_by_path(tree, keys):
    """Access a nested object in tree by sequence of keys."""
    return reduce(getitem, keys, tree)

QFVS/parse_config.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import os
import h5py
import numpy as np
import torch
import transformers
from semantic_evaluation import calculate_semantic_matching, load_videos_tag
import pdb
import json
from collections import OrderedDict
from tqdm import tqdm

from torch.utils.data import Dataset, DataLoader
from utils_QFVS import load_pickle
import math
import torch.nn as nn
import torch.nn.init as init
from pathlib import Path
from model.model_summary import Transformer_Encoder


def weights_init(m):
    init.xavier_uniform_(m.weight, gain=np.sqrt(2.0))
    if m.bias is not None:
        init.constant_(m.bias, 0.1)


class Feature_Dataset(Dataset):
    def __init__(self, feature_list):

        self.feature_list = feature_list

    def __len__(self):
        return len(self.feature_list)

    def __getitem__(self, idx):

        feature = self.feature_list[idx]

        return feature['seg_len'], feature['fusedfeat_concept1_batched_list'], feature['fusedfeat_concept2_batched_list'], feature['fusedfeat_oracle_batched_list'], feature['concept1_GT'], feature['concept2_GT'], feature['mask_GT'], feature['oracle_summary']


class Runner_summary_transformer():
    def __init__(self,config, last_bias, train_videos, test_video, cuda_base, device_ids, feature_dict):
        self.config=config
        self.last_bias=last_bias
        self.train_videos = train_videos
        self.test_video = [test_video] #convert int to list
        self.cuda_base = cuda_base
        self.device_ids = device_ids
        self.feature_dict = feature_dict
        self._build_dataloader()
        self._bulid_model()
        self._build_optimizer()
        self.max_f1=0
        self.max_p=0
        self.max_r=0

    def _bulid_model(self):

        self.model_summary = Transformer_Encoder(self.config['model_dim'], self.config['nhead'], self.config['num_layers'])

    def _build_dataset(self):

        self.feature_list = []
        for video_id in self.train_videos:
            self.feature_list += self.feature_dict[str(video_id)]
        
        return Feature_Dataset(self.feature_list)

    def _build_dataloader(self):
        dataset=self._build_dataset()
        self.dataloader=DataLoader(dataset, batch_size=self.config["batch_size"], shuffle=True, num_workers=self.config["num_workers"])
        print(len(self.dataloader))

    def _build_optimizer(self):
        param_weights = []
        param_biases = []
        for param in self.model_summary.parameters():
            if param.ndim == 1:
                param_biases.append(param)
            else:
                param_weights.append(param)
        parameters = [{'params': param_weights}, {'params': param_biases}]
        self.optimizer = torch.optim.AdamW(parameters, lr=self.config["lr"], weight_decay=0.1)

    def output(self):
        print(" max_p = ",self.max_p," max_r = ",self.max_r," max_f1 = ",self.max_f1)

    def adjust_learning_rate(self, config, optimizer, loader, step):
        max_steps = config["epoch"] * len(loader)
        warmup_steps = 2 * len(loader)
        base_lr = 0.00001
        if step < warmup_steps:
            lr = base_lr * step / warmup_steps
        else:
            step -= warmup_steps
            max_steps -= warmup_steps
            q = 0.5 * (1 + math.cos(math.pi * step / max_steps))
            end_lr = base_lr * 0.001
            lr = base_lr * q + end_lr * (1 - q)
        optimizer.param_groups[0]['lr'] = lr * config["learning_rate_weights"]
        optimizer.param_groups[1]['lr'] = lr * config["learning_rate_biases"]


    def train(self):

        device = torch.device(self.cuda_base if torch.cuda.is_available() else 'cpu')
        criterion=torch.nn.BCEWithLogitsLoss()
        self.model_summary.summ_head.apply(weights_init)
        torch.nn.init.normal_(self.model_summary.summ_head.weight, mean=0.0, std=0.1)
        torch.nn.init.constant_(self.model_summary.summ_head.bias, self.last_bias)
        
        self.model_summary = self.model_summary.to(device)
        self.model_summary = nn.DataParallel(self.model_summary, device_ids = self.device_ids)
 
        self.model_summary.train()
        scaler = torch.cuda.amp.GradScaler()

        self.evaluate(self.test_video,self.config["top_percent"],0)

        for epoch in range(self.config["epoch"]):

            step = epoch*len(self.dataloader)
            batch_count=0
            
            for _, (seg_len, fusedfeat_concept1_batched_list, fusedfeat_concept2_batched_list, fusedfeat_oracle_batched_list, concept1_GT, concept2_GT, mask_GT, oracle_summary) in enumerate(tqdm(self.dataloader)):
                
                train_num=seg_len.shape[0]
                batch_count+=1

                self.optimizer.zero_grad()

                mask=torch.zeros(train_num,self.config["max_segment_num"],self.config["max_frame_num"],dtype=torch.bool).cuda()
                for i in range(train_num):
                    for j in range(len(seg_len[i])):
                        for k in range(seg_len[i][j]):
                            mask[i][j][k]=1


                with torch.cuda.amp.autocast():
                    
                    concept1_score = self.model_summary(fusedfeat_concept1_batched_list, seg_len)
                    concept2_score = self.model_summary(fusedfeat_concept2_batched_list, seg_len)
                    oracle_score = self.model_summary(fusedfeat_oracle_batched_list, seg_len)
                    
                    loss=torch.zeros(1).to(device)

                    for i in range(train_num):
                        concept1_score_tmp=concept1_score[i].masked_select(mask[i]).unsqueeze(0)
                        concept2_score_tmp=concept2_score[i].masked_select(mask[i]).unsqueeze(0)
                        concept1_GT_tmp=concept1_GT[i].masked_select(mask_GT[i]).unsqueeze(0)
                        concept2_GT_tmp=concept2_GT[i].masked_select(mask_GT[i]).unsqueeze(0)

                        score_tmp = oracle_score[i].masked_select(mask[i]).unsqueeze(0)
                        oracle_summary_sample = oracle_summary[i][0:score_tmp.shape[1]].unsqueeze(0)

                        loss_concept1=criterion(concept1_score_tmp,concept1_GT_tmp.to(device))
                        loss_concept2=criterion(concept2_score_tmp,concept2_GT_tmp.to(device))
                        loss_oracle=criterion(score_tmp, oracle_summary_sample.to(device))
                        loss+=loss_concept1+loss_concept2+loss_oracle

                self.adjust_learning_rate(self.config, self.optimizer, self.dataloader, step)
                scaler.scale(loss).backward()
                scaler.step(self.optimizer)
                scaler.update()
                step+=1


            self.evaluate(self.test_video,self.config["top_percent"],epoch)

    def evaluate(self,test_video,top_percent,epoch):

        f1_sum=0
        p_sum=0
        r_sum=0

        evaluation_num=0


        test_feature_list = []
        for video_id in test_video:
            test_feature_list += self.feature_dict[str(video_id)]

        test_dataset = Feature_Dataset(test_feature_list)
        test_dataloader=DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=0)
        self.model_summary.eval()
        
        evaluation_num=len(test_dataloader)

        for _, (seg_len, fusedfeat_concept1_batched_list, fusedfeat_concept2_batched_list, fusedfeat_oracle_batched_list, concept1_GT, concept2_GT, mask_GT, oracle_summary) in enumerate(test_dataloader):

            summaries_GT = torch.nonzero(oracle_summary[0:1,:], as_tuple=True)
            summaries_GT = summaries_GT[1].numpy().tolist()

            mask=torch.zeros(1,self.config["max_segment_num"],self.config["max_frame_num"],dtype=torch.bool).cuda()
            for i in range(1):
                for j in range(len(seg_len[i])):
                    for k in range(seg_len[i][j]):
                        mask[i][j][k]=1

            score = self.model_summary(fusedfeat_oracle_batched_list, seg_len)

            score=score.masked_select(mask)

            _,top_index=score.topk(int(score.shape[0]*top_percent))


            video_shots_tag = load_videos_tag(mat_path="./Tags.mat")
            p, r, f1 = calculate_semantic_matching(list(top_index.cpu().numpy()), summaries_GT, video_shots_tag, video_id=video_id-1)

            f1_sum+=f1
            r_sum+=r
            p_sum+=p

        if f1_sum/evaluation_num>self.max_f1:
            self.max_f1=f1_sum/evaluation_num
            self.max_p=p_sum/evaluation_num
            self.max_r=r_sum/evaluation_num

        print("p = ",p_sum/evaluation_num," r = ",r_sum/evaluation_num," f1 = ",f1_sum/evaluation_num)



QFVS/runner_train.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


from .cpd_auto import cpd_auto

QFVS/segment/__init__.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import numpy as np
from .cpd_nonlin import cpd_nonlin

def cpd_auto(K, ncp, vmax, desc_rate=1, **kwargs):
    """Main interface
    
    Detect change points automatically selecting their number
        K       - kernel between each pair of frames in video
        ncp     - maximum ncp
        vmax    - special parameter
    Optional arguments:
        lmin     - minimum segment length
        lmax     - maximum segment length
        desc_rate - rate of descriptor sampling (vmax always corresponds to 1x)

    Note:
        - cps are always calculated in subsampled coordinates irrespective to
            desc_rate
        - lmin and m should be in agreement
    ---
    Returns: (cps, costs)
        cps   - best selected change-points
        costs - costs for 0,1,2,...,m change-points
        
    Memory requirement: ~ (3*N*N + N*ncp)*4 bytes ~= 16 * N^2 bytes
    That is 1,6 Gb for the N=10000.
    """
    m = ncp
    (_, scores) = cpd_nonlin(K, m, backtrack=False, **kwargs)
    # print("scores ",scores)
    
    N = K.shape[0]
    N2 = N*desc_rate  # length of the video before subsampling
    
    penalties = np.zeros(m+1)
    # Prevent division by zero (in case of 0 changes)
    ncp = np.arange(1, m+1)
    penalties[1:] = (vmax*ncp/(2.0*N2))*(np.log(float(N2)/ncp)+1)
    
    costs = scores/float(N) + penalties
    m_best = np.argmin(costs)
    # print("cost ",costs)
    # print("m_best ",m_best)
    (cps, scores2) = cpd_nonlin(K, m_best, **kwargs)

    return (cps, costs)
    

# ------------------------------------------------------------------------------
# Extra functions (currently not used)

def estimate_vmax(K_stable):
    """K_stable - kernel between all frames of a stable segment"""
    n = K_stable.shape[0]
    vmax = np.trace(centering(K_stable)/n)
    return vmax


def centering(K):
    """Apply kernel centering"""
    mean_rows = np.mean(K, 1)[:, np.newaxis]
    return K - mean_rows - mean_rows.T + np.mean(mean_rows)


def eval_score(K, cps):
    """ Evaluate unnormalized empirical score
        (sum of kernelized scatters) for the given change-points """
    N = K.shape[0]
    cps = [0] + list(cps) + [N]
    V1 = 0
    V2 = 0
    for i in range(len(cps)-1):
        K_sub = K[cps[i]:cps[i+1], :][:, cps[i]:cps[i+1]]
        V1 += np.sum(np.diag(K_sub))
        V2 += np.sum(K_sub) / float(cps[i+1] - cps[i])
    return (V1 - V2)


def eval_cost(K, cps, score, vmax):
    """ Evaluate cost function for automatic number of change points selection
    K      - kernel between all frames
    cps    - selected change-points
    score  - unnormalized empirical score (sum of kernelized scatters)
    vmax   - vmax parameter"""
    
    N = K.shape[0]
    penalty = (vmax*len(cps)/(2.0*N))*(np.log(float(N)/len(cps))+1)
    return score/float(N) + penalty


QFVS/segment/cpd_auto.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import numpy as np

def calc_scatters(K):
    n = K.shape[0]
    K1 = np.cumsum([0] + list(np.diag(K)))
    K2 = np.zeros((n+1, n+1))
    K2[1:, 1:] = np.cumsum(np.cumsum(K, 0), 1) # TODO: use the fact that K - symmetric

    scatters = np.zeros((n, n))
    
    
    for i in range(n):
        for j in range(i, n):
            scatters[i,j] = K1[j+1] - K1[i] - (K2[j+1,j+1]+K2[i,i]-K2[j+1,i]-K2[i,j+1])/(j-i+1)
    return scatters

def cpd_nonlin(K, ncp, lmin=1, lmax=100000, backtrack=True, verbose=True,
    out_scatters=None):
    """ Change point detection with dynamic programming
    K - square kernel matrix 
    ncp - number of change points to detect (ncp >= 0)
    lmin - minimal length of a segment
    lmax - maximal length of a segment
    backtrack - when False - only evaluate objective scores (to save memory)
    
    Returns: (cps, obj)
        cps - detected array of change points: mean is thought to be constant on [ cps[i], cps[i+1] )    
        obj_vals - values of the objective function for 0..m changepoints
        
    """   
    m = int(ncp)  # prevent numpy.int64
    
    (n, n1) = K.shape
    assert(n == n1), "Kernel matrix awaited."    
    
    assert(n >= (m + 1)*lmin)
    assert(n <= (m + 1)*lmax)
    assert(lmax >= lmin >= 1)
    
    if verbose:
        #print "n =", n
        print("Precomputing scatters...")
    J = calc_scatters(K)
    
    if out_scatters != None:
        out_scatters[0] = J

    if verbose:
        print("Inferring best change points...")
    I = 1e101*np.ones((m+1, n+1))
    I[0, lmin:lmax] = J[0, lmin-1:lmax-1]

    if backtrack:
        p = np.zeros((m+1, n+1), dtype=int)
    else:
        p = np.zeros((1,1), dtype=int)
        
    
    for k in range(1, m+1):
        for l in range((k+1)*lmin, n+1):
            I[k, l] = 1e100
            for t in range(max(k*lmin,l-lmax), l-lmin+1):
                c = I[k-1, t] + J[t, l-1]
                if (c < I[k, l]):
                    I[k, l] = c
                    if (backtrack == 1):
                        p[k, l] = t
    
    
    # Collect change points
    cps = np.zeros(m, dtype=int)
    
    if backtrack:
        cur = n
        for k in range(m, 0, -1):
            cps[k-1] = p[k, cur]
            cur = cps[k-1]

    scores = I[:, n].copy() 
    scores[scores > 1e99] = np.inf
    return cps, scores
    


QFVS/segment/cpd_nonlin.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import networkx as nx
from sklearn.metrics import pairwise_distances
import numpy as np
import scipy.io

def process_video_mat(video_mat):
    result = []
    for shot_vec in video_mat:
        shot_vec= shot_vec[0][0]
        result.append(shot_vec)
    result = np.array(result)
    return result


def process_mat(mat):
    videos = mat['Tags'][0]
    result = []
    for video_mat in videos:
        video_mat = video_mat[0]
        video_data = process_video_mat(video_mat)
        result.append(video_data)
    return result


def load_videos_tag(mat_path="./data/ute_query/Tags.mat"):
    mat = scipy.io.loadmat(mat_path)
    videos_tag = process_mat(mat)
    return videos_tag

def semantic_iou(a, b):
    intersection = a * b
    intersection_num = sum(intersection)
    union = a + b
    union[union>0] = 1
    union_num = sum(union)
    if union_num != 0:
        return intersection_num / union_num
    else:
        return 0


def build_graph_from_pariwise_weights(weight_matrix):
    B = nx.Graph()
    bottom_nodes = list(map(lambda x: "b-{}".format(x), list(range(weight_matrix.shape[0]))))
    top_nodes = list(map(lambda x: "t-{}".format(x), list(range(weight_matrix.shape[1]))))
    edges = []
    for i in range(weight_matrix.shape[0]):
        for j in range(weight_matrix.shape[1]):
            weight = weight_matrix[i][j]
            edges.append(("b-{}".format(i), "t-{}".format(j), weight))
    B.add_weighted_edges_from(edges)
    return B


def calculate_semantic_matching(machine_summary, gt_summary, video_shots_tag, video_id):
    video_shots_tag = video_shots_tag[video_id]
    machine_summary_mat = video_shots_tag[machine_summary]
    gt_summary_mat = video_shots_tag[gt_summary]
    weights = pairwise_distances(machine_summary_mat, gt_summary_mat, metric=semantic_iou)
    B = build_graph_from_pariwise_weights(weights)
    matching_edges = nx.algorithms.matching.max_weight_matching(B)
    sum_weights = 0
    i = 0
    for edge in matching_edges:
        edge_data = B.get_edge_data(edge[0], edge[1])
        sum_weights += edge_data['weight']
        i += 1
    precision = sum_weights / machine_summary_mat.shape[0]
    recall = sum_weights / gt_summary_mat.shape[0]
    f1 = 2 * precision * recall / (precision + recall)
    return precision, recall, f1


if __name__=='__main__':
    video_shots_tag = load_videos_tag(mat_path="./Tags.mat")
    machine_summary = [1, 39, 99, 31, 778]
    gt_summary = [1, 34, 101, 29, 774]
    print(calculate_semantic_matching(machine_summary, gt_summary, video_shots_tag, video_id=0))

QFVS/semantic_evaluation.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

from .util import *

QFVS/utils/__init__.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import numbers
from typing import List, Tuple

import torch
from torch import Tensor
from torchvision.transforms import functional_pil as F_pil, functional_tensor as F_t
from torchvision.transforms.functional import center_crop, crop


def _get_image_size(img: Tensor) -> List[int]:
    """Returns image size as [w, h]
    """
    if isinstance(img, torch.Tensor):
        return F_t._get_image_size(img)

    return F_pil._get_image_size(img)

def center_plus_four_crops(img: Tensor, size: List[int],
                              margin_h: int, margin_w: int) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
    """Crop the given image into four tiled borders and the central crop.
    """

    if isinstance(size, numbers.Number):
        size = (int(size), int(size))
    elif isinstance(size, (tuple, list)) and len(size) == 1:
        size = (size[0], size[0])

    if len(size) != 2:
        raise ValueError("Please provide only two dimensions (h, w) for size.")

    image_width, image_height = _get_image_size(img)

    crop_height, crop_width = size

    if crop_width > image_width or crop_height > image_height:
        msg = "Requested crop size {} is bigger than input size {}"
        raise ValueError(msg.format(size, (image_height, image_width)))

    if crop_width + margin_w > image_width:
        msg = "Requested margin size {} + input {} is bigger than input size {}"
        raise ValueError(msg.format((margin_h, margin_w), size, (image_height, image_width)))

    #vertical_border_height = image_height - crop_height
    #horizontal_border_height = image_width - crop_width

    #x1 = horizontal_border_height // 2
    x11 = (image_width - crop_width - 2 * margin_w) // 2
    x12 = x11 + margin_w
    x21 = x12 + crop_width
    x22 = x21 + margin_w

    y11 = (image_height - crop_height - 2 * margin_h) // 2
    y12 = y11 + margin_h
    y21 = y12 + crop_height
    y22 = y21 + margin_h

    tl = crop(img, y11, x11, margin_h, margin_w + crop_width)
    tr = crop(img, y11, x21, margin_h + crop_height, margin_w)
    bl = crop(img, y12, x11, margin_h + crop_height, margin_w)
    br = crop(img, y21, x12, margin_h, margin_w + crop_width)
    center = center_crop(img, [crop_height, crop_width])

    return tl, tr, bl, br, center



def center_plus_twohori_crops(img: Tensor, size: List[int],
                              margin_w: int) -> Tuple[Tensor, Tensor, Tensor]:
    """Crop the given image into four tiled borders and the central crop.
    """

    if isinstance(size, numbers.Number):
        size = (int(size), int(size))
    elif isinstance(size, (tuple, list)) and len(size) == 1:
        size = (size[0], size[0])

    if len(size) != 2:
        raise ValueError("Please provide only two dimensions (h, w) for size.")

    image_width, image_height = _get_image_size(img)

    crop_height, crop_width = size

    if crop_width > image_width or crop_height > image_height:
        msg = "Requested crop size {} is bigger than input size {}"
        raise ValueError(msg.format(size, (image_height, image_width)))

    if crop_width + margin_w > image_width :
        msg = "Requested margin size {} + input {} is bigger than input size {}"
        raise ValueError(msg.format((0, margin_w), size, (image_height, image_width)))

    # vertical_border_height = image_height - crop_height
    # horizontal_border_height = image_width - crop_width

    # x1 = horizontal_border_height // 2
    x11 = (image_width - crop_width - 2 * margin_w) // 2
    x12 = x11 + margin_w
    x21 = x12 + crop_width

    y11 = (image_height - crop_height) // 2

    left = crop(img, y11, x11, crop_height, margin_w)
    right = crop(img, y11, x21, crop_height, margin_w)
    center = center_crop(img, [crop_height, crop_width])

    return left, right, center

from torch import nn
class TwoHoriCrop(nn.Module):
    def __init__(self, size, margin_w):
        super().__init__()
        self.size = size
        self.margin_w = margin_w

    def forward(self, x):
        return center_plus_twohori_crops(x, self.size, self.margin_w)

if __name__ == "__main__":
    from PIL import Image

    img = Image.open('visualisations/guitar.png')
    crops = center_plus_four_crops(img, [336, 336], 112, 112)
    order = ['tl', 'tr', 'bl', 'br', 'center']

    for idx, subimg in zip(order, crops):
        subimg.save(f'visualisations/guitar_{idx}.png')

    crops = center_plus_twohori_crops(img, [448, 448], 112)
    order = ['left', 'right', 'center2']

    for idx, subimg in zip(order, crops):
        subimg.save(f'visualisations/guitar_{idx}.png')

QFVS/utils/custom_transforms.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import os

import dominate
from dominate.tags import a, attr, br, h3, img, meta, p, source, span, table, td, tr, video


class HTML:
    """This HTML class allows us to save images and write texts into a single HTML file.

     It consists of functions such as <add_header> (add a text header to the HTML file),
     <add_images> (add a row of images to the HTML file), and <save> (save the HTML to the disk).
     It is based on Python library 'dominate', a Python library for creating and
     manipulating HTML documents using a DOM API.
    """

    def __init__(self, web_dir, title, refresh=0):
        """Initialize the HTML classes

        Parameters:
            web_dir (str) -- a directory that stores the webpage. HTML file will be
            created at <web_dir>/index.html; images will be saved at <web_dir/images/
            title (str)   -- the webpage name
            reflect (int) -- how often the website refresh itself; if 0; no refreshing
        """
        self.title = title
        self.web_dir = web_dir
        self.img_dir = os.path.join(self.web_dir, "images")
        if not os.path.exists(self.web_dir):
            os.makedirs(self.web_dir)
        if not os.path.exists(self.img_dir):
            os.makedirs(self.img_dir)

        self.doc = dominate.document(title=title)
        if refresh > 0:
            with self.doc.head:
                meta(http_equiv="refresh", content=str(refresh))

    def get_image_dir(self):
        """Return the directory that stores images"""
        return self.img_dir

    def add_header(self, text):
        """Insert a header to the HTML file

        Parameters:
            text (str) -- the header text
        """
        with self.doc:
            h3(text)

    def add_videos(self, vids, txts, links, width=400, hidden_tag="hidden"):
        """add images to the HTML file

        Parameters:
            vids (str list)   -- a list of image paths
            txts (str list)  -- a list of image names shown on the website
            links (str list) --  a list of hyperref links; when you click an image,
            it will redirect you to a new page
        """
        self.t = table(border=1, style="table-layout: fixed;")  # Insert a table
        self.doc.add(self.t)
        colors = ["red", "blue", "gold", "salman"]
        with self.t:
            with tr():
                for vid, txt, link in zip(vids, txts, links):
                    td_style = "word-wrap: break-word; width:{}px".format(width)
                    with td(style=td_style, halign="center", valign="top"):
                        with p():
                            vid_path = str(vid)
                            if vid_path == hidden_tag:
                                p_style = "font-weight: bold; width:{}px;"
                                p_style = p_style.format(width * 3)
                                p("hidden video", style=p_style)
                            else:
                                with a(href=str(link)):
                                    with video():
                                        attr(controls="controls")
                                        source(src=vid_path, type="video/mp4")
                            br()
                            rows = txt.split("<br>")
                            for idx, row in enumerate(rows):
                                color = colors[idx % len(colors)]
                                bold_tag = "<b>"
                                if not row.startswith(bold_tag):
                                    s_style = "color:{};".format(color)
                                else:
                                    s_style = "color:black; font-weight: bold;"
                                    row = row[len(bold_tag):]
                                span(row, style=s_style)

    def add_images(self, ims, txts, links, width=400):
        """add images to the HTML file

        Parameters:
            ims (str list)   -- a list of image paths
            txts (str list)  -- a list of image names shown on the website
            links (str list) --  a list of hyperref links; when you click an image,
            it will redirect you to a new page
        """
        self.t = table(border=1, style="table-layout: fixed;")  # Insert a table
        self.doc.add(self.t)
        colors = ["red", "blue", "gold", "salman"]
        with self.t:
            with tr():
                for im, txt, link in zip(ims, txts, links):
                    td_style = "word-wrap: break-word;"
                    with td(style=td_style, halign="center", valign="top"):
                        with p():
                            with a(href=link):
                                img(
                                    style="width:%dpx" % width,
                                    src=im,
                                )
                            br()
                            rows = txt.split("<br>")
                            for idx, row in enumerate(rows):
                                color = colors[idx % len(colors)]
                                bold_tag = "<b>"
                                if not row.startswith(bold_tag):
                                    s_style = "color:{};".format(color)
                                else:
                                    s_style = "color:black; font-weight: bold;"
                                    row = row[len(bold_tag):]
                                span(row, style=s_style)

    def save(self):
        """save the current content to the HMTL file"""
        html_file = "%s/index.html" % self.web_dir
        f = open(html_file, "wt")
        f.write(self.doc.render())
        f.close()


if __name__ == "__main__":  # we show an example usage here.
    html = HTML("web/", "test_html")
    html.add_header("hello world")

    ims, txts, links = [], [], []
    for n in range(4):
        ims.append("image_%d.png" % n)
        txts.append("text_%d" % n)
        links.append("image_%d.png" % n)
    html.add_images(ims, txts, links)
    html.save()

QFVS/utils/html.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import numpy as np


def calculate_mAP(sim_mat, relevancy_matrix):
    """
    Computes the mean average precision according to the following formula of
    average precision:
    \frac{\sum_{k=1}^n p(k) x rel(k)}{num_rel_docs}

    where p(k) is the precision at k, rel(k) is an indicator function
    determining whether the kth returned item is relevant or not and
    num_rel_docs is the number of relevant items to find within the search.

    The mean average precision is the mean of the average precision for each
    query item (i.e row in the matrix)

    This function takes in two parameters:
        - sim_mat: a NxM matrix which represents the similarity between two
        modalities (with modality 1 being of size N and modality 2 of size M).
        - relevancy_matrix: an NxM matrix which represents the relevancy between two
        modalities of items (with modality 1 being of size N and modality 2 of
        size M).
    """
    #Find the order of the items in modality 2 according to modality 1
    ranked_order = (-sim_mat).argsort()
    ranked_sim_mat = sim_mat[np.arange(sim_mat.shape[0])[:, None], ranked_order]
    #re-order the relevancy matrix to accommodate the proposals
    ranked_rel_mat = relevancy_matrix[np.arange(relevancy_matrix.shape[0])[:, None], ranked_order]

    #find the number of relevant items found at each k
    cumulative_rel_mat = np.cumsum(ranked_rel_mat, axis=1)
    #Mask this ensuring that it is non zero if the kth term is 1 (rel(k) above)
    cumulative_rel_mat[ranked_rel_mat != 1] = 0
    #find the divisor for p(k)
    divisor = np.arange(ranked_rel_mat.shape[1]) + 1

    #find the number of relevant docs per query item
    number_rel_docs = np.sum(ranked_rel_mat==1, axis=1)

    #find the average precision per query, within np.sum finds p(k) * rel(k)
    avg_precision = np.sum(cumulative_rel_mat / divisor, axis=1) / number_rel_docs
    mAP = np.mean(avg_precision)
    return mAP

QFVS/utils/mAP.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import numpy as np

def calculate_DCG(similarity_matrix, relevancy_matrix, k_counts):
    """
    Calculates the Discounted Cumulative Gain (DCG) between two modalities for
    the first modality.
    DCG = \sum_{i=1}^k \frac{rel_i}{log_2(i + 1)}
    i.e. the sum of the k relevant retrievals which is calculated as the scaled
    relevancy for the ith item. The scale is designed such that early
    retrievals are more important than later retrievals.
    Params:
        - similarity_matrix: matrix of size n1 x n2 where n1 is the number of
          items in the first modality and n2 is the number of items in the
          second modality. The [ith,jth] element is the predicted similarity
          between the ith item from the first modality and the jth item from
          the second modality.
        - relevancy_matrix: matrix of size n1 x n2 (see similarity_matrix
          above). The [ith, jth] element is the semantic relevancy between the
          ith item from the first modality and the jth item from the second
          modality.
        - k_counts: matrix of size n1 x n2 (see similarity_matrix above) which
          includes information on which items to use to calculate the DCG for
          (see calculate_k_counts for more info on this matrix).
    Returns:
        - The DCG for each item in the first modality, a n1 length vector.
    """
    x_sz, y_sz = similarity_matrix.shape
    ranks = np.argsort(similarity_matrix)[:, ::-1]
    #Create vector of size (n,) where n is the length of the last dimension in
    #similarity matrix
    #This vector is of the form log(i+1)
    logs = np.log2(np.arange(y_sz) + 2)
    #Convert logs into the divisor for the DCG calculation, of size similarity
    #matrix
    divisors = np.repeat(np.expand_dims(logs, axis=0), x_sz, axis=0)

    #mask out the sorted relevancy matrix to only use the first k relevant
    #retrievals for each item.
    columns = np.repeat(np.expand_dims(np.arange(x_sz), axis=1), y_sz, axis=1)
    numerators = relevancy_matrix[columns, ranks] * k_counts
    #Calculate the final DCG score (note that this isn't expected to sum to 1)
    return np.sum(numerators / divisors, axis=1)

def calculate_k_counts(relevancy_matrix):
    """
    Works out the maximum number of allowed retrievals when working out the
    Discounted Cumulative Gain. For each query the DCG only uses the first k
    items retrieved which constitute the k relevant items for that query
    (otherwise the nDCG scores can be deceptively high for bad rankings).
    Params:
        - relevancy_matrix: matrix of size n1 x n2 where n1 is the number of
          items in the first modality and n2 is the number of items in the
          second modality.  The [ith, jth] element is the semantic relevancy
          between the ith item from the first modality and the jth item from
          the second modality.
    Returns:
        - Matrix of size n1 x n2 (see relevancy matrix for more info). This is
          created as a mask such that if the [ith, jth] element is 1 it
          represents a valid item to use for the calculation of DCG for the
          ith item after sorting. For example, if relevancy matrix of:
        [[1, 0.5, 0],
          [0, 0  , 1]]
          is given, then the k_counts matrix will be:
        [[1, 1, 0],
         [1, 0, 0]]
         i.e. the first row has 2 non-zero items, so the first two retrieved
         items should be used in the calculation. In the second row there is
         only 1 relevant item, therefore only the first retrieved item should
         be used for the DCG calculation.
    """
    return (np.sort(relevancy_matrix)[:, ::-1] > 0).astype(int)


def calculate_IDCG(relevancy_matrix, k_counts):
    """
    Calculates the Ideal Discounted Cumulative Gain (IDCG) which is the value
    of the Discounted Cumulative Gain (DCG) for a perfect retrieval, i.e. the
    items in the second modality were retrieved in order of their descending
    relevancy.
    Params:
        - relevancy_matrix: matrix of size n1 x n2 where n1 is the number of
          items in the first modality and n2 is the number of items in the
          second modality. The [ith, jth] element is the semantic relevancy
          between the ith item from the first modality and the jth item from
          the second modality.
        - k_counts: matrix of size n1 x n2 (see similarity_matrix above) which
          includes information on which items to use to calculate the DCG for
          (see calculate_k_counts for more info on this matrix).
    """
    return calculate_DCG(relevancy_matrix, relevancy_matrix, k_counts)

def calculate_nDCG(similarity_matrix, relevancy_matrix, k_counts=None, IDCG=None, reduction='mean'):
    """
    Calculates the normalised Discounted Cumulative Gain (nDCG) between two
    modalities for the first modality using the Discounted Cumulative Gain
    (DCG) and the Ideal Discounted Cumulative Gain (IDCG).

    nDCG = \frac{DCG}{IDCG}

    Params:
        - similarity_matrix: matrix of size n1 x n2 where n1 is the number of
          items in the first modality and n2 is the number of items in the second
          modality. The [ith,jth] element is the predicted similarity between
          the ith item from the first modality and the jth item from the second
          modality.
        - relevancy_matrix: matrix of size n1 x n2 (see similarity_matrix
          above). The [ith, jth] element is the semantic relevancy between the
          ith item from the first modality and the jth item from the second
          modality.
        - k_counts: optional parameter: matrix of size n1 x n2 (see
          similarity_matrix above) which includes information on which items to
          use to calculate the DCG for (see calculate_k_counts for more info on
          this matrix). This will be calculated using calculate_IDCG if not
          present, but should be pre-processed for efficiency.
        - IDCG: Optional parameter which includes the pre-processed Ideal
          Discounted Cumulative Gain (IDCG). This is a vector of size n1 (see
          similarity_matrix above) which contains the IDCG value for each item
          from the first modality. This will be calculated using calculate_IDCG
          if not present, but should be pre-processed for efficiency.
        - reduction: what to use to reduce the different nDCG scores. By
          default this applies np.mean across all different queries.
    Returns:
        - The nDCG values for the first modality.
    """
    if k_counts is None:
        k_counts = calculate_k_counts(relevancy_matrix)
    DCG = calculate_DCG(similarity_matrix, relevancy_matrix, k_counts)
    if IDCG is None:
        IDCG = calculate_IDCG(relevancy_matrix, k_counts)
    if reduction == 'mean':
        return np.mean(DCG / IDCG)
    elif reduction is None:
        return DCG / IDCG


if __name__ == '__main__':
    sim_matrix = np.array([
        [1.0, 0.7, 0.4, 0.0],
        [0.3, 0.9, 0.6, 0.1],
        [0.2, 0.5, 0.8, 0.4]
    ])
    rel_matrix = np.array([
        [1.0, 0.5, 0.25, 0.0],
        [0.0, 1.0, 0.4, 0.0],
        [0.5, 0.3, 1.0, 0.0]
    ])
    k_counts = np.array([
        [1, 1, 1, 0],
        [1, 1, 0, 0],
        [1, 1, 1, 0]
    ])
    assert (k_counts == calculate_k_counts(rel_matrix)).all()
    nDCG = calculate_nDCG(sim_matrix, rel_matrix, k_counts)
    assert nDCG == 0.9371789900735429
    DCG = calculate_DCG(sim_matrix, rel_matrix, k_counts)
    IDCG = calculate_IDCG(rel_matrix, k_counts)

    assert nDCG == np.mean(DCG / IDCG)

    pre_nDCG = calculate_nDCG(sim_matrix, rel_matrix, k_counts, IDCG=IDCG)
    assert pre_nDCG == nDCG

    post_mean_nDCG = calculate_nDCG(sim_matrix, rel_matrix, k_counts, IDCG=IDCG, reduction=None)
    assert np.mean(post_mean_nDCG) == pre_nDCG

QFVS/utils/nDCG.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import functools
import json
import os
import socket
import time
from collections import OrderedDict
from datetime import datetime
from itertools import repeat
from pathlib import Path

import humanize
import numpy as np
import psutil


def replace_nested_dict_item(obj, key, replace_value):
    for k, v in obj.items():
        if isinstance(v, dict):
            obj[k] = replace_nested_dict_item(v, key, replace_value)
    if key in obj:
        obj[key] = replace_value
    return obj


def state_dict_data_parallel_fix(load_state_dict, curr_state_dict):
    load_keys = list(load_state_dict.keys())
    curr_keys = list(curr_state_dict.keys())

    redo_dp = False
    undo_dp = False
    if not curr_keys[0].startswith('module.') and load_keys[0].startswith('module.'):   # this
        undo_dp = True
    elif curr_keys[0].startswith('module.') and not load_keys[0].startswith('module.'):
        redo_dp = True

    if undo_dp: # this
        from collections import OrderedDict
        new_state_dict = OrderedDict()
        for k, v in load_state_dict.items():
            name = k[7:]  # remove `module.`
            new_state_dict[name] = v
        # load params
    elif redo_dp:
        from collections import OrderedDict
        new_state_dict = OrderedDict()
        for k, v in load_state_dict.items():
            name = 'module.' + k  # remove `module.`
            new_state_dict[name] = v
    else:
        new_state_dict = load_state_dict
    return new_state_dict

def print_numpy(x, val=True, shp=False):
    """Print the mean, min, max, median, std, and size of a numpy array
    Parameters:
        val (bool) -- if print the values of the numpy array
        shp (bool) -- if print the shape of the numpy array
    """
    x = x.astype(np.float64)
    if shp:
        print('shape,', x.shape)
    if val:
        x = x.flatten()
        print('mean = %3.3f, min = %3.3f, max = %3.3f, median = %3.3f, std=%3.3f' % (
            np.mean(x), np.min(x), np.max(x), np.median(x), np.std(x)))


def mkdirs(paths):
    """create empty directories if they don't exist
    Parameters:
        paths (str list) -- a list of directory paths
    """
    if isinstance(paths, list) and not isinstance(paths, str):
        for path in paths:
            mkdir(path)
    else:
        mkdir(paths)


def mkdir(path):
    """create a single empty directory if it didn't exist
    Parameters:
        path (str) -- a single directory path
    """
    if not os.path.exists(path):
        os.makedirs(path)

def read_json(fname):
    with fname.open('rt') as handle:
        return json.load(handle, object_hook=OrderedDict)

def write_json(content, fname):
    with fname.open('wt') as handle:
        json.dump(content, handle, indent=4, sort_keys=False)

def inf_loop(data_loader):
    ''' wrapper function for endless data loader. '''
    for loader in repeat(data_loader):
        yield from loader

def memory_summary():
    vmem = psutil.virtual_memory()
    msg = (
        f">>> Currently using {vmem.percent}% of system memory "
        f"{humanize.naturalsize(vmem.used)}/{humanize.naturalsize(vmem.available)}"
    )
    print(msg)

@functools.lru_cache(maxsize=64, typed=False)
def memcache(path):
    suffix = Path(path).suffix
    print(f"loading features >>>", end=" ")
    tic = time.time()
    if suffix == ".npy":
        res = np_loader(path)
    else:
        raise ValueError(f"unknown suffix: {suffix} for path {path}")
    print(f"[Total: {time.time() - tic:.1f}s] ({socket.gethostname() + ':' + str(path)})")
    return res

def np_loader(np_path, l2norm=False):
    with open(np_path, "rb") as f:
        data = np.load(f, encoding="latin1", allow_pickle=True)
    if isinstance(data, np.ndarray) and data.size == 1:
        data = data[()]  # handle numpy dict storage convnetion
    if l2norm:
        print("L2 normalizing features")
        if isinstance(data, dict):
            for key in data:
                feats_ = data[key]
                feats_ = feats_ / max(np.linalg.norm(feats_), 1E-6)
                data[key] = feats_
        elif data.ndim == 2:
            data_norm = np.linalg.norm(data, axis=1)
            data = data / np.maximum(data_norm.reshape(-1, 1), 1E-6)
        else:
            raise ValueError("unexpected data format {}".format(type(data)))
    return data


class Timer:
    def __init__(self):
        self.cache = datetime.now()

    def check(self):
        now = datetime.now()
        duration = now - self.cache
        self.cache = now
        return duration.total_seconds()

    def reset(self):
        self.cache = datetime.now()

QFVS/utils/util.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import random

import cv2
import numpy as np
import torch


def load_frames_from_video_path(path, num_frames, sample='rand'):
    cap = cv2.VideoCapture(path)
    assert (cap.isOpened())
    vlen = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    acc_samples = min(num_frames, vlen)
    intervals = np.linspace(start=0, stop=vlen, num=acc_samples + 1).astype(int)
    ranges = []
    for idx, interv in enumerate(intervals[:-1]):
        ranges.append((interv, intervals[idx + 1] - 1))
    if sample == 'rand':
        frame_idxs = [random.choice(range(x[0], x[1])) for x in ranges]
    elif sample == 'uniform':
        frame_idxs = [(x[0] + x[1]) // 2 for x in ranges]
    else:
        raise NotImplementedError

    frames = []
    for index in frame_idxs:
        cap.set(cv2.CAP_PROP_POS_FRAMES, index)
        ret, frame = cap.read()
        if ret:
            cv2.imwrite(f'images/{index}.jpg', frame)
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frame = torch.from_numpy(frame)
            # (H x W x C) to (C x H x W)
            frame = frame.permute(2, 0, 1)
            frames.append(frame)
        else:
            raise ValueError

    frames = torch.stack(frames).float() / 255
    cap.release()
    return frames, frame_idxs

QFVS/utils/video.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import os
import cv2
import shutil
import json
import sys
import multiprocessing
from multiprocessing import Pool
import subprocess
import pandas as pd
from csv import reader, writer

video_dir = './dataset/ego4d_256_multiprocessing'
output_dir = './dataset/ego4d_chunked_multiprocessing'

save_dir = os.path.join(output_dir)
if not os.path.exists(save_dir):
    os.makedirs(save_dir)

dur_limit = 600

def video2segments(infos):
    global count
    index, uid, dur = infos[0], infos[1], infos[2]
    input_path = os.path.join(video_dir, uid + '.mp4')

    output_uid_dir = os.path.join(output_dir, uid)
    if not os.path.exists(output_uid_dir):
        os.makedirs(output_uid_dir)

    #if index % num_partition != partition:
    #    return

    assert os.path.exists(input_path)

    cap = cv2.VideoCapture(input_path)
    rate = cap.get(5)
    frame_num = cap.get(7)
    duration = frame_num / rate

    if duration <= dur_limit:
        shutil.copyfile(input_path, os.path.join(output_uid_dir, '0.mp4'))
        return

    num_seg = duration // dur_limit

    s1_time = 0;
    s2_time = dur_limit
    num_finished = 0
    while num_finished <= num_seg:
        output_mp4_path = os.path.join(output_uid_dir, str(num_finished) + '.mp4')

        cmd = 'ffmpeg -y -i {} -ss {} -to {} -async 1 {}'.format(input_path, s1_time, s2_time,
                                                                                 output_mp4_path)
        # print(cmd)
        subprocess.call(cmd, shell=True)

        # Update for next steps
        s1_time = s2_time
        s2_time = min(s1_time + dur_limit, duration)
        num_finished += 1
    return

if __name__ == "__main__":
    with open('./dataset/ego4d/manifest.csv', 'r') as csv:
        csv_reader = list(reader(csv))[1:]

    downloaded = os.listdir('./dataset/ego4d')

    uid_list = []
    infos_list = []
    num_valid = 0
    for i_item, item in enumerate(csv_reader):
        uid, dur = item[0], float(item[2])
        existed = uid + '.mp4' in downloaded

        if not existed:
            continue

        uid_list.append(uid)
        infos_list.append([num_valid, uid, dur])
        num_valid += 1

    # for infos in infos_list:
    #     video2segments(infos)

    pool = Pool(multiprocessing.cpu_count())
    pool.map(video2segments, tuple(infos_list))

QFVS/utils/video_chunk.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import os
import time
import sys
import subprocess
import multiprocessing
from multiprocessing import Pool, Value

folder_path = './dataset/ego4d'
output_path = './dataset/ego4d_256_multiprocessing'

def videos_resize(videoinfos):
    global count

    videoid, videoname = videoinfos

    if os.path.exists(os.path.join(output_path, videoname)):
        print(f'{videoname} is resized.')
        return

    inname = folder_path + '/' + videoname
    outname = output_path + '/' + videoname

    cmd = "ffmpeg -y -i {} -filter:v scale=\"trunc(oh*a/2)*2:256\" -c:a copy {}".format( inname, outname)
    subprocess.call(cmd, shell=True)

    return

if __name__ == "__main__":
    file_list = []
    mp4_list = [item for item in os.listdir(folder_path) if item.endswith('.mp4')]
    for id, video in enumerate(mp4_list):
        file_list.append([id, video])

    pool = Pool(multiprocessing.cpu_count()) 
    pool.map(videos_resize, tuple(file_list))

    # for file in file_list:
    #    videos_resize(file)

QFVS/utils/video_resize.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import matplotlib
import numpy as np
import torch

matplotlib.use('Agg')


def visualise_path(pred, target, window):
    """
    :param pred: (P, 2) Tensor where P is the number of predictions, and 2 is the (i,j) coordinate
    :param target: (T, 2) Tensor where T is the number of targets, and 2 is the (i,j) coordinate
    :param dims: (H, W) tup/list the desired height and width of matrix (should be >= to max(i), max(j))
    :param assignment_method: Method of assignment (dtw, minimum etc.)
    :return: image, visualisation of path prediction and target.
    """
    tp = torch.Tensor((64, 191, 64))
    fp = torch.Tensor((191, 64, 64))
    gt = torch.Tensor((102, 153, 255))

    grid = torch.ones_like(window).unsqueeze(0).repeat(3, 1, 1) * 255
    inf = 130 * torch.ones_like(grid)
    grid = torch.where(torch.isnan(window), inf, grid)

    clip_idxs = [t[0] for t in target]
    local_idxs = np.unique(np.array(clip_idxs)).tolist()

    for t in target:
        local_idx = local_idxs.index(t[0])
        grid[:, local_idx, t[1]] = gt

    for p in pred:
        local_idx = local_idxs.index(p[0])
        if (grid[:, local_idx, p[1]] == gt).all():
            grid[:, local_idx, p[1]] = tp
        else:
            grid[:, local_idx, p[1]] = fp

    return grid / 255


def batch_path_vis(pred_dict, target, window):
    grids = []

    window = window.cpu()
    for key, pred in pred_dict.items():
        tmp_window = window
        if key == 'min_dist':
            tmp_window = torch.zeros_like(window)
        grids.append(visualise_path(pred, target, tmp_window))

    return torch.stack(grids)


if __name__ == "__main__":
    pred = [[1, 1], [2, 4]]
    gt = [[1, 1], [3, 4]]
    window = torch.zeros((5, 6))
    visualise_path(pred, gt, window)

QFVS/utils/visualisation.py
facebookresearch/EgoVLPv2
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import json
import pickle

def load_json(filename):
    with open(filename, encoding='utf8') as f:
        return json.load(f)

def load_pickle(filename):
    with open(filename,'rb') as f:
        return pickle.load(f)

def save_pickle(object,filename):
    with open(filename, 'wb') as f:
        pickle.dump(object,f)

QFVS/utils_QFVS.py
facebookresearch/EgoVLPv2