A Multi-Task Evaluation Benchmark for Audio-Visual Representation Models
"""
Custom class for loading audio-visual data
Modified from https://github.com/s3prl/s3prl/blob/main/s3prl/downstream/example/dataset.py
"""
import random
import os

import torch
import torch.nn as nn
from torch.utils.data.dataset import Dataset

# Example parameters
AUDIO_SAMPLE_RATE = 44100
VIDEO_FRAME_RATE = 30
MIN_SEC = 5
MAX_SEC = 20
DATASET_SIZE = 200
HEIGHT = 224
WIDTH = 224


class RandomDataset(Dataset):
    def __init__(self, preprocess=None , preprocess_audio=None, preprocess_video=None, **kwargs):
        """
        Your dataset should take two preprocessing transform functions,
        preprocess_audio and preprocess_video as input.

        These two functions will be defined by the upstream models, and
        will transform raw waveform & video frames into the desired
        format of the upstream model.

        They take two arguments, the input audio/video Tensor, and the
        audio sample rate/video frame rate, respectively.

        Optionally, if you wish to obtain raw data for testing purposes,
        you may also specify these functions to be None, and return the
        raw data when the functions are not defined.
        """
        self.class_num = 48
        self.audio_sample_rates = [AUDIO_SAMPLE_RATE] * len(self)
        self.video_frame_rates = [VIDEO_FRAME_RATE] * len(self)
        self.preprocess = preprocess
        self.preprocess_audio = preprocess_audio
        self.preprocess_video = preprocess_video

        self.upstream_name = kwargs['upstream']
        self.upstream_feature_selection = kwargs['upstream_feature_selection']
        self.pooled_features_path = kwargs['pooled_features_path']

    def get_rates(self, idx):
        """
        Return the audio sample rate and video frame rate of the idx-th video.
        (Datasets may contain data with different sample rates)
        """
        return self.audio_sample_rates[idx], self.video_frame_rates[idx]

    def __getitem__(self, idx):
        length = random.randint(MIN_SEC, MAX_SEC)
        audio_samples = length * AUDIO_SAMPLE_RATE
        video_samples = length * VIDEO_FRAME_RATE
        audio_sr, video_fps = self.get_rates(idx)
        # You may use the following function to read video data:
        # frames, wav = torchvision.io.read_video(path, pts_unit="sec", output_format="TCHW")
        wav = torch.randn(audio_samples)
        frames = torch.ones(
            video_samples, 3, random.randint(50, HEIGHT), random.randint(50, WIDTH), dtype=torch.uint8
        )

        # Run preprocessing only if features are not precomputed
        fname = "path to your video"
        basename = fname.rsplit('/')[-1].rsplit('.')[0]
        label = random.randint(0, self.class_num - 1)

        # Directly load pooled features if exist, 
        # skipping video loading and preprocessing
        if self.pooled_features_path:
            pooled_feature_path = f"{self.pooled_features_path}/{self.upstream_name}_{self.upstream_feature_selection}/{basename}_pooled.pt"
            if os.path.exists(pooled_feature_path):
                pooled_feature = torch.load(pooled_feature_path)
                return pooled_feature, pooled_feature, label, True

        feature_path = f"/work/b07901163/features/{self.upstream_name}/{basename}.pt"
        if os.path.exists(feature_path):
            processed_wav, processed_frames = torch.load(feature_path)
        else:
            if self.preprocess is not None:
                processed_frames, processed_wav = self.preprocess(frames, wav, video_fps, audio_sr)
            else:    
                if self.preprocess_audio is not None:
                    processed_wav = self.preprocess_audio(wav, audio_sr)
                else:
                    processed_wav = wav
                if self.preprocess_video is not None:
                    processed_frames = self.preprocess_video(frames, video_fps)
                else:
                    processed_frames = frames
            # Uncomment the next line
            # torch.save([processed_wav, processed_frames], feature_path)

        return processed_wav, processed_frames, label, basename

    def __len__(self):
        return DATASET_SIZE

    def collate_fn(self, samples):
        wavs, videos, *others = zip(*samples)
        # Concise way of doing:
        # wavs, videos, labels = [], [], []
        # for wav, frames, label in samples:
        #     wavs.append(wav)
        #     videos.append(frames)
        #     labels.append(label)
        return wavs, videos, *others

downstream_tasks/example/dataset.py
roger-tseng/av-superb
"""
Custom class for training/testing code for downstream audio-visual tasks
Modified from https://github.com/s3prl/s3prl/blob/main/s3prl/downstream/example/expert.py
"""
import math
import os
import random

import torch
import torch.nn as nn
from torch.distributed import is_initialized
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import DataLoader, Dataset, DistributedSampler

from .dataset import RandomDataset
from .model import Model


def get_ddp_sampler(dataset: Dataset, epoch: int):
    """
    This function will create a DistributedSampler if DDP is initialized,
    and will just return None if DDP is not initialized.
    """
    if is_initialized():
        sampler = DistributedSampler(dataset)
        sampler.set_epoch(epoch)
    else:
        sampler = None
    return sampler


class DownstreamExpert(nn.Module):
    """
    Used to handle downstream-specific operations
    eg. downstream forward, metric computation, contents to log
    """

    def __init__(
        self,
        preprocess,
        preprocess_audio,
        preprocess_video,
        upstream_dim,
        downstream_expert,
        expdir,
        **kwargs,
    ):
        """
        Your dataset should take two preprocessing transform functions,
        preprocess_audio and preprocess_video as input.

        These two functions will be defined by the upstream models, and
        will transform raw waveform & video frames into the desired
        format of the upstream model.

        They take two arguments, the input audio/video Tensor, and the
        audio sample rate/video frame rate, respectively.

        Optionally, if you wish to obtain raw data for testing purposes,
        you may also specify these functions to be None, and return the
        raw data when the functions are not defined.
        Args:
            preprocess_audio: function
                Defined by specified upstream model, transforms raw waveform into
                desired input format.
                Takes two arguments, input audio Tensor, and audio sample rate.

            preprocess_video: function
                Defined by specified upstream model, transforms raw video frames
                into desired input format.
                Takes two arguments, input video Tensor, and video frame rate.

            upstream_dim: int
                Different upstream models will give different representation dimension
                You might want to first project them to the same dimension

            downstream_expert: dict
                The 'downstream_expert' field specified in your downstream config file
                eg. downstream/example/config.yaml

            expdir: string
                The expdir from command-line argument, you should save all results into
                this directory, like some logging files.

            **kwargs: dict
                All the arguments specified by the argparser in run_downstream.py
                and all the other fields in config.yaml, in case you need it.

                Note. Feel free to add new argument for __init__ as long as it is
                a command-line argument or a config field. You can check the constructor
                code in downstream/runner.py
        """

        super(DownstreamExpert, self).__init__()
        self.upstream_dim = upstream_dim
        self.datarc = downstream_expert["datarc"]  # config for dataset
        self.modelrc = downstream_expert["modelrc"]  # config for model

        self.train_dataset = RandomDataset(
            preprocess, 
            preprocess_audio, 
            preprocess_video,
            upstream=kwargs['upstream'],
            pooled_features_path=kwargs['pooled_features_path'],
            upstream_feature_selection=kwargs['upstream_feature_selection'],
            **self.datarc
        )
        self.dev_dataset = RandomDataset(
            preprocess, 
            preprocess_audio, 
            preprocess_video,
            upstream=kwargs['upstream'],
            pooled_features_path=kwargs['pooled_features_path'],
            upstream_feature_selection=kwargs['upstream_feature_selection'],
            **self.datarc
        )
        self.test_dataset = RandomDataset(
            preprocess, 
            preprocess_audio, 
            preprocess_video,
            upstream=kwargs['upstream'],
            pooled_features_path=kwargs['pooled_features_path'],
            upstream_feature_selection=kwargs['upstream_feature_selection'],
            **self.datarc
        )

        self.connector = nn.Linear(upstream_dim, self.modelrc["input_dim"])
        self.model = Model(
            output_class_num=self.train_dataset.class_num, **self.modelrc
        )
        self.objective = nn.CrossEntropyLoss()
        self.register_buffer("best_score", torch.zeros(1))

    # Interface
    def get_dataloader(self, split, epoch: int = 0):
        """
        Args:
            split: string
                'train'
                    will always be called before the training loop

                'dev', 'test', or more
                    defined by the 'eval_dataloaders' field in your downstream config
                    these will be called before the evaluation loops during the training loop

        Return:
            a torch.utils.data.DataLoader returning each batch in the format of:

            [(wav1,vid1), (wav2,vid2), ...], your_other_contents1, your_other_contents2, ...
        """

        if split == "train":
            return self._get_train_dataloader(self.train_dataset, epoch)
        elif split == "dev":
            return self._get_eval_dataloader(self.dev_dataset)
        elif split == "test":
            return self._get_eval_dataloader(self.test_dataset)

    def _get_train_dataloader(self, dataset, epoch: int):
        sampler = get_ddp_sampler(dataset, epoch)
        return DataLoader(
            dataset,
            batch_size=self.datarc["train_batch_size"],
            shuffle=(sampler is None),
            sampler=sampler,
            num_workers=self.datarc["num_workers"],
            collate_fn=dataset.collate_fn,
        )

    def _get_eval_dataloader(self, dataset):
        return DataLoader(
            dataset,
            batch_size=self.datarc["eval_batch_size"],
            shuffle=False,
            num_workers=self.datarc["num_workers"],
            collate_fn=dataset.collate_fn,
        )

    # Interface
    def forward(self, split, features, labels, basenames, records, **kwargs):
        """
        Args:
            split: string
                'train'
                    when the forward is inside the training loop

                'dev', 'test' or more
                    when the forward is inside the evaluation loop

            features:
                list of unpadded features [feat1, feat2, ...]
                each feat is in torch.FloatTensor and already
                put in the device assigned by command-line args

            your_other_contents1, ... :
                in the order defined by your dataloader (dataset + collate_fn)
                these are all in cpu, and you can move them to the same device
                as features

            records:
                defaultdict(list), by appending contents into records,
                these contents can be averaged and logged on Tensorboard
                later by self.log_records (also customized by you)

                Note1. downstream/runner.py will call self.log_records
                    1. every `log_step` during training
                    2. once after evalute the whole dev/test dataloader

                Note2. `log_step` is defined in your downstream config
                eg. downstream/example/config.yaml

        Return:
            loss:
                the loss to be optimized, should not be detached
                a single scalar in torch.FloatTensor
        """
        features = pad_sequence(features, batch_first=True)
        features = self.connector(features)
        predicted = self.model(features)

        utterance_labels = labels
        labels = torch.LongTensor(utterance_labels).to(features.device)
        loss = self.objective(predicted, labels)

        predicted_classid = predicted.max(dim=-1).indices

        records["loss"].append(loss.item())
        records["acc"] += (predicted_classid == labels).view(-1).cpu().float().tolist()

        return loss

    # interface
    def log_records(
        self, split, records, logger, global_step, batch_ids, total_batch_num, **kwargs
    ):
        """
        Args:
            split: string
                'train':
                    records and batchids contain contents for `log_step` batches
                    `log_step` is defined in your downstream config
                    eg. downstream/example/config.yaml

                'dev', 'test' or more:
                    records and batchids contain contents for the entire evaluation dataset

            records:
                defaultdict(list), contents already prepared by self.forward

            logger:
                Tensorboard SummaryWriter
                please use f'{your_task_name}/{split}-{key}' as key name to log your contents,
                preventing conflict with the logging of other tasks

            global_step:
                The global_step when training, which is helpful for Tensorboard logging

            batch_ids:
                The batches contained in records when enumerating over the dataloader

            total_batch_num:
                The total amount of batches in the dataloader

        Return:
            a list of string
                Each string is a filename we wish to use to save the current model
                according to the evaluation result, like the best.ckpt on the dev set
                You can return nothing or an empty list when no need to save the checkpoint
        """
        save_names = []
        for key, values in records.items():
            average = torch.FloatTensor(values).mean().item()
            logger.add_scalar(
                f"example/{split}-{key}", average, global_step=global_step
            )
            if split == "dev" and key == "acc" and average > self.best_score:
                self.best_score = torch.ones(1) * average
                save_names.append(f"{split}-best.ckpt")
        return save_names

downstream_tasks/example/expert.py
roger-tseng/av-superb
"""
Small probing model for downstream audio-visual tasks
Modified from https://github.com/s3prl/s3prl/blob/main/s3prl/downstream/example/model.py
"""
import torch
import torch.nn as nn


class Model(nn.Module):
    def __init__(self, input_dim, output_class_num, **kwargs):
        super(Model, self).__init__()
        self.linear = nn.Linear(input_dim, output_class_num)

    def forward(self, features):
        pooled = features.mean(dim=1)
        predicted = self.linear(pooled)
        return predicted

downstream_tasks/example/model.py
roger-tseng/av-superb
# Author: Yuan Tseng
# Hub for custom model classes

# Modified from S3PRL 
# (Authors: Leo Yang, Andy T. Liu and S3PRL team, https://github.com/s3prl/s3prl/blob/main/s3prl/hub.py)

from upstream_models.example.hubconf import *
from upstream_models.replai.hubconf import *
from upstream_models.vhubert.hubconf import *
from upstream_models.avbert.hubconf import *
from upstream_models.mavil.hubconf import *
from upstream_models.hubert.hubconf import *
from upstream_models.fbank.hubconf import *
from upstream_models.hog.hubconf import *


def options(only_registered_ckpt: bool = False):
    all_options = []
    for name, value in globals().items():
        torch_hubconf_policy = not name.startswith("_") and callable(value)
        if torch_hubconf_policy and name != "options":
            if only_registered_ckpt and (
                name.endswith("_local")
                or name.endswith("_url")
                or name.endswith("_gdriveid")
                or name.endswith("_custom")
            ):
                continue
            all_options.append(name)

    return all_options

hub.py
roger-tseng/av-superb
# Author: Yuan Tseng
# Custom classes for interfacing upstream models and downstream tasks

# Modified from S3PRL
# (Authors: Leo Yang, Andy T. Liu and S3PRL team, https://github.com/s3prl/s3prl/blob/main/s3prl/upstream/interfaces.py)

import random
import sys
from typing import Callable, Dict, List, Tuple, Union

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor

from utils.helper import show

TOLERABLE_SEQLEN_DIFF = 5
AUDIO_SAMPLE_RATE = 16000
VIDEO_SAMPLE_RATE = 16
HEIGHT = 112
WIDTH = 112
MIN_SEC = 2
MAX_SEC = 4


class Hook:
    def __init__(self, module_path, transform, unique_identifier=None):
        self.module_path = module_path
        self.transform = transform
        self.unique_identifier = unique_identifier or module_path
        self.handler = None

        assert isinstance(self.module_path, str)
        assert callable(self.transform)
        assert isinstance(self.unique_identifier, str)


class initHook(type):
    def __call__(cls, *args, **kwargs):
        instance = super().__call__(*args, **kwargs)
        for hook in instance.hooks:
            if hook.handler is None:
                instance._register_hook_handler(hook)
        return instance


class UpstreamBase(nn.Module, metaclass=initHook):
    def __init__(
        self,
        hooks: List[Tuple] = None,
        hook_postprocess: Callable[
            [List[Tuple[str, Tensor]]], List[Tuple[str, Tensor]]
        ] = None,
        **kwargs,
    ):
        """
        Args:
            hooks: each Tuple is an argument list for the Hook initializer
        """
        super().__init__()
        self.hooks: List[Hook] = [Hook(*hook) for hook in hooks] if hooks else []
        self.hook_postprocess = hook_postprocess
        self._hook_hiddens: List[Tuple(str, Tensor)] = []

    def remove_all_hooks(self):
        for hook in self.hooks:
            hook.handler.remove()
        self.hooks.clear()

    def remove_hook(self, unique_identifier: str):
        updated_hooks = []
        for hook in self.hooks:
            if hook.unique_identifier == unique_identifier:
                hook.handler.remove()
            else:
                updated_hooks.append(hook)
        self.hooks = updated_hooks

    def add_hook(self, *args, **kwargs):
        hook = Hook(*args, **kwargs)
        self._register_hook_handler(hook)
        self.hooks.append(hook)

    def _register_hook_handler(self, hook: Hook):
        module = eval(hook.module_path)
        if not isinstance(module, nn.Module):
            show(
                f"[UpstreamBase] - {hook.module_path} is not a valid nn.Module. Skip.",
                file=sys.stderr,
            )
            return

        if callable(hook.handler):
            show(
                f"[UpstreamBase] - Existing hook handler for {hook.unique_identifier} is found. Remove the existing one.",
                file=sys.stderr,
            )
            hook.handler.remove()

        def generate_hook_handler(hiddens: List, hook: Hook):
            def hook_handler(self, input, output):
                hiddens.append((hook.unique_identifier, hook.transform(input, output)))

            return hook_handler

        hook.handler = module.register_forward_hook(
            generate_hook_handler(self._hook_hiddens, hook)
        )

    def __call__(self, wavs: List[Tensor], *args, **kwargs):
        self._hook_hiddens.clear()

        result = super().__call__(wavs, *args, **kwargs) or {}
        assert isinstance(result, dict)

        if len(self._hook_hiddens) > 0:
            hook_hiddens = self._hook_hiddens.copy()
            self._hook_hiddens.clear()

            if callable(self.hook_postprocess):
                hook_hiddens = self.hook_postprocess(hook_hiddens)

            names, hiddens = zip(*hook_hiddens)
            if "fusion" in names[0]:
                key = "fusion_feats"
            elif "video" in names[0]:
                key = "video_feats"
            elif "audio" in names[0]:
                key = "audio_feats"

            if (
                result.get("_hidden_states_info") is not None
                or result.get(key) is not None
            ):
                show(
                    f"[UpstreamBase] - If there are registered hooks, '_hidden_states_info' and '{key}' "
                    "are reserved and should not be included in child class's return dict.",
                    file=sys.stderr,
                )
                raise ValueError

            result["_hidden_states_info"], result[key] = names, hiddens

        return result


class Featurizer(nn.Module):
    def __init__(
        self,
        upstream: UpstreamBase,
        feature_selection: str = "hidden_states",
        upstream_device: str = "cuda",
        layer_selection: int = None,
        normalize: bool = False,
        **kwargs,
    ):
        super().__init__()
        self.name = "Featurizer"

        upstream.eval()
        length = random.randint(MIN_SEC, MAX_SEC)
        audio_samples = length * AUDIO_SAMPLE_RATE
        video_samples = length * VIDEO_SAMPLE_RATE
        paired_wavs = [
            (
                torch.randn(audio_samples),
                torch.ones(video_samples, 3, HEIGHT, WIDTH, dtype=torch.uint8),
            )
            for i in range(2)
        ]

        with torch.no_grad():
            if hasattr(upstream, "preprocess"):

                paired = [
                    (
                        upstream.preprocess(video, audio, VIDEO_SAMPLE_RATE, AUDIO_SAMPLE_RATE)
                    )
                    for audio, video in paired_wavs
                ]
                paired_input = [
                    (
                        audio.to(upstream_device),
                        video.to(upstream_device),
                    )
                    for video, audio in paired
                ]
                paired_features = upstream(paired_input)
            elif hasattr(upstream, "preprocess_audio") and hasattr(
                upstream, "preprocess_video"
            ):
                paired_input = [
                    (
                        upstream.preprocess_audio(audio, AUDIO_SAMPLE_RATE).to(
                            upstream_device
                        ),
                        upstream.preprocess_video(video, VIDEO_SAMPLE_RATE).to(
                            upstream_device
                        ),
                    )
                    for audio, video in paired_wavs
                ]
                paired_features = upstream(paired_input)
            else:
                show(
                    f"[{self.name}] - Error: Your upstream model does not implement its preprocessing functions."
                    f" Please follow upstream_models/example/expert.py to add your preprocess_audio and preprocess_video functions for your upstream.",
                    file=sys.stderr,
                )
                raise NotImplementedError
        if feature_selection not in paired_features:
            if "hidden_states" in paired_features:
                show(
                    f"[{self.name}] - Warning: {feature_selection} is not a supported args.upstream_feature_selection."
                    f' Using "hidden_states" as the default key.',
                    file=sys.stderr,
                )
                feature_selection = "hidden_states"
            else:
                show(
                    f"[{self.name}] - Error: {feature_selection} is not a supported args.upstream_feature_selection."
                    f' The default key "hidden_states" is also not supported.'
                    f" Please specify -s with the following options: {list(paired_wavs.keys())}",
                    file=sys.stderr,
                )
                raise ValueError
        self.feature_selection = feature_selection
        self.layer_selection = layer_selection
        self.normalize = normalize

        feature = self._select_feature(paired_features)
        if isinstance(feature, (list, tuple)):
            self.layer_num = len(feature)
            show(
                f"[{self.name}] - Take a list of {self.layer_num} features and weighted sum them.",
                file=sys.stderr,
            )
            self.weights = nn.Parameter(torch.zeros(self.layer_num))
            feature = self._weighted_sum([f.cpu() for f in feature])
        else:
            feature = feature.cpu()

        self.output_dim = feature.size(-1)
        if hasattr(upstream, "get_downsample_rates"):
            self.downsample_rate = upstream.get_downsample_rates(feature_selection)
            show(
                f"[{self.name}] - The selected feature {feature_selection}'s downsample rate is {self.downsample_rate}",
                file=sys.stderr,
            )
        else:
            self.downsample_rate = round(
                max(len(wav[0]) for wav in paired_wavs) / feature.size(1)
            )
            # TODO: add back in downsample rates for padding
            # show(
            #     f"[{self.name}] - Warning: The provided upstream does not give statis downsample rate"
            #     ' by the "get_downsample_rates" interface (see upstream/example/expert.py).'
            #     " The downsample rate is calculated dynamically basing on the shape of the"
            #     f" input waveforms v.s. the output features: {self.downsample_rate}",
            #     file=sys.stderr,
            # )

    def _select_feature(self, features):
        feature = features.get(self.feature_selection)

        if isinstance(feature, dict):
            feature = list(feature.values())

        if isinstance(feature, (list, tuple)) and len(feature) == 1:
            feature = feature[0]

        if isinstance(feature, (list, tuple)) and isinstance(self.layer_selection, int):
            feature = feature[self.layer_selection]

        return feature

    def _weighted_sum(self, feature):
        assert self.layer_num == len(feature), (
            "If you run into this error, there is a great chance"
            " you are finetuning the upstream with wav2vec2's transformer blocks"
            " in weighted-sum mode (default), including wav2vec2, hubert, and decoar2."
            " These models use the layerdrop technique which causes the different number"
            " of layer forwards between different model forwards, resulting in different"
            " number of hidden states for different model forwards. Hence, finetuning"
            " these upstreams is essentially incompatible with weight-sum mode unless"
            " you turn off the layerdrop option in fairseq. See:"
            " https://github.com/pytorch/fairseq/blob/f6abcc2a67328bee8b15c596bb626ce2d720aae6/fairseq/models/wav2vec/wav2vec2.py#L857"
            " However, since finetuning upstreams will backward the gradient through all layers"
            " which serves the same functionality as weighted-sum: all layers can be used for different"
            " downstream tasks. Hence instead of finetuning upstream with weighted-sum, we suggest to"
            " follow the more common setting: finetuning upstream with the last layer. Please use the"
            " following options: --upstream_trainable --upstream_feature_selection last_hidden_state."
            " Or: -f -s last_hidden_state"
        )
        stacked_feature = torch.stack(feature, dim=0)

        if self.normalize:
            stacked_feature = F.layer_norm(
                stacked_feature, (stacked_feature.shape[-1],)
            )

        _, *origin_shape = stacked_feature.shape
        stacked_feature = stacked_feature.view(self.layer_num, -1)
        norm_weights = F.softmax(self.weights, dim=-1)
        weighted_feature = (norm_weights.unsqueeze(-1) * stacked_feature).sum(dim=0)
        weighted_feature = weighted_feature.view(*origin_shape)

        return weighted_feature

    def tolist(self, paired_wavs: List[Tuple[Tensor, Tensor]], paired_feature: Tensor):
        assert paired_feature.dim() == 3, "(batch_size, max_seq_len, feat_dim)"
        # TODO: check if this is needed
        # feature_len = [round(len(wav[0]) / self.downsample_rate) for wav in paired_wavs]
        # length_diff = abs(
        #     paired_feature.size(1)
        #     - round(max([len(wav[0]) for wav in paired_wavs]) / self.downsample_rate)
        # )
        # assert (
        #     length_diff < TOLERABLE_SEQLEN_DIFF
        # ), f"{length_diff} >= {TOLERABLE_SEQLEN_DIFF}, {paired_feature.size(1)}, {max([len(wav[0]) for wav in paired_wavs])}"
        # feature = [f[:l] for f, l in zip(paired_feature, feature_len)]
        feature = [f for f in paired_feature]
        return feature

    def forward(
        self,
        paired_wavs: List[Tuple[Tensor, Tensor]],
        paired_features: Dict[str, Union[Tensor, List[Tensor], Dict[str, Tensor]]],
    ):
        feature = self._select_feature(paired_features)
        if isinstance(feature, (list, tuple)):
            feature = self._weighted_sum(feature)

        return self.tolist(paired_wavs, feature)

interfaces.py
roger-tseng/av-superb
# Author: Yuan Tseng
# Main script for running experiments

# Modified from S3PRL 
# (Authors: Leo Yang, Andy T. Liu and S3PRL team, https://github.com/s3prl/s3prl/blob/main/s3prl/run_downstream.py)

import argparse
import glob
import logging
import os
import random
from argparse import Namespace

import numpy as np
import torch
import torchaudio
import yaml
from torch.distributed import get_world_size, is_initialized

import hub
from runner import Runner
from utils.helper import (
    backup,
    get_time_tag,
    hack_isinstance,
    is_leader_process,
    override,
)


def get_downstream_args():
    parser = argparse.ArgumentParser()

    # train or test for this experiment
    parser.add_argument(
        "-m", "--mode", choices=["train", "evaluate", "inference"], required=True
    )
    parser.add_argument("-t", "--evaluate_split", default="test")
    parser.add_argument(
        "-o",
        "--override",
        help="Used to override args and config, this is at the highest priority",
    )

    # distributed training
    parser.add_argument(
        "--backend", default="nccl", help="The backend for distributed training"
    )
    parser.add_argument(
        "--local_rank",
        type=int,
        help=f"The GPU id this process should use while distributed training. \
                               None when not launched by torch.distributed.launch",
    )

    # use a ckpt as the experiment initialization
    # if set, all the args and config below this line will be overwrited by the ckpt
    # if a directory is specified, the latest ckpt will be used by default
    parser.add_argument(
        "-e",
        "--past_exp",
        metavar="{CKPT_PATH,CKPT_DIR}",
        help="Resume training from a checkpoint",
    )

    # only load the parameters in the checkpoint without overwriting arguments and config, this is for evaluation
    parser.add_argument(
        "-i",
        "--init_ckpt",
        metavar="CKPT_PATH",
        help="Load the checkpoint for evaluation",
    )

    # configuration for the experiment, including runner and downstream
    parser.add_argument(
        "-c",
        "--config",
        help="The yaml file for configuring the whole experiment except the upstream model",
    )

    # downstream settings
    parser.add_argument(
        "-d",
        "--downstream",
        help="\
        Typically downstream dataset need manual preparation.\
        Please check downstream/README.md for details",
    )
    parser.add_argument(
        "-v", "--downstream_variant", help="Downstream vairants given the same expert"
    )

    # upstream settings
    parser.add_argument(
        "--hub",
        default="torch",
        choices=["torch", "huggingface"],
        help="The model Hub used to retrieve the upstream model.",
    )

    upstreams = [attr for attr in dir(hub) if attr[0] != "_"]
    parser.add_argument(
        "-u",
        "--upstream",
        help=""
        'Upstreams with "_local" or "_url" postfix need local ckpt (-k) or config file (-g). '
        "Other upstreams download two files on-the-fly and cache them, so just -u is enough and -k/-g are not needed. "
        "Please check upstream/README.md for details. "
        f"Available options in S3PRL: {upstreams}. "
        "If you do not see your upstream model here, check hub.py to see if your upstream folder is correctly imported.",
    )
    parser.add_argument(
        "-k",
        "--upstream_ckpt",
        metavar="{PATH,URL,GOOGLE_DRIVE_ID}",
        help="Only set when the specified upstream need it",
    )
    parser.add_argument(
        "-g",
        "--upstream_model_config",
        help="The config file for constructing the pretrained model",
    )
    parser.add_argument(
        "-r",
        "--upstream_refresh",
        action="store_true",
        help="Re-download cached ckpts for on-the-fly upstream variants",
    )
    parser.add_argument(
        "-f",
        "--upstream_trainable",
        action="store_true",
        help="Fine-tune, set upstream.train(). Default is upstream.eval()",
    )
    parser.add_argument(
        "-s",
        "--upstream_feature_selection",
        default="fusion_feats",
        help="Specify the layer to be extracted as the representation",
    )
    parser.add_argument(
        "-l",
        "--upstream_layer_selection",
        type=int,
        help="Select a specific layer for the features selected by -s",
    )
    parser.add_argument(
        "--upstream_feature_normalize",
        action="store_true",
        help="Specify whether to normalize hidden features before weighted sum",
    )

    # experiment directory, choose one to specify
    # expname uses the default root directory: result/downstream
    parser.add_argument(
        "-n", "--expname", help="Save experiment at result/downstream/expname"
    )
    parser.add_argument("-p", "--expdir", help="Save experiment at expdir")
    parser.add_argument(
        "-a",
        "--auto_resume",
        action="store_true",
        help="Auto-resume if the expdir contains checkpoints",
    )

    # options
    parser.add_argument("--seed", default=1337, type=int)
    parser.add_argument("--device", default="cuda", help="model.to(device)")
    parser.add_argument(
        "--cache_dir", help="The cache directory for pretrained model downloading"
    )
    parser.add_argument("--verbose", action="store_true", help="Print model infomation")
    parser.add_argument("--disable_cudnn", action="store_true", help="Disable CUDNN")
    # Path to where to save the features
    parser.add_argument('--pooled_features_path', type=str)
    args = parser.parse_args()
    backup_files = []

    if args.expdir is None:
        args.expdir = f"result/downstream/{args.expname}"

    if args.auto_resume:
        if os.path.isdir(args.expdir):
            ckpt_pths = glob.glob(f"{args.expdir}/states-*.ckpt")
            if len(ckpt_pths) > 0:
                args.past_exp = args.expdir

    if args.past_exp:
        # determine checkpoint path
        if os.path.isdir(args.past_exp):
            ckpt_pths = glob.glob(f"{args.past_exp}/states-*.ckpt")
            assert len(ckpt_pths) > 0
            ckpt_pths = sorted(
                ckpt_pths, key=lambda pth: int(pth.split("-")[-1].split(".")[0])
            )
            ckpt_pth = ckpt_pths[-1]
        else:
            ckpt_pth = args.past_exp

        print(f"[Runner] - Resume from {ckpt_pth}")

        # load checkpoint
        ckpt = torch.load(ckpt_pth, map_location="cpu")

        def update_args(old, new, preserve_list=None):
            out_dict = vars(old)
            new_dict = vars(new)
            for key in list(new_dict.keys()):
                if key in preserve_list:
                    new_dict.pop(key)
            out_dict.update(new_dict)
            return Namespace(**out_dict)

        # overwrite args
        cannot_overwrite_args = [
            "mode",
            "evaluate_split",
            "override",
            "backend",
            "local_rank",
            "past_exp",
            "device",
        ]
        args = update_args(args, ckpt["Args"], preserve_list=cannot_overwrite_args)
        os.makedirs(args.expdir, exist_ok=True)
        args.init_ckpt = ckpt_pth
        config = ckpt["Config"]

    else:
        print("[Runner] - Start a new experiment")
        os.makedirs(args.expdir, exist_ok=True)

        if args.config is None:
            args.config = f"./downstream_tasks/{args.downstream}/config.yaml"
        with open(args.config, "r") as file:
            config = yaml.load(file, Loader=yaml.FullLoader)

        if args.upstream_model_config is not None and os.path.isfile(
            args.upstream_model_config
        ):
            backup_files.append(args.upstream_model_config)

    if args.override is not None and args.override.lower() != "none":
        override(args.override, args, config)
        os.makedirs(args.expdir, exist_ok=True)

    return args, config, backup_files


def main():
    logging.basicConfig(level=logging.INFO)

    torch.multiprocessing.set_sharing_strategy("file_system")
    torchaudio.set_audio_backend("sox_io")
    hack_isinstance()

    # get config and arguments
    args, config, backup_files = get_downstream_args()
    if args.cache_dir is not None:
        torch.hub.set_dir(args.cache_dir)

    # When torch.distributed.launch is used
    if args.local_rank is not None:
        torch.cuda.set_device(args.local_rank)
        torch.distributed.init_process_group(args.backend)

    if args.mode == "train" and args.past_exp:
        ckpt = torch.load(args.init_ckpt, map_location="cpu")

        now_use_ddp = is_initialized()
        original_use_ddp = ckpt["Args"].local_rank is not None
        assert now_use_ddp == original_use_ddp, f"{now_use_ddp} != {original_use_ddp}"

        if now_use_ddp:
            now_world = get_world_size()
            original_world = ckpt["WorldSize"]
            assert now_world == original_world, f"{now_world} != {original_world}"

    # Save command
    if is_leader_process():
        with open(
            os.path.join(args.expdir, f"args_{get_time_tag()}.yaml"), "w"
        ) as file:
            yaml.dump(vars(args), file)

        with open(
            os.path.join(args.expdir, f"config_{get_time_tag()}.yaml"), "w"
        ) as file:
            yaml.dump(config, file)

        for file in backup_files:
            backup(file, args.expdir)

    # Fix seed and make backends deterministic
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(args.seed)
    if args.disable_cudnn:
        torch.backends.cudnn.enabled = False
    else:
        torch.backends.cudnn.enabled = True
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False

    runner = Runner(args, config)
    eval(f"runner.{args.mode}")()


if __name__ == "__main__":
    main()

run_downstream.py
roger-tseng/av-superb
# Author: Yuan Tseng
# Train/eval loop

# Modified from S3PRL 
# (Authors: Leo Yang, Andy T. Liu and S3PRL team, https://github.com/s3prl/s3prl/blob/main/s3prl/downstream/runner.py)

import glob
import importlib
import math
import os
import random
import shutil
import sys
import tempfile
import uuid
from pathlib import Path

import numpy as np
import torch
import torchaudio
from tensorboardX import SummaryWriter
from torch.distributed import get_rank, get_world_size, is_initialized
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DistributedSampler
from tqdm import tqdm

import hub
from interfaces import Featurizer
from utils.helper import defaultdict, get_model_state, is_leader_process, show
from utils.optimizers import get_optimizer
from utils.schedulers import get_scheduler

SAMPLE_RATE = 16000


class ModelEntry:
    def __init__(self, model, name, trainable, interfaces):
        self.model = model
        self.name = name
        self.trainable = trainable
        self.interfaces = interfaces


class Runner:
    """
    Used to handle high-level concepts of a ML experiment
    eg. training loop, evaluation loop, upstream propagation, optimization, logging, checkpoint saving
    """

    def __init__(self, args, config):
        self.args = args
        self.config = config
        self.init_ckpt = (
            torch.load(self.args.init_ckpt, map_location="cpu")
            if self.args.init_ckpt
            else {}
        )

        self.upstream = self._get_upstream()
        self.featurizer = self._get_featurizer()
        self.downstream = self._get_downstream(
            self.upstream.model.preprocess if hasattr(self.upstream.model, "preprocess") else None ,self.upstream.model.preprocess_audio, self.upstream.model.preprocess_video
        )
        self.all_entries = [self.upstream, self.featurizer, self.downstream]

    def _load_weight(self, model, name):
        init_weight = self.init_ckpt.get(name)
        if init_weight:
            show(f"[Runner] - Loading {name} weights from the previous experiment")
            model.load_state_dict(init_weight)

    def _init_model(self, model, name, trainable, interfaces=None):
        for interface in interfaces or []:
            assert hasattr(model, interface), interface

        self._load_weight(model, name)

        if (
            is_initialized()
            and trainable
            and any((p.requires_grad for p in model.parameters()))
        ):
            model = DDP(
                model, device_ids=[self.args.local_rank], find_unused_parameters=True
            )
            for interface in interfaces or []:
                setattr(model, interface, getattr(model.module, interface))

        return ModelEntry(model, name, trainable, interfaces)

    def _get_upstream(self):
        Upstream = getattr(hub, self.args.upstream)
        ckpt_path = self.args.upstream_ckpt
        upstream_refresh = self.args.upstream_refresh

        if is_initialized() and get_rank() > 0:
            torch.distributed.barrier()
            upstream_refresh = False

        model = Upstream(
            ckpt=ckpt_path,
            model_config=self.args.upstream_model_config,
            refresh=upstream_refresh,
        ).to(self.args.device)

        if is_initialized() and get_rank() == 0:
            torch.distributed.barrier()

        return self._init_model(
            model=model,
            name="Upstream",
            trainable=self.args.upstream_trainable,
            interfaces=["preprocess_audio", "preprocess_video"],
        )

    def _get_featurizer(self):
        model = Featurizer(
            upstream=self.upstream.model,
            feature_selection=self.args.upstream_feature_selection,
            layer_selection=self.args.upstream_layer_selection,
            upstream_device=self.args.device,
            normalize=self.args.upstream_feature_normalize,
        ).to(self.args.device)

        return self._init_model(
            model=model,
            name="Featurizer",
            trainable=True,
            interfaces=["output_dim", "downsample_rate"],
        )

    def _get_downstream(self, preprocess, preprocess_audio, preprocess_video):
        expert = importlib.import_module(
            f"downstream_tasks.{self.args.downstream}.expert"
        )
        Downstream = getattr(expert, "DownstreamExpert")

        model = Downstream(
            preprocess=preprocess,
            preprocess_audio=preprocess_audio,
            preprocess_video=preprocess_video,
            upstream_dim=self.featurizer.model.output_dim,
            upstream_rate=self.featurizer.model.downsample_rate,
            **self.config,
            **vars(self.args),
        ).to(self.args.device)

        return self._init_model(
            model=model,
            name="Downstream",
            trainable=True,
            interfaces=["get_dataloader", "log_records"],
        )

    def _get_optimizer(self, model_params):
        optimizer = get_optimizer(
            model_params, self.config["runner"]["total_steps"], self.config["optimizer"]
        )
        self._load_weight(optimizer, "Optimizer")
        return optimizer

    def _get_scheduler(self, optimizer):
        scheduler = get_scheduler(
            optimizer, self.config["runner"]["total_steps"], self.config["scheduler"]
        )
        self._load_weight(scheduler, "Scheduler")
        return scheduler

    def train(self):
        # trainable parameters and train/eval mode
        trainable_models = []
        trainable_paras = []
        for entry in self.all_entries:
            if entry.trainable:
                entry.model.train()
                trainable_models.append(entry.model)
                trainable_paras += list(entry.model.parameters())
            else:
                entry.model.eval()

        # optimizer
        optimizer = self._get_optimizer(trainable_models)

        # scheduler
        scheduler = None
        if self.config.get("scheduler"):
            scheduler = self._get_scheduler(optimizer)

        # progress bar
        tqdm_file = sys.stderr if is_leader_process() else open(os.devnull, "w")
        pbar = tqdm(
            total=self.config["runner"]["total_steps"],
            dynamic_ncols=True,
            desc="overall",
            file=tqdm_file,
        )
        init_step = self.init_ckpt.get("Step")
        if init_step:
            pbar.n = init_step

        # Tensorboard logging
        if is_leader_process():
            logger = SummaryWriter(self.args.expdir)

        batch_ids = []
        backward_steps = 0
        records = defaultdict(list)
        epoch = self.init_ckpt.get("Epoch", 0)
        train_split = self.config["runner"].get("train_dataloader", "train")
        while pbar.n < pbar.total:
            try:
                dataloader = self.downstream.model.get_dataloader(
                    train_split, epoch=epoch
                )
            except TypeError as e:
                if "unexpected keyword argument 'epoch'" in str(e):
                    dataloader = self.downstream.model.get_dataloader(train_split)
                    if hasattr(dataloader, "sampler") and isinstance(
                        dataloader.sampler, DistributedSampler
                    ):
                        dataloader.sampler.set_epoch(epoch)
                else:
                    raise

            for batch_id, (wavs, frames, *others) in enumerate(
                tqdm(dataloader, dynamic_ncols=True, desc="train", file=tqdm_file)
            ):
                # try/except block for forward/backward
                try:
                    if pbar.n >= pbar.total:
                        break
                    global_step = pbar.n + 1

                    assert len(wavs) == len(frames)
                    if self.args.pooled_features_path and all(i == True for i in others[-1]):
                        source = None
                        features = dict()
                        # "wavs" is overloaded into saved features here
                        # can be list of Tensors, or list of list of Tensors
                        if isinstance(wavs[0], (list, tuple)):
                            features[self.args.upstream_feature_selection] = [torch.stack(layer).to(self.args.device) for layer in zip(*wavs)]
                        else:
                            features[self.args.upstream_feature_selection] = torch.stack(wavs).to(self.args.device)
                    else:
                        source = [
                            (
                                wav.float().to(self.args.device),
                                frame.float().to(self.args.device),
                            )
                            for wav, frame in zip(wavs, frames)
                        ]
                        if self.upstream.trainable:
                            features = self.upstream.model(source)
                        else:
                            with torch.no_grad():
                                features = self.upstream.model(source)
                        if self.args.pooled_features_path:
                            if batch_id == 0:
                                for feature_selection in features.keys():
                                    os.makedirs(f"{self.args.pooled_features_path}/{self.args.upstream}_{feature_selection}", exist_ok=True)

                            show(f"[Runner] - Save mean-pooled features of batch no. {batch_id}")
                            assert isinstance(others[-1][0], str)
                            with torch.no_grad():
                                for key, feature in features.items():
                                    if key[0] == '_':
                                        continue

                                    if isinstance(feature, (list, tuple)):
                                        feature = [layer.mean(dim=1, keepdim=True) for layer in feature]
                                    else:
                                        feature = feature.mean(dim=1, keepdim=True)

                                    for i, names_k in enumerate(others[-1]):
                                        if isinstance(feature, (list, tuple)):
                                            save_target = [f[i].detach().cpu() for f in feature]
                                        else:
                                            save_target = feature[i].detach().cpu()
                                        torch.save(save_target, f"{self.args.pooled_features_path}/{self.args.upstream}_{key}/{names_k}_pooled.pt")

                    features = self.featurizer.model(source, features)

                    loss = self.downstream.model(
                        train_split,
                        features,
                        *others,
                        records=records,
                    )
                    batch_ids.append(batch_id)

                    gradient_accumulate_steps = self.config["runner"].get(
                        "gradient_accumulate_steps"
                    )
                    (loss / gradient_accumulate_steps).backward()
                    del loss

                except RuntimeError as e:
                    if "CUDA out of memory" in str(e):
                        print(f"[Runner] - CUDA out of memory at step {global_step}")
                        if is_initialized():
                            raise
                        with torch.cuda.device(self.args.device):
                            torch.cuda.empty_cache()
                        optimizer.zero_grad()
                        continue
                    else:
                        raise

                # whether to accumulate gradient
                backward_steps += 1
                if backward_steps % gradient_accumulate_steps > 0:
                    continue

                # gradient clipping
                grad_norm = torch.nn.utils.clip_grad_norm_(
                    trainable_paras, self.config["runner"]["gradient_clipping"]
                )

                # optimize
                if math.isnan(grad_norm):
                    print(f"[Runner] - grad norm is NaN at step {global_step}")
                else:
                    optimizer.step()
                optimizer.zero_grad()

                # adjust learning rate
                if scheduler:
                    scheduler.step()

                if not is_leader_process():
                    batch_ids = []
                    records = defaultdict(list)
                    continue

                # logging
                if global_step % self.config["runner"]["log_step"] == 0:
                    self.downstream.model.log_records(
                        train_split,
                        records=records,
                        logger=logger,
                        global_step=global_step,
                        batch_ids=batch_ids,
                        total_batch_num=len(dataloader),
                    )
                    batch_ids = []
                    records = defaultdict(list)

                # evaluation and save checkpoint
                save_names = []

                if global_step % self.config["runner"]["eval_step"] == 0:
                    for split in self.config["runner"]["eval_dataloaders"]:
                        save_names += self.evaluate(split, logger, global_step)

                if global_step % self.config["runner"]["save_step"] == 0:

                    def check_ckpt_num(directory):
                        max_keep = self.config["runner"]["max_keep"]
                        ckpt_pths = glob.glob(f"{directory}/states-*.ckpt")
                        if len(ckpt_pths) >= max_keep:
                            ckpt_pths = sorted(
                                ckpt_pths,
                                key=lambda pth: int(pth.split("-")[-1].split(".")[0]),
                            )
                            for ckpt_pth in ckpt_pths[: len(ckpt_pths) - max_keep + 1]:
                                os.remove(ckpt_pth)

                    check_ckpt_num(self.args.expdir)
                    save_names.append(f"states-{global_step}.ckpt")

                if len(save_names) > 0:
                    all_states = {
                        "Optimizer": optimizer.state_dict(),
                        "Step": global_step,
                        "Epoch": epoch,
                        "Args": self.args,
                        "Config": self.config,
                    }

                    for entry in self.all_entries:
                        if entry.trainable:
                            all_states[entry.name] = get_model_state(entry.model)

                    if scheduler:
                        all_states["Scheduler"] = scheduler.state_dict()

                    if is_initialized():
                        all_states["WorldSize"] = get_world_size()

                    save_paths = [
                        os.path.join(self.args.expdir, name) for name in save_names
                    ]
                    tqdm.write(f"[Runner] - Save the checkpoint to:")
                    for i, path in enumerate(save_paths):
                        tqdm.write(f"{i + 1}. {path}")
                        torch.save(all_states, path)

                pbar.update(1)
            epoch += 1

        pbar.close()

        if is_leader_process():
            logger.close()

    def evaluate(self, split=None, logger=None, global_step=0):
        """evaluate function will always be called on a single process even during distributed training"""

        # When this member function is called directly by command line
        not_during_training = split is None and logger is None and global_step == 0
        if not_during_training:
            split = self.args.evaluate_split
            tempdir = tempfile.mkdtemp()
            logger = SummaryWriter(tempdir)

        # fix seed to guarantee the same evaluation protocol across steps
        random.seed(self.args.seed)
        np.random.seed(self.args.seed)
        torch.manual_seed(self.args.seed)
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(self.args.seed)
            with torch.cuda.device(self.args.device):
                torch.cuda.empty_cache()

        # record original train/eval states and set all models to eval
        trainings = []
        for entry in self.all_entries:
            trainings.append(entry.model.training)
            entry.model.eval()

        # prepare data
        dataloader = self.downstream.model.get_dataloader(split)
        evaluate_ratio = float(self.config["runner"].get("evaluate_ratio", 1))
        evaluate_steps = round(len(dataloader) * evaluate_ratio)

        batch_ids = []
        records = defaultdict(list)
        for batch_id, (wavs, frames, *others) in enumerate(
            tqdm(dataloader, dynamic_ncols=True, desc=split, total=evaluate_steps)
        ):
            if batch_id > evaluate_steps:
                break

            assert len(wavs) == len(frames)
            if self.args.pooled_features_path and all(i == True for i in others[-1]):
                source = None
                features = dict()
                # "wavs" is overloaded into saved features here
                # can be list of Tensors, or list of list of Tensors
                if isinstance(wavs[0], (list, tuple)):
                    features[self.args.upstream_feature_selection] = [torch.stack(layer).to(self.args.device) for layer in zip(*wavs)]
                else:
                    features[self.args.upstream_feature_selection] = torch.stack(wavs).to(self.args.device)
            else:
                source = [
                    (
                        wav.float().to(self.args.device),
                        frame.float().to(self.args.device),
                    )
                    for wav, frame in zip(wavs, frames)
                ]
                with torch.no_grad():
                    features = self.upstream.model(source)
                if self.args.pooled_features_path:
                    show(f"[Runner] - Save mean-pooled features of batch no. {batch_id}")
                    assert isinstance(others[-1][0], str)
                    with torch.no_grad():
                        for key, feature in features.items():

                            if key[0] == '_':
                                continue

                            if isinstance(feature, (list, tuple)):
                                feature = [layer.mean(dim=1, keepdim=True) for layer in feature]
                            else:
                                feature = feature.mean(dim=1, keepdim=True)

                            for i, names_k in enumerate(others[-1]):
                                if isinstance(feature, (list, tuple)):
                                    save_target = [f[i].detach().cpu() for f in feature]
                                else:
                                    save_target = feature[i].detach().cpu()
                                torch.save(save_target, f"{self.args.pooled_features_path}/{self.args.upstream}_{key}/{names_k}_pooled.pt")


            with torch.no_grad():
                features = self.featurizer.model(source, features)
                self.downstream.model(
                    split,
                    features,
                    *others,
                    records=records,
                    batch_id=batch_id,
                )
                batch_ids.append(batch_id)

        save_names = self.downstream.model.log_records(
            split,
            records=records,
            logger=logger,
            global_step=global_step,
            batch_ids=batch_ids,
            total_batch_num=len(dataloader),
        )
        batch_ids = []
        records = defaultdict(list)

        # prepare back to training
        if torch.cuda.is_available():
            with torch.cuda.device(self.args.device):
                torch.cuda.empty_cache()

        for entry, training in zip(self.all_entries, trainings):
            if training:
                entry.model.train()

        if not_during_training:
            logger.close()
            shutil.rmtree(tempdir)

        return [] if type(save_names) is not list else save_names

    def inference(self):
        raise NotImplementedError("not updated to audio-visual models")
        filepath = Path(self.args.evaluate_split)
        assert filepath.is_file(), filepath
        filename = filepath.stem

        if hasattr(self.downstream.model, "load_audio"):
            wav = self.downstream.model.load_audio(filepath)
        else:
            wav, sr = torchaudio.load(str(filepath))
            assert sr == SAMPLE_RATE, sr
        wavs = [wav.view(-1).to(self.args.device)]

        for entry in self.all_entries:
            entry.model.eval()

        with torch.no_grad():
            features = self.upstream.model(wavs)
            features = self.featurizer.model(wavs, features)
            self.downstream.model.inference(features, [filename])

runner.py
roger-tseng/av-superb
import os
import argparse
from datetime import datetime
import torch
from pathlib import Path
"""Configs."""
from fvcore.common.config import CfgNode
import warnings

project_dir = str(Path(__file__).resolve().parent.parent)
dataset_root = os.path.join(project_dir, 'datasets')
output_root = os.path.join(project_dir, 'runs')


# -----------------------------------------------------------------------------
# Config definition
# -----------------------------------------------------------------------------
_C = CfgNode()


# ---------------------------------------------------------------------------- #
# Batch norm options
# ---------------------------------------------------------------------------- #
_C.BN = CfgNode()

# Weight decay value that applies on BN.
_C.BN.WEIGHT_DECAY = 0.0


# ---------------------------------------------------------------------------- #
# Pretraining options.
# ---------------------------------------------------------------------------- #
_C.PRETRAIN = CfgNode()

# Path to the checkpoint to load the pretrained weight.
_C.PRETRAIN.CHECKPOINT_FILE_PATH = ""


# ---------------------------------------------------------------------------- #
# Training options.
# ---------------------------------------------------------------------------- #
_C.TRAIN = CfgNode()

# If True Train the model, else skip training.
_C.TRAIN.ENABLE = True

# Dataset.
_C.TRAIN.DATASET = "KineticsSounds"

# Dataset split.
_C.TRAIN.DATASET_SPLIT = 1

# Number of samples to sample from a data file.
_C.TRAIN.NUM_SAMPLES = 10

# Total mini-batch size.
_C.TRAIN.BATCH_SIZE = 5

# Evaluate model on validation data every eval period epochs.
_C.TRAIN.EVAL_PERIOD = 1

# Evaluate model on test data every eval period epochs.
_C.TRAIN.TEST_PERIOD = 1

# Save model checkpoint every checkpoint period epochs.
_C.TRAIN.SAVE_EVERY_EPOCH = 30

# Path to the checkpoint to load the initial weight.
_C.TRAIN.CHECKPOINT_FILE_PATH = ""

# PREEMPTIBLE
_C.TRAIN.PREEMPTIBLE = False


# ---------------------------------------------------------------------------- #
# Validation options
# ---------------------------------------------------------------------------- #
_C.VAL = CfgNode()

# If True validate the model, else skip the testing.
_C.VAL.ENABLE = False

# Dataset for validation.
_C.VAL.DATASET = "KineticsSounds"

# Dataset split.
_C.VAL.DATASET_SPLIT = 1


# ---------------------------------------------------------------------------- #
# Testing options
# ---------------------------------------------------------------------------- #
_C.TEST = CfgNode()

# If True test the model, else skip the testing.
_C.TEST.ENABLE = False

# Dataset for testing.
_C.TEST.DATASET = "KineticsSounds"

# Dataset split.
_C.TEST.DATASET_SPLIT = 1

# Total mini-batch size
_C.TEST.BATCH_SIZE = 5

# Path to the checkpoint to load the initial weight.
_C.TEST.CHECKPOINT_FILE_PATH = ""

# Number of samples to sample from a data file.
_C.TEST.NUM_SAMPLES = 30

# Number of samples to sample from a video uniformly for aggregating the
# prediction results.
_C.TEST.NUM_ENSEMBLE_VIEWS = 10

# Number of crops to sample from a frame spatially for aggregating the
# prediction results.
_C.TEST.NUM_SPATIAL_CROPS = 3


# -----------------------------------------------------------------------------
# Visual options
# -----------------------------------------------------------------------------
_C.VIS = CfgNode()

# Visual Model architecture.
_C.VIS.ARCH = "resnet"

# Visual Model name.
_C.VIS.MODEL_NAME = "ResNet"


# -----------------------------------------------------------------------------
# ResNet options
# -----------------------------------------------------------------------------
_C.RESNET = CfgNode()

# Transformation function.
_C.RESNET.TRANS_FUNC = "bottleneck_transform"

# Number of groups. 1 for ResNet, and larger than 1 for ResNeXt).
_C.RESNET.NUM_GROUPS = 1

# Width of each group (64 -> ResNet; 4 -> ResNeXt).
_C.RESNET.WIDTH_PER_GROUP = 64

# Apply relu in a inplace manner.
_C.RESNET.INPLACE_RELU = True

# Apply stride to 1x1 conv.
_C.RESNET.STRIDE_1X1 = False

# Number of weight layers.
_C.RESNET.DEPTH = 50

# If the current block has more than NUM_BLOCK_TEMP_KERNEL blocks, use temporal
# kernel of 1 for the rest of the blocks.
_C.RESNET.NUM_BLOCK_TEMP_KERNEL = [[3], [4], [6], [3]]

# Size of stride on different res stages.
_C.RESNET.SPATIAL_STRIDES = [[1], [2], [2], [2]]

# Size of dilation on different res stages.
_C.RESNET.SPATIAL_DILATIONS = [[1], [1], [1], [1]]


# -----------------------------------------------------------------------------
# Audio options
# -----------------------------------------------------------------------------
_C.AUD = CfgNode()

# Audio Model architecture.
_C.AUD.ARCH = "resnet"

# Audio Model name
_C.AUD.MODEL_NAME = "AudioResNet"


# -----------------------------------------------------------------------------
# AudioResNet options
# -----------------------------------------------------------------------------
_C.AUDIO_RESNET = CfgNode()

# Transformation function.
_C.AUDIO_RESNET.TRANS_FUNC = "bottleneck_transform"

# Number of groups. 1 for ResNet, and larger than 1 for ResNeXt).
_C.AUDIO_RESNET.NUM_GROUPS = 1

# Width of each group (32 -> ResNet; 4 -> ResNeXt).
_C.AUDIO_RESNET.WIDTH_PER_GROUP = 32

# Apply relu in a inplace manner.
_C.AUDIO_RESNET.INPLACE_RELU = True

# Apply stride to 1x1 conv.
_C.AUDIO_RESNET.STRIDE_1X1 = False

# Number of weight layers.
_C.AUDIO_RESNET.DEPTH = 50

# Size of stride on different res stages.
_C.AUDIO_RESNET.STRIDES = [2, 2, 2, 2]

# Size of dilation on different res stages.
_C.AUDIO_RESNET.DILATIONS = [1, 1, 1, 1]


# -----------------------------------------------------------------------------
# Model options
# -----------------------------------------------------------------------------

_C.MODEL = CfgNode()

# Downstream task.
_C.MODEL.TASK = "MultimodalSequenceClassify"

# Modality fusion strategy for downstream task.
_C.MODEL.DOWNSTREAM_FUSION = "concat"

# The std to initialize the fc layer(s).
_C.MODEL.FC_INIT_STD = 0.01

#  If true, initialize the gamma of the final BN of each block of ResNet to zero.
_C.MODEL.ZERO_INIT_FINAL_BN = True

# Epsilon value for normalization layers.
_C.MODEL.EPSILON = 1e-5

# Momentum value for normalization layers.
_C.MODEL.MOMENTUM = 0.1

# The number of classes to predict for the model.
_C.MODEL.NUM_CLASSES = 32

# Dropout rate.
_C.MODEL.DROPOUT_RATE = 0.5

# If True, average-pool feature maps from convolutional models and feed
# them into another models (e.g. Transformers).
# If False, average-pool feature maps from convolutional models and directly
# perform prediction.
_C.MODEL.POOLING = True

# Loss function.
_C.MODEL.LOSS_FUNC = "cross_entropy"

# Margin value for MultiMarginLoss.
_C.MODEL.MARGIN = 0.1

# Activation layer for the output head.
_C.MODEL.HEAD_ACT = "softmax"

# Transformer architecture.
_C.MODEL.ARCH = "visual_audio"

# Transformer modality fusion strategy.
_C.MODEL.FUSION = "mid"

# Whether to use Transformers.
_C.MODEL.USE_TRANSFORMER = True


# -----------------------------------------------------------------------------
# Transformer options
# -----------------------------------------------------------------------------

_C.TRANSFORMER = CfgNode()

# The ropout ratio for the attention probabilities
_C.TRANSFORMER.ATTENTION_PROBS_DROPOUT_PROB = 0.1

# The non-linear activation function in the encoder.
_C.TRANSFORMER.HIDDEN_ACT = "gelu"

# The dropout probability for all fully connected layers in the encoder.
_C.TRANSFORMER.HIDDEN_DROPOUT_PROB = 0.1

# Size of the encoder layers.
_C.TRANSFORMER.HIDDEN_SIZE = 768

# The stddev of the normal initializer for initializing weight matrices.
_C.TRANSFORMER.INITIALIZER_RANGE = 0.02

# The size of the "intermediate" (i.e., feed-forward) layer in the encoder.
_C.TRANSFORMER.INTERMEDIATE_SIZE = 3072

# The maximum sequence length
_C.TRANSFORMER.MAX_POSITION_EMBEDDINGS = 22

# Number of attention heads for each attention layer in the encoder.
_C.TRANSFORMER.NUM_ATTENTION_HEADS = 12

# Number of hidden layers in the encoder.
_C.TRANSFORMER.NUM_HIDDEN_LAYERS = 6

# Number of group for the hidden layers, parameters in the same group are shared.
_C.TRANSFORMER.NUM_HIDDEN_GROUPS = 1

# Network structure type.
_C.TRANSFORMER.NET_STRUCTURE_TYPE = 0

# Gap size.
_C.TRANSFORMER.GAP_SIZE = 0

# Number of memory blocks.
_C.TRANSFORMER.NUM_MEMORY_BLOCKS = 0

# Number of inner repetition of attention and ffn.
_C.TRANSFORMER.INNER_GROUP_NUM = 1

# The scale to apply.
_C.TRANSFORMER.DOWN_SCALE_FACTOR = 1

# The vocabular size (number of modalities) of the `token_type_ids`
_C.TRANSFORMER.TYPE_VOCAB_SIZE = 2

# Whether to output hidden states of all layers or not.
_C.TRANSFORMER.OUTPUT_HIDDEN_STATES = True

# Whether to use mean pooling for positional encoding.
_C.TRANSFORMER.USE_MEAN_POOLING = False

# Parameter sharing across Transformer layers.
_C.TRANSFORMER.SHARING_ACROSS_LAYERS = True

# Parameter sharing across modalities.
_C.TRANSFORMER.SHARING_ACROSS_MODELS = True

# Perform low-rank decomposition.
_C.TRANSFORMER.USE_DECOMPOSITION = True

# Low-rank decomposition size
_C.TRANSFORMER.ORTHOGONAL_SIZE = 128

# Method to impose the orthogonalization constraint.
_C.TRANSFORMER.ORTHOGONAL_MODE = "static"

# Method to initialize orthogonal matrices.
_C.TRANSFORMER.ORTHOGONAL_INIT = "cayley"

# Hyperparameter for imposing the orthogonalization constraint.
_C.TRANSFORMER.ORTHOGONAL_K = 100


# -----------------------------------------------------------------------------
# Data options
# -----------------------------------------------------------------------------
_C.DATA = CfgNode()

# The spatial crop size of the input clip.
_C.DATA.CROP_SIZE = 112

# The number of frames to sample from the input clip.
_C.DATA.NUM_FRAMES = 32

# The frame sampling rate of the input clip. (interval between two sampled frames)
_C.DATA.SAMPLING_RATE = 2

# Input videos may have different fps, convert it to the target video fps before
# frame sampling.
_C.DATA.TARGET_FPS = 30

# List of input frame channel dimensions.
_C.DATA.INPUT_CHANNEL_NUM = [3]

# The mean value of the video raw pixels across the R G B channels.
_C.DATA.MEAN = [0.45, 0.45, 0.45]

# The std value of the video raw pixels across the R G B channels.
_C.DATA.STD = [0.225, 0.225, 0.225]

# The spatial crop size for training
_C.DATA.TRAIN_CROP_SIZE = 112

# The spatial augmentation jitter scales for training.
_C.DATA.TRAIN_JITTER_SCALES = [128, 160]

# The spatial crop size for testing.
_C.DATA.TEST_CROP_SIZE = 128

# Input audio clip duration (sec)
_C.DATA.CLIP_DURATION = 2

# Input audios may have different sampling rate, convert it to the target audio
# sampling rate.
_C.DATA.TARGET_AUDIO_RATE = 44100

# Number of mel bins for log-mel-scaled spectrograms.
_C.DATA.AUDIO_FREQUENCY = 80

# Time dimension for log-mel-scaled spectrograms.
_C.DATA.AUDIO_TIME = 128

# The audio frequency masking droput rate.
_C.DATA.FREQUENCY_MASK_RATE = 0.05

# The audio temporal masking dropout rate.
_C.DATA.TIME_MASK_RATE = 0.05

# Method to perform the ensemble, options include "sum" and "max".
_C.DATA.ENSEMBLE_METHOD = "sum"

# List of data augmentations.
_C.DATA.TRANSFORMATIONS = ["resize_crop", "flip", "color_normalize"]

# Sequene length of Transformer inputs.
_C.DATA.SEQUENCE_LENGTH = 10


# ---------------------------------------------------------------------------- #
# Optimizer options
# ---------------------------------------------------------------------------- #
_C.SOLVER = CfgNode()

# Base learning rate.
_C.SOLVER.BASE_LR = 0.001

# Base finetuning learning rate
_C.SOLVER.FINETUNE_LR = 0.0001

# Evaluation protocol, options include "linear_eval" and "finetune".
_C.SOLVER.PROTOCOL = "finetune"

# Learning rate policy (see utils/lr_policy.py for options and examples).
_C.SOLVER.LR_POLICY = "cosine"

# Maximal number of epochs.
_C.SOLVER.MAX_EPOCH = 10

# Maximal number of steps.
_C.SOLVER.NUM_STEPS = -1

# Momentum.
_C.SOLVER.MOMENTUM = 0.9

# Momentum dampening.
_C.SOLVER.DAMPENING = 0.0

# Nesterov momentum.
_C.SOLVER.NESTEROV = True

# Use AMSGrad
_C.SOLVER.USE_AMSGRAD = True

# L2 regularization.
_C.SOLVER.WEIGHT_DECAY = 5e-6

# L2 regularization for finetuning.
_C.SOLVER.FINETUNE_WEIGHT_DECAY = 5e-6

# The start learning rate of the warm up.
_C.SOLVER.WARMUP_START_LR = 0.0

# The start finetuning learning rate of the warm up.
_C.SOLVER.WARMUP_START_FINETUNE_LR = 0.0

# Gradually warm up the lrs over this number of steps.
_C.SOLVER.WARMUP_STEPS = -1

# Gradually warm up the lrs over this number of epochs.
_C.SOLVER.WARMUP_EPOCHS = 0

# Gradually warm up the lrs over the first proportion of total steps.
_C.SOLVER.WARMUP_PROPORTION = 0.0

# Optimization method.
_C.SOLVER.OPTIMIZING_METHOD = "adam"


# ---------------------------------------------------------------------------- #
# Misc options
# ---------------------------------------------------------------------------- #

# Number of GPUs to use per machine (applies to both training and testing).
_C.NUM_GPUS = -1

# Number of machine to use for the job.
_C.NUM_SHARDS = 1

# The index of the current machine.
_C.SHARD_ID = 0

# Data basedir.
_C.DATASET_ROOT = ""

# Data dir.
_C.DATASET_DIR = ""

# Output basedir.
_C.OUTPUT_ROOT = ""

# Checkpoints dir.
_C.SAVE_DIR = ""

# Log dir.
_C.LOG_DIR = ""

# Note that non-determinism may still be present due to non-deterministic
# operator implementations in GPU operator libraries.
_C.RNG_SEED = -1

# Log period in iters.
_C.LOG_PERIOD = 50

# Distributed init method.
_C.DIST_INIT_METHOD = "tcp://localhost:9999"

# Distributed backend.
_C.DIST_BACKEND = "nccl"


# ---------------------------------------------------------------------------- #
# Common train/test data loader options
# ---------------------------------------------------------------------------- #
_C.DATA_LOADER = CfgNode()

# Number of data loader workers per training process.
_C.DATA_LOADER.NUM_WORKERS = 10

# Load data to pinned host memory.
_C.DATA_LOADER.PIN_MEMORY = True

def get_cfg():
    """
    Get a copy of the default config.
    """
    config = _C.clone()

    return config

upstream_models/avbert/avbert/config.py
roger-tseng/av-superb
"""PyTorch ALBERT model. """


from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import copy
import json
import math
import torch
import torch.nn as nn
from io import open

from .bert import (ACT2FN, BaseConfig, prune_linear_layer,
                         BertEmbeddings, BertSelfAttention, BaseModel)

from .factor_linear import FactorLinear


class AlbertConfig(BaseConfig):
    """Configuration for `AlbertModel`.

    The default settings match the configuration of model `albert_xxlarge`.
    """

    def __init__(self,
                 hidden_size=4096,
                 num_hidden_layers=12,
                 num_hidden_groups=1,
                 num_attention_heads=64,
                 intermediate_size=16384,
                 inner_group_num=1,
                 hidden_act="gelu_new",
                 hidden_dropout_prob=0,
                 attention_probs_dropout_prob=0,
                 max_position_embeddings=512,
                 type_vocab_size=2,
                 initializer_range=0.02,
                 layer_norm_eps=1e-12, **kwargs):
        """Constructs AlbertConfig.

        Args:
            hidden_size: Size of the encoder layers and the pooler layer.
            num_hidden_layers: Number of hidden layers in the Transformer encoder.
            num_hidden_groups: Number of group for the hidden layers, parameters in
                the same group are shared.
            num_attention_heads: Number of attention heads for each attention layer in
                the Transformer encoder.
            intermediate_size: The size of the "intermediate" (i.e., feed-forward)
                layer in the Transformer encoder.
            inner_group_num: int, number of inner repetition of attention and ffn.
            down_scale_factor: float, the scale to apply
            hidden_act: The non-linear activation function (function or string) in the
                encoder and pooler.
            hidden_dropout_prob: The dropout probability for all fully connected
                layers in the embeddings, encoder, and pooler.
            attention_probs_dropout_prob: The dropout ratio for the attention
                probabilities.
            max_position_embeddings: The maximum sequence length that this model might
                ever be used with. Typically set this to something large just in case
                (e.g., 512 or 1024 or 2048).
            type_vocab_size: The vocabulary size (number of modalities) of the `token_type_ids` passed into
                `AlbertModel`.
            initializer_range: The stdev of the truncated_normal_initializer for
                initializing all weight matrices.
        """
        super(AlbertConfig, self).__init__(**kwargs)
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.num_hidden_groups = num_hidden_groups
        self.num_attention_heads = num_attention_heads
        self.inner_group_num = inner_group_num
        self.hidden_act = hidden_act
        self.intermediate_size = intermediate_size
        self.hidden_dropout_prob = hidden_dropout_prob
        self.attention_probs_dropout_prob = attention_probs_dropout_prob
        self.max_position_embeddings = max_position_embeddings
        self.type_vocab_size = type_vocab_size
        self.initializer_range = initializer_range
        self.layer_norm_eps = layer_norm_eps


class AlbertEmbeddings(BertEmbeddings):
    """
    Construct the embeddings from word, position and token_type embeddings.
    """
    def __init__(self, config, use_type_embeddings=False):
        super(AlbertEmbeddings, self).__init__(config, use_type_embeddings)

        self.LayerNorm = torch.nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)


class AlbertAttention(BertSelfAttention):
    def __init__(self, config):
        super(AlbertAttention, self).__init__(config)

        self.output_attentions = config.output_attentions
        self.num_attention_heads = config.num_attention_heads
        self.hidden_size = config.hidden_size
        self.attention_head_size = config.hidden_size // config.num_attention_heads
        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
        self.use_decomposition = config.use_decomposition
        # Low-rank decomposition
        if config.use_decomposition:
            self.factor_dense = FactorLinear(
                config.num_modality_groups,
                config.hidden_size,
                config.orthogonal_size,
                config.hidden_size,
                config.orthogonal_mode,
                config.orthogonal_init,
                config.orthogonal_k,
            )
        else:
            self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.pruned_heads = set()

    def prune_heads(self, heads):
        if len(heads) == 0:
            return
        mask = torch.ones(self.num_attention_heads, self.attention_head_size)
        heads = set(heads) - self.pruned_heads  # Convert to set and emove already pruned heads
        for head in heads:
            # Compute how many pruned heads are before the head and move the index accordingly
            head = head - sum(1 if h < head else 0 for h in self.pruned_heads)
            mask[head] = 0
        mask = mask.view(-1).contiguous().eq(1)
        index = torch.arange(len(mask))[mask].long()

        # Prune linear layers
        self.query = prune_linear_layer(self.query, index)
        self.key = prune_linear_layer(self.key, index)
        self.value = prune_linear_layer(self.value, index)
        self.dense = prune_linear_layer(self.dense, index, dim=1)

        # Update hyper params and store pruned heads
        self.num_attention_heads = self.num_attention_heads - len(heads)
        self.all_head_size = self.attention_head_size * self.num_attention_heads
        self.pruned_heads = self.pruned_heads.union(heads)

    def forward(self, input_ids, attention_mask=None, head_mask=None, modality_idx=0):
        # Low-rank decomposition
        if self.use_decomposition:
            mixed_query_layer = self.factor_query(input_ids, modality_idx)
            mixed_key_layer = self.factor_key(input_ids, modality_idx)
            mixed_value_layer = self.factor_value(input_ids, modality_idx)
        else:
            mixed_query_layer = self.query(input_ids)
            mixed_key_layer = self.key(input_ids)
            mixed_value_layer = self.value(input_ids)

        query_layer = self.transpose_for_scores(mixed_query_layer)
        key_layer = self.transpose_for_scores(mixed_key_layer)
        value_layer = self.transpose_for_scores(mixed_value_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))
        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 BertModel 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,)
        reshaped_context_layer = context_layer.view(*new_context_layer_shape)

        # Low-rank decomposition
        if self.use_decomposition:
            projected_context_layer = self.factor_dense(
                reshaped_context_layer,
                modality_idx,
            )
        else:
            w = (
                self.dense.weight.t()
                .view(self.num_attention_heads, self.attention_head_size, self.hidden_size)
                .to(context_layer.dtype)
            )
            b = self.dense.bias.to(context_layer.dtype)
            projected_context_layer = torch.einsum("bfnd,ndh->bfh", context_layer, w) + b
        projected_context_layer_dropout = self.dropout(projected_context_layer)
        layernormed_context_layer = self.LayerNorm(input_ids + projected_context_layer_dropout)
        return (layernormed_context_layer, attention_probs) if self.output_attentions else (layernormed_context_layer,)


class AlbertLayer(nn.Module):
    def __init__(self, config):
        super(AlbertLayer, self).__init__()

        self.config = config
        self.full_layer_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.attention = AlbertAttention(config)
        self.use_decomposition = config.use_decomposition
        # Low-rank decomposition
        if config.use_decomposition:
            self.factor_ffn = FactorLinear(
                config.num_modality_groups,
                config.hidden_size,
                config.orthogonal_size,
                config.intermediate_size,
                config.orthogonal_mode,
                config.orthogonal_init,
                config.orthogonal_k,
            )
        else:
            self.ffn = nn.Linear(config.hidden_size, config.intermediate_size)

        # Low-rank decomposition
        if config.use_decomposition:
            self.factor_ffn_output = FactorLinear(
                config.num_modality_groups,
                config.intermediate_size,
                config.orthogonal_size,
                config.hidden_size,
                config.orthogonal_mode,
                config.orthogonal_init,
                config.orthogonal_k,
            )
        else:
            self.ffn_output = nn.Linear(config.intermediate_size, config.hidden_size)
        self.activation = ACT2FN[config.hidden_act]

    def forward(self, hidden_states, attention_mask=None, head_mask=None, modality_idx=0):
        attention_output = self.attention(hidden_states, attention_mask, head_mask, modality_idx)
        # Low-rank decomposition
        if self.use_decomposition:
            ffn_output = self.factor_ffn(attention_output[0], modality_idx)
        else:
            ffn_output = self.ffn(attention_output[0])
        ffn_output = self.activation(ffn_output)
        # Low-rank decomposition
        if self.use_decomposition:
            ffn_output = self.factor_ffn_output(ffn_output, modality_idx)
        else:
            ffn_output = self.ffn_output(ffn_output)
        hidden_states = self.full_layer_layer_norm(ffn_output + attention_output[0])

        return (hidden_states,) + attention_output[1:]  # add attentions if we output them


class AlbertLayerGroup(nn.Module):
    def __init__(self, config):
        super(AlbertLayerGroup, self).__init__()

        self.output_attentions = config.output_attentions
        self.output_hidden_states = config.output_hidden_states
        self.albert_layers = nn.ModuleList([AlbertLayer(config) for _ in range(config.inner_group_num)])

    def forward(self, hidden_states, attention_mask=None, head_mask=None, modality_idx=0):
        layer_hidden_states = ()
        layer_attentions = ()

        for layer_index, albert_layer in enumerate(self.albert_layers):
            layer_output = albert_layer(hidden_states, attention_mask, head_mask[layer_index], modality_idx)
            hidden_states = layer_output[0]

            if self.output_attentions:
                layer_attentions = layer_attentions + (layer_output[1],)

            if self.output_hidden_states:
                layer_hidden_states = layer_hidden_states + (hidden_states,)

        outputs = (hidden_states,)
        if self.output_hidden_states:
            outputs = outputs + (layer_hidden_states,)
        if self.output_attentions:
            outputs = outputs + (layer_attentions,)
        return outputs  # last-layer hidden state, (layer hidden states), (layer attentions)


class AlbertTransformer(nn.Module):
    def __init__(self, config):
        super(AlbertTransformer, self).__init__()

        self.config = config
        self.output_attentions = config.output_attentions
        self.output_hidden_states = config.output_hidden_states
        self.albert_layer_groups = nn.ModuleList([AlbertLayerGroup(config) for _ in range(config.num_hidden_groups)])

    def forward(self, hidden_states, attention_mask=None, head_mask=None, modality_idx=0):
        all_attentions = ()

        if self.output_hidden_states:
            all_hidden_states = (hidden_states,)

        for i in range(self.config.num_hidden_layers):
            # Number of layers in a hidden group
            layers_per_group = int(self.config.num_hidden_layers / self.config.num_hidden_groups)

            # Index of the hidden group
            group_idx = int(i / (self.config.num_hidden_layers / self.config.num_hidden_groups))

            # Index of the layer inside the group
            layer_idx = int(i - group_idx * layers_per_group)

            layer_group_output = self.albert_layer_groups[group_idx](hidden_states, attention_mask, head_mask[group_idx*layers_per_group:(group_idx+1)*layers_per_group], modality_idx)
            hidden_states = layer_group_output[0]
            # NOTE: should prbly return 1st element instead of 0th, but behavious is identical for one layer

            if self.output_attentions:
                all_attentions = all_attentions + layer_group_output[-1]

            if self.output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)


        outputs = (hidden_states,)
        if self.output_hidden_states:
            outputs = outputs + (all_hidden_states,)
        if self.output_attentions:
            outputs = outputs + (all_attentions,)
        return outputs  # last-layer hidden state, (all hidden states), (all attentions)


class Albert(BaseModel):
    r"""
    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
        **last_hidden_state**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length, hidden_size)``
            Sequence of hidden-states at the output of the last layer of the model.
        **pooler_output**: ``torch.FloatTensor`` of shape ``(batch_size, hidden_size)``
            Last layer hidden-state of the first token of the sequence (classification token)
            further processed by a Linear layer and a Tanh activation function. The Linear
            layer weights are trained from the next sentence prediction (classification)
            objective during Bert pretraining. This output is usually *not* a good summary
            of the semantic content of the input, you're often better with averaging or pooling
            the sequence of hidden-states for the whole input sequence.
        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
            list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)
            of shape ``(batch_size, sequence_length, hidden_size)``:
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        **attentions**: (`optional`, returned when ``config.output_attentions=True``)
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.
    """

    config_class = AlbertConfig
    base_model_prefix = "albert"

    def _init_weights(self, module):
        """ Initialize the weights.
        """
        if isinstance(module, (nn.Linear, nn.Embedding)):
            # 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 isinstance(module, (nn.Linear)) and module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)

    def __init__(self, config):
        super(Albert, self).__init__(config)

        self.config = config
        self.encoder = AlbertTransformer(config)

    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}
            ALBERT has a different architecture in that its layers are shared across groups, which then has inner groups.
            If an ALBERT model has 12 hidden layers and 2 hidden groups, with two inner groups, there
            is a total of 4 different layers.

            These layers are flattened: the indices [0,1] correspond to the two inner groups of the first hidden layer,
            while [2,3] correspond to the two inner groups of the second hidden layer.

            Any layer with in index other than [0,1,2,3] will result in an error.
            See base class PreTrainedModel for more information about head pruning
        """
        for layer, heads in heads_to_prune.items():
            group_idx = int(layer / self.config.inner_group_num)
            inner_group_idx = int(layer - group_idx * self.config.inner_group_num)
            self.encoder.albert_layer_groups[group_idx].albert_layers[inner_group_idx].attention.prune_heads(heads)

    def forward(self, embedding_output, attention_mask=None, head_mask=None, modality_idx=0):

        input_shape = embedding_output.size()[:-1]
        device = embedding_output.device

        if attention_mask is None:
            attention_mask = torch.ones(input_shape, device=device)

        extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
        extended_attention_mask = extended_attention_mask.to(dtype=next(self.parameters()).dtype) # fp16 compatibility
        extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
        if head_mask is not None:
            if head_mask.dim() == 1:
                head_mask = head_mask.unsqueeze(0).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
                head_mask = head_mask.expand(self.config.num_hidden_layers, -1, -1, -1, -1)
            elif head_mask.dim() == 2:
                head_mask = head_mask.unsqueeze(1).unsqueeze(-1).unsqueeze(-1)  # We can specify head_mask for each layer
            head_mask = head_mask.to(dtype=next(self.parameters()).dtype) # switch to fload if need + fp16 compatibility
        else:
            head_mask = [None] * self.config.num_hidden_layers

        encoder_outputs = self.encoder(embedding_output,
                                       extended_attention_mask,
                                       head_mask=head_mask,
                                       modality_idx=modality_idx)

        sequence_output = encoder_outputs[0]

        outputs = (sequence_output, ) + encoder_outputs[1:]  # add hidden_states and attentions if they are here
        return outputs

upstream_models/avbert/avbert/models/albert.py
roger-tseng/av-superb
"""Audio models."""

import torch
import torch.nn as nn
import math

from . import audio_head_helper, audio_resnet_helper, audio_stem_helper
from .build import MODEL_REGISTRY

# Number of blocks for different stages given the model depth.
_MODEL_STAGE_DEPTH = {18: (2, 2, 2, 2), 34: (3, 4, 6, 3), 50: (3, 4, 6, 3), 101: (3, 4, 23, 3)}


@MODEL_REGISTRY.register()
class AudioResNet(nn.Module):
    """
    Audio Resnet model builder.
    """

    def __init__(self, cfg):
        """
        Args:
            cfg (CfgNode): model building configs, details are in the
                comments of the config file.
        """
        super(AudioResNet, self).__init__()
        self.pooling = cfg.MODEL.POOLING
        self._construct_network(cfg)

    def _compute_dim_in(
        self,
        idx,
        trans_func,
        width_per_group,
    ):
        """
        Compute the input dimension of each convolutional stage.
        args:
            idx (int): the index of the convolutional stage.
            trans_func (string): transform function to be used to contrusct each
                ResBlock.
            width_per_group (int): width of each group.
        returns:
            dim_in (int): the input dimension.
        """
        if trans_func == 'basic_transform':
            factor = 1 if idx == 0 else 2 ** (idx - 1)
        elif trans_func == 'bottleneck_transform':
            factor = 1 if idx == 0 else 2 * (2 ** idx)
        else:
            raise NotImplementedError(
                "Does not support {} transfomration".format(trans_func)
            )

        dim_in = width_per_group * factor
        return dim_in

    def _compute_dim_out(
        self,
        idx,
        trans_func,
        width_per_group,
    ):
        """
        Compute the output dimension of each convolutional stage.
        args:
            idx (int): the index of the convolutional stage.
            trans_func (string): transform function to be used to contrusct each
                ResBlock.
            width_per_group (int): width of each group.
        returns:
            dim_out (int): the output dimension.
        """
        if trans_func == 'basic_transform':
            factor = 2 ** idx
        elif trans_func == 'bottleneck_transform':
            factor = 4 * (2 ** idx)
        else:
            raise NotImplementedError(
                "Does not support {} transfomration".format(trans_func)
            )

        dim_out = width_per_group * factor
        return dim_out

    def _construct_network(self, cfg):
        """
        Builds a ResNet model.

        Args:
            cfg (CfgNode): model building configs, details are in the
                comments of the config file.
        """
        assert cfg.AUDIO_RESNET.DEPTH in _MODEL_STAGE_DEPTH.keys()

        (d2, d3, d4, d5) = _MODEL_STAGE_DEPTH[cfg.AUDIO_RESNET.DEPTH]

        num_groups = cfg.AUDIO_RESNET.NUM_GROUPS
        width_per_group = cfg.AUDIO_RESNET.WIDTH_PER_GROUP
        dim_inner = num_groups * width_per_group

        self.s1 = audio_stem_helper.AudioModelStem(
            dim_in=1,
            dim_out=width_per_group,
            kernel=[9, 9],
            stride=[1, 1],
            padding=[4, 4],
            separable=True,
            eps=cfg.MODEL.EPSILON,
            bn_mmt=cfg.MODEL.MOMENTUM,
        )

        dim_in_l = [
            self._compute_dim_in(
                i,
                cfg.AUDIO_RESNET.TRANS_FUNC,
                width_per_group
            )
            for i in range(4)
        ]

        dim_out_l = [
            self._compute_dim_out(
                i,
                cfg.AUDIO_RESNET.TRANS_FUNC,
                width_per_group
            )
            for i in range(4)
        ]

        self.s2 = audio_resnet_helper.ResStage(
            dim_in=dim_in_l[0],
            dim_out=dim_out_l[0],
            dim_inner=dim_inner,
            stride=cfg.AUDIO_RESNET.STRIDES[0],
            num_blocks=d2,
            num_groups=num_groups,
            trans_func_name=cfg.AUDIO_RESNET.TRANS_FUNC,
            stride_1x1=cfg.AUDIO_RESNET.STRIDE_1X1,
            inplace_relu=cfg.AUDIO_RESNET.INPLACE_RELU,
            dilation=cfg.AUDIO_RESNET.DILATIONS[0],
            separable=True,
            eps=cfg.MODEL.EPSILON,
            bn_mmt=cfg.MODEL.MOMENTUM,
        )

        self.s3 = audio_resnet_helper.ResStage(
            dim_in=dim_in_l[1],
            dim_out=dim_out_l[1],
            dim_inner=dim_inner * 2,
            stride=cfg.AUDIO_RESNET.STRIDES[1],
            num_blocks=d3,
            num_groups=num_groups,
            trans_func_name=cfg.AUDIO_RESNET.TRANS_FUNC,
            stride_1x1=cfg.AUDIO_RESNET.STRIDE_1X1,
            inplace_relu=cfg.AUDIO_RESNET.INPLACE_RELU,
            dilation=cfg.AUDIO_RESNET.DILATIONS[1],
            separable=True,
            eps=cfg.MODEL.EPSILON,
            bn_mmt=cfg.MODEL.MOMENTUM,
        )

        self.s4 = audio_resnet_helper.ResStage(
            dim_in=dim_in_l[2],
            dim_out=dim_out_l[2],
            dim_inner=dim_inner * 4,
            stride=cfg.AUDIO_RESNET.STRIDES[2],
            num_blocks=d4,
            num_groups=num_groups,
            trans_func_name=cfg.AUDIO_RESNET.TRANS_FUNC,
            stride_1x1=cfg.AUDIO_RESNET.STRIDE_1X1,
            inplace_relu=cfg.AUDIO_RESNET.INPLACE_RELU,
            dilation=cfg.AUDIO_RESNET.DILATIONS[2],
            separable=False,
            eps=cfg.MODEL.EPSILON,
            bn_mmt=cfg.MODEL.MOMENTUM,
        )

        self.s5 = audio_resnet_helper.ResStage(
            dim_in=dim_in_l[3],
            dim_out=dim_out_l[3],
            dim_inner=dim_inner * 8,
            stride=cfg.AUDIO_RESNET.STRIDES[3],
            num_blocks=d5,
            num_groups=num_groups,
            trans_func_name=cfg.AUDIO_RESNET.TRANS_FUNC,
            stride_1x1=cfg.AUDIO_RESNET.STRIDE_1X1,
            inplace_relu=cfg.AUDIO_RESNET.INPLACE_RELU,
            dilation=cfg.AUDIO_RESNET.DILATIONS[3],
            separable=False,
            eps=cfg.MODEL.EPSILON,
            bn_mmt=cfg.MODEL.MOMENTUM,
        )

        if self.pooling:
            self.head = audio_head_helper.ResNetPoolingHead(
                pool_size=[
                    math.ceil(cfg.DATA.AUDIO_FREQUENCY / 16),
                    math.ceil(cfg.DATA.AUDIO_TIME / 16),
                ],
            )
        else:
            self.head = audio_head_helper.ResNetBasicHead(
                dim_in=dim_out_l[3],
                num_classes=cfg.MODEL.NUM_CLASSES,
                pool_size=[
                    math.ceil(cfg.DATA.AUDIO_FREQUENCY / 16),
                    math.ceil(cfg.DATA.AUDIO_TIME / 16),
                ],
                dropout_rate=cfg.MODEL.DROPOUT_RATE,
                act_func=cfg.MODEL.HEAD_ACT,
                fusion=cfg.MODEL.DOWNSTREAM_FUSION,
            )

        self.output_size = dim_out_l[3]

    def get_feature_map(self, x):
        x = self.s1(x)
        x = self.s2(x)
        x = self.s3(x)
        x = self.s4(x)
        x = self.s5(x)
        return x

    def get_logit(self, feature_map):
        return self.head(feature_map)

    def forward(self, x):
        x = self.s1(x)
        x = self.s2(x)
        x = self.s3(x)
        x = self.s4(x)
        x = self.s5(x)
        x = self.head(x)
        return x

upstream_models/avbert/avbert/models/audio_model_builder.py
roger-tseng/av-superb
import copy
import numpy as np

import torch
import torch.nn as nn

from .build import MODEL_REGISTRY
from .encoding_helper import SummaryEncoding
from .bert import (
    BertConfig,
    BertEmbeddings,
    Bert
)
from .albert import (
    AlbertConfig,
    AlbertEmbeddings,
    Albert
)


def _identity(x, attention_mask=None, modality_idx=0):
    return [x]


@MODEL_REGISTRY.register()
class AVBert(nn.Module):
    """
    Multimodal Transformer model builder.
    It builds a multimodal Transformer backbone that takes audio and visual
    sequences as the input.
    """
    def __init__(self, cfg):
        super(AVBert, self).__init__()
        self.cfg = cfg
        # Shring across layers
        transformer_cfg_func = \
            AlbertConfig if cfg.TRANSFORMER.SHARING_ACROSS_LAYERS else BertConfig
        transformer_embeddings = \
            AlbertEmbeddings if cfg.TRANSFORMER.SHARING_ACROSS_LAYERS else \
            BertEmbeddings
        transformer_model = \
            Albert if cfg.TRANSFORMER.SHARING_ACROSS_LAYERS else Bert
        transformer_cfg = transformer_cfg_func.from_dict(
            {key.lower(): value for key, value in cfg.TRANSFORMER.items()}
        )

        self.num_modalities = 2 if 'visual_audio' in cfg.MODEL.ARCH else 1
        self.idx2modality = []
        if 'visual' in cfg.MODEL.ARCH:
            self.visual_conv = MODEL_REGISTRY.get(cfg.VIS.MODEL_NAME)(cfg)
            self.visual_fc = nn.Linear(
                self.visual_conv.head.output_size, transformer_cfg.hidden_size
            )
            self.visual_norm = nn.LayerNorm(
                transformer_cfg.hidden_size, transformer_cfg.layer_norm_eps
            )
            self.idx2modality.append('visual')
        if 'audio' in cfg.MODEL.ARCH:
            self.audio_conv = MODEL_REGISTRY.get(cfg.AUD.MODEL_NAME)(cfg)
            self.audio_fc = nn.Linear(
                self.audio_conv.output_size, transformer_cfg.hidden_size
            )
            self.audio_norm = nn.LayerNorm(
                transformer_cfg.hidden_size, transformer_cfg.layer_norm_eps
            )
            self.idx2modality.append('audio')

        self.modality2idx = \
            {value: idx for idx, value in enumerate(self.idx2modality)}

        # BOS embeddings
        self.summary_encoding = SummaryEncoding(
            self.idx2modality,
            transformer_cfg.hidden_size,
            transformer_cfg.use_mean_pooling,
            cfg.DATA.SEQUENCE_LENGTH,
            transformer_cfg.layer_norm_eps
        )

        # Modality fusion strategies
        if cfg.MODEL.FUSION == 'early':
            assert 'visual_audio' in cfg.MODEL.ARCH
            for i in range(self.num_modalities):
                setattr(
                    self,
                    f"single_{self.idx2modality[i]}_embeddings",
                    lambda x: x
                )
                setattr(
                    self,
                    f"single_{self.idx2modality[i]}_transformer",
                    _identity
                )
            transformer_cfg.num_modality_groups = 1
            self.multi_embeddings = \
                transformer_embeddings(transformer_cfg, True)
            self.multi_transformer = transformer_model(transformer_cfg)
        elif cfg.MODEL.FUSION == 'mid':
            assert 'visual_audio' in cfg.MODEL.ARCH
            transformer_cfg.num_modality_groups = self.num_modalities + 1
            for i in range(self.num_modalities):
                self.add_module(
                    f"single_{self.idx2modality[i]}_embeddings",
                    transformer_embeddings(transformer_cfg, False)
                )
            self.multi_embeddings = \
                transformer_embeddings(transformer_cfg, True)
            # Sharing across modalities
            if cfg.TRANSFORMER.SHARING_ACROSS_MODELS:
                self.transformer = transformer_model(transformer_cfg)
                for i in range(self.num_modalities):
                    setattr(
                        self,
                        f"single_{self.idx2modality[i]}_transformer",
                        self.transformer
                    )
                self.multi_transformer = self.transformer
            else:
                for i in range(self.num_modalities):
                    self.add_module(
                        f"single_{self.idx2modality[i]}_transformer",
                        transformer_model(transformer_cfg)
                    )
                self.multi_transformer = \
                    transformer_model(transformer_cfg)
        elif cfg.MODEL.FUSION == 'late':
            transformer_cfg.num_modality_groups = self.num_modalities
            for i in range(self.num_modalities):
                self.add_module(
                    f"single_{self.idx2modality[i]}_embeddings",
                    transformer_embeddings(transformer_cfg, False)
                )
            # Sharing across modalities
            if cfg.TRANSFORMER.SHARING_ACROSS_MODELS:
                self.transformer = transformer_model(transformer_cfg)
                for i in range(self.num_modalities):
                    setattr(
                        self,
                        f"single_{self.idx2modality[i]}_transformer",
                        self.transformer
                    )
            else:
                for i in range(self.num_modalities):
                    self.add_module(
                        f"single_{self.idx2modality[i]}_transformer",
                        transformer_model(transformer_cfg)
                    )
            self.multi_embeddings = None
            self.multi_transformer = None
        else:
            raise NotImplementedError(
                "Does not support {} fusion".format(cfg.MODEL.FUSION)
            )

    def forward(
        self,
        visual_seq=None,
        audio_seq=None,
    ):
        batch_size, seqlen = \
            visual_seq[0].size()[:2] if visual_seq is not None else \
            audio_seq.size()[:2]

        # ConvNet
        conv_outputs = []
        _conv_outputs = []
        idx2modality = []
        if 'visual' in self.cfg.MODEL.ARCH and visual_seq is not None:
            nchannels, _, H, W = visual_seq[0].size()[2:]
            vconv_repr = self.visual_conv.get_feature_map(
                [
                    t.view(
                        batch_size * seqlen,
                        nchannels,
                        -1,
                        H,
                        W,
                    )
                    for t in visual_seq
                ]
            )
            conv_outputs.append(vconv_repr)
            _conv_outputs.append(
                self.visual_conv.get_logit(vconv_repr).view(batch_size, seqlen, -1)
            )
            idx2modality.append('visual')
        if 'audio' in self.cfg.MODEL.ARCH and audio_seq is not None:
            nchannels, frequency, time = audio_seq.size()[2:]
            aconv_repr = self.audio_conv.get_feature_map(
                audio_seq.view(
                    batch_size * seqlen,
                    nchannels,
                    frequency,
                    time,
                )
            )
            conv_outputs.append(aconv_repr)
            _conv_outputs.append(
                self.audio_conv.get_logit(aconv_repr).view(batch_size, seqlen, -1)
            )
            idx2modality.append('audio')
        # assert len({len(_conv_outputs), self.num_modalities}) == 1
        modality2idx = {value: idx for idx, value in enumerate(idx2modality)}
        num_modalities = len(idx2modality)

        # Transformers
        # Single modality
        # Projection and normalization
        # Prepend summary embeddings
        single_inputs = []
        for idx in range(num_modalities):
            fc = getattr(self, f"{idx2modality[idx]}_fc")
            norm = getattr(self, f"{idx2modality[idx]}_norm")
            conv_repr_prj = norm(fc(_conv_outputs[idx]))
            _idx = self.modality2idx[idx2modality[idx]]
            single_inputs.append(
                self.summary_encoding(conv_repr_prj, _idx)
            )

        att_mask = torch.ones(
            batch_size, 1 + seqlen,
            dtype=torch.long,
            device=_conv_outputs[0].device,
        )

        single_outputs = []
        single_hiddens = []
        for idx in range(num_modalities):
            s_embeddings = getattr(
                self,
                f"single_{idx2modality[idx]}_embeddings"
            )
            s_transformer = getattr(
                self,
                f"single_{idx2modality[idx]}_transformer"
            )
            _idx = self.modality2idx[idx2modality[idx]]
            if self.cfg.TRANSFORMER.OUTPUT_HIDDEN_STATES:
                single_hidden = s_transformer(
                    s_embeddings(single_inputs[idx]),
                    attention_mask=att_mask,
                    modality_idx=_idx
                )[1] # hidden_states
                single_hiddens.append(single_hidden)
            single_output = s_transformer(
                s_embeddings(single_inputs[idx]),
                attention_mask=att_mask,
                modality_idx=_idx
            )[0] # last layer
            single_outputs.append(single_output)


        # Multi modality
        multi_output = None
        if self.multi_transformer is not None and num_modalities == 2:
            multi_input = torch.cat(single_outputs, dim=1)
            token_type_ids = torch.cat(
                [
                    torch.zeros(
                        batch_size,
                        1 + seqlen,
                        dtype=torch.long,
                        device=multi_input.device
                    ),
                    torch.ones(
                        batch_size,
                        1 + seqlen,
                        dtype=torch.long,
                        device=multi_input.device
                    )
                ],
                dim=1
            )
            input_shape = multi_input.size()[:-1]
            position_ids = torch.cat(
                [
                    torch.arange(
                        1 + seqlen,
                        dtype=torch.long,
                        device=multi_input.device
                    ),
                    torch.arange(
                        1 + seqlen,
                        dtype=torch.long,
                        device=multi_input.device
                    )
                ]
            )
            position_ids = position_ids.unsqueeze(0).expand(input_shape)

            multi_attention_mask = torch.cat(
                [att_mask, att_mask],
                dim=1,
            )
            if self.cfg.TRANSFORMER.OUTPUT_HIDDEN_STATES:
                multi_hidden = self.multi_transformer(
                    self.multi_embeddings(
                        multi_input,
                        token_type_ids=token_type_ids,
                        position_ids=position_ids
                    ),
                    attention_mask=multi_attention_mask,
                    modality_idx=self.num_modalities,
                )[1] # hidden_states
                return conv_outputs, single_hiddens, multi_hidden
            else:
                multi_output = self.multi_transformer(
                    self.multi_embeddings(
                        multi_input,
                        token_type_ids=token_type_ids,
                        position_ids=position_ids
                    ),
                    attention_mask=multi_attention_mask,
                    modality_idx=self.num_modalities,
                )[0] # last layer
                return conv_outputs, single_outputs, multi_output
upstream_models/avbert/avbert/models/avbert.py
roger-tseng/av-superb
"""PyTorch BERT model. """

from __future__ import absolute_import, division, print_function, unicode_literals

import copy
import json
import math
from io import open

import torch
from torch import nn

from .factor_linear import FactorLinear


def gelu(x):
    """ Original Implementation of the gelu activation function in Google Bert repo when initially created.
        For information: OpenAI GPT's gelu is slightly different (and gives slightly different results):
        0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))
        Also see https://arxiv.org/abs/1606.08415
    """
    return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0)))


def gelu_new(x):
    """ Implementation of the gelu activation function currently in Google Bert repo (identical to OpenAI GPT).
        Also see https://arxiv.org/abs/1606.08415
    """
    return 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))


def swish(x):
    return x * torch.sigmoid(x)


def mish(x):
    return x * torch.tanh(nn.functional.softplus(x))


def identity(x):
    return x


ACT2FN = {"identity": identity, "gelu": gelu, "relu": torch.nn.functional.relu, "swish": swish, "gelu_new": gelu_new, "mish": mish}


BertLayerNorm = torch.nn.LayerNorm


class BaseConfig(object):
    r""" Base class for all configuration classes.
        Handles a few parameters common to all models' configurations as well as methods for loading/downloading/saving configurations.

        Note:
            A configuration file can be loaded and saved to disk. Loading the configuration file and using this file to initialize a model does **not** load the model weights.
            It only affects the model's configuration.

        Class attributes (overridden by derived classes):
            - ``pretrained_config_archive_map``: a python ``dict`` of with `short-cut-names` (string) as keys and `url` (string) of associated pretrained model configurations as values.

        Parameters:
            ``finetuning_task``: string, default `None`. Name of the task used to fine-tune the model. This can be used when converting from an original (TensorFlow or PyTorch) checkpoint.
            ``num_labels``: integer, default `2`. Number of classes to use when the model is a classification model (sequences/tokens)
            ``output_attentions``: boolean, default `False`. Should the model returns attentions weights.
            ``output_hidden_states``: string, default `False`. Should the model returns all hidden-states.
            ``torchscript``: string, default `False`. Is the model used with Torchscript.
    """
    def __init__(self, **kwargs):
        # Attributes with defaults
        self.output_attentions = kwargs.pop('output_attentions', False)
        self.output_hidden_states = kwargs.pop('output_hidden_states', False)
        self.output_past = kwargs.pop('output_past', True)  # Not used by all models
        self.torchscript = kwargs.pop('torchscript', False)  # Only used by PyTorch models
        self.use_bfloat16 = kwargs.pop('use_bfloat16', False)
        self.pruned_heads = kwargs.pop('pruned_heads', {})
        self.is_decoder = kwargs.pop('is_decoder', False)
        self.use_mean_pooling = kwargs.pop('use_mean_pooling', False)

        # Fine-tuning task arguments
        self.finetuning_task = kwargs.pop('finetuning_task', None)
        self.num_labels = kwargs.pop('num_labels', 2)
        self.id2label = kwargs.pop('id2label', {i: 'LABEL_{}'.format(i) for i in range(self.num_labels)})
        self.id2label = dict((int(key), value) for key, value in self.id2label.items())
        self.label2id = kwargs.pop('label2id', dict(zip(self.id2label.values(), self.id2label.keys())))
        self.label2id = dict((key, int(value)) for key, value in self.label2id.items())

        # Additional attributes without default values
        for key, value in kwargs.items():
            try:
                setattr(self, key, value)
            except AttributeError as err:
                raise err

    @classmethod
    def from_dict(cls, json_object):
        """Constructs a `Config` from a Python dictionary of parameters."""
        return cls(**json_object)

    @classmethod
    def from_json_file(cls, json_file):
        """Constructs a `Config` from a json file of parameters."""
        with open(json_file, "r", encoding='utf-8') as reader:
            text = reader.read()
        dict_obj = json.loads(text)
        return cls(**dict_obj)

    def to_dict(self):
        """Serializes this instance to a Python dictionary."""
        output = copy.deepcopy(self.__dict__)
        return output

    def to_json_string(self):
        """Serializes this instance to a JSON string."""
        return json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n"

    def to_json_file(self, json_file_path):
        """ Save this instance to a json file."""
        with open(json_file_path, "w", encoding='utf-8') as writer:
            writer.write(self.to_json_string())

    def __eq__(self, other):
        return self.__dict__ == other.__dict__

    def __repr__(self):
        return str(self.to_json_string())


class BertConfig(BaseConfig):
    r"""
        :class:`BertConfig` is the configuration class to store the configuration of a
        `BertModel`.


        Arguments:
            hidden_size: Size of the encoder layers and the pooler layer.
            num_hidden_layers: Number of hidden layers in the Transformer encoder.
            num_attention_heads: Number of attention heads for each attention layer in
                the Transformer encoder.
            intermediate_size: The size of the "intermediate" (i.e., feed-forward)
                layer in the Transformer encoder.
            hidden_act: The non-linear activation function (function or string) in the
                encoder and pooler. If string, "gelu", "relu", "swish" and "gelu_new" are supported.
            hidden_dropout_prob: The dropout probabilitiy for all fully connected
                layers in the embeddings, encoder, and pooler.
            attention_probs_dropout_prob: The dropout ratio for the attention
                probabilities.
            max_position_embeddings: The maximum sequence length that this model might
                ever be used with. Typically set this to something large just in case
                (e.g., 512 or 1024 or 2048).
            type_vocab_size: The vocabulary size (number of modalities) of the `token_type_ids` passed into
                `BertModel`.
            initializer_range: The sttdev of the truncated_normal_initializer for
                initializing all weight matrices.
            layer_norm_eps: The epsilon used by LayerNorm.
    """
    def __init__(self,
                 hidden_size=768,
                 num_hidden_layers=12,
                 num_attention_heads=12,
                 intermediate_size=3072,
                 hidden_act="gelu",
                 hidden_dropout_prob=0.1,
                 attention_probs_dropout_prob=0.1,
                 max_position_embeddings=512,
                 type_vocab_size=2,
                 initializer_range=0.02,
                 layer_norm_eps=1e-12,
                 **kwargs):
        super(BertConfig, self).__init__(**kwargs)
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.hidden_act = hidden_act
        self.intermediate_size = intermediate_size
        self.hidden_dropout_prob = hidden_dropout_prob
        self.attention_probs_dropout_prob = attention_probs_dropout_prob
        self.max_position_embeddings = max_position_embeddings
        self.type_vocab_size = type_vocab_size
        self.initializer_range = initializer_range
        self.layer_norm_eps = layer_norm_eps


BertLayerNorm = torch.nn.LayerNorm


class BertEmbeddings(nn.Module):
    """Construct the embeddings from word, position and token_type embeddings.
    """
    def __init__(self, config, use_type_embeddings=False):
        super(BertEmbeddings, self).__init__()
        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
        self.use_type_embeddings = use_type_embeddings
        if use_type_embeddings:
            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 = BertLayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, inputs_embeds=None, token_type_ids=None, position_ids=None):
        input_shape = inputs_embeds.size()[:-1]

        seq_length = input_shape[1]
        device = inputs_embeds.device
        if position_ids is None:
            position_ids = torch.arange(seq_length, dtype=torch.long, device=device)
            position_ids = position_ids.unsqueeze(0).expand(input_shape)
        if self.use_type_embeddings and token_type_ids is None:
            token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)

        position_embeddings = self.position_embeddings(position_ids)

        embeddings = inputs_embeds + position_embeddings
        if self.use_type_embeddings:
            token_type_embeddings = self.token_type_embeddings(token_type_ids)
            embeddings += token_type_embeddings
        embeddings = self.LayerNorm(embeddings)
        embeddings = self.dropout(embeddings)
        return embeddings


def prune_linear_layer(layer, index, dim=0):
    """ Prune a linear layer (a model parameters) to keep only entries in index.
        Return the pruned layer as a new layer with requires_grad=True.
        Used to remove heads.
    """
    index = index.to(layer.weight.device)
    W = layer.weight.index_select(dim, index).clone().detach()
    if layer.bias is not None:
        if dim == 1:
            b = layer.bias.clone().detach()
        else:
            b = layer.bias[index].clone().detach()
    new_size = list(layer.weight.size())
    new_size[dim] = len(index)
    new_layer = nn.Linear(new_size[1], new_size[0], bias=layer.bias is not None).to(layer.weight.device)
    new_layer.weight.requires_grad = False
    new_layer.weight.copy_(W.contiguous())
    new_layer.weight.requires_grad = True
    if layer.bias is not None:
        new_layer.bias.requires_grad = False
        new_layer.bias.copy_(b.contiguous())
        new_layer.bias.requires_grad = True
    return new_layer


class BertSelfAttention(nn.Module):
    def __init__(self, config):
        super(BertSelfAttention, self).__init__()
        if config.hidden_size % config.num_attention_heads != 0:
            raise ValueError(
                "The hidden size (%d) is not a multiple of the number of attention "
                "heads (%d)" % (config.hidden_size, config.num_attention_heads))
        self.output_attentions = config.output_attentions

        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
        # Low-rank decomposition
        self.use_decomposition = config.use_decomposition

        if config.use_decomposition:
            self.factor_query = FactorLinear(
                config.num_modality_groups,
                config.hidden_size,
                config.orthogonal_size,
                self.all_head_size,
                config.orthogonal_mode,
                config.orthogonal_init,
                config.orthogonal_k,
            )

            self.factor_key = FactorLinear(
                config.num_modality_groups,
                config.hidden_size,
                config.orthogonal_size,
                self.all_head_size,
                config.orthogonal_mode,
                config.orthogonal_init,
                config.orthogonal_k,
            )

            self.factor_value = FactorLinear(
                config.num_modality_groups,
                config.hidden_size,
                config.orthogonal_size,
                self.all_head_size,
                config.orthogonal_mode,
                config.orthogonal_init,
                config.orthogonal_k,
            )
        else:
            self.query = nn.Linear(config.hidden_size, self.all_head_size)
            self.key = nn.Linear(config.hidden_size, self.all_head_size)
            self.value = nn.Linear(config.hidden_size, self.all_head_size)

        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)

    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, modality_idx=0):
        # Low-rank decomposition
        if self.use_decomposition:
            mixed_query_layer = self.factor_query(hidden_states, modality_idx)

            if encoder_hidden_states is not None:
                mixed_key_layer = self.factor_key(encoder_hidden_states, modality_idx)
                mixed_value_layer = self.factor_value(encoder_hidden_states, modality_idx)
                attention_mask = encoder_attention_mask
            else:
                mixed_key_layer = self.factor_key(hidden_states, modality_idx)
                mixed_value_layer = self.factor_value(hidden_states, modality_idx)
        else:
            mixed_query_layer = self.query(hidden_states)

            if encoder_hidden_states is not None:
                mixed_key_layer = self.key(encoder_hidden_states)
                mixed_value_layer = self.value(encoder_hidden_states)
                attention_mask = encoder_attention_mask
            else:
                mixed_key_layer = self.key(hidden_states)
                mixed_value_layer = self.value(hidden_states)

        query_layer = self.transpose_for_scores(mixed_query_layer)
        key_layer = self.transpose_for_scores(mixed_key_layer)
        value_layer = self.transpose_for_scores(mixed_value_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))
        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 BertModel 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 self.output_attentions else (context_layer,)
        return outputs


class BertSelfOutput(nn.Module):
    def __init__(self, config):
        super(BertSelfOutput, self).__init__()
        self.use_decomposition = config.use_decomposition
        # Low-rank decomposition
        if config.use_decomposition:
            self.factor_dense = FactorLinear(
                config.num_modality_groups,
                config.hidden_size,
                config.orthogonal_size,
                config.hidden_size,
                config.orthogonal_mode,
                config.orthogonal_init,
                config.orthogonal_k,
            )
        else:
            self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        self.LayerNorm = BertLayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, hidden_states, input_tensor, modality_idx=0):
        # Low-rank decomposition
        if self.use_decomposition:
            hidden_states = self.factor_dense(hidden_states, modality_idx)
        else:
            hidden_states = self.dense(hidden_states)
        hidden_states = self.dropout(hidden_states)
        hidden_states = self.LayerNorm(hidden_states + input_tensor)
        return hidden_states


class BertAttention(nn.Module):
    def __init__(self, config):
        super(BertAttention, self).__init__()
        self.self = BertSelfAttention(config)
        self.output = BertSelfOutput(config)
        self.pruned_heads = set()

    def prune_heads(self, heads):
        if len(heads) == 0:
            return
        mask = torch.ones(self.self.num_attention_heads, self.self.attention_head_size)
        heads = set(heads) - self.pruned_heads  # Convert to set and remove already pruned heads
        for head in heads:
            # Compute how many pruned heads are before the head and move the index accordingly
            head = head - sum(1 if h < head else 0 for h in self.pruned_heads)
            mask[head] = 0
        mask = mask.view(-1).contiguous().eq(1)
        index = torch.arange(len(mask))[mask].long()

        # 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, modality_idx=0):
        self_outputs = self.self(hidden_states, attention_mask, head_mask, encoder_hidden_states, encoder_attention_mask, modality_idx)
        attention_output = self.output(self_outputs[0], hidden_states, modality_idx)
        outputs = (attention_output,) + self_outputs[1:]  # add attentions if we output them
        return outputs


class BertIntermediate(nn.Module):
    def __init__(self, config):
        super(BertIntermediate, self).__init__()
        self.use_decomposition = config.use_decomposition
        # Low-rank decomposition
        if config.use_decomposition:
            self.factor_dense = FactorLinear(
                config.num_modality_groups,
                config.hidden_size,
                config.orthogonal_size,
                config.intermediate_size,
                config.orthogonal_mode,
                config.orthogonal_init,
                config.orthogonal_k,
            )
        else:
            self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
        if isinstance(config.hidden_act, str) or (sys.version_info[0] == 2 and isinstance(config.hidden_act, unicode)):
            self.intermediate_act_fn = ACT2FN[config.hidden_act]
        else:
            self.intermediate_act_fn = config.hidden_act

    def forward(self, hidden_states, modality_idx=0):
        # Low-rank decomposition
        if self.use_decomposition:
            hidden_states = self.factor_dense(hidden_states, modality_idx)
        else:
            hidden_states = self.dense(hidden_states)
        hidden_states = self.intermediate_act_fn(hidden_states)
        return hidden_states


class BertOutput(nn.Module):
    def __init__(self, config):
        super(BertOutput, self).__init__()
        self.use_decomposition = config.use_decomposition
        # Low-rank decomposition
        if config.use_decomposition:
            self.factor_dense = FactorLinear(
                config.num_modality_groups,
                config.intermediate_size,
                config.orthogonal_size,
                config.hidden_size,
                config.orthogonal_mode,
                config.orthogonal_init,
                config.orthogonal_k,
            )
        else:
            self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
        self.LayerNorm = BertLayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, hidden_states, input_tensor, modality_idx=0):
        # Low-rank decomposition
        if self.use_decomposition:
            hidden_states = self.factor_dense(hidden_states, modality_idx)
        else:
            hidden_states = self.dense(hidden_states)
        hidden_states = self.dropout(hidden_states)
        hidden_states = self.LayerNorm(hidden_states + input_tensor)
        return hidden_states


class BertLayer(nn.Module):
    def __init__(self, config):
        super(BertLayer, self).__init__()
        self.attention = BertAttention(config)
        self.is_decoder = config.is_decoder
        if self.is_decoder:
            self.crossattention = BertAttention(config)
        self.intermediate = BertIntermediate(config)
        self.output = BertOutput(config)

    def forward(self, hidden_states, attention_mask=None, head_mask=None, encoder_hidden_states=None, encoder_attention_mask=None, modality_idx=0):
        self_attention_outputs = self.attention(hidden_states, attention_mask, head_mask, modality_idx=modality_idx)
        attention_output = self_attention_outputs[0]
        outputs = self_attention_outputs[1:]  # add self attentions if we output attention weights

        if self.is_decoder and encoder_hidden_states is not None:
            cross_attention_outputs = self.crossattention(attention_output, attention_mask, head_mask, encoder_hidden_states, encoder_attention_mask, modality_idx=modality_idx)
            attention_output = cross_attention_outputs[0]
            outputs = outputs + cross_attention_outputs[1:]  # add cross attentions if we output attention weights

        intermediate_output = self.intermediate(attention_output, modality_idx)
        layer_output = self.output(intermediate_output, attention_output, modality_idx)
        outputs = (layer_output,) + outputs
        return outputs


class BertEncoder(nn.Module):
    def __init__(self, config):
        super(BertEncoder, self).__init__()
        self.output_attentions = config.output_attentions
        self.output_hidden_states = config.output_hidden_states
        self.layer = nn.ModuleList([BertLayer(config) for _ 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, modality_idx=0):
        all_hidden_states = ()
        all_attentions = ()
        for i, layer_module in enumerate(self.layer):
            if self.output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

            layer_outputs = layer_module(hidden_states, attention_mask, head_mask[i], encoder_hidden_states, encoder_attention_mask, modality_idx)
            hidden_states = layer_outputs[0]

            if self.output_attentions:
                all_attentions = all_attentions + (layer_outputs[1],)

        # Add last layer
        if self.output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        outputs = (hidden_states,)
        if self.output_hidden_states:
            outputs = outputs + (all_hidden_states,)
        if self.output_attentions:
            outputs = outputs + (all_attentions,)
        return outputs  # last-layer hidden state, (all hidden states), (all attentions)


class BertPooler(nn.Module):
    def __init__(self, config):
        super(BertPooler, self).__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(pooled_output)
        return pooled_output


class BaseModel(nn.Module):
    r""" Base class for all models.

        :class:`~transformers.PreTrainedModel` takes care of storing the configuration of the models and handles methods for loading/downloading/saving models
        as well as a few methods common to all models to (i) resize the input embeddings and (ii) prune heads in the self-attention heads.

        Class attributes (overridden by derived classes):
            - ``config_class``: a class derived from :class:`~transformers.PretrainedConfig` to use as configuration class for this model architecture.
            - ``pretrained_model_archive_map``: a python ``dict`` of with `short-cut-names` (string) as keys and `url` (string) of associated pretrained weights as values.
            - ``load_tf_weights``: a python ``method`` for loading a TensorFlow checkpoint in a PyTorch model, taking as arguments:

                - ``model``: an instance of the relevant subclass of :class:`~transformers.PreTrainedModel`,
                - ``config``: an instance of the relevant subclass of :class:`~transformers.PretrainedConfig`,
                - ``path``: a path (string) to the TensorFlow checkpoint.

            - ``base_model_prefix``: a string indicating the attribute associated to the base model in derived classes of the same architecture adding modules on top of the base model.
    """
    config_class = None
    base_model_prefix = ""

    def __init__(self, config, *inputs, **kwargs):
        super(BaseModel, self).__init__()
        if not isinstance(config, BaseConfig):
            raise ValueError(
                "Parameter config in `{}(config)` should be an instance of class `BaseConfig`. "
            )
        # Save config in model
        self.config = config

    @property
    def base_model(self):
        return getattr(self, self.base_model_prefix, self)

    def get_output_embeddings(self):
        """ Get model's output embeddings
            Return None if the model doesn't have output embeddings
        """
        return None  # Overwrite for models with output embeddings

    def init_weights(self):
        """ Initialize and prunes weights if needed. """
        # Initialize weights
        self.apply(self._init_weights)

        # Prune heads if needed
        if self.config.pruned_heads:
            self.prune_heads(self.config.pruned_heads)

    def prune_heads(self, heads_to_prune):
        """ Prunes heads of the base model.

            Arguments:

                heads_to_prune: dict with keys being selected layer indices (`int`) and associated values being the list of heads to prune in said layer (list of `int`).
                E.g. {1: [0, 2], 2: [2, 3]} will prune heads 0 and 2 on layer 1 and heads 2 and 3 on layer 2.
        """
        # save new sets of pruned heads as union of previously stored pruned heads and newly pruned heads
        for layer, heads in heads_to_prune.items():
            union_heads = set(self.config.pruned_heads.get(layer, [])) | set(heads)
            self.config.pruned_heads[layer] = list(union_heads)  # Unfortunately we have to store it as list for JSON

        self.base_model._prune_heads(heads_to_prune)


class Bert(BaseModel):
    r"""
    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
        **last_hidden_state**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length, hidden_size)``
            Sequence of hidden-states at the output of the last layer of the model.
        **pooler_output**: ``torch.FloatTensor`` of shape ``(batch_size, hidden_size)``
            Last layer hidden-state of the first token of the sequence (classification token)
            further processed by a Linear layer and a Tanh activation function. The Linear
            layer weights are trained from the next sentence prediction (classification)
            objective during Bert pretraining. This output is usually *not* a good summary
            of the semantic content of the input, you're often better with averaging or pooling
            the sequence of hidden-states for the whole input sequence.
        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
            list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)
            of shape ``(batch_size, sequence_length, hidden_size)``:
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        **attentions**: (`optional`, returned when ``config.output_attentions=True``)
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.

    Examples::

        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
        model = BertModel.from_pretrained('bert-base-uncased')
        input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0)  # Batch size 1
        outputs = model(input_ids)
        last_hidden_states = outputs[0]  # The last hidden-state is the first element of the output tuple

    """
    config_class = BertConfig
    base_model_prefix = "bert"

    def _init_weights(self, module):
        """ Initialize the weights """
        if isinstance(module, (nn.Linear, nn.Embedding)):
            # 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)
        elif isinstance(module, BertLayerNorm):
            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_()

    def __init__(self, config):
        super(Bert, self).__init__(config)
        self.config = config

        self.encoder = BertEncoder(config)
        # self.pooler = BertPooler(config)

        # 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)

    def forward(self, embedding_output, attention_mask=None,
                head_mask=None, encoder_hidden_states=None, encoder_attention_mask=None,
                modality_idx=0):
        """ Forward pass on the Model.

        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
        `is_decoder` argument of the configuration set to `True`; an
        `encoder_hidden_states` is expected as an input to the forward pass.

        .. _`Attention is all you need`:
            https://arxiv.org/abs/1706.03762

        """

        input_shape = embedding_output.size()[:-1]
        device = embedding_output.device

        if attention_mask is None:
            attention_mask = torch.ones(input_shape, 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.
        if attention_mask.dim() == 3:
            extended_attention_mask = attention_mask[:, None, :, :]
        elif attention_mask.dim() == 2:
            # Provided a padding mask of dimensions [batch_size, seq_length]
            # - if the model is a decoder, apply a causal mask in addition to the padding mask
            # - if the model is an encoder, make the mask broadcastable to [batch_size, num_heads, seq_length, seq_length]
            if self.config.is_decoder:
                batch_size, seq_length = input_shape
                seq_ids = torch.arange(seq_length, device=device)
                causal_mask = seq_ids[None, None, :].repeat(batch_size, seq_length, 1) <= seq_ids[None, :, None]
                causal_mask = causal_mask.to(torch.long)  # not converting to long will cause errors with pytorch version < 1.3
                extended_attention_mask = causal_mask[:, None, :, :] * attention_mask[:, None, None, :]
            else:
                extended_attention_mask = attention_mask[:, None, None, :]
        else:
            raise ValueError("Wrong shape for input_ids (shape {}) or attention_mask (shape {})".format(input_shape, attention_mask.shape))

        # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
        # masked positions, this operation will create a tensor which is 0.0 for
        # positions we want to attend and -10000.0 for masked positions.
        # Since we are adding it to the raw scores before the softmax, this is
        # effectively the same as removing these entirely.
        extended_attention_mask = extended_attention_mask.to(dtype=next(self.parameters()).dtype)  # fp16 compatibility
        extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0

        # If a 2D ou 3D attention mask is provided for the cross-attention
        # we need to make broadcastabe to [batch_size, num_heads, seq_length, seq_length]
        if self.config.is_decoder and 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)

            if encoder_attention_mask.dim() == 3:
                encoder_extended_attention_mask = encoder_attention_mask[:, None, :, :]
            elif encoder_attention_mask.dim() == 2:
                encoder_extended_attention_mask = encoder_attention_mask[:, None, None, :]
            else:
                raise ValueError("Wrong shape for encoder_hidden_shape (shape {}) or encoder_attention_mask (shape {})".format(encoder_hidden_shape,
                                                                                                                               encoder_attention_mask.shape))

            encoder_extended_attention_mask = encoder_extended_attention_mask.to(dtype=next(self.parameters()).dtype)  # fp16 compatibility
            encoder_extended_attention_mask = (1.0 - encoder_extended_attention_mask) * -10000.0
        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]
        if head_mask is not None:
            if head_mask.dim() == 1:
                head_mask = head_mask.unsqueeze(0).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
                head_mask = head_mask.expand(self.config.num_hidden_layers, -1, -1, -1, -1)
            elif head_mask.dim() == 2:
                head_mask = head_mask.unsqueeze(1).unsqueeze(-1).unsqueeze(-1)  # We can specify head_mask for each layer
            head_mask = head_mask.to(dtype=next(self.parameters()).dtype)  # switch to fload if need + fp16 compatibility
        else:
            head_mask = [None] * self.config.num_hidden_layers

        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,
                                       modality_idx=modality_idx)
        sequence_output = encoder_outputs[0]

        outputs = (sequence_output, ) + encoder_outputs[1:]  # add hidden_states and attentions if they are here
        return outputs  # sequence_output, pooled_output, (hidden_states), (attentions)

upstream_models/avbert/avbert/models/bert.py
roger-tseng/av-superb
import math

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

from .head_helper import ResNetBasicHead as VisualHead
from .audio_head_helper import ResNetBasicHead as AudioHead
from .video_model_builder import _POOL1
from ..utils.weight_init_helper import init_weights
from .build import MODEL_REGISTRY
from .avbert import AVBert


def get_visual_dim_in(
    trans_func,
    width_per_group,
    pooling,
):
    """
    Compute the input dimension to the VisualClassifyHead.
    args:
        trans_func (string): transform function to be used to contrusct each
            ResBlock.
        width_per_group (int): width of each group.
        pooling (bool): Whether to use the output of the Transformer or
            the ConvNet.
    returns:
        dim_in (int or list): the input dimension.
    """
    if trans_func == "basic_transform":
        factor = 2 ** 3
    elif trans_func == "bottleneck_transform":
        factor = 4 * (2 ** 3)
    else:
        raise NotImplementedError(
            "Does not support {} transfomration".format(trans_func)
        )
    dim_in = [width_per_group * factor]

    if pooling:
        dim_in = sum(dim_in)

    return dim_in


def get_audio_dim_in(
    trans_func,
    width_per_group,
):
    """
    Compute the input dimension to the AudioClassifyHead.
    args:
        trans_func (string): transform function to be used to contrusct each
            ResBlock.
        width_per_group (int): width of each group.
    returns:
        dim_in (int): the input dimension.
    """
    if trans_func == "basic_transform":
        factor = 2 ** 3
    elif trans_func == "bottleneck_transform":
        factor = 4 * (2 ** 3)
    else:
        raise NotImplementedError(
            "Does not support {} transfomration".format(trans_func)
        )
    dim_in = width_per_group * factor
    return dim_in


def get_visual_pool_size(
    vis_arch,
    num_frames,
    crop_size,
):
    """
    Compute the pooling size used in VisualClassifyHead.
    args:
        vis_arch (string): the architecture of the visual conv net.
        num_frames (int): number of frames per clip.
        crop_size (int): spatial size of frames.
    returns:
        pool_size (list): list of p the kernel sizes of spatial tempoeral
            poolings, temporal pool kernel size, height pool kernel size,
            width pool kernel size in order.
    """
    _pool_size = _POOL1[vis_arch]
    _num_frames = num_frames // 2
    pool_size = [
        [
            _num_frames // _pool_size[0][0],
            math.ceil(crop_size / 32) // _pool_size[0][1],
            math.ceil(crop_size / 32) // _pool_size[0][2],
        ]
    ]

    return pool_size


def get_audio_pool_size(
    audio_frequency,
    audio_time,
):
    """
    Compute the pooling size used in AudioClassifyHead.
    args:
        audio_frequency (int): frequency dimension of the audio clip.
        audio_time (int): time dimension of the audio clip.
    returns:
        pool_size (list): list of the kernel sizes of an avg pooling,
            frequency pool kernel size, time pool kernel size in order.
    """
    pool_size = [
        math.ceil(audio_frequency / 16),
        math.ceil(audio_time / 16),
    ]

    return pool_size


class ConcatHead(nn.Module):
    """
    Concatenation head for the conv net outputs.
    This layer takes as input the concateneation of the outputs of audio and
    visual convolutional nets and performs a fully-connected projection.
    """
    def __init__(
        self,
        dim_in,
        num_classes,
        dropout_rate=0.0,
        act_func="softmax",
    ):
        """
        Args:
            dim_in (int): the channel dimension of the input.
            num_classes (int): the channel dimension of the output.
            dropout_rate (float): dropout rate. If equal to 0.0, perform no
                dropout.
            act_func (string): activation function to use. 'softmax': applies
                softmax on the output. 'sigmoid': applies sigmoid on the output.
        """
        super(ConcatHead, self).__init__()
        if dropout_rate > 0.0:
            self.dropout = nn.Dropout(dropout_rate)
        self.projection = nn.Linear(dim_in, num_classes, bias=True)

        if act_func == "softmax":
            self.act = nn.Softmax(dim=-1)
        elif act_func == "sigmoid":
            self.act = nn.Sigmoid()
        else:
            raise NotImplementedError(
                "{} is not supported as an activation"
                "function.".format(act_func)
            )

    def forward(self, x):
        if hasattr(self, "dropout"):
            x = self.dropout(x)
        x = self.projection(x)
        if not self.training:
            x = self.act(x)
        return x


class ConcatSeqHead(nn.Module):
    """
    Concatenation head for the multimodal Transformer outputs.
    This layer takes as input the Transformer outputs from audio and visual
    sequences, and performs a fully-connected projection.
    """
    def __init__(
        self,
        dim_in,
        dim_hidden,
        num_classes,
        dropout_rate=0.0,
        eps=1e-5,
        act_func="softmax",
    ):
        """
        Args:
            dim_in (int): the channel dimension of the input.
            dim_hidden (int): the channel dimension of the intermediate output.
            num_classes (int): the channel dimension of the output.
            dropout_rate (float): dropout rate. If equal to 0.0, perform no
                dropout.
            eps (float): epsilon for normalization.
            act_func (string): activation function to use. 'softmax': applies
                softmax on the output. 'sigmoid': applies sigmoid on the output.
        """
        super(ConcatSeqHead, self).__init__()
        if dropout_rate > 0.0:
            self.dropout1 = nn.Dropout(dropout_rate)
            self.dropout2 = nn.Dropout(dropout_rate)
        self.fc1 = nn.Linear(dim_in, dim_hidden, bias=False)
        self.fc2 = nn.Linear(dim_in, dim_hidden, bias=False)
        self.ln = nn.LayerNorm(dim_hidden, eps=eps)
        self.projection = nn.Linear(dim_hidden, num_classes)

        if act_func == "softmax":
            self.act = nn.Softmax(dim=-1)
        elif act_func == "sigmoid":
            self.act = nn.Sigmoid()
        else:
            raise NotImplementedError(
                "{} is not supported as an activation"
                "function.".format(act_func)
            )

    def forward(self, x, y):
        if hasattr(self, "dropout1"):
            x = self.dropout1(x)
            y = self.dropout2(y)
        x = self.fc1(x)
        y = self.fc2(y)
        out = self.ln(x + y)
        out = self.projection(out)
        if not self.training:
            out = self.act(out)
        return out


class ClassifyHead(nn.Module):
    """
    Classification head.
    For linear evaluation, only this classification head will be trained.
    """
    def __init__(
        self,
        cfg,
    ):
        """
        Args:
            cfg (CfgNode): model building configs, details are in the
                comments of the config file.
        """
        super(ClassifyHead, self).__init__()
        self.cfg = cfg
        if cfg.MODEL.TASK in ["VisualClassify"]:
            visual_dim_in = get_visual_dim_in(
                cfg.RESNET.TRANS_FUNC,
                cfg.RESNET.WIDTH_PER_GROUP,
                False,
            )
            visual_pool_size = get_visual_pool_size(
                cfg.VIS.ARCH,
                cfg.DATA.NUM_FRAMES,
                cfg.DATA.CROP_SIZE,
            )
            self.visual_head = VisualHead(
                dim_in=visual_dim_in,
                num_classes=cfg.MODEL.NUM_CLASSES,
                pool_size=visual_pool_size,
                dropout_rate=cfg.MODEL.DROPOUT_RATE,
                act_func=cfg.MODEL.HEAD_ACT,
                fusion=cfg.MODEL.DOWNSTREAM_FUSION,
            )
        elif cfg.MODEL.TASK in ["AudioClassify"]:
            audio_dim_in = get_audio_dim_in(
                cfg.AUDIO_RESNET.TRANS_FUNC,
                cfg.AUDIO_RESNET.WIDTH_PER_GROUP,
            )
            audio_pool_size = get_audio_pool_size(
                cfg.DATA.AUDIO_FREQUENCY,
                cfg.DATA.AUDIO_TIME,
            )
            self.audio_head = AudioHead(
                dim_in=audio_dim_in,
                num_classes=cfg.MODEL.NUM_CLASSES,
                pool_size=audio_pool_size,
                dropout_rate=cfg.MODEL.DROPOUT_RATE,
                act_func=cfg.MODEL.HEAD_ACT,
                fusion=cfg.MODEL.DOWNSTREAM_FUSION,
            )
        elif cfg.MODEL.TASK in ["MultimodalSequenceClassify"]:
            visual_pool_size = get_visual_pool_size(
                cfg.VIS.ARCH,
                cfg.DATA.NUM_FRAMES,
                cfg.DATA.CROP_SIZE,
            )
            audio_dim_in = get_audio_dim_in(
                cfg.AUDIO_RESNET.TRANS_FUNC,
                cfg.AUDIO_RESNET.WIDTH_PER_GROUP,
            )
            audio_pool_size = get_audio_pool_size(
                cfg.DATA.AUDIO_FREQUENCY,
                cfg.DATA.AUDIO_TIME,
            )
            visual_dim_in = get_visual_dim_in(
                cfg.RESNET.TRANS_FUNC,
                cfg.RESNET.WIDTH_PER_GROUP,
                True,
            )
            self.concat_head = ConcatHead(
                visual_dim_in + audio_dim_in,
                cfg.MODEL.NUM_CLASSES,
                cfg.MODEL.DROPOUT_RATE,
                cfg.MODEL.HEAD_ACT,
            )
            self.concat_seq_head = ConcatSeqHead(
                2 * cfg.TRANSFORMER.HIDDEN_SIZE,
                cfg.TRANSFORMER.HIDDEN_SIZE,
                cfg.MODEL.NUM_CLASSES,
                cfg.MODEL.DROPOUT_RATE,
                cfg.MODEL.EPSILON,
                cfg.MODEL.HEAD_ACT,
            )

        init_weights(
            self,
            cfg.MODEL.FC_INIT_STD,
            cfg.MODEL.ZERO_INIT_FINAL_BN,
        )

    def visual_forward(self, visual_feature_map):
        return [self.visual_head(visual_feature_map)]

    def audio_forward(self, audio_feature_map):
        return [self.audio_head(audio_feature_map)]

    def multimodal_seq_forward(self, v_f, a_f, m_t):
        batch_size = m_t.shape[0]
        seqlen = (m_t.shape[1] - 2) // 2
        mv = torch.cat(
            (m_t[:, 0].unsqueeze(1).expand(-1, seqlen, -1), m_t[:, 1:1+seqlen]),
            dim=-1,
        )
        ma = torch.cat(
            (m_t[:, 1+seqlen].unsqueeze(1).expand(-1, seqlen, -1), m_t[:, 2+seqlen:]),
            dim=-1,
        )
        mv = mv.view(batch_size * seqlen, -1)
        ma = ma.view(batch_size * seqlen, -1)
        visual_logit, audio_logit = v_f, a_f
        clip_logit = self.concat_head(
            torch.cat((visual_logit, audio_logit), dim=-1)
        )
        seq_logit = self.concat_seq_head(mv, ma)
        multimodal_logit = [clip_logit, seq_logit]

        return multimodal_logit

    def forward(self, feature_maps):
        if self.cfg.MODEL.TASK in ["VisualClassify"]:
            return self.visual_forward(*feature_maps)
        elif self.cfg.MODEL.TASK in ["AudioClassify"]:
            return self.audio_forward(*feature_maps)
        elif self.cfg.MODEL.TASK in ["MultimodalSequenceClassify"]:
            return self.multimodal_seq_forward(*feature_maps)


@MODEL_REGISTRY.register()
class VisualClassify(nn.Module):
    """
    Visual classifier
    """
    def __init__(self, cfg):
        """
        Args:
            cfg (CfgNode): model building configs, details are in the
                comments of the config file.
        """
        super(VisualClassify, self).__init__()
        self.cfg = cfg
        self.visual_conv = MODEL_REGISTRY.get(cfg.VIS.MODEL_NAME)(cfg)
        assert cfg.MODEL.DOWNSTREAM_FUSION == "late", \
            f"Visual Classifier needs late fusion"
        init_weights(
            self,
            cfg.MODEL.FC_INIT_STD,
            cfg.MODEL.ZERO_INIT_FINAL_BN,
        )

    def forward(self, visual_clip, protocol):
        if protocol == "linear_eval":
            return (self.visual_conv.get_feature_map(visual_clip), )
        else:
            return [self.visual_conv(visual_clip)]


@MODEL_REGISTRY.register()
class AudioClassify(nn.Module):
    """
    Audio classifier
    """
    def __init__(self, cfg):
        """
        Args:
            cfg (CfgNode): model building configs, details are in the
                comments of the config file.
        """
        super(AudioClassify, self).__init__()
        self.cfg = cfg
        self.audio_conv = MODEL_REGISTRY.get(cfg.AUD.MODEL_NAME)(cfg)
        assert cfg.MODEL.DOWNSTREAM_FUSION == "late", \
            f"Audio Classifier needs late fusion"
        init_weights(
            self,
            cfg.MODEL.FC_INIT_STD,
            cfg.MODEL.ZERO_INIT_FINAL_BN,
        )

    def forward(self, audio_clip, protocol):
        if protocol == "linear_eval":
            return (self.audio_conv.get_feature_map(audio_clip), )
        else:
            return [self.audio_conv(audio_clip)]


@MODEL_REGISTRY.register()
class MultimodalSequenceClassify(nn.Module):
    """
    Multimodal sequence classifier
    """
    def __init__(self, cfg):
        """
        Args:
            cfg (CfgNode): model building configs, details are in the
                comments of the config file.
        """
        super(MultimodalSequenceClassify, self).__init__()
        self.cfg = cfg
        self.avbert = AVBert(cfg)
        if self.cfg.SOLVER.PROTOCOL == "finetune":
            visual_pool_size = get_visual_pool_size(
                cfg.VIS.ARCH,
                cfg.DATA.NUM_FRAMES,
                cfg.DATA.CROP_SIZE,
            )
            audio_dim_in = get_audio_dim_in(
                cfg.AUDIO_RESNET.TRANS_FUNC,
                cfg.AUDIO_RESNET.WIDTH_PER_GROUP,
            )
            audio_pool_size = get_audio_pool_size(
                cfg.DATA.AUDIO_FREQUENCY,
                cfg.DATA.AUDIO_TIME,
            )
            visual_dim_in = get_visual_dim_in(
                cfg.RESNET.TRANS_FUNC,
                cfg.RESNET.WIDTH_PER_GROUP,
                True,
            )
            self.concat_head = ConcatHead(
                visual_dim_in + audio_dim_in,
                cfg.MODEL.NUM_CLASSES,
                cfg.MODEL.DROPOUT_RATE,
                cfg.MODEL.HEAD_ACT,
            )
            self.concat_seq_head = ConcatSeqHead(
                2 * cfg.TRANSFORMER.HIDDEN_SIZE,
                cfg.TRANSFORMER.HIDDEN_SIZE,
                cfg.MODEL.NUM_CLASSES,
                cfg.MODEL.DROPOUT_RATE,
                cfg.MODEL.EPSILON,
                cfg.MODEL.HEAD_ACT,
            )

        init_weights(
            self,
            cfg.MODEL.FC_INIT_STD,
            cfg.MODEL.ZERO_INIT_FINAL_BN,
        )

    def multimodal_seq_forward(self, v_f, a_f, m_t):
        batch_size = m_t.shape[0]
        seqlen = (m_t.shape[1] - 2) // 2
        mv = torch.cat(
            (m_t[:, 0].unsqueeze(1).expand(-1, seqlen, -1), m_t[:, 1:1+seqlen]),
            dim=-1,
        )
        ma = torch.cat(
            (m_t[:, 1+seqlen].unsqueeze(1).expand(-1, seqlen, -1), m_t[:, 2+seqlen:]),
            dim=-1,
        )
        mv = mv.view(batch_size * seqlen, -1)
        ma = ma.view(batch_size * seqlen, -1)
        visual_logit, audio_logit = v_f, a_f
        clip_logit = self.concat_head(
            torch.cat((visual_logit, audio_logit), dim=-1)
        )
        seq_logit = self.concat_seq_head(mv, ma)
        multimodal_logit = [clip_logit, seq_logit]

        return multimodal_logit

    def forward(self, visual_seq, audio_seq, protocol):
        conv_outputs, single_outputs, multi_output = self.avbert(
            visual_seq=visual_seq, audio_seq=audio_seq
        )
        feature_maps = (conv_outputs[0], conv_outputs[1], multi_output)
        feature_maps = (
            self.avbert.visual_conv.head(feature_maps[0]),
            self.avbert.audio_conv.head(feature_maps[1]),
            feature_maps[2],
        )
        if protocol == "linear_eval":
            return feature_maps
        else:
            return self.multimodal_seq_forward(*feature_maps)

upstream_models/avbert/avbert/models/classify.py
roger-tseng/av-superb
import torch
import torch.nn as nn


class PositionalEncoding(nn.Module):
    """
    Implement the PE function.
    Ashish Vaswani et al.
    "Attention Is All You Need."
    """
    def __init__(self, d_model, max_len=5000):
        super(PositionalEncoding, self).__init__()
        # Compute the positional encodings once in log space.
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1).to(torch.float)
        div_term = torch.exp(
            torch.arange(0, d_model, 2).to(torch.float) *
            -(math.log(10000.0) / d_model)
        )
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)

    def forward(self, x):
        x_pos = x + self.pe[:, :x.size(1)]
        return x_pos


class SummaryEncoding(nn.Module):
    """
    Implement BOS Encoding.
    """
    def __init__(
        self,
        modalities,
        d_model,
        use_mean_pooling=False,
        max_len=5000,
        layer_norm_eps=1e-12
    ):
        super(SummaryEncoding, self).__init__()
        self.modalities = modalities
        self.bos_embeddings = nn.Embedding(len(modalities), d_model)
        self.use_mean_pooling = use_mean_pooling
        if use_mean_pooling:
            self.positional_encoding = PositionalEncoding(
                d_model,
                max_len
            )
            for modality in modalities:
                self.add_module(
                    f"bos_{modality}_norm",
                    nn.LayerNorm(d_model, layer_norm_eps)
                )

    def forward(self, x, modality_idx):
        batch_size = x.size(0)
        bos = self.bos_embeddings(
            torch.full(
                (batch_size, 1),
                modality_idx,
                dtype=torch.long,
                device=x.device
            )
        )
        if self.use_mean_pooling:
            x_pos = self.positional_encoding(x)
            x_pooled = x_pos.mean(dim=1, keepdim=True)
            bos = getattr(self, f"bos_{self.modalities[modality_idx]}_norm")(
                bos + x_pooled
            )
        x = torch.cat([bos, x], dim=1)
        return x

upstream_models/avbert/avbert/models/encoding_helper.py
roger-tseng/av-superb
"""
Adaptation of expm and expm_frechet in numpy for torch
"""

#
# Authors: Travis Oliphant, March 2002
#          Anthony Scopatz, August 2012 (Sparse Updates)
#          Jake Vanderplas, August 2012 (Sparse Updates)
#

from __future__ import division, print_function, absolute_import

import math

import numpy as np

import torch
import scipy.special

def _onenorm_matrix_power_nnm(A, p):
    """
    Compute the 1-norm of a non-negative integer power of a non-negative matrix.

    Parameters
    ----------
    A : a square ndarray or matrix or sparse matrix
        Input matrix with non-negative entries.
    p : non-negative integer
        The power to which the matrix is to be raised.

    Returns
    -------
    out : float
        The 1-norm of the matrix power p of A.

    """
    # check input
    if int(p) != p or p < 0:
        raise ValueError('expected non-negative integer p')
    p = int(p)
    if len(A.shape) != 2 or A.shape[0] != A.shape[1]:
        raise ValueError('expected A to be like a square matrix')

    # Explicitly make a column vector so that this works when A is a
    # numpy matrix (in addition to ndarray and sparse matrix).
    v = torch.ones((A.shape[0], 1), dtype=A.dtype, device=A.device)
    M = A.t()
    for _ in range(p):
        v = M.mm(v)
    return torch.max(v).item()


def _onenorm(A):
    return torch.norm(A, 1).item()


def _ident_like(A):
    return torch.eye(A.shape[0], A.shape[1], dtype=A.dtype, device=A.device)

class _ExpmPadeHelper(object):
    """
    Help lazily evaluate a matrix exponential.

    The idea is to not do more work than we need for high expm precision,
    so we lazily compute matrix powers and store or precompute
    other properties of the matrix.

    """
    def __init__(self, A):
        """
        Initialize the object.

        Parameters
        ----------
        A : a dense or sparse square numpy matrix or ndarray
            The matrix to be exponentiated.
        """
        self.A = A
        self._A2 = None
        self._A4 = None
        self._A6 = None
        self._A8 = None
        self._A10 = None
        self._d4_exact = None
        self._d6_exact = None
        self._d8_exact = None
        self._d10_exact = None
        self._d4_approx = None
        self._d6_approx = None
        self._d8_approx = None
        self._d10_approx = None
        self.ident = _ident_like(A)

    @property
    def A2(self):
        if self._A2 is None:
            self._A2 = self.A.mm(self.A)
        return self._A2

    @property
    def A4(self):
        if self._A4 is None:
            self._A4 = self.A2.mm(self.A2)
        return self._A4

    @property
    def A6(self):
        if self._A6 is None:
            self._A6 = self.A4.mm(self.A2)
        return self._A6

    @property
    def A8(self):
        if self._A8 is None:
            self._A8 = self.A6.mm(self.A2)
        return self._A8

    @property
    def A10(self):
        if self._A10 is None:
            self._A10 = self.A4.mm(self.A6)
        return self._A10

    @property
    def d4_tight(self):
        if self._d4_exact is None:
            self._d4_exact = _onenorm(self.A4)**(1/4.)
        return self._d4_exact

    @property
    def d6_tight(self):
        if self._d6_exact is None:
            self._d6_exact = _onenorm(self.A6)**(1/6.)
        return self._d6_exact

    @property
    def d8_tight(self):
        if self._d8_exact is None:
            self._d8_exact = _onenorm(self.A8)**(1/8.)
        return self._d8_exact

    @property
    def d10_tight(self):
        if self._d10_exact is None:
            self._d10_exact = _onenorm(self.A10)**(1/10.)
        return self._d10_exact

    @property
    def d4_loose(self):
        return self.d4_tight

    @property
    def d6_loose(self):
        return self.d6_tight

    @property
    def d8_loose(self):
        return self.d8_tight

    @property
    def d10_loose(self):
        return self.d10_tight

    def pade3(self):
        b = (120., 60., 12., 1.)
        U = self.A.mm(b[3]*self.A2 + b[1]*self.ident)
        V = b[2]*self.A2 + b[0]*self.ident
        return U, V

    def pade5(self):
        b = (30240., 15120., 3360., 420., 30., 1.)
        U = self.A.mm(b[5]*self.A4 + b[3]*self.A2 + b[1]*self.ident)
        V = b[4]*self.A4 + b[2]*self.A2 + b[0]*self.ident
        return U, V

    def pade7_scaled(self, s):
        b = (17297280., 8648640., 1995840., 277200., 25200., 1512., 56., 1.)

        B = self.A * 2**-s
        B2 = self.A2 * 2**(-2*s)
        B4 = self.A4 * 2**(-4*s)
        B6 = self.A6 * 2**(-6*s)

        U = B.mm(b[7]*B6 + b[5]*B4 + b[3]*B2 + b[1]*self.ident)
        V = b[6]*B6 + b[4]*B4 + b[2]*B2 + b[0]*self.ident
        return U, V


def expm32(A):
    """
    Compute the matrix exponential using Pade approximation.

    Parameters
    ----------
    A : (M,M) array_like or sparse matrix
        2D Array or Matrix (sparse or dense) to be exponentiated

    Returns
    -------
    expA : (M,M) ndarray
        Matrix exponential of `A`

    Notes
    -----
    This is algorithm (6.1) which is a simplification of algorithm (5.1).

    .. versionadded:: 0.12.0

    References
    ----------
    .. [1] Awad H. Al-Mohy and Nicholas J. Higham (2009)
           "A New Scaling and Squaring Algorithm for the Matrix Exponential."
           SIAM Journal on Matrix Analysis and Applications.
           31 (3). pp. 970-989. ISSN 1095-7162

    """
    return _expm(A)


def _expm(A):
    # Core of expm, separated to allow testing exact and approximate
    # algorithms.

    # Avoid indiscriminate asarray() to allow sparse or other strange arrays.
    if len(A.shape) != 2 or A.shape[0] != A.shape[1]:
        raise ValueError('expected a square matrix')

    # Trivial case
    if A.shape == (1, 1):
        return torch.exp(A)

    # Track functions of A to help compute the matrix exponential.
    h = _ExpmPadeHelper(A)

    # Try Pade order 3.
    eta_1 = max(h.d4_loose, h.d6_loose)
    theta3 = 4.2587300348979312e-001
    if eta_1 < theta3 and _ell(h.A, 3) == 0:
        U, V = h.pade3()
        return _solve_P_Q(U, V)

    # Try Pade order 5.
    eta_2 = max(h.d4_tight, h.d6_loose)
    theta5 = 1.8801526985337688e+000
    if eta_2 < theta5 and _ell(h.A, 5) == 0:
        U, V = h.pade5()
        return _solve_P_Q(U, V)

    theta_7 = 3.9257248464332842e+000
    eta_3 = max(h.d6_tight, h.d8_loose)
    s = max(int(np.ceil(np.log2(eta_3 / theta_7))), 0)

    s += _ell(2**-s * h.A, 7)
    U, V = h.pade7_scaled(s)
    X = _solve_P_Q(U, V)
    return torch.matrix_power(X, 2**s)


def _solve_P_Q(U, V):
    P = U + V
    Q = -U + V
    # return torch.solve(P, Q)[0]
    return torch.linalg.solve(Q, P)

def _ell(A, m):
    """
    A helper function for expm_2009.

    Parameters
    ----------
    A : linear operator
        A linear operator whose norm of power we care about.
    m : int
        The power of the linear operator

    Returns
    -------
    value : int
        A value related to a bound.

    """
    if len(A.shape) != 2 or A.shape[0] != A.shape[1]:
        raise ValueError('expected A to be like a square matrix')

    p = 2*m + 1

    # The c_i are explained in (2.2) and (2.6) of the 2005 expm paper.
    # They are coefficients of terms of a generating function series expansion.
    choose_2p_p = scipy.special.comb(2*p, p, exact=True)
    abs_c_recip = float(choose_2p_p * math.factorial(2*p + 1))

    # This is explained after Eq. (1.2) of the 2009 expm paper.
    # It is the "unit roundoff" of IEEE double precision arithmetic.
    u = 2.**-24

    # Compute the one-norm of matrix power p of abs(A).
    A_abs_onenorm = _onenorm_matrix_power_nnm(abs(A), p)

    # Treat zero norm as a special case.
    if not A_abs_onenorm:
        return 0

    alpha = A_abs_onenorm / (_onenorm(A) * abs_c_recip)
    return max(int(np.ceil(np.log2(alpha/u) / (2 * m))), 0)

def differential(f, A, E):
    """ Computes the differential of f at A when acting on E:  (df)_A(E) """
    n = A.size(0)
    M = torch.zeros(2*n, 2*n, dtype=A.dtype, device=A.device, requires_grad=False)
    M[:n, :n] = A
    M[n:, n:] = A
    M[:n, n:] = E
    return f(M)[:n, n:]

upstream_models/avbert/avbert/models/expRNN/expm32.py
roger-tseng/av-superb
import torch
import torch.nn as nn

from .parametrization import Parametrization


class Orthogonal(Parametrization):
    """ Class that implements optimization restricted to the Stiefel manifold """
    def __init__(self, input_size, output_size, initializer_skew, mode, param):
        """
        mode: "static" or a tuple such that:
                mode[0] == "dynamic"
                mode[1]: int, K, the number of steps after which we should change the basis of the dyn triv
                mode[2]: int, M, the number of changes of basis after which we should project back onto the manifold the basis. This is particularly helpful for small values of K.

        param: A parametrization of in terms of skew-symmetyric matrices
        """
        max_size = max(input_size, output_size)
        A = torch.empty(max_size, max_size)
        # base = torch.empty(input_size, output_size)
        base = torch.empty(max_size, max_size)
        super(Orthogonal, self).__init__(A, base, mode)
        self.input_size = input_size
        self.output_size = output_size
        self.param = param
        self.init_A = initializer_skew
        self.init_base = nn.init.eye_

        self.reset_parameters()

    def reset_parameters(self):
        self.init_A(self.A)
        self.init_base(self.base)

    def forward(self, input):
        return input.matmul(self.B)

    def retraction(self, A, base):
        # This could be any parametrization of a tangent space
        A = A.triu(diagonal=1)
        A = A - A.t()
        B = base.mm(self.param(A))
        if self.input_size != self.output_size:
            B = B[:self.input_size, :self.output_size]
        return B

    def project(self, base):
        try:
            # Compute the projection using the thin SVD decomposition
            U, _, V = torch.svd(base, some=True)
            return U.mm(V.t())
        except RuntimeError:
            # If the svd does not converge, fallback to the (thin) QR decomposition
            x = base
            if base.size(0) < base.size(1):
                x = base.t()
            ret = torch.qr(x, some=True).Q
            if base.size(0) < base.size(1):
                ret = ret.t()
            return ret


class modrelu(nn.Module):
    def __init__(self, features):
        # For now we just support square layers
        super(modrelu, self).__init__()
        self.features = features
        self.b = nn.Parameter(torch.Tensor(self.features))
        self.reset_parameters()

    def reset_parameters(self):
        self.b.data.uniform_(-0.01, 0.01)

    def forward(self, inputs):
        norm = torch.abs(inputs)
        biased_norm = norm + self.b
        magnitude = nn.functional.relu(biased_norm)
        phase = torch.sign(inputs)

        return phase * magnitude


class OrthogonalRNN(nn.Module):
    def __init__(self, input_size, hidden_size, initializer_skew, mode, param):
        super(OrthogonalRNN, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.recurrent_kernel = Orthogonal(hidden_size, hidden_size, initializer_skew, mode, param=param)
        self.input_kernel = nn.Linear(in_features=self.input_size, out_features=self.hidden_size, bias=False)
        self.nonlinearity = modrelu(hidden_size)

        self.reset_parameters()

    def reset_parameters(self):
        nn.init.kaiming_normal_(self.input_kernel.weight.data, nonlinearity="relu")

    def default_hidden(self, input):
        return input.new_zeros(input.size(0), self.hidden_size, requires_grad=False)

    def forward(self, input, hidden):
        input = self.input_kernel(input)
        hidden = self.recurrent_kernel(hidden)
        out = input + hidden
        out = self.nonlinearity(out)

        return out, out

upstream_models/avbert/avbert/models/expRNN/orthogonal.py
roger-tseng/av-superb
import torch
from torch import nn
import torch.nn.functional as F

from .expRNN.orthogonal import Orthogonal
from .expRNN.initialization import henaff_init_, cayley_init_
from .expRNN.trivializations import cayley_map, expm


class FactorLinear(nn.Module):
    def __init__(
        self,
        num_modality_groups,
        input_size,
        orthogonal_size,
        output_size,
        mode,
        init,
        K="100",
    ):
        super(FactorLinear, self).__init__()
        for i in range(num_modality_groups):
            self.add_module(
                "ms{}_linear".format(i),
                nn.Linear(input_size, orthogonal_size, bias=False)
            )

        if init == "cayley":
            _init = cayley_init_
        elif init == "henaff":
            _init = henaff_init_

        if K != "infty":
            K = int(K)

        if mode == "static":
            _mode = "static"
            param = expm
        elif mode == "dtriv":
            _mode = ("dynamic", K, 100)
            param = expm
        elif mode == "cayley":
            _mode = "static"
            param = cayley_map

        for i in range(num_modality_groups):
            self.add_module(
                "orthogonal_ms{}_linear".format(i),
                Orthogonal(
                    orthogonal_size,
                    orthogonal_size,
                    _init,
                    _mode,
                    param,
                )
            )

        self.s_linear = nn.Linear(orthogonal_size, output_size)

    def forward(self, x, modality_idx):
        x = getattr(self, "ms{}_linear".format(modality_idx))(x)
        x = F.normalize(x, dim=-1)
        x = getattr(self, "orthogonal_ms{}_linear".format(modality_idx))(x)
        x = self.s_linear(x)
        return x

upstream_models/avbert/avbert/models/factor_linear.py
roger-tseng/av-superb
"""Visual Conv models."""

import torch
import torch.nn as nn
import math

from . import head_helper, resnet_helper, stem_helper
from .build import MODEL_REGISTRY

# Number of blocks for different stages given the model depth.
_MODEL_STAGE_DEPTH = {18: (2, 2, 2, 2), 34: (3, 4, 6, 3), 50: (3, 4, 6, 3), 101: (3, 4, 23, 3)}

# Basis of temporal kernel sizes for each of the stage.
_TEMPORAL_KERNEL_BASIS = {
    "resnet": [
        [[5]],  # conv1 temporal kernel.
        [[1]],  # res2 temporal kernel.
        [[1]],  # res3 temporal kernel.
        [[3]],  # res4 temporal kernel.
        [[3]],  # res5 temporal kernel.
    ],
}

_POOL1 = {
    "resnet": [[1, 1, 1]],
}


@MODEL_REGISTRY.register()
class ResNet(nn.Module):
    """
    ResNet model builder. It builds a ResNet like network backbone.
    """

    def __init__(self, cfg):
        """
        Args:
            cfg (CfgNode): model building configs, details are in the
                comments of the config file.
        """
        super(ResNet, self).__init__()
        self.pooling = cfg.MODEL.POOLING
        self.num_pathways = 1
        self._construct_network(cfg)

    def _compute_dim_in(
        self,
        idx,
        trans_func,
        width_per_group,
    ):
        """
        Compute the input dimension of each convolutional stage.
        args:
            idx (int): the index of the convolutional stage.
            trans_func (string): transform function to be used to contrusct each
                ResBlock.
            width_per_group (int): width of each group.
        returns:
            dim_in (list): list containing the input dimension.
        """
        if trans_func == 'basic_transform':
            factor = 1 if idx == 0 else 2 ** (idx - 1)
        elif trans_func == 'bottleneck_transform':
            factor = 1 if idx == 0 else 2 * (2 ** idx)
        else:
            raise NotImplementedError(
                "Does not support {} transfomration".format(trans_func)
            )

        dim_in = [width_per_group * factor]
        return dim_in

    def _compute_dim_out(
        self,
        idx,
        trans_func,
        width_per_group,
    ):
        """
        Compute the output dimension of each convolutional stage.
        args:
            idx (int): the index of the convolutional stage.
            trans_func (string): transform function to be used to contrusct each
                ResBlock.
            width_per_group (int): width of each group.
        returns:
            dim_out (list): list containing the output dimension.
        """
        if trans_func == 'basic_transform':
            factor = 2 ** idx
        elif trans_func == 'bottleneck_transform':
            factor = 4 * (2 ** idx)
        else:
            raise NotImplementedError(
                "Does not support {} transfomration".format(trans_func)
            )

        dim_out = [width_per_group * factor]
        return dim_out

    def _construct_network(self, cfg):
        """
        Builds a ResNet model.

        Args:
            cfg (CfgNode): model building configs, details are in the
                comments of the config file.
        """
        assert cfg.VIS.ARCH in _POOL1.keys()
        pool_size = _POOL1[cfg.VIS.ARCH]
        assert len({len(pool_size), self.num_pathways}) == 1
        assert cfg.RESNET.DEPTH in _MODEL_STAGE_DEPTH.keys()

        (d2, d3, d4, d5) = _MODEL_STAGE_DEPTH[cfg.RESNET.DEPTH]

        num_groups = cfg.RESNET.NUM_GROUPS
        width_per_group = cfg.RESNET.WIDTH_PER_GROUP
        dim_inner = num_groups * width_per_group

        temp_kernel = _TEMPORAL_KERNEL_BASIS[cfg.VIS.ARCH]

        self.s1 = stem_helper.VideoModelStem(
            dim_in=cfg.DATA.INPUT_CHANNEL_NUM,
            dim_out=[width_per_group],
            kernel=[temp_kernel[0][0] + [7, 7]],
            stride=[[2, 2, 2]],
            padding=[[temp_kernel[0][0][0] // 2, 3, 3]],
            eps=cfg.MODEL.EPSILON,
            bn_mmt=cfg.MODEL.MOMENTUM,
        )

        dim_in_l = [
            self._compute_dim_in(
                i,
                cfg.RESNET.TRANS_FUNC,
                width_per_group
            )
            for i in range(4)
        ]

        dim_out_l = [
            self._compute_dim_out(
                i,
                cfg.RESNET.TRANS_FUNC,
                width_per_group
            )
            for i in range(4)
        ]

        self.s2 = resnet_helper.ResStage(
            dim_in=dim_in_l[0],
            dim_out=dim_out_l[0],
            dim_inner=[dim_inner],
            temp_kernel_sizes=temp_kernel[1],
            stride=cfg.RESNET.SPATIAL_STRIDES[0],
            num_blocks=[d2],
            num_groups=[num_groups],
            num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[0],
            trans_func_name=cfg.RESNET.TRANS_FUNC,
            stride_1x1=cfg.RESNET.STRIDE_1X1,
            inplace_relu=cfg.RESNET.INPLACE_RELU,
            dilation=cfg.RESNET.SPATIAL_DILATIONS[0],
            eps=cfg.MODEL.EPSILON,
            bn_mmt=cfg.MODEL.MOMENTUM,
        )

        for pathway in range(self.num_pathways):
            pool = nn.MaxPool3d(
                kernel_size=pool_size[pathway],
                stride=pool_size[pathway],
                padding=[0, 0, 0],
            )
            self.add_module("pathway{}_pool".format(pathway), pool)

        self.s3 = resnet_helper.ResStage(
            dim_in=dim_in_l[1],
            dim_out=dim_out_l[1],
            dim_inner=[dim_inner * 2],
            temp_kernel_sizes=temp_kernel[2],
            stride=cfg.RESNET.SPATIAL_STRIDES[1],
            num_blocks=[d3],
            num_groups=[num_groups],
            num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[1],
            trans_func_name=cfg.RESNET.TRANS_FUNC,
            stride_1x1=cfg.RESNET.STRIDE_1X1,
            inplace_relu=cfg.RESNET.INPLACE_RELU,
            dilation=cfg.RESNET.SPATIAL_DILATIONS[1],
            eps=cfg.MODEL.EPSILON,
            bn_mmt=cfg.MODEL.MOMENTUM,
        )

        self.s4 = resnet_helper.ResStage(
            dim_in=dim_in_l[2],
            dim_out=dim_out_l[2],
            dim_inner=[dim_inner * 4],
            temp_kernel_sizes=temp_kernel[3],
            stride=cfg.RESNET.SPATIAL_STRIDES[2],
            num_blocks=[d4],
            num_groups=[num_groups],
            num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[2],
            trans_func_name=cfg.RESNET.TRANS_FUNC,
            stride_1x1=cfg.RESNET.STRIDE_1X1,
            inplace_relu=cfg.RESNET.INPLACE_RELU,
            dilation=cfg.RESNET.SPATIAL_DILATIONS[2],
            eps=cfg.MODEL.EPSILON,
            bn_mmt=cfg.MODEL.MOMENTUM,
        )

        self.s5 = resnet_helper.ResStage(
            dim_in=dim_in_l[3],
            dim_out=dim_out_l[3],
            dim_inner=[dim_inner * 8],
            temp_kernel_sizes=temp_kernel[4],
            stride=cfg.RESNET.SPATIAL_STRIDES[3],
            num_blocks=[d5],
            num_groups=[num_groups],
            num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[3],
            trans_func_name=cfg.RESNET.TRANS_FUNC,
            stride_1x1=cfg.RESNET.STRIDE_1X1,
            inplace_relu=cfg.RESNET.INPLACE_RELU,
            dilation=cfg.RESNET.SPATIAL_DILATIONS[3],
            eps=cfg.MODEL.EPSILON,
            bn_mmt=cfg.MODEL.MOMENTUM,
        )

        _num_frames = cfg.DATA.NUM_FRAMES // 2

        if self.pooling:
            self.head = head_helper.ResNetPoolingHead(
                dim_in=dim_out_l[3],
                pool_size=[
                    [
                        _num_frames // pool_size[0][0],
                        math.ceil(cfg.DATA.CROP_SIZE / 32) // pool_size[0][1],
                        math.ceil(cfg.DATA.CROP_SIZE / 32) // pool_size[0][2],
                    ]
                ],
            )
        else:
            self.head = head_helper.ResNetBasicHead(
                dim_in=dim_out_l[3],
                num_classes=cfg.MODEL.NUM_CLASSES,
                pool_size=[
                    [
                        _num_frames // pool_size[0][0],
                        math.ceil(cfg.DATA.CROP_SIZE / 32) // pool_size[0][1],
                        math.ceil(cfg.DATA.CROP_SIZE / 32) // pool_size[0][2],
                    ]
                ],
                dropout_rate=cfg.MODEL.DROPOUT_RATE,
                act_func=cfg.MODEL.HEAD_ACT,
                fusion=cfg.MODEL.DOWNSTREAM_FUSION,
            )

    def get_feature_map(self, x):
        x = self.s1(x)
        x = self.s2(x)
        for pathway in range(self.num_pathways):
            pool = getattr(self, "pathway{}_pool".format(pathway))
            x[pathway] = pool(x[pathway])
        x = self.s3(x)
        x = self.s4(x)
        x = self.s5(x)
        return x

    def get_logit(self, feature_map):
        return self.head(feature_map)

    def forward(self, x):
        x = self.s1(x)
        x = self.s2(x)
        for pathway in range(self.num_pathways):
            pool = getattr(self, "pathway{}_pool".format(pathway))
            x[pathway] = pool(x[pathway])
        x = self.s3(x)
        x = self.s4(x)
        x = self.s5(x)
        x = self.head(x)
        return x

upstream_models/avbert/avbert/models/video_model_builder.py
roger-tseng/av-superb
import os
import shutil
from collections import OrderedDict

import torch

from . import logging as logging

logger = logging.get_logger(__name__)


def load_checkpoint(model, state_dict, data_parallel=False):
    """
    Load the trained weights from the checkpoint.
    Args:
        model (model): model to load the weights from the checkpoint.
        state_dict (OrderedDict): checkpoint.
        data_parallel (bool): if true, model is wrapped by
            torch.nn.parallel.DistributedDataParallel.
    """
    ms = model.module if data_parallel else model
    ms.load_state_dict(state_dict)


def load_finetune_checkpoint(model, state_dict, data_parallel=False, use_trans=False):
    """
    Load the pretrained weights from the checkpoint.
    Args:
        model (model): model to load the weights from the checkpoint.
        state_dict (OrderedDict): checkpoint.
        data_parallel (bool): if true, model is wrapped by
            torch.nn.parallel.DistributedDataParallel.
        use_trans (bool): if true, load the Transformer weights.
    """
    ms = model.module if data_parallel else model
    model_dict = ms.state_dict()
    partial_dict = OrderedDict()
    if use_trans:
        for key in state_dict.keys():
            if "_B" not in key and "head" not in key:
                partial_dict[key] = state_dict[key]
    else:
        for key in state_dict.keys():
            if 'visual_conv' in key and 'head' not in key:
                partial_dict[key[7:]] = state_dict[key]
            if 'audio_conv' in key and 'head' not in key:
                partial_dict[key[7:]] = state_dict[key]

    update_dict = {k: v for k, v in partial_dict.items() if k in model_dict}
    missing_keys, unexpected_keys = ms.load_state_dict(update_dict, strict=True)
    return missing_keys, unexpected_keys

def save_checkpoint(state, is_best=False, filename='checkpoint.pyth'):
    """
    Save the model weights to the checkpoint.
    Args:
        state (Dict): model states
        is_best (bool): whether the model has achieved the best performance so far.
        filename (str): path to the checkpoint to save.
    """
    torch.save(state, filename)
    if is_best:
        shutil.copyfile(filename, 'model_best.pyth')

upstream_models/avbert/avbert/utils/checkpoint.py
roger-tseng/av-superb
"""Logging."""

import builtins
import decimal
import logging
import sys
import simplejson

from . import distributed as du


def _suppress_print():
    """
    Suppresses printing from the current process.
    """

    def print_pass(*objects, sep=" ", end="\n", file=sys.stdout, flush=False):
        pass

    builtins.print = print_pass


def setup_logging(log_path=None):
    """
    Sets up the logging for multiple processes. Only enable the logging for the
    master process, and suppress logging for the non-master processes.
    """
    # Set up logging format.
    _FORMAT = "[%(asctime)s: %(levelname)s: %(filename)s: %(lineno)4d]: %(message)s"

    if du.is_master_proc():
        # Enable logging for the master process.
        # logging.root.handlers = []
        handlers = [logging.FileHandler(log_path),
                    logging.StreamHandler()] if log_path is not None else \
            [logging.StreamHandler()]
        logging.basicConfig(
            level=logging.INFO, datefmt='%m/%d/%Y %H:%M:%S',
            format=_FORMAT, handlers=handlers
        )
    else:
        # Suppress logging for non-master processes.
        _suppress_print()


def get_logger(name):
    """
    Retrieve the logger with the specified name or, if name is None, return a
    logger which is the root logger of the hierarchy.
    Args:
        name (string): name of the logger.
    """
    return logging.getLogger(name)


def log_json_stats(stats):
    """
    Logs json stats.
    Args:
        stats (dict): a dictionary of statistical information to log.
    """
    stats = {
        k: decimal.Decimal("{:.6f}".format(v)) if isinstance(v, float) else v
        for k, v in stats.items()
    }
    json_stats = simplejson.dumps(stats, sort_keys=True, use_decimal=True)
    logger = get_logger(__name__)
    logger.info("json_stats: {:s}".format(json_stats))

upstream_models/avbert/avbert/utils/logging.py
roger-tseng/av-superb
"""Meters."""
import datetime
import numpy as np
import os
from collections import defaultdict, deque
import torch
from fvcore.common.timer import Timer

from . import logging as logging
from . import metrics as metrics

logger = logging.get_logger(__name__)


class ScalarMeter(object):
    """
    A scalar meter uses a deque to track a series of scaler values with a given
    window size. It supports calculating the median and average values of the
    window, and also supports calculating the global average.
    """

    def __init__(self, window_size):
        """
        Args:
            window_size (int): size of the max length of the deque.
        """
        self.deque = deque(maxlen=window_size)
        self.total = 0.0
        self.count = 0

    def reset(self):
        """
        Reset the deque.
        """
        self.deque.clear()
        self.total = 0.0
        self.count = 0

    def add_value(self, value):
        """
        Add a new scalar value to the deque.
        """
        self.deque.append(value)
        self.count += 1
        self.total += value

    def get_win_median(self):
        """
        Calculate the current median value of the deque.
        """
        return np.median(self.deque)

    def get_win_avg(self):
        """
        Calculate the current average value of the deque.
        """
        return np.mean(self.deque)

    def get_global_avg(self):
        """
        Calculate the global mean value.
        """
        return self.total / self.count


class ClassifyTrainMeter(object):
    """
    Measure classfication training stats.
    """

    def __init__(self, writer, epoch_iters, num_epochs, num_iters, cfg):
        """
        Args:
            writer (SummaryWriter): tensorboard summary writer.
            epoch_iters (int): the overall number of iterations of one epoch.
            num_epochs (int): the overall number of epochs
            num_iters (int): the overall number of iterations
            cfg (CfgNode): configs.
        """
        self._cfg = cfg
        self.writer = writer
        self.epoch_iters = epoch_iters
        self.NUM_EPOCHS = int(num_epochs)
        self.NUM_ITERS = int(num_iters)
        self.iter_timer = Timer()
        self.loss = ScalarMeter(cfg.LOG_PERIOD)
        self.loss_total = 0.0
        self.lr = None
        # Current minibatch accs (smoothed over a window).
        self.mb_top1_acc = ScalarMeter(cfg.LOG_PERIOD)
        self.mb_top5_acc = ScalarMeter(cfg.LOG_PERIOD)
        # Number of correctly classified examples.
        self.num_top1_acc = 0
        self.num_top5_acc = 0
        self.num_samples = 0

    def reset(self):
        """
        Reset the Meter.
        """
        self.loss.reset()
        self.loss_total = 0.0
        self.lr = None
        self.mb_top1_acc.reset()
        self.mb_top5_acc.reset()
        self.num_top1_acc = 0
        self.num_top5_acc = 0
        self.num_samples = 0

    def iter_tic(self):
        """
        Start to record time.
        """
        self.iter_timer.reset()

    def iter_toc(self):
        """
        Stop to record time.
        """
        self.iter_timer.pause()

    def update_stats(self, top1_acc, top5_acc, loss, lr, mb_size):
        """
        Update the current stats.
        Args:
            top1_acc (float): top1 accuracy.
            top5_acc (float): top5 accuracy.
            loss (float): loss value.
            lr (float): learning rate.
            mb_size (int): mini batch size.
        """
        self.loss.add_value(loss)
        self.lr = lr
        self.loss_total += loss * mb_size
        self.num_samples += mb_size

        # Current minibatch stats
        self.mb_top1_acc.add_value(top1_acc)
        self.mb_top5_acc.add_value(top5_acc)
        # Aggregate stats
        self.num_top1_acc += top1_acc * mb_size
        self.num_top5_acc += top5_acc * mb_size

    @property
    def use_plot(self):
        """
        Whether to log the stats on tensorboard.
        """
        return self.writer is not None

    def log_iter_stats(self, cur_epoch, cur_iter, global_iter):
        """
        log the stats of the current iteration.
        Args:
            cur_epoch (int): the number of current epoch.
            cur_iter (int): the number of current iteration in current epoch.
            global_iter (int): the number of current iteration in training phase
        """
        if (cur_iter + 1) % self._cfg.LOG_PERIOD != 0:
            return
        eta_sec = self.iter_timer.seconds() * (
            self.NUM_ITERS - (cur_epoch * self.epoch_iters + cur_iter + 1)
        )
        eta = str(datetime.timedelta(seconds=int(eta_sec)))
        stats = {
            "_type": "train_iter",
            "global_iter": "{}/{}".format(global_iter, self.NUM_ITERS),
            "epoch": "{}/{}".format(cur_epoch + 1, self.NUM_EPOCHS),
            "iter": "{}/{}".format(cur_iter + 1, self.epoch_iters),
            "time_diff": self.iter_timer.seconds(),
            "eta": eta,
            "loss": self.loss.get_win_median(),
            "lr": self.lr,
        }
        stats["top1_acc"] = self.mb_top1_acc.get_win_median()
        stats["top5_acc"] = self.mb_top5_acc.get_win_median()
        logging.log_json_stats(stats)
        if self.use_plot:
            self.writer.add_scalar(
                "Train/iter_loss",
                stats["loss"],
                global_iter,
            )
            self.writer.add_scalar(
                "Train/iter_top1_acc",
                stats["top1_acc"],
                global_iter,
            )
            self.writer.add_scalar(
                "Train/iter_top5_acc",
                stats["top5_acc"],
                global_iter,
            )

    def log_epoch_stats(self, cur_epoch, global_iter):
        """
        Log the stats of the current epoch.
        Args:
            cur_epoch (int): the number of current epoch.
            global_iter (int): the number of current iteration in training phase
        """
        eta_sec = self.iter_timer.seconds() * (
            self.NUM_ITERS - (cur_epoch + 1) * self.epoch_iters
        )
        eta = str(datetime.timedelta(seconds=int(eta_sec)))
        stats = {
            "_type": "train_epoch",
            "global_iter": "{}/{}".format(global_iter, self.NUM_ITERS),
            "epoch": "{}/{}".format(cur_epoch + 1, self.NUM_EPOCHS),
            "time_diff": self.iter_timer.seconds(),
            "eta": eta,
            "lr": self.lr,
        }
        top1_acc = self.num_top1_acc / self.num_samples
        top5_acc = self.num_top5_acc / self.num_samples
        avg_loss = self.loss_total / self.num_samples
        stats["top1_acc"] = top1_acc
        stats["top5_acc"] = top5_acc
        stats["loss"] = avg_loss
        logging.log_json_stats(stats)
        if self.use_plot:
            self.writer.add_scalar(
                "Train/epoch_loss",
                stats["loss"],
                cur_epoch + 1,
            )
            self.writer.add_scalar(
                "Train/epoch_top1_acc",
                stats["top1_acc"],
                cur_epoch + 1,
            )
            self.writer.add_scalar(
                "Train/epoch_top5_acc",
                stats["top5_acc"],
                cur_epoch + 1,
            )


class ClassifyValMeter(object):
    """
    Measures classification validation stats.
    """

    def __init__(self, writer, epoch_iters, num_epochs, cfg):
        """
        Args:
            writer (SummaryWriter): tensorboard summary writer.
            epoch_iters (int): the overall number of iterations of one epoch.
            num_epochs (int): the overall number of epochs
            cfg (CfgNode): configs.
        """
        self._cfg = cfg
        self.writer = writer
        self.epoch_iters = epoch_iters
        self.NUM_EPOCHS = int(num_epochs)
        self.iter_timer = Timer()
        # Current minibatch accs (smoothed over a window).
        self.mb_top1_acc = ScalarMeter(cfg.LOG_PERIOD)
        self.mb_top5_acc = ScalarMeter(cfg.LOG_PERIOD)
        # Max accs (over the full val set).
        self.max_top1_acc = 0.0
        self.max_top5_acc = 0.0
        # Number of correctly classified examples.
        self.num_top1_acc = 0
        self.num_top5_acc = 0
        self.num_samples = 0

    def reset(self):
        """
        Reset the Meter.
        """
        self.iter_timer.reset()
        self.mb_top1_acc.reset()
        self.mb_top5_acc.reset()
        self.num_top1_acc = 0
        self.num_top5_acc = 0
        self.num_samples = 0

    def iter_tic(self):
        """
        Start to record time.
        """
        self.iter_timer.reset()

    def iter_toc(self):
        """
        Stop to record time.
        """
        self.iter_timer.pause()

    def update_stats(self, top1_acc, top5_acc, mb_size):
        """
        Update the current stats.
        Args:
            top1_acc (float): top1 accuracy.
            top5_acc (float): top5 accuracy.
            mb_size (int): mini batch size.
        """
        self.mb_top1_acc.add_value(top1_acc)
        self.mb_top5_acc.add_value(top5_acc)
        self.num_top1_acc += top1_acc * mb_size
        self.num_top5_acc += top5_acc * mb_size
        self.num_samples += mb_size

    @property
    def use_plot(self):
        """
        Whether to log the stats on tensorboard.
        """
        return self.writer is not None

    def log_iter_stats(self, cur_epoch, cur_iter):
        """
        log the stats of the current iteration.
        Args:
            cur_epoch (int): the number of current epoch.
            cur_iter (int): the number of current iteration.
        """
        if (cur_iter + 1) % self._cfg.LOG_PERIOD != 0:
            return
        eta_sec = self.iter_timer.seconds() * (self.epoch_iters - cur_iter - 1)
        eta = str(datetime.timedelta(seconds=int(eta_sec)))
        stats = {
            "_type": "val_iter",
            "epoch": "{}/{}".format(cur_epoch + 1, self.NUM_EPOCHS),
            "iter": "{}/{}".format(cur_iter + 1, self.epoch_iters),
            "time_diff": self.iter_timer.seconds(),
            "eta": eta,
        }
        stats["top1_acc"] = self.mb_top1_acc.get_win_median()
        stats["top5_acc"] = self.mb_top5_acc.get_win_median()
        logging.log_json_stats(stats)

    def log_epoch_stats(self, cur_epoch):
        """
        Log the stats of the current epoch.
        Args:
            cur_epoch (int): the number of current epoch.
        """
        stats = {
            "_type": "val_epoch",
            "epoch": "{}/{}".format(cur_epoch + 1, self.NUM_EPOCHS),
            "time_diff": self.iter_timer.seconds(),
        }
        top1_acc = self.num_top1_acc / self.num_samples
        top5_acc = self.num_top5_acc / self.num_samples
        self.max_top1_acc = max(self.max_top1_acc, top1_acc)
        self.max_top5_acc = max(self.max_top5_acc, top5_acc)

        stats["top1_acc"] = top1_acc
        stats["top5_acc"] = top5_acc
        stats["max_top1_acc"] = self.max_top1_acc
        stats["max_top5_acc"] = self.max_top5_acc
        logging.log_json_stats(stats)
        if self.use_plot:
            self.writer.add_scalar(
                "Val/epoch_top1_acc",
                stats["top1_acc"],
                cur_epoch + 1,
            )
            self.writer.add_scalar(
                "Val/epoch_top5_acc",
                stats["top5_acc"],
                cur_epoch + 1,
            )


class ClassifyTestMeter(object):
    """
    Perform the multi-view ensemble for testing: each data point with an unique index
    will be sampled with multiple clips, and the predictions of the clips will
    be aggregated to produce the final prediction for the data point.
    The accuracy is calculated with the given ground truth labels.
    """

    def __init__(
        self,
        num_points,
        num_clips,
        num_cls,
        overall_iters,
        ensemble_method="sum",
        log_period=1,
    ):
        """
        Construct tensors to store the predictions and labels. Expect to get
        num_clips predictions from each data point, and calculate the metrics on
        num_points data points.
        Args:
            num_points (int): number of data points to test.
            num_clips (int): number of clips sampled from each data point for
                aggregating the final prediction for the data point.
            num_cls (int): number of classes for each prediction.
            overall_iters (int): overall iterations for testing.
            ensemble_method (str): method to perform the ensemble, options
                include "sum", and "max".
            log_period (int): log period.
        """

        self.iter_timer = Timer()
        self.num_clips = num_clips
        self.overall_iters = overall_iters
        self.ensemble_method = ensemble_method
        self.log_period = log_period
        # Initialize tensors.
        self.point_preds = torch.zeros((num_points, num_cls))

        self.point_labels = torch.zeros((num_points)).long()
        self.clip_count = torch.zeros((num_points)).long()
        self.view_count = torch.zeros((num_points, num_clips)).long()
        # Reset metric.
        self.reset()

    def reset(self):
        """
        Reset the metric.
        """
        self.clip_count.zero_()
        self.view_count.zero_()
        self.point_preds.zero_()
        self.point_labels.zero_()

    def update_stats(self, preds, labels, clip_ids):
        """
        Collect the predictions from the current batch and perform on-the-flight
        summation as ensemble.
        Args:
            preds (tensor): predictions from the current batch. Dimension is
                N x C where N is the batch size and C is the channel size
                (num_cls).
            labels (tensor): the corresponding labels of the current batch.
                Dimension is N.
            clip_ids (tensor): clip indexes of the current batch, dimension is
                N.
        """
        for ind in range(preds.shape[0]):
            point_id = int(clip_ids[ind]) // self.num_clips
            clip_id = int(clip_ids[ind]) % self.num_clips
            if self.view_count[point_id, clip_id] == 1:
                continue
            if self.point_labels[point_id].sum() > 0:
                assert torch.equal(
                    self.point_labels[point_id].type(torch.FloatTensor),
                    labels[ind].type(torch.FloatTensor),
                )
            self.point_labels[point_id] = labels[ind]
            if self.ensemble_method == "sum":
                self.point_preds[point_id] += preds[ind]
            elif self.ensemble_method == "max":
                self.point_preds[point_id] = torch.max(
                    self.point_preds[point_id], preds[ind]
                )
            else:
                raise NotImplementedError(
                    "Ensemble Method {} is not supported".format(
                        self.ensemble_method
                    )
                )
            self.clip_count[point_id] += 1
            self.view_count[point_id, clip_id] = 1

    def log_iter_stats(self, cur_iter):
        """
        Log the stats.
        Args:
            cur_iter (int): the current iteration of testing.
        """
        if (cur_iter + 1) % self.log_period != 0:
            return
        eta_sec = self.iter_timer.seconds() * (self.overall_iters - cur_iter - 1)
        eta = str(datetime.timedelta(seconds=int(eta_sec)))
        stats = {
            "split": "test_iter",
            "cur_iter": "{}/{}".format(cur_iter + 1, self.overall_iters),
            "eta": eta,
            "time_diff": self.iter_timer.seconds(),
        }
        logging.log_json_stats(stats)

    def iter_tic(self):
        """
        Start to record time.
        """
        self.iter_timer.reset()

    def iter_toc(self):
        """
        Stop to record time.
        """
        self.iter_timer.pause()

    def finalize_metrics(self, ks=(1, 5), cur_epoch=None, num_epochs=None, writer=None):
        """
        Calculate and log the final ensembled metrics.
        ks (tuple): list of top-k values for topk_accuracies. For example,
            ks = (1, 5) correspods to top-1 and top-5 accuracy.
        """
        if not all(self.clip_count == self.num_clips):
            logger.warning(
                "clip count {} ~= num clips {}".format(
                    ", ".join(
                        [
                            "{}: {}".format(i, k)
                            for i, k in enumerate(self.clip_count.tolist())
                        ]
                    ),
                    self.num_clips,
                )
            )

        stats = {"split": "test_final"}
        if cur_epoch is not None and num_epochs is not None:
            stats['epoch'] = "{}/{}".format(cur_epoch + 1, num_epochs)
        num_topks_correct = metrics.topks_correct(
            self.point_preds, self.point_labels, ks
        )
        topks = [
            (x / self.point_preds.size(0)) * 100.0
            for x in num_topks_correct
        ]
        assert len({len(ks), len(topks)}) == 1
        for k, topk in zip(ks, topks):
            stats["top{}_acc".format(k)] = "{:.{prec}f}".format(
                topk, prec=2
            )
        logging.log_json_stats(stats)
        if writer is not None and cur_epoch is not None and num_epochs is not None:
            for k, topk in zip(ks, topks):
                writer.add_scalar(
                    f"Test/epoch_top{k}_acc",
                    topk.item(),
                    cur_epoch + 1,
                )

        return stats

upstream_models/avbert/avbert/utils/meters.py
roger-tseng/av-superb
import math
from datetime import datetime

from . import logging as logging

logger = logging.get_logger(__name__)


def check_nan_losses(loss):
    """
    Determine whether the loss is NaN (not a number).
    Args:
        loss (loss): loss to check whether is NaN.
    """
    if math.isnan(loss):
        raise RuntimeError("ERROR: Got NaN losses {}".format(datetime.now()))


def is_eval_epoch(eval_period, cur_epoch, num_epochs):
    """
    Determine if the model should be evaluated at the current epoch.
    """
    return (
        cur_epoch + 1
    ) % eval_period == 0 or cur_epoch + 1 == num_epochs

upstream_models/avbert/avbert/utils/misc.py
roger-tseng/av-superb
import math
from collections import OrderedDict
from typing import Dict, List, Tuple, Union

import torch
import torch.nn as nn
import torchaudio
import torchaudio.transforms as aT
import torchvision
from torch import Tensor
from torch.nn.utils.rnn import pad_sequence
from .preprocess_function import get_audio_seq, get_visual_seq, resample
from .avbert.utils import checkpoint as cu
from .avbert.config import get_cfg
from .avbert.models.avbert import AVBert


class UpstreamExpert(nn.Module):
    def __init__(self, ckpt: str = None, model_config: str = None, **kwargs):
        """
        Args:
            ckpt:
                checkpoint path for loading pretrained weights.

            model_config:
                config path for your model.
        """
        super().__init__()

        self.cfg = get_cfg()
        self.feature_concat_axis = kwargs['feature_concat_axis']

        # NOTE: Encoders should return (batch_size, seq_len, hidden_dims)
        
        weights = torch.load(
            ckpt, map_location='cpu'
        )['state_dict']
        weights = OrderedDict(
            {
                k.split('.', 1)[1]: weights[k]
                for k in weights.keys()
                if k.startswith("avbert.")
            }
        )

        self.multi_encoder = AVBert(self.cfg)
        missing_keys, unexpected_keys = cu.load_finetune_checkpoint(
            self.multi_encoder,
            weights,
            self.cfg.NUM_GPUS > 1,
            self.cfg.MODEL.USE_TRANSFORMER,
        )
        assert len(missing_keys) == 0 and len(unexpected_keys) == 0

    def preprocess_video(self, video, video_frame_rate):
        """
        Replace this function to preprocess videos into your input format
        video: (video_length, video_channels, height, width), where video_channels is usually 3 for RGB or 1 for greyscale
        """

        return video

    def preprocess_audio(self, audio, audio_sample_rate):
        """
        Replace this function to preprocess audio waveforms into your input format
        audio: (audio_channels, audio_length), where audio_channels is usually 1 or 2
        """

        return audio
    
    def preprocess(self, video, audio, video_frame_rate, audio_sample_rate):
        
        video = video.permute(0, 2, 3 ,1)

        visual_seqs = []

        num_frames = (
            self.cfg.DATA.NUM_FRAMES *
            self.cfg.DATA.SAMPLING_RATE *
            video_frame_rate /
            self.cfg.DATA.TARGET_FPS
        )

        waveform_size = int(
            self.cfg.DATA.TARGET_AUDIO_RATE *
            self.cfg.DATA.NUM_FRAMES *
            self.cfg.DATA.SAMPLING_RATE /
            self.cfg.DATA.TARGET_FPS
        )

        visual_delta = max(video.size(0) - num_frames, 0)
        for spatial_sample_index in range(self.cfg.TEST.NUM_SPATIAL_CROPS):
            visual_start_idx = [
                visual_delta * temporal_sample_index / (self.cfg.TEST.NUM_ENSEMBLE_VIEWS - 1)
                for temporal_sample_index in range(self.cfg.TEST.NUM_ENSEMBLE_VIEWS)
            ]
            visual_end_idx = [s + num_frames - 1 for s in visual_start_idx]
            visual_seq = get_visual_seq(
                video,
                visual_start_idx,
                visual_end_idx,
                self.cfg.DATA.NUM_FRAMES,
                self.cfg.DATA.TEST_CROP_SIZE,
                spatial_sample_index,
            )
            visual_seqs.append(visual_seq)

        visual_seqs = torch.stack(visual_seqs)

        if len(audio.shape) == 1:
            waveform = audio.unsqueeze(0)
        else:
            waveform = audio.mean(dim=0, keepdim=True)

        waveform = resample(
            waveform,
            audio_sample_rate,
            self.cfg.DATA.TARGET_AUDIO_RATE,
            use_mono=True,
        )

        audio_delta = max(waveform.size(-1) - waveform_size, 0)
        audio_start_idx = [
            int(audio_delta * (idx / visual_delta)) if visual_delta != 0 else 0
            for idx in visual_start_idx
        ]
        audio_end_idx = [s + waveform_size for s in audio_start_idx]

        audio_seq = get_audio_seq(
            waveform,
            audio_start_idx,
            audio_end_idx,
            self.cfg.DATA.TARGET_AUDIO_RATE,
            self.cfg.DATA.AUDIO_FREQUENCY,
            self.cfg.DATA.AUDIO_TIME,
        )
        # visual: 10 1-sec clips, 3 views per clip, 32 frames per view, 3 x 128 x 128 frame
        return visual_seqs, audio_seq


    def forward(
        self, source: List[Tuple[Tensor, Tensor]]
    ) -> Dict[str, Union[Tensor, List[Tensor]]]:
        """
        Replace this function run a forward pass with your model
        source: list of audio-video Tensor tuples
                [(wav1,vid1), (wav2,vid2), ...]
                in your input format
        """
        audio, video = zip(*source)

        # Collate audio and video into batch
        audios = pad_sequence(audio, batch_first=True)
        videos = torch.stack(video)

        videos = videos.transpose(0, 1).contiguous()

        _ , single_hiddens_0, multi_hidden_0 = self.multi_encoder(
            visual_seq=[videos[0]], audio_seq=audios
        )

        _ , single_hiddens_1, multi_hidden_1 = self.multi_encoder(
            visual_seq=[videos[1]], audio_seq=audios
        )

        _ , single_hiddens_2, multi_hidden_2 = self.multi_encoder(
            visual_seq=[videos[2]], audio_seq=audios
        )

        # Concat features of three views along time axis (B x 3L x 768) or along hidden dim axis (B x L x 3*768)
        audio_feats = single_hiddens_0[1]
        if self.feature_concat_axis == 'time':
            video_feats = []
            for layer_0, layer_1, layer_2 in zip(single_hiddens_0[0], single_hiddens_1[0], single_hiddens_2[0]):
                stacked = torch.stack((layer_0, layer_1, layer_2), dim = 2)
                video_feats.append(torch.flatten(stacked, start_dim=1, end_dim=2))
            fusion_feats = []
            for layer_0, layer_1, layer_2 in zip(multi_hidden_0, multi_hidden_1, multi_hidden_2):
                stacked = torch.stack((layer_0, layer_1, layer_2), dim = 2)
                fusion_feats.append(torch.flatten(stacked, start_dim=1, end_dim=2))
        elif self.feature_concat_axis == 'hidden':
            video_feats = tuple(torch.cat((layer_0, layer_1, layer_2), dim = 2) for layer_0, layer_1, layer_2 in zip(single_hiddens_0[0], single_hiddens_1[0], single_hiddens_2[0]))
            fusion_feats = tuple(torch.cat((layer_0, layer_1, layer_2), dim = 2) for layer_0, layer_1, layer_2 in zip(multi_hidden_0, multi_hidden_1, multi_hidden_2))
        else:
            raise NotImplementedError


        # Return intermediate layer representations for potential layer-wise experiments
        # Dict should contain three items, with keys as listed below:
        # video_feats: features that only use visual modality as input
        # audio_feats: features that only use auditory modality as input
        # fusion_feats: features that consider both modalities
        # Each item should be a list of features that are of the same shape
        return {
            "video_feats": video_feats,
            "audio_feats": audio_feats,
            "fusion_feats": fusion_feats,
        }

upstream_models/avbert/expert.py
roger-tseng/av-superb
import os

from utils.download import _urls_to_filepaths

from .expert import UpstreamExpert as _UpstreamExpert

def google_large_file_link(file_id):
    return f"https://drive.google.com/u/0/uc?id={file_id}&export=download&confirm=t&uuid=b73dc818-c9b0-49e8-a488-13b829fdbb7e&at=ANzk5s4TfW9_erYs4M9PH7KYuKzy:1681055195563"

def avbert_local(ckpt, *args, **kwargs):
    """
    The model from local ckpt
        ckpt (str): PATH
    """
    assert os.path.isfile(ckpt)
    return _UpstreamExpert(ckpt, *args, **kwargs)


def avbert_url(ckpt, refresh=False, *args, **kwargs):
    """
    The model from URL
        ckpt (str): URL
    """
    return avbert_local(_urls_to_filepaths(ckpt, refresh=refresh, agent="gdown"), *args, **kwargs)

def avbert_time(refresh=False, *args, **kwargs):
    """
    Interleave features of three video crops along time axis
        refresh (bool): whether to download ckpt/config again if existed
    """
    kwargs["ckpt"] = google_large_file_link("1T55DuJq0OYSdXVnDlOrgPSTwffVKsiFY")
    kwargs["feature_concat_axis"] = 'time'
    return avbert_url(refresh=refresh, *args, **kwargs)

def avbert_hidden(refresh=False, *args, **kwargs):
    """
    Concat features of three video crops along hidden dim axis
        refresh (bool): whether to download ckpt/config again if existed
    """
    kwargs["ckpt"] = google_large_file_link("1T55DuJq0OYSdXVnDlOrgPSTwffVKsiFY")
    kwargs["feature_concat_axis"] = 'hidden'
    return avbert_url(refresh=refresh, *args, **kwargs)

def avbert(refresh=False, *args, **kwargs):
    """
    Interleave features along time axis by default
    """
    return avbert_time(refresh=refresh, *args, **kwargs)

upstream_models/avbert/hubconf.py
roger-tseng/av-superb
import torch
import torchaudio
import math
import numpy as np

def get_visual_seq(
    frames,
    start_idx,
    end_idx,
    video_num_frames,
    crop_size,
    spatial_sample_index,
):
    """
    Sample a sequence of clips from the input video, and apply
    visual transformations.
    args:
        frames (tensor): a tensor of video frames, dimension is
            `num frames` x `height` x `width` x `channel`.
        start_idx (list): list of the index of the start frame of each clip.
        end_idx (list): list of the index of the end frame of each clip.
        min_scale (int): the minimal size of scaling.
        max_scale (int): the maximal size of scaling.
        crop_size (int): the size of height and width used to crop the
            frames.
        spatial_sample_index (int): if -1, perform random spatial sampling.
            If 0, 1, or 2, perform left, center, right crop if width is
            larger than height, and perform top, center, buttom crop if
            height is larger than width.
    returns:
        clip_seq (tensor): a sequence of sampled frames. The dimension is
            `sequence length` x `channel` x `num frames` x `height` x `width`.
    """
    # Temporal sampling.
    clip_seq = []
    for s, e in zip(start_idx, end_idx):
        clip_seq.append(
            temporal_sampling(
                frames,
                s,
                e,
                video_num_frames,
            )
        )
    clip_seq = torch.stack(clip_seq, dim=0)

    # Convert frames of the uint type in the range [0, 255] to
    # a torch.FloatTensor in the range [0.0, 1.0]
    clip_seq = clip_seq.float()
    clip_seq = clip_seq / 255.0

    # S T H W C -> SxT C H W
    clip_seq = clip_seq.view(
        clip_seq.shape[0] * clip_seq.shape[1],
        clip_seq.shape[2],
        clip_seq.shape[3],
        clip_seq.shape[4],
    )
    clip_seq = clip_seq.permute(0, 3, 1, 2)
    # Visual transformations.
    clip_seq = apply_visual_transform(
        clip_seq,
        spatial_idx=spatial_sample_index,
        crop_size=crop_size,
    )
    # SxT C H W -> S C T H W
    clip_seq = clip_seq.reshape(
        len(start_idx),
        video_num_frames,
        clip_seq.shape[1],
        clip_seq.shape[2],
        clip_seq.shape[3],
    )
    clip_seq = clip_seq.transpose(1, 2).contiguous()

    return clip_seq

def apply_visual_transform(
    frames,
    spatial_idx,
    crop_size = 128,
):
    """
    Apply visual data transformations on the given video frames.
    Data transformations include `resize_crop`, `flip` and `color_normalize`.
    Args:
        cfg (CfgNode): configs.
        frames (tensor): frames of images sampled from the video. The
            dimension is `num frames` x `channel` x `height` x `width`.
        spatial_idx (int): if -1, perform random spatial sampling. If 0, 1,
            or 2, perform left, center, right crop if width is larger than
            height, and perform top, center, buttom crop if height is larger
            than width.
        min_scale (int): the minimal size of scaling.
        max_scale (int): the maximal size of scaling.
        crop_size (int): the size of height and width used to crop the
            frames.
    Returns:
        frames (tensor): spatially sampled frames.
    """
    assert spatial_idx in [0, 1, 2]
    # The testing is deterministic and no jitter should be performed.
    # min_scale, max_scale, and crop_size are expect to be the same.

    frames = short_side_scale_jitter(frames, crop_size)
    frames = uniform_crop(frames, crop_size, spatial_idx)
    frames = color_normalization(frames, (0.45, 0.45, 0.45), (0.225, 0.225, 0.225))
    return frames

def short_side_scale_jitter(
    images, crop_size
):
    """
    Perform a spatial short scale jittering on the given images.
    Args:
        images (tensor): images to perform scale jitter. Dimension is
            `num frames` x `channel` x `height` x `width`.
        min_size (int): the minimal size to scale the frames.
        max_size (int): the maximal size to scale the frames.
    Returns:
        (tensor): the scaled images with dimension of
            `num frames` x `channel` x `new height` x `new width`.
    """
    size = crop_size

    height = images.shape[2]
    width = images.shape[3]
    if (width <= height and width == size) or (
        height <= width and height == size
    ):
        return images
    new_width = size
    new_height = size
    if width < height:
        new_height = int(math.floor((float(height) / width) * size))
    else:
        new_width = int(math.floor((float(width) / height) * size))

    return torch.nn.functional.interpolate(
        images,
        size=(new_height, new_width),
        mode="nearest",
        # align_corners=False,
    )

def uniform_crop(images, size, spatial_idx):
    """
    Perform uniform spatial sampling on the images.
    args:
        images (tensor): images to perform uniform crop. The dimension is
            `num frames` x `channel` x `height` x `width`.
        size (int): size of height and weight to crop the images.
        spatial_idx (int): 0, 1, or 2 for left, center, and right crop if width
            is larger than height. or 0, 1, or 2 for top, center, and bottom
            crop if height is larger than width.
    returns:
        cropped (tensor): images with dimension of
            `num frames` x `channel` x `size` x `size`.
    """
    assert spatial_idx in [0, 1, 2]
    height = images.shape[2]
    width = images.shape[3]

    y_offset = int(math.ceil((height - size) / 2))
    x_offset = int(math.ceil((width - size) / 2))

    if height > width:
        if spatial_idx == 0:
            y_offset = 0
        elif spatial_idx == 2:
            y_offset = height - size
    else:
        if spatial_idx == 0:
            x_offset = 0
        elif spatial_idx == 2:
            x_offset = width - size
    cropped = images[
        :, :, y_offset : y_offset + size, x_offset : x_offset + size
    ]

    return cropped

def color_normalization(images, mean, stddev):
    """
    Perform color nomration on the given images.
    Args:
        images (tensor): images to perform color normalization. Dimension is
            `num frames` x `channel` x `height` x `width`.
        mean (list): mean values for normalization.
        stddev (list): standard deviations for normalization.

    Returns:
        out_images (tensor): the noramlized images, the dimension is
            `num frames` x `channel` x `height` x `width`.
    """
    assert len(mean) == images.shape[1], "channel mean not computed properly"
    assert (
        len(stddev) == images.shape[1]
    ), "channel stddev not computed properly"

    out_images = torch.zeros_like(images)
    for idx in range(len(mean)):
        out_images[:, idx] = (images[:, idx] - mean[idx]) / stddev[idx]

    return out_images

def temporal_sampling(frames, start_idx, end_idx, num_samples):
    """
    Given the start and end frame index, sample num_samples frames between
    the start and end with equal interval.
    Args:
        frames (tensor): a tensor of video frames, dimension is
            `num video frames` x `channel` x `height` x `width`.
        start_idx (int): the index of the start frame.
        end_idx (int): the index of the end frame.
        num_samples (int): number of frames to sample.
    Returns:
        frames (tersor): a tensor of temporal sampled video frames, dimension is
            `num clip frames` x `channel` x `height` x `width`.
    """
    index = torch.linspace(start_idx, end_idx, num_samples)
    index = torch.clamp(index, 0, frames.shape[0] - 1).long()
    frames = torch.index_select(frames, 0, index)
    return frames

def get_audio_seq(
    waveform,
    start_idx,
    end_idx,
    audio_fps,
    frequency,
    time
):
    """
    Sample a sequence of clips from the input audio, and apply
    audio transformations.
    args:
        waveform (tensor): a tensor of audio waveform, dimension is
            `channel` x `time`.
        start_idx (list): list of the start index.
        end_idx (list): list of the end index (not inclusive).
        apply_transform (bool): whether to apply transformations.
    returns:
        (tensor): a sequence of log-mel-scaled spectrogram with dimension of
            `sequence length` x `channel` x `frequency` x `time`.
    """
    audio_seq = []
    for s, e in zip(start_idx, end_idx):
        # Temporal sampling.
        waveform_view = waveform[:, s:e]
        # Convert it to log-mel-scaled spectrogram.
        log_mel_spectrogram = get_log_mel_spectrogram(
            waveform_view,
            audio_fps,
            frequency,
            time
        )
        audio_seq.append(log_mel_spectrogram)
    # S x C x F x T
    audio_seq = torch.stack(audio_seq, dim=0)

    return audio_seq

def resample(waveform, orig_freq, new_freq, use_mono=True):
    """
    Resample the input waveform to ``new_freq``.
    args:
        waveform (tensor): waveform to perform resampling. The dimension is
            `channel` x `frequency` x `width`.
        `orig_freq` (int): original sampling rate of `waveform`.
        `new_freq` (int): target sampling rate of `waveform`.
        `use_mono` (bool): If True, first convert `waveform` to a monophonic signal.
    returns:
         (tensor): waveform with dimension of
            `channel` x `time`.
    """
    if waveform.size(0) != 1 and use_mono:
        waveform = waveform.mean(0, keepdim=True)

    if orig_freq != new_freq:
        waveform = torchaudio.transforms.Resample(
            orig_freq, new_freq,
        )(waveform)

    return waveform

def get_log_mel_spectrogram(
    waveform,
    audio_fps,
    frequency,
    time,
):
    """
    Convert the input waveform to log-mel-scaled spectrogram.
    args:
        waveform (tensor): input waveform. The dimension is
            `channel` x `time.`
        `audio_fps` (int): sampling rate of `waveform`.
        `frequency` (int): target frequecy dimension (number of mel bins).
        `time` (int): target time dimension.
    returns:
        (tensor): log-mel-scaled spectrogram with dimension of
            `channel` x `frequency` x `time`.
    """
    w = waveform.size(-1)
    n_fft = 2 * (math.floor(w / time) + 1)
    mel_spectrogram = torchaudio.transforms.MelSpectrogram(
        audio_fps, n_fft=n_fft, n_mels=frequency,
    )(waveform)
    log_mel_spectrogram = torch.log(1e-6 + mel_spectrogram)
    _nchannels, _frequency, _time = log_mel_spectrogram.size()
    assert _frequency == frequency, \
        f"frequency {_frequency} must be {frequency}"
    if _time != time:
        t = torch.zeros(
            _nchannels,
            frequency,
            time,
            dtype=log_mel_spectrogram.dtype,
        )
        min_time = min(time, _time)
        t[:, :, :min_time] = log_mel_spectrogram[:, :, :min_time]
        log_mel_spectrogram = t

    return log_mel_spectrogram

upstream_models/avbert/preprocess_function.py
roger-tseng/av-superb
"""
Custom class for loading audio-visual model and extract features
Modified from https://github.com/s3prl/s3prl/blob/main/s3prl/upstream/example/expert.py
"""
from typing import Dict, List, Tuple, Union

import torch
import torch.nn as nn
import torchaudio
import torchaudio.transforms as aT
import torchvision
from torch import Tensor
from torch.nn.utils.rnn import pad_sequence

HIDDEN_DIM = 8


class UpstreamExpert(nn.Module):
    def __init__(self, ckpt: str = None, model_config: str = None, **kwargs):
        """
        Args:
            ckpt:
                checkpoint path for loading pretrained weights.

            model_config:
                config path for your model.
        """
        super().__init__()

        self.model1 = nn.Linear(2, HIDDEN_DIM)
        self.model2 = nn.Linear(HIDDEN_DIM, HIDDEN_DIM)

        self.audio_sample_rate = 16000
        self.melspec_transform = [aT.MelSpectrogram(self.audio_sample_rate, n_mels=80)]

        self.video_frame_size = (224, 224)
        self.video_frame_rate = 25

        # NOTE: Encoders should return (batch_size, seq_len, hidden_dims)
        self.audio_encoder = lambda x: x.reshape(x.size(0), 1, -1).mean(
            dim=-1, keepdim=True
        )
        self.video_encoder = lambda x: x.float().reshape(x.size(0), 1, -1).mean(
            dim=-1, keepdim=True
        )

    def preprocess_video(self, video, video_frame_rate):
        """
        Replace this function to preprocess videos into your input format
        video: (video_length, video_channels, height, width), where video_channels is usually 3 for RGB or 1 for greyscale
        """
        # Resize video frames
        video_frames = []
        for frame in video:
            video_frames.append(
                torchvision.transforms.functional.resize(
                    frame, self.video_frame_size, antialias=False
                )
            )
        video = torch.stack(video_frames)

        # Resample video
        # (from https://github.com/pytorch/vision/blob/5b07d6c9c6c14cf88fc545415d63021456874744/torchvision/datasets/video_utils.py#L278)
        step = float(video_frame_rate) / self.video_frame_rate
        if step.is_integer():
            # optimization: if step is integer, don't need to perform
            # advanced indexing
            step = int(step)
            idxs = slice(None, None, step)
        else:
            num_frames = int(len(video) / step)
            idxs = torch.arange(num_frames, dtype=torch.float32) * step
            idxs = idxs.floor().to(torch.int64)
        video = video[idxs]

        # Other preprocessing steps (i.e. cropping, flipping, etc.)
        # e.g. take first three frames to ensure all videos have same size
        video = video[:3]
        return video.float()

    def preprocess_audio(self, audio, audio_sample_rate):
        """
        Replace this function to preprocess audio waveforms into your input format
        audio: (audio_channels, audio_length), where audio_channels is usually 1 or 2
        """
        # Resample audio
        if audio_sample_rate != self.audio_sample_rate:
            audio = torchaudio.functional.resample(
                audio, audio_sample_rate, self.audio_sample_rate
            )

        # Other preprocessing steps (e.g. trimming, transform to melspectrogram etc.)
        mel_specgram = self.melspec_transform[0](audio).transpose(0, 1)

        return mel_specgram

    def forward(
        self, source: List[Tuple[Tensor, Tensor]]
    ) -> Dict[str, Union[Tensor, List[Tensor]]]:
        """
        Replace this function run a forward pass with your model
        source: list of audio-video Tensor tuples
                [(wav1,vid1), (wav2,vid2), ...]
                in your input format
        """
        audio, video = zip(*source)

        # Collate audio and video into batch
        audios = pad_sequence(audio, batch_first=True)
        videos = torch.stack(video)

        # Run through audio and video encoders
        audio_feats = self.audio_encoder(audios)
        video_feats = self.video_encoder(videos)

        layer1 = self.model1(torch.cat((audio_feats, video_feats), dim=-1))

        layer2 = self.model2(layer1)

        # Return intermediate layer representations for potential layer-wise experiments
        # Dict should contain three items, with keys as listed below:
        # video_feats: features that only use visual modality as input
        # audio_feats: features that only use auditory modality as input
        # fusion_feats: features that consider both modalities
        # Each item should be a list of features that are of the same shape
        return {
            "video_feats": [video_feats],
            "audio_feats": [audio_feats],
            "fusion_feats": [layer1, layer2],
        }

upstream_models/example/expert.py
roger-tseng/av-superb
# Author: Yuan Tseng
# Modified from S3PRL
# (Authors: Leo Yang, Andy T. Liu and S3PRL team, https://github.com/s3prl/s3prl/blob/main/s3prl/upstream/baseline/expert.py

# -*- coding: utf-8 -*- #
"""*********************************************************************************************"""
#   FileName     [ upstream/baseline/extracter.py ]
#   Synopsis     [ the baseline feature extracter with torchaudio.compliance.kaldi as backend ]
#   Author       [ S3PRL ]
#   Copyright    [ Copyleft(c), Speech Lab, NTU, Taiwan ]
"""*********************************************************************************************"""

import yaml
from torch.nn.utils.rnn import pad_sequence

from interfaces import UpstreamBase
from .extracter import get_extracter
import torchaudio
# from .preprocessor import get_preprocessor

SAMPLE_RATE = 16000

###################
# UPSTREAM EXPERT #
###################
class UpstreamExpert(UpstreamBase):
    """
    Extract baseline features from wavforms by torchaudio.compliance.kaldi or torchaudio preprocessor
    Support: fbank
    """

    def __init__(self, model_config, **kwargs):
        super().__init__(**kwargs)

        self.audio_sample_rate = 16000
        
        with open(model_config, "r") as file:
            self.config = yaml.load(file, Loader=yaml.FullLoader)

        if "kaldi" in self.config:
            self.extracter, self.output_dim, frame_shift = get_extracter(self.config)
            self.downsample_rate = round(frame_shift * SAMPLE_RATE / 1000)
        else:
            raise NotImplementedError
            # self.extracter, self.output_dim, _ = get_preprocessor(
            #     self.config, process_input_only=True
            # )
            # self.downsample_rate = round(
            #     self.config.get("hop_ms", 10) * SAMPLE_RATE / 1000
            # )

    def _extractor_forward(self, wavs):
        feats = []
        for wav in wavs:
            feats.append(self.extracter(wav))
        return feats

    def get_downsample_rates(self, key: str) -> int:
        return self.downsample_rate

    def preprocess_video(self, video, video_frame_rate):
        return video[0][0][0]
    
    def preprocess_audio(self, audio, audio_sample_rate):
        """
        Replace this function to preprocess audio waveforms into your input format
        audio: (audio_channels, audio_length), where audio_channels is usually 1 or 2
        """
        if len(audio.shape) == 2:
            audio = audio.mean(dim=0)

        # Resample audio
        if audio_sample_rate != self.audio_sample_rate:
            audio = torchaudio.functional.resample(
                audio, audio_sample_rate, self.audio_sample_rate
            )

        return audio
    
    def forward(self, source):

        wavs, video = zip(*source)

        if "kaldi" in self.config:
            feats = self._extractor_forward(wavs)
        else:
            raise NotImplementedError

        padded_feats = pad_sequence(feats, batch_first=True)
        return {
            "audio_feats": [padded_feats],
        }
upstream_models/fbank/expert.py
roger-tseng/av-superb
# Author: Yuan Tseng
# Modified from S3PRL
# (Authors: Leo Yang, Andy T. Liu and S3PRL team, https://github.com/s3prl/s3prl/blob/main/s3prl/upstream/baseline/extracter.py)

# -*- coding: utf-8 -*- #
"""*********************************************************************************************"""
#   FileName     [ upstream/baseline/extracter.py ]
#   Synopsis     [ the baseline feature extracter with torchaudio.compliance.kaldi as backend ]
#   Author       [ S3PRL ]
#   Copyright    [ Copyleft(c), Speech Lab, NTU, Taiwan ]
"""*********************************************************************************************"""

###############
# IMPORTATION #
###############
import copy
from collections import namedtuple

# -------------#
import torch
import torch.nn as nn
import torch.nn.functional as F

# -------------#
import torchaudio
from torchaudio import transforms

############
# CONSTANT #
############
SAMPLE_RATE = 16000
EXAMPLE_SEC = 5


def get_extracter(config):
    transforms = [
        ExtractAudioFeature(**config.get("kaldi", {})),
        Delta(**config.get("delta", {})),
        CMVN(**config.get("cmvn", {})),
    ]
    extracter = nn.Sequential(*transforms)
    output_dim = extracter(torch.randn(EXAMPLE_SEC * SAMPLE_RATE)).size(-1)

    return extracter, output_dim, extracter[0].frame_shift


class ExtractAudioFeature(nn.Module):
    def __init__(self, feat_type="fbank", **kwargs):
        super(ExtractAudioFeature, self).__init__()
        self.extract_fn = eval(f"torchaudio.compliance.kaldi.{feat_type}")
        self.kwargs = kwargs[feat_type]
        self.frame_shift = self.kwargs.get("frame_shift", 10.0)

    def forward(self, waveform):
        # waveform: (time, )
        x = self.extract_fn(
            waveform.view(1, -1), sample_frequency=SAMPLE_RATE, **self.kwargs
        )
        # x: (feat_seqlen, feat_dim)
        return x


class Delta(nn.Module):
    def __init__(self, order=2, **kwargs):
        super(Delta, self).__init__()
        self.order = order
        self.compute_delta = transforms.ComputeDeltas(**kwargs)

    def forward(self, x):
        # x: (feat_seqlen, feat_dim)
        feats = [x]
        for o in range(self.order):
            feat = feats[-1].transpose(0, 1).unsqueeze(0)
            delta = self.compute_delta(feat)
            feats.append(delta.squeeze(0).transpose(0, 1))
        x = torch.cat(feats, dim=-1)
        # x: (feat_seqlen, feat_dim)
        return x


class CMVN(nn.Module):
    def __init__(self, use_cmvn, eps=1e-10):
        super(CMVN, self).__init__()
        self.eps = eps
        self.use_cmvn = use_cmvn

    def forward(self, x):
        # x: (feat_seqlen, feat_dim)
        if self.use_cmvn:
            x = (x - x.mean(dim=0, keepdim=True)) / (
                self.eps + x.std(dim=0, keepdim=True)
            )
        return x
upstream_models/fbank/extracter.py
roger-tseng/av-superb
# Author: Yuan Tseng
# Modified from S3PRL
# (Authors: Leo Yang, Andy T. Liu and S3PRL team, https://github.com/s3prl/s3prl/blob/main/s3prl/upstream/baseline/hubconf.py)

# -*- coding: utf-8 -*- #
"""*********************************************************************************************"""
#   FileName     [ upstream/baseline/hubconf.py ]
#   Synopsis     [ the baseline torch hubconf ]
#   Author       [ S3PRL ]
#   Copyright    [ Copyleft(c), Speech Lab, NTU, Taiwan ]
"""*********************************************************************************************"""

###############
# IMPORTATION #
###############
import os

# -------------#
from .expert import UpstreamExpert as _UpstreamExpert


def baseline_local(model_config, *args, **kwargs):
    """
    Baseline feature
        model_config: PATH
    """
    assert os.path.isfile(model_config)
    return _UpstreamExpert(model_config, *args, **kwargs)

def fbank(*args, **kwargs):
    """
    Baseline feature - Fbank, or Mel-scale spectrogram
    """
    kwargs["model_config"] = os.path.join(os.path.dirname(__file__), "fbank.yaml")
    return baseline_local(*args, **kwargs)

upstream_models/fbank/hubconf.py
roger-tseng/av-superb
import yaml
import torch
import torch.nn as nn
import torchaudio
import torchaudio.transforms as aT
import torchvision
from torch import Tensor
from torch.nn.utils.rnn import pad_sequence

from interfaces import UpstreamBase

from skimage.feature import hog

###################
# UPSTREAM EXPERT #
###################
class UpstreamExpert(UpstreamBase):

    def __init__(self, model_config, **kwargs):
        super().__init__(**kwargs)

        self.video_frame_size = (128, 64)
        self.video_frame_rate = 25

        with open(model_config, "r") as file:
            self.config = yaml.load(file, Loader=yaml.FullLoader)

    def preprocess_video(self, video, video_frame_rate):
        
        # Resize video frames
        video_frames = []
        for frame in video:
            video_frames.append(
                torchvision.transforms.functional.resize(
                    frame, self.video_frame_size, antialias=False
                )
            )
        video = torch.stack(video_frames)

        # Resample video
        # (from https://github.com/pytorch/vision/blob/5b07d6c9c6c14cf88fc545415d63021456874744/torchvision/datasets/video_utils.py#L278)
        step = float(video_frame_rate) / self.video_frame_rate
        if step.is_integer():
            # optimization: if step is integer, don't need to perform
            # advanced indexing
            step = int(step)
            idxs = slice(None, None, step)
        else:
            num_frames = int(len(video) / step)
            idxs = torch.arange(num_frames, dtype=torch.float32) * step
            idxs = idxs.floor().to(torch.int64)
        video = video[idxs]

        if video.shape[1] == 3:
            video = 0.2989 * video[:, 0] + 0.587 * video[:, 1] + 0.114 * video[:, 2]
        else:
            video = video.mean(dim=1)

        return video
    
    def preprocess_audio(self, audio, audio_sample_rate):
        """
        Replace this function to preprocess audio waveforms into your input format
        audio: (audio_channels, audio_length), where audio_channels is usually 1 or 2
        """
        if len(audio.shape) == 2:
            audio = audio[0]

        return audio[0]
    
    def forward(self, source):

        audio, video = zip(*source)

        # TODO: add here, try out hyperparams that work decently
        # 
        # ( Add moving object detection if useful, as shown in Section 3.1: 
        # https://www.mdpi.com/1424-8220/20/24/7299 )
        #

        videos = pad_sequence(video, batch_first = True)

        video_feats = []
        for video in videos:
            video_feat = []
            for frame in video:
                feat = hog(frame.cpu(), orientations=8, pixels_per_cell=(16, 16), cells_per_block=(1, 1))
                feat = torch.from_numpy(feat)
                video_feat.append(feat)
            video_feat = torch.stack(video_feat)
            video_feats.append(video_feat)
        video_feats = torch.stack(video_feats).cuda()

        return {
            "video_feats": [video_feats],
            "audio_feats": [],
            "fusion_feats": [],
        }

upstream_models/hog/expert.py
roger-tseng/av-superb
# -*- coding: utf-8 -*- #
"""*********************************************************************************************"""
#   FileName     [ Modified from https://github.com/s3prl/s3prl/blob/main/s3prl/upstream/baseline/hubconf.py ]
#   Synopsis     [ the baseline torch hubconf ]
#   Author       [ S3PRL (Leo Yang) ]
#   Copyright    [ Copyleft(c), Speech Lab, NTU, Taiwan ]
"""*********************************************************************************************"""


###############
# IMPORTATION #
###############
import os

# -------------#
from .expert import UpstreamExpert as _UpstreamExpert


def baseline_local(model_config, *args, **kwargs):
    """
    Baseline feature
        model_config: PATH
    """
    assert os.path.isfile(model_config)
    return _UpstreamExpert(model_config, *args, **kwargs)

def hog(*args, **kwargs):
    """
    Baseline feature - Histogram of gradients
    """
    kwargs["model_config"] = os.path.join(os.path.dirname(__file__), "hog.yaml")
    return baseline_local(*args, **kwargs)

upstream_models/hog/hubconf.py
roger-tseng/av-superb
# Author: Yuan Tseng
# Modified from S3PRL
# (https://github.com/s3prl/s3prl/blob/main/s3prl/upstream/hubert/expert.py

# Copyright (c) Facebook, Inc. All Rights Reserved

# -*- coding: utf-8 -*- #
"""*********************************************************************************************"""
#   FileName     [ upstream/hubert/expert.py ]
#   Synopsis     [ the HuBERT wrapper ]
#   Author       [ Kushal Lakhotia ]
"""*********************************************************************************************"""

import logging

import torch
import torch.nn.functional as F
from torch.nn.utils.rnn import pad_sequence
import torchaudio

import fairseq
from interfaces import UpstreamBase

logger = logging.getLogger(__name__)


class UpstreamExpert(UpstreamBase):
    def __init__(self, ckpt, **kwargs):
        super().__init__(**kwargs)
    
        model, cfg, task = fairseq.checkpoint_utils.load_model_ensemble_and_task([ckpt])
        self.model = model[0]
        self.task = task

        self.model.feature_grad_mult = 0.0
        self.model.encoder.layerdrop = 0.0
        self.audio_sample_rate = 16000

        if len(self.hooks) == 0:
            module_name = "self.model.encoder.layers"
            for module_id in range(len(eval(module_name))):
                self.add_hook(
                    f"{module_name}[{module_id}]",
                    lambda input, output: input[0].transpose(0, 1),
                )
            self.add_hook("self.model.encoder", lambda input, output: output[0])

            def postprocess(xs):
                names, hiddens = zip(*xs)
                names = [name + "_audio" for name in names]
                unpad_len = min([hidden.size(1) for hidden in hiddens])
                hiddens = [hidden[:, :unpad_len, :] for hidden in hiddens]
                return list(zip(names, hiddens))

            self.hook_postprocess = postprocess

    def get_downsample_rates(self, key: str) -> int:
        return 320

    def preprocess_video(self, video, video_frame_rate):
        return video[0][0][0]
    
    def preprocess_audio(self, audio, audio_sample_rate):
        """
        Replace this function to preprocess audio waveforms into your input format
        audio: (audio_channels, audio_length), where audio_channels is usually 1 or 2
        """
        if len(audio.shape) == 2:
            audio = audio.mean(dim=0)

        # Resample audio
        if audio_sample_rate != self.audio_sample_rate:
            audio = torchaudio.functional.resample(
                audio, audio_sample_rate, self.audio_sample_rate
            )

        return audio
    
    def forward(self, source):

        wavs, video = zip(*source)

        if self.task.cfg.normalize:
            wavs = [F.layer_norm(wav, wav.shape) for wav in wavs]

        device = wavs[0].device
        wav_lengths = torch.LongTensor([len(wav) for wav in wavs]).to(device)
        wav_padding_mask = ~torch.lt(
            torch.arange(max(wav_lengths)).unsqueeze(0).to(device),
            wav_lengths.unsqueeze(1),
        )
        padded_wav = pad_sequence(wavs, batch_first=True)

        features, feat_padding_mask = self.model.extract_features(
            padded_wav,
            padding_mask=wav_padding_mask,
            mask=None,
        )

        # This forward function only does the model forward
        # The return dict is then handled by UpstreamBase's hooks



upstream_models/hubert/expert.py
roger-tseng/av-superb
# Copyright (c) Facebook, Inc. All Rights Reserved

# -*- coding: utf-8 -*- #
"""*********************************************************************************************"""
#   FileName     [ upstream/hubert/hubconf.py ]
#   Synopsis     [ the HuBERT torch hubconf ]
#   Author       [ S3PRL / Kushal Lakhotia ]
"""*********************************************************************************************"""


import logging
import os
import time
from pathlib import Path

from filelock import FileLock

from utils.download import _urls_to_filepaths

from .expert import UpstreamExpert as _UpstreamExpert

logger = logging.getLogger(__name__)

NEW_ENOUGH_SECS = 2.0


def hubert_custom(
    ckpt: str,
    legacy: bool = False,
    refresh: bool = False,
    **kwargs,
):

    if ckpt.startswith("http"):
        ckpt = _urls_to_filepaths(ckpt, refresh=refresh)

    assert os.path.isfile(ckpt)
    return _UpstreamExpert(ckpt, **kwargs)


def hubert_local(*args, **kwargs):
    return hubert_custom(*args, **kwargs)


def hubert_url(*args, **kwargs):
    return hubert_custom(*args, **kwargs)


def hubert(refresh=False, *args, **kwargs):
    """
    The default model - Base
        refresh (bool): whether to download ckpt/config again if existed
    """
    return hubert_base(refresh=refresh, *args, **kwargs)


def hubert_base(refresh=False, legacy=True, **kwargs):
    """
    The Base model
        refresh (bool): whether to download ckpt/config again if existed
    """
    kwargs["ckpt"] = "https://dl.fbaipublicfiles.com/hubert/hubert_base_ls960.pt"
    return hubert_custom(refresh=refresh, legacy=legacy, **kwargs)


def hubert_large_ll60k(refresh=False, legacy=False, **kwargs):
    """
    The Large model
        refresh (bool): whether to download ckpt/config again if existed
    """
    kwargs["ckpt"] = "https://dl.fbaipublicfiles.com/hubert/hubert_large_ll60k.pt"
    return hubert_custom(refresh=refresh, legacy=legacy, **kwargs)


def hubert_base_robust_mgr(refresh=False, legacy=False, **kwargs):
    """
    The Base model, continually trained with Libri 960 hr with Musan noise, Gaussian noise and Reverberation.
        refresh (bool): whether to download ckpt/config again if existed
    """
    kwargs[
        "ckpt"
    ] = "https://huggingface.co/kphuang68/HuBERT_base_robust_mgr/resolve/main/HuBERT_base_robust_mgr_best_loss_2.7821.pt"
    return hubert_custom(refresh=refresh, legacy=legacy, **kwargs)

upstream_models/hubert/hubconf.py
roger-tseng/av-superb
from typing import Dict, List, Tuple, Union

import numpy as np
import torch
import torch.nn as nn
import torchaudio

from torch import Tensor
from torch.nn.utils.rnn import pad_sequence

from torchvision.transforms import Compose, Resize 
from torchvision.transforms._transforms_video import ToTensorVideo, NormalizeVideo

from . import models_vitmm

from .util.patch_embed import PatchEmbed_new
from timm.models.layers import to_2tuple

class UpstreamExpert(nn.Module):
    def __init__(self, ckpt: str = None, model_config: str = None, **kwargs):
        """
        Args:
            ckpt:
                checkpoint path for loading pretrained weights.

            model_config:
                config path for your model.
        """
        super().__init__()

        av_fusion = True
        ckpt = torch.load(ckpt, map_location='cpu')
        # Fine-tuned model uses three fusion layers
        self.ft = True if any('blocks_av.2' in key for key in ckpt['model'].keys()) else False
        self.model = models_vitmm.vitmm_base_patch16(
                    num_classes=527,
                    drop_path_rate=0.1,
                    global_pool=True,
                    mask_2d=True,
                    av_fusion=av_fusion, 
                    depth_av=3 if av_fusion and self.ft else 2 if av_fusion else 0,
                    n_frm=8, # 8 frames per video
                    pos_train=False,
                    ft=self.ft,
                )

        img_size = (1024, 128) # 1024, 128
        in_chans = 1
        emb_dim = 768
        self.model.patch_embed = PatchEmbed_new(img_size=img_size, patch_size=(16,16), in_chans=in_chans, embed_dim=emb_dim, stride=16) # no overlap. stride=img_size=16
        num_patches = self.model.patch_embed.num_patches
        self.model.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, emb_dim), requires_grad=False)  # fixed sin-cos embedding
        
        checkpoint_model = ckpt['model']
        self.model.load_state_dict(checkpoint_model, strict=False)

        self.audio_conf = {'sample_rate': 16000,
                    'num_mel_bins': 128,
                    'target_length': 1024,  # crop/pad spectogram to roughly 10 seconds
                    # 'mean': -4.2677393,     # precomputed on audioset
                    # 'std': 4.5689974,       # precomputed on audioset
                    'noise': False,
                    }
        
        self.video_frame_size = (224, 224)
        self.video_len = 4                  # MAViL takes 4 sec clips
        self.video_frame_rate = 2           # at 2 frames per second

        self.video_transform = Compose([
            ToTensorVideo(),
            Resize((224,224)),
            NormalizeVideo((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
        ])

    def preprocess_video(self, video, video_frame_rate):
        
        # 1. TCHW -> THWC and repeat until four sec. if too short
        video = torch.permute(video, (0,2,3,1))
        n_target_frm = int(self.video_len * video_frame_rate)
        if video.shape[0] < n_target_frm:
            n_repeats = n_target_frm//video.shape[0] + 1
            video = video.repeat(n_repeats,1,1,1)

        # 2. Resample video
        # (from https://github.com/pytorch/vision/blob/5b07d6c9c6c14cf88fc545415d63021456874744/torchvision/datasets/video_utils.py#L278)
        step = float(video_frame_rate) / self.video_frame_rate
        if step.is_integer():
            # optimization: if step is integer, don't need to perform
            # advanced indexing
            step = int(step)
            idxs = slice(None, None, step)
        else:
            num_frames = int(len(video) / step)
            idxs = torch.arange(num_frames, dtype=torch.float32) * step
            idxs = idxs.floor().to(torch.int64)
        video = video[idxs]

        # 3. Crop to 4 seconds
        video = video[:self.video_len * self.video_frame_rate] # 8 frames per clip

        preprocessed_video = self.video_transform(video)
        return preprocessed_video
    
    def preprocess_audio(self, audio, audio_sample_rate, fbank_mean=None, fbank_std=None):
        if len(audio.shape) == 2:
            audio = audio.mean(dim=0)
        waveform = audio.unsqueeze(0)
        waveform = waveform - waveform.mean()

        # Audio resampling done by fbank

        # 498x128, 998x128, 1000x128
        fbank = torchaudio.compliance.kaldi.fbank(waveform, htk_compat=True, sample_frequency=audio_sample_rate, use_energy=False,
                                                  window_type='hanning', num_mel_bins=self.audio_conf.get('num_mel_bins'), dither=0.0, frame_shift=10)
        target_length = self.audio_conf.get('target_length')
        n_frames = fbank.shape[0]

        p = target_length - n_frames

        # cut and pad
        if p > 0:
            m = torch.nn.ZeroPad2d((0, 0, 0, p))
            fbank = m(fbank)
        elif p < 0:
            fbank = fbank[0:target_length, :]

        # fbank = fbank.transpose(0,1).unsqueeze(0) # 1, 128, 1024 (...,freq,time)
        # fbank = torch.transpose(fbank.squeeze(), 0, 1) # time, freq

        # normalize fbank (allow precomputed mean/std)
        # on-the-fly normalization causes 6% mAP drop on subset of Audioset
        if fbank_mean is not None and fbank_std is not None:
            fbank = (fbank - fbank_mean) / (fbank_std * 2) 
        else:
            fbank = (fbank - fbank.mean()) / (fbank.std() * 2) 

        return fbank.unsqueeze(0)

    def forward(
        self, source: List[Tuple[Tensor, Tensor]]
    ) -> Dict[str, Union[Tensor, List[Tensor]]]:
        """
        Replace this function run a forward pass with your model
        source: list of audio-video Tensor tuples
                [(wav1,vid1), (wav2,vid2), ...]
                in your input format
        """
        audio, video = zip(*source)
        x = torch.stack(audio, dim=0)
        v = torch.stack(video, dim=0)
        v = v[:,:,:self.model.n_frm,:,:]

        B = x.shape[0]

        # audio
        x = self.model.patch_embed(x)
        x = x + self.model.pos_embed[:, 1:, :]
        cls_token = self.model.cls_token + self.model.pos_embed[:, :1, :]
        cls_tokens = cls_token.expand(B, -1, -1)  # cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.model.pos_drop(x)

        # audio encoding
        audio_seq_feats = []
        audio_pooled_feats = []
        for blk in self.model.blocks:
            audio_seq_feats.append(x[:, 1:, :]) # drop CLS token
            audio_pooled_feats.append(x[:, 1:, :].mean(dim=1, keepdim=True)) # drop CLS token and pool
            x = blk(x)

        audio_seq_feats.append(x[:, 1:, :]) # drop CLS token
        audio_pooled_feats.append(x[:, 1:, :].mean(dim=1, keepdim=True)) # drop CLS token and pool

        # video
        v = self.model.patch_embed_v(v)
        v = v + self.model.pos_embed_v[:, 1:, :]
        cls_token_v = self.model.cls_token_v + self.model.pos_embed_v[:, :1, :]
        cls_tokens_v = cls_token_v.expand(B, -1, -1)
        v = torch.cat((cls_tokens_v, v), dim=1)
        v = self.model.pos_drop(v)

        # video encoding
        video_seq_feats = []
        video_pooled_feats = []
        for blk in self.model.blocks_v:
            video_seq_feats.append(v[:, 1:, :]) # drop CLS token
            video_pooled_feats.append(v[:, 1:, :].mean(dim=1, keepdim=True)) # drop CLS token and pool
            v = blk(v)
        
        video_seq_feats.append(v[:, 1:, :]) # drop CLS token
        video_pooled_feats.append(v[:, 1:, :].mean(dim=1, keepdim=True)) # drop CLS token and pool


        if self.model.av_fusion:
            x_len = x.shape[1]
            xv = torch.cat((x,v), dim=1)

            fusion_seq_feats = []
            fusion_pooled_feats = []
            for blk in self.model.blocks_av:
                
                fusion_seq_feats.append(xv) 
                x = xv[:, 1:x_len, :].mean(dim=1, keepdim=True)  # global pool without cls token
                v = xv[:, x_len+1:,:].mean(dim=1, keepdim=True) # global pool without cls token
                fusion_pooled_feats.append(torch.cat((x,v),dim=2))

                xv =  blk(xv)

            fusion_seq_feats.append(xv) 
            x = xv[:, 1:x_len, :].mean(dim=1, keepdim=True)  # global pool without cls token
            v = xv[:, x_len+1:,:].mean(dim=1, keepdim=True) # global pool without cls token
            if self.ft:
                fusion_pooled_feats.append(self.model.fc_norm_av(torch.cat((x,v),dim=2)))
            else:
                fusion_pooled_feats.append(torch.cat((x,v),dim=2))

        # Return intermediate layer representations for potential layer-wise experiments
        # Dict should contain three items, with keys as listed below:
        # video_feats: features that only use visual modality as input
        # audio_feats: features that only use auditory modality as input
        # fusion_feats: features that consider both modalities
        # Each item should be a list of features that are of the same shape
        return {
            "video_feats": video_pooled_feats,
            "audio_feats": audio_pooled_feats,
            "fusion_feats": fusion_pooled_feats,
            "_video_seq_feats": video_seq_feats,
            "_audio_seq_feats": audio_seq_feats,
            "_fusion_seq_feats": fusion_seq_feats,
        }

upstream_models/mavil/expert.py
roger-tseng/av-superb
import os

from utils.download import _urls_to_filepaths

from .expert import UpstreamExpert as _UpstreamExpert

def google_large_file_link(file_id):
    return f"https://drive.google.com/u/0/uc?id={file_id}&export=download&confirm=t&uuid=b73dc818-c9b0-49e8-a488-13b829fdbb7e&at=ANzk5s4TfW9_erYs4M9PH7KYuKzy:1681055195563"

def mavil_local(ckpt, *args, **kwargs):
    """
    The model from local ckpt
        ckpt (str): PATH
    """
    assert os.path.isfile(ckpt)
    return _UpstreamExpert(ckpt, *args, **kwargs)

def mavil_url(ckpt, refresh=False, *args, **kwargs):
    """
    The model from URL
        ckpt (str): URL
    """
    return mavil_local(_urls_to_filepaths(ckpt, refresh=refresh, agent="gdown"), *args, **kwargs)

def mavil_base(refresh=False, *args, **kwargs):
    """
    The pretrained model
        refresh (bool): whether to download ckpt/config again if existed
    """
    kwargs["ckpt"] = google_large_file_link("160pJQDQGlNwIb5xWlSc73A713I4uXCtg")
    return mavil_url(refresh=refresh, *args, **kwargs)

upstream_models/mavil/hubconf.py
roger-tseng/av-superb
# 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.
# --------------------------------------------------------
# References:
# timm: https://github.com/rwightman/pytorch-image-models/tree/master/timm
# DeiT: https://github.com/facebookresearch/deit
# --------------------------------------------------------

from functools import partial

import torch
import torch.nn as nn
import torchaudio
import numpy as np
import timm.models.vision_transformer
from timm.models.vision_transformer import PatchEmbed, Block
from .util.patch_embed import PatchEmbed_new, PatchEmbed3D_new
from .util.pos_embed import get_3d_sincos_pos_embed


class VisionTransformerMM(timm.models.vision_transformer.VisionTransformer):
    """ Mulitmodal Vision Transformer with support for global average pooling (assume loading both audio and video)
    """
    def __init__(self, global_pool=False, mask_2d=True, av_fusion=False, n_frm=8, depth_av=0, pos_train=False, ft=False, roll_mag_aug=False, dataset="audioset", **kwargs):
        super(VisionTransformerMM, self).__init__(**kwargs)
        self.av_fusion=av_fusion
        self.global_pool = global_pool
        if self.global_pool:
            norm_layer = kwargs['norm_layer']
            embed_dim = kwargs['embed_dim']

            if not self.av_fusion:
                self.fc_norm_a = norm_layer(embed_dim)
                self.fc_norm_v = norm_layer(embed_dim)
                self.head_a = nn.Linear(embed_dim, self.num_classes)
                self.head_v = nn.Linear(embed_dim, self.num_classes)

            del self.norm
            # must deleete otherwise distributed 
        else:
            assert(0) # not implemented yet
        self.mask_2d = mask_2d

        # video branch
        num_heads = kwargs['num_heads']
        depth = kwargs['depth']
        mlp_ratio = kwargs['mlp_ratio']

        patch_size=16
        self.temporal_kernel = 2
        self.temporal_stride = 2
        self.n_frm = n_frm

        self.patch_embed_v = PatchEmbed3D_new(video_size=(n_frm, 224, 224), 
                                        patch_size=(self.temporal_kernel, patch_size, patch_size),
                                        stride=(self.temporal_stride, patch_size, patch_size),
                                        in_chans=3, embed_dim=embed_dim)
        
        self.num_patches_v = self.patch_embed_v.num_patches
        self.cls_token_v = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.pos_embed_v = nn.Parameter(torch.zeros(1, self.num_patches_v + 1, embed_dim), requires_grad=pos_train)  
        self.blocks_v = nn.ModuleList([Block(embed_dim, num_heads, mlp_ratio, qkv_bias=True, qk_scale=None, norm_layer=norm_layer) for i in range(depth)])

        self.av_fusion=av_fusion
        self.depth_av=depth_av
        if self.av_fusion:
            self.blocks_av = nn.ModuleList([
                Block(embed_dim, num_heads, mlp_ratio, qkv_bias=True, qk_scale=None, norm_layer=norm_layer)
            for i in range(self.depth_av)])
            if ft:
                self.fc_norm_av = norm_layer(2*embed_dim)
                self.head_av = nn.Linear(2*embed_dim, self.num_classes)

        del self.head
        # must deleete otherwise distributed complains
        self.initialize_weights()


    def initialize_weights(self):
        pass

    def no_weight_decay(self):
        return {'pos_embed', 'cls_token', 'pos_embed_v', 'cls_token_v'}

    def random_masking(self, x, mask_ratio):
        """
        Perform per-sample random masking by per-sample shuffling.
        Per-sample shuffling is done by argsort random noise.
        x: [N, L, D], sequence
        """
        N, L, D = x.shape  # batch, length, dim
        len_keep = int(L * (1 - mask_ratio))
        
        noise = torch.rand(N, L, device=x.device)  # noise in [0, 1]
        
        # sort noise for each sample
        ids_shuffle = torch.argsort(noise, dim=1)  # ascend: small is keep, large is remove
        ids_restore = torch.argsort(ids_shuffle, dim=1)

        # keep the first subset
        ids_keep = ids_shuffle[:, :len_keep]
        x_masked = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D))

        # generate the binary mask: 0 is keep, 1 is remove
        mask = torch.ones([N, L], device=x.device)
        mask[:, :len_keep] = 0
        # unshuffle to get the binary mask
        mask = torch.gather(mask, dim=1, index=ids_restore)

        return x_masked, mask, ids_restore

    def random_masking_2d(self, x, mask_t_prob, mask_f_prob):
        """
        2D: Spectrogram (msking t and f under mask_t_prob and mask_f_prob)
        Perform per-sample random masking by per-sample shuffling.
        Per-sample shuffling is done by argsort random noise.
        x: [N, L, D], sequence
        """
        
        N, L, D = x.shape  # batch, length, dim
        # ## for AS 
        T=64
        F=8
        
        # mask T
        x = x.reshape(N, T, F, D)
        len_keep_T = int(T * (1 - mask_t_prob))
        noise = torch.rand(N, T, device=x.device)  # noise in [0, 1]
        # sort noise for each sample
        ids_shuffle = torch.argsort(noise, dim=1)  # ascend: small is keep, large is remove
        ids_keep = ids_shuffle[:, :len_keep_T]
        index = ids_keep.unsqueeze(-1).unsqueeze(-1).repeat(1, 1, F, D)
        x = torch.gather(x, dim=1, index=index) # N, len_keep_T(T'), F, D

        # mask F
        x = x.permute(0,2,1,3) # N T' F D => N F T' D
        len_keep_F = int(F * (1 - mask_f_prob))
        noise = torch.rand(N, F, device=x.device)  # noise in [0, 1]
        # sort noise for each sample
        ids_shuffle = torch.argsort(noise, dim=1)  # ascend: small is keep, large is remove
        ids_keep = ids_shuffle[:, :len_keep_F]
        index = ids_keep.unsqueeze(-1).unsqueeze(-1).repeat(1, 1, len_keep_T, D)
        x_masked = torch.gather(x, dim=1, index=index)
        x_masked = x_masked.permute(0,2,1,3) # N F' T' D => N T' F' D 
        x_masked = x_masked.reshape(N,len_keep_F*len_keep_T,D)
            
        return x_masked, None, None


    def forward_features_a(self, x):
        B = x.shape[0]

        # audio
        x = self.patch_embed(x)
        x = x + self.pos_embed[:, 1:, :]
        cls_token = self.cls_token + self.pos_embed[:, :1, :]
        cls_tokens = cls_token.expand(B, -1, -1)  # cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.pos_drop(x)        

        for blk in self.blocks:
            x = blk(x)

        if self.global_pool:
            x = x[:, 1:, :].mean(dim=1)  # global pool without cls token
            outcome = self.fc_norm_a(x)
        else:
            x = self.norm(x)
            outcome = x[:, 0]
        return outcome     

    def forward_features_av(self, x, v):
        B = x.shape[0]

        # audio
        x = self.patch_embed(x)
        x = x + self.pos_embed[:, 1:, :]
        cls_token = self.cls_token + self.pos_embed[:, :1, :]
        cls_tokens = cls_token.expand(B, -1, -1)  # cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.pos_drop(x)

        # audio encoding
        for blk in self.blocks:
            x = blk(x)

        # video
        v = self.patch_embed_v(v)
        v = v + self.pos_embed_v[:, 1:, :]
        cls_token_v = self.cls_token_v + self.pos_embed_v[:, :1, :]
        cls_tokens_v = cls_token_v.expand(B, -1, -1)
        v = torch.cat((cls_tokens_v, v), dim=1)
        v = self.pos_drop(v)

        # video encoding
        for blk in self.blocks_v:
            v = blk(v)

        if self.av_fusion:
            xv = torch.cat((x,v), dim=1)
            for blk in self.blocks_av:
                xv =  blk(xv)        
            x_len = x.shape[1]
            x = xv[:, 1:x_len, :].mean(dim=1)  # global pool without cls token
            v = xv[:, x_len+1:,:].mean(dim=1) # global pool without cls token
            outcome_av=self.fc_norm_av(torch.cat((x,v),dim=1))
            return outcome_av, outcome_av
        else:
            x = x[:, 1:, :].mean(dim=1)  # global pool without cls token
            v = v[:, 1:, :].mean(dim=1)  # global pool without cls token

            outcome_a = self.fc_norm_a(x)
            outcome_v = self.fc_norm_v(v)
        return outcome_a, outcome_v

    def forward_features_av_mask(self, x, v, mask_t_prob, mask_f_prob, mask_v_prob):
        B = x.shape[0] #4,1,1024,128
        # audio
        x = self.patch_embed(x) # 4, 512, 768
        x = x + self.pos_embed[:, 1:, :]
        if self.random_masking_2d:
            x, mask, ids_restore = self.random_masking_2d(x, mask_t_prob, mask_f_prob)
        else:
            x, mask, ids_restore = self.random_masking(x, mask_t_prob)
        cls_token = self.cls_token + self.pos_embed[:, :1, :]
        cls_tokens = cls_token.expand(B, -1, -1)
        x = torch.cat((cls_tokens, x), dim=1)        
        x = self.pos_drop(x)

        for blk in self.blocks:
            x = blk(x)        
        
        # video
        v = self.patch_embed_v(v)
        v = v + self.pos_embed_v[:, 1:, :]
        v, mask, ids_restore = self.random_masking(v, mask_v_prob)
        cls_token_v = self.cls_token_v + self.pos_embed_v[:, :1, :]
        cls_tokens_v = cls_token_v.expand(B, -1, -1)
        v = torch.cat((cls_tokens_v, v), dim=1)
        v = self.pos_drop(v)

        for blk in self.blocks_v:
            v = blk(v)

        # av fusion
        if self.av_fusion:
            xv = torch.cat((x,v), dim=1)
            for blk in self.blocks_av:
                xv = blk(xv)
            x_len = x.shape[1]
            x = xv[:, 1:x_len, :].mean(dim=1)  # global pool without cls token
            v = xv[:, x_len+1:,:].mean(dim=1)
            outcome_av=self.fc_norm_av(torch.cat((x,v),dim=1))
            return outcome_av, outcome_av
        else:
            x = x[:, 1:, :].mean(dim=1)  # global pool without cls token
            v = v[:, 1:, :].mean(dim=1)  # global pool without cls token
            outcome_a = self.fc_norm_a(x)
            outcome_v = self.fc_norm_v(v)
            return outcome_a, outcome_v


    # overwrite original timm
    #def forward(self, x, v, mask_t_prob=0.0, mask_f_prob=0.0, mask_v_prob=0.0):
    def forward(self, source, mask_t_prob=0.0, mask_f_prob=0.0, mask_v_prob=0.0):
        x, v = [], []
        for audio, video in source:
            x.append(audio.unsqueeze(0))
            v.append(video.unsqueeze(0))
        x = torch.cat(x, dim=0)
        v = torch.cat(v, dim=0)
        # x = x.to(device, non_blocking=True)
        # v = v.to(device, non_blocking=True)

        v = v[:,:,:self.n_frm,:,:]
        # add 10/23 for test
        if mask_t_prob > 0.0 or mask_f_prob > 0.0 or mask_v_prob > 0.0:
            x, v = self.forward_features_av_mask(x, v, mask_t_prob=mask_t_prob, mask_f_prob=mask_f_prob, mask_v_prob=mask_v_prob)
        else:
            x, v = self.forward_features_av(x, v)
        if self.av_fusion:
            x = self.head_av(x)
            return x, x.detach()
        else:
            x = self.head_a(x)
            v = self.head_v(v)
            return x, v #x.detach()


def vitmm_small_patch16(**kwargs):
    model = VisionTransformerMM(
        patch_size=16, embed_dim=384, depth=12, num_heads=6, mlp_ratio=4, qkv_bias=True,
        norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)        
    return model

def vitmm_base_patch16(**kwargs):
    model = VisionTransformerMM(
        patch_size=16, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4, qkv_bias=True,
        norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
    return model


def vitmm_large_patch16(**kwargs):
    model = VisionTransformerMM(
        patch_size=16, embed_dim=1024, depth=24, num_heads=16, mlp_ratio=4, qkv_bias=True,
        norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
    return model


def vitmm_huge_patch14(**kwargs):
    model = VisionTransformerMM(
        patch_size=14, embed_dim=1280, depth=32, num_heads=16, mlp_ratio=4, qkv_bias=True,
        norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
    return model

upstream_models/mavil/models_vitmm.py
roger-tseng/av-superb
# 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.
# --------------------------------------------------------
# References:
# ELECTRA https://github.com/google-research/electra
# BEiT: https://github.com/microsoft/unilm/tree/master/beit
# --------------------------------------------------------

import json


def param_groups_lrd(model, weight_decay=0.05, no_weight_decay_list=[], layer_decay=.75):
    """
    Parameter groups for layer-wise lr decay
    Following BEiT: https://github.com/microsoft/unilm/blob/master/beit/optim_factory.py#L58
    """
    param_group_names = {}
    param_groups = {}

    num_layers = len(model.blocks) + 1

    layer_scales = list(layer_decay ** (num_layers - i) for i in range(num_layers + 1))

    for n, p in model.named_parameters():
        if not p.requires_grad:
            continue

        # no decay: all 1D parameters and model specific ones
        if p.ndim == 1 or n in no_weight_decay_list:
            g_decay = "no_decay"
            this_decay = 0.
        else:
            g_decay = "decay"
            this_decay = weight_decay
            
        layer_id = get_layer_id_for_vit(n, num_layers)
        group_name = "layer_%d_%s" % (layer_id, g_decay)

        if group_name not in param_group_names:
            this_scale = layer_scales[layer_id]

            param_group_names[group_name] = {
                "lr_scale": this_scale,
                "weight_decay": this_decay,
                "params": [],
            }
            param_groups[group_name] = {
                "lr_scale": this_scale,
                "weight_decay": this_decay,
                "params": [],
            }

        param_group_names[group_name]["params"].append(n)
        param_groups[group_name]["params"].append(p)

    # print("parameter groups: \n%s" % json.dumps(param_group_names, indent=2))

    return list(param_groups.values())


vision_list=['patch_embed_v.proj.weight','patch_embed_v.proj.bias',
'blocks_v.0.norm1.weight','blocks_v.0.norm1.bias','blocks_v.0.attn.qkv.weight','blocks_v.0.attn.qkv.bias','blocks_v.0.attn.proj.weight','blocks_v.0.attn.proj.bias','blocks_v.0.norm2.weight','blocks_v.0.norm2.bias','blocks_v.0.mlp.fc1.weight','blocks_v.0.mlp.fc1.bias','blocks_v.0.mlp.fc2.weight','blocks_v.0.mlp.fc2.bias',
'blocks_v.1.norm1.weight','blocks_v.1.norm1.bias','blocks_v.1.attn.qkv.weight','blocks_v.1.attn.qkv.bias','blocks_v.1.attn.proj.weight','blocks_v.1.attn.proj.bias','blocks_v.1.norm2.weight','blocks_v.1.norm2.bias','blocks_v.1.mlp.fc1.weight','blocks_v.1.mlp.fc1.bias','blocks_v.1.mlp.fc2.weight','blocks_v.1.mlp.fc2.bias',
'blocks_v.2.norm1.weight','blocks_v.2.norm1.bias','blocks_v.2.attn.qkv.weight','blocks_v.2.attn.qkv.bias','blocks_v.2.attn.proj.weight','blocks_v.2.attn.proj.bias','blocks_v.2.norm2.weight','blocks_v.2.norm2.bias','blocks_v.2.mlp.fc1.weight','blocks_v.2.mlp.fc1.bias','blocks_v.2.mlp.fc2.weight','blocks_v.2.mlp.fc2.bias',
'blocks_v.3.norm1.weight','blocks_v.3.norm1.bias','blocks_v.3.attn.qkv.weight','blocks_v.3.attn.qkv.bias','blocks_v.3.attn.proj.weight','blocks_v.3.attn.proj.bias','blocks_v.3.norm2.weight','blocks_v.3.norm2.bias','blocks_v.3.mlp.fc1.weight','blocks_v.3.mlp.fc1.bias','blocks_v.3.mlp.fc2.weight','blocks_v.3.mlp.fc2.bias',
'blocks_v.4.norm1.weight','blocks_v.4.norm1.bias','blocks_v.4.attn.qkv.weight','blocks_v.4.attn.qkv.bias','blocks_v.4.attn.proj.weight','blocks_v.4.attn.proj.bias','blocks_v.4.norm2.weight','blocks_v.4.norm2.bias','blocks_v.4.mlp.fc1.weight','blocks_v.4.mlp.fc1.bias','blocks_v.4.mlp.fc2.weight','blocks_v.4.mlp.fc2.bias',
'blocks_v.5.norm1.weight','blocks_v.5.norm1.bias','blocks_v.5.attn.qkv.weight','blocks_v.5.attn.qkv.bias','blocks_v.5.attn.proj.weight','blocks_v.5.attn.proj.bias','blocks_v.5.norm2.weight','blocks_v.5.norm2.bias','blocks_v.5.mlp.fc1.weight','blocks_v.5.mlp.fc1.bias','blocks_v.5.mlp.fc2.weight','blocks_v.5.mlp.fc2.bias',
'blocks_v.6.norm1.weight','blocks_v.6.norm1.bias','blocks_v.6.attn.qkv.weight','blocks_v.6.attn.qkv.bias','blocks_v.6.attn.proj.weight','blocks_v.6.attn.proj.bias','blocks_v.6.norm2.weight','blocks_v.6.norm2.bias','blocks_v.6.mlp.fc1.weight','blocks_v.6.mlp.fc1.bias','blocks_v.6.mlp.fc2.weight','blocks_v.6.mlp.fc2.bias',
'blocks_v.7.norm1.weight','blocks_v.7.norm1.bias','blocks_v.7.attn.qkv.weight','blocks_v.7.attn.qkv.bias','blocks_v.7.attn.proj.weight','blocks_v.7.attn.proj.bias','blocks_v.7.norm2.weight','blocks_v.7.norm2.bias','blocks_v.7.mlp.fc1.weight','blocks_v.7.mlp.fc1.bias','blocks_v.7.mlp.fc2.weight','blocks_v.7.mlp.fc2.bias',
'blocks_v.8.norm1.weight','blocks_v.8.norm1.bias','blocks_v.8.attn.qkv.weight','blocks_v.8.attn.qkv.bias','blocks_v.8.attn.proj.weight','blocks_v.8.attn.proj.bias','blocks_v.8.norm2.weight','blocks_v.8.norm2.bias','blocks_v.8.mlp.fc1.weight','blocks_v.8.mlp.fc1.bias','blocks_v.8.mlp.fc2.weight','blocks_v.8.mlp.fc2.bias',
'blocks_v.9.norm1.weight','blocks_v.9.norm1.bias','blocks_v.9.attn.qkv.weight','blocks_v.9.attn.qkv.bias','blocks_v.9.attn.proj.weight','blocks_v.9.attn.proj.bias','blocks_v.9.norm2.weight','blocks_v.9.norm2.bias','blocks_v.9.mlp.fc1.weight','blocks_v.9.mlp.fc1.bias','blocks_v.9.mlp.fc2.weight','blocks_v.9.mlp.fc2.bias',
'blocks_v.10.norm1.weight','blocks_v.10.norm1.bias','blocks_v.10.attn.qkv.weight','blocks_v.10.attn.qkv.bias','blocks_v.10.attn.proj.weight','blocks_v.10.attn.proj.bias','blocks_v.10.norm2.weight','blocks_v.10.norm2.bias','blocks_v.10.mlp.fc1.weight','blocks_v.10.mlp.fc1.bias','blocks_v.10.mlp.fc2.weight','blocks_v.10.mlp.fc2.bias',
'blocks_v.11.norm1.weight','blocks_v.11.norm1.bias','blocks_v.11.attn.qkv.weight','blocks_v.11.attn.qkv.bias','blocks_v.11.attn.proj.weight','blocks_v.11.attn.proj.bias','blocks_v.11.norm2.weight','blocks_v.11.norm2.bias','blocks_v.11.mlp.fc1.weight','blocks_v.11.mlp.fc1.bias','blocks_v.11.mlp.fc2.weight','blocks_v.11.mlp.fc2.bias',
'blocks_v.12.norm1.weight','blocks_v.12.norm1.bias','blocks_v.12.attn.qkv.weight','blocks_v.12.attn.qkv.bias','blocks_v.12.attn.proj.weight','blocks_v.12.attn.proj.bias','blocks_v.12.norm2.weight','blocks_v.12.norm2.bias','blocks_v.12.mlp.fc1.weight','blocks_v.12.mlp.fc1.bias','blocks_v.12.mlp.fc2.weight','blocks_v.12.mlp.fc2.bias',
'blocks_v.13.norm1.weight','blocks_v.13.norm1.bias','blocks_v.13.attn.qkv.weight','blocks_v.13.attn.qkv.bias','blocks_v.13.attn.proj.weight','blocks_v.13.attn.proj.bias','blocks_v.13.norm2.weight','blocks_v.13.norm2.bias','blocks_v.13.mlp.fc1.weight','blocks_v.13.mlp.fc1.bias','blocks_v.13.mlp.fc2.weight','blocks_v.13.mlp.fc2.bias',
'blocks_v.14.norm1.weight','blocks_v.14.norm1.bias','blocks_v.14.attn.qkv.weight','blocks_v.14.attn.qkv.bias','blocks_v.14.attn.proj.weight','blocks_v.14.attn.proj.bias','blocks_v.14.norm2.weight','blocks_v.14.norm2.bias','blocks_v.14.mlp.fc1.weight','blocks_v.14.mlp.fc1.bias','blocks_v.14.mlp.fc2.weight','blocks_v.14.mlp.fc2.bias',
'blocks_v.15.norm1.weight','blocks_v.15.norm1.bias','blocks_v.15.attn.qkv.weight','blocks_v.15.attn.qkv.bias','blocks_v.15.attn.proj.weight','blocks_v.15.attn.proj.bias','blocks_v.15.norm2.weight','blocks_v.15.norm2.bias','blocks_v.15.mlp.fc1.weight','blocks_v.15.mlp.fc1.bias','blocks_v.15.mlp.fc2.weight','blocks_v.15.mlp.fc2.bias',
'blocks_v.16.norm1.weight','blocks_v.16.norm1.bias','blocks_v.16.attn.qkv.weight','blocks_v.16.attn.qkv.bias','blocks_v.16.attn.proj.weight','blocks_v.16.attn.proj.bias','blocks_v.16.norm2.weight','blocks_v.16.norm2.bias','blocks_v.16.mlp.fc1.weight','blocks_v.16.mlp.fc1.bias','blocks_v.16.mlp.fc2.weight','blocks_v.16.mlp.fc2.bias',
'blocks_v.17.norm1.weight','blocks_v.17.norm1.bias','blocks_v.17.attn.qkv.weight','blocks_v.17.attn.qkv.bias','blocks_v.17.attn.proj.weight','blocks_v.17.attn.proj.bias','blocks_v.17.norm2.weight','blocks_v.17.norm2.bias','blocks_v.17.mlp.fc1.weight','blocks_v.17.mlp.fc1.bias','blocks_v.17.mlp.fc2.weight','blocks_v.17.mlp.fc2.bias',
'blocks_v.18.norm1.weight','blocks_v.18.norm1.bias','blocks_v.18.attn.qkv.weight','blocks_v.18.attn.qkv.bias','blocks_v.18.attn.proj.weight','blocks_v.18.attn.proj.bias','blocks_v.18.norm2.weight','blocks_v.18.norm2.bias','blocks_v.18.mlp.fc1.weight','blocks_v.18.mlp.fc1.bias','blocks_v.18.mlp.fc2.weight','blocks_v.18.mlp.fc2.bias',
'blocks_v.19.norm1.weight','blocks_v.19.norm1.bias','blocks_v.19.attn.qkv.weight','blocks_v.19.attn.qkv.bias','blocks_v.19.attn.proj.weight','blocks_v.19.attn.proj.bias','blocks_v.19.norm2.weight','blocks_v.19.norm2.bias','blocks_v.19.mlp.fc1.weight','blocks_v.19.mlp.fc1.bias','blocks_v.19.mlp.fc2.weight','blocks_v.19.mlp.fc2.bias',
'blocks_v.20.norm1.weight','blocks_v.20.norm1.bias','blocks_v.20.attn.qkv.weight','blocks_v.20.attn.qkv.bias','blocks_v.20.attn.proj.weight','blocks_v.20.attn.proj.bias','blocks_v.20.norm2.weight','blocks_v.20.norm2.bias','blocks_v.20.mlp.fc1.weight','blocks_v.20.mlp.fc1.bias','blocks_v.20.mlp.fc2.weight','blocks_v.20.mlp.fc2.bias',
'blocks_v.21.norm1.weight','blocks_v.21.norm1.bias','blocks_v.21.attn.qkv.weight','blocks_v.21.attn.qkv.bias','blocks_v.21.attn.proj.weight','blocks_v.21.attn.proj.bias','blocks_v.21.norm2.weight','blocks_v.21.norm2.bias','blocks_v.21.mlp.fc1.weight','blocks_v.21.mlp.fc1.bias','blocks_v.21.mlp.fc2.weight','blocks_v.21.mlp.fc2.bias',
'blocks_v.22.norm1.weight','blocks_v.22.norm1.bias','blocks_v.22.attn.qkv.weight','blocks_v.22.attn.qkv.bias','blocks_v.22.attn.proj.weight','blocks_v.22.attn.proj.bias','blocks_v.22.norm2.weight','blocks_v.22.norm2.bias','blocks_v.22.mlp.fc1.weight','blocks_v.22.mlp.fc1.bias','blocks_v.22.mlp.fc2.weight','blocks_v.22.mlp.fc2.bias',
'blocks_v.23.norm1.weight','blocks_v.23.norm1.bias','blocks_v.23.attn.qkv.weight','blocks_v.23.attn.qkv.bias','blocks_v.23.attn.proj.weight','blocks_v.23.attn.proj.bias','blocks_v.23.norm2.weight','blocks_v.23.norm2.bias','blocks_v.23.mlp.fc1.weight','blocks_v.23.mlp.fc1.bias','blocks_v.23.mlp.fc2.weight','blocks_v.23.mlp.fc2.bias',
]
audio_video_list=[
'patch_embed_v.proj.weight','patch_embed_v.proj.bias', 'pos_emb_v'
'blocks_v.0.norm1.weight','blocks_v.0.norm1.bias','blocks_v.0.attn.qkv.weight','blocks_v.0.attn.qkv.bias','blocks_v.0.attn.proj.weight','blocks_v.0.attn.proj.bias','blocks_v.0.norm2.weight','blocks_v.0.norm2.bias','blocks_v.0.mlp.fc1.weight','blocks_v.0.mlp.fc1.bias','blocks_v.0.mlp.fc2.weight','blocks_v.0.mlp.fc2.bias',
'blocks_v.1.norm1.weight','blocks_v.1.norm1.bias','blocks_v.1.attn.qkv.weight','blocks_v.1.attn.qkv.bias','blocks_v.1.attn.proj.weight','blocks_v.1.attn.proj.bias','blocks_v.1.norm2.weight','blocks_v.1.norm2.bias','blocks_v.1.mlp.fc1.weight','blocks_v.1.mlp.fc1.bias','blocks_v.1.mlp.fc2.weight','blocks_v.1.mlp.fc2.bias',
'blocks_v.2.norm1.weight','blocks_v.2.norm1.bias','blocks_v.2.attn.qkv.weight','blocks_v.2.attn.qkv.bias','blocks_v.2.attn.proj.weight','blocks_v.2.attn.proj.bias','blocks_v.2.norm2.weight','blocks_v.2.norm2.bias','blocks_v.2.mlp.fc1.weight','blocks_v.2.mlp.fc1.bias','blocks_v.2.mlp.fc2.weight','blocks_v.2.mlp.fc2.bias',
'blocks_v.3.norm1.weight','blocks_v.3.norm1.bias','blocks_v.3.attn.qkv.weight','blocks_v.3.attn.qkv.bias','blocks_v.3.attn.proj.weight','blocks_v.3.attn.proj.bias','blocks_v.3.norm2.weight','blocks_v.3.norm2.bias','blocks_v.3.mlp.fc1.weight','blocks_v.3.mlp.fc1.bias','blocks_v.3.mlp.fc2.weight','blocks_v.3.mlp.fc2.bias',
'blocks_v.4.norm1.weight','blocks_v.4.norm1.bias','blocks_v.4.attn.qkv.weight','blocks_v.4.attn.qkv.bias','blocks_v.4.attn.proj.weight','blocks_v.4.attn.proj.bias','blocks_v.4.norm2.weight','blocks_v.4.norm2.bias','blocks_v.4.mlp.fc1.weight','blocks_v.4.mlp.fc1.bias','blocks_v.4.mlp.fc2.weight','blocks_v.4.mlp.fc2.bias',
'blocks_v.5.norm1.weight','blocks_v.5.norm1.bias','blocks_v.5.attn.qkv.weight','blocks_v.5.attn.qkv.bias','blocks_v.5.attn.proj.weight','blocks_v.5.attn.proj.bias','blocks_v.5.norm2.weight','blocks_v.5.norm2.bias','blocks_v.5.mlp.fc1.weight','blocks_v.5.mlp.fc1.bias','blocks_v.5.mlp.fc2.weight','blocks_v.5.mlp.fc2.bias',
'blocks_v.6.norm1.weight','blocks_v.6.norm1.bias','blocks_v.6.attn.qkv.weight','blocks_v.6.attn.qkv.bias','blocks_v.6.attn.proj.weight','blocks_v.6.attn.proj.bias','blocks_v.6.norm2.weight','blocks_v.6.norm2.bias','blocks_v.6.mlp.fc1.weight','blocks_v.6.mlp.fc1.bias','blocks_v.6.mlp.fc2.weight','blocks_v.6.mlp.fc2.bias',
'blocks_v.7.norm1.weight','blocks_v.7.norm1.bias','blocks_v.7.attn.qkv.weight','blocks_v.7.attn.qkv.bias','blocks_v.7.attn.proj.weight','blocks_v.7.attn.proj.bias','blocks_v.7.norm2.weight','blocks_v.7.norm2.bias','blocks_v.7.mlp.fc1.weight','blocks_v.7.mlp.fc1.bias','blocks_v.7.mlp.fc2.weight','blocks_v.7.mlp.fc2.bias',
'blocks_v.8.norm1.weight','blocks_v.8.norm1.bias','blocks_v.8.attn.qkv.weight','blocks_v.8.attn.qkv.bias','blocks_v.8.attn.proj.weight','blocks_v.8.attn.proj.bias','blocks_v.8.norm2.weight','blocks_v.8.norm2.bias','blocks_v.8.mlp.fc1.weight','blocks_v.8.mlp.fc1.bias','blocks_v.8.mlp.fc2.weight','blocks_v.8.mlp.fc2.bias',
'blocks_v.9.norm1.weight','blocks_v.9.norm1.bias','blocks_v.9.attn.qkv.weight','blocks_v.9.attn.qkv.bias','blocks_v.9.attn.proj.weight','blocks_v.9.attn.proj.bias','blocks_v.9.norm2.weight','blocks_v.9.norm2.bias','blocks_v.9.mlp.fc1.weight','blocks_v.9.mlp.fc1.bias','blocks_v.9.mlp.fc2.weight','blocks_v.9.mlp.fc2.bias',
'blocks_v.10.norm1.weight','blocks_v.10.norm1.bias','blocks_v.10.attn.qkv.weight','blocks_v.10.attn.qkv.bias','blocks_v.10.attn.proj.weight','blocks_v.10.attn.proj.bias','blocks_v.10.norm2.weight','blocks_v.10.norm2.bias','blocks_v.10.mlp.fc1.weight','blocks_v.10.mlp.fc1.bias','blocks_v.10.mlp.fc2.weight','blocks_v.10.mlp.fc2.bias',
'blocks_v.11.norm1.weight','blocks_v.11.norm1.bias','blocks_v.11.attn.qkv.weight','blocks_v.11.attn.qkv.bias','blocks_v.11.attn.proj.weight','blocks_v.11.attn.proj.bias','blocks_v.11.norm2.weight','blocks_v.11.norm2.bias','blocks_v.11.mlp.fc1.weight','blocks_v.11.mlp.fc1.bias','blocks_v.11.mlp.fc2.weight','blocks_v.11.mlp.fc2.bias'
'blocks_v.12.norm1.weight','blocks_v.12.norm1.bias','blocks_v.12.attn.qkv.weight','blocks_v.12.attn.qkv.bias','blocks_v.12.attn.proj.weight','blocks_v.12.attn.proj.bias','blocks_v.12.norm2.weight','blocks_v.12.norm2.bias','blocks_v.12.mlp.fc1.weight','blocks_v.12.mlp.fc1.bias','blocks_v.12.mlp.fc2.weight','blocks_v.12.mlp.fc2.bias',
'blocks_v.13.norm1.weight','blocks_v.13.norm1.bias','blocks_v.13.attn.qkv.weight','blocks_v.13.attn.qkv.bias','blocks_v.13.attn.proj.weight','blocks_v.13.attn.proj.bias','blocks_v.13.norm2.weight','blocks_v.13.norm2.bias','blocks_v.13.mlp.fc1.weight','blocks_v.13.mlp.fc1.bias','blocks_v.13.mlp.fc2.weight','blocks_v.13.mlp.fc2.bias',
'blocks_v.14.norm1.weight','blocks_v.14.norm1.bias','blocks_v.14.attn.qkv.weight','blocks_v.14.attn.qkv.bias','blocks_v.14.attn.proj.weight','blocks_v.14.attn.proj.bias','blocks_v.14.norm2.weight','blocks_v.14.norm2.bias','blocks_v.14.mlp.fc1.weight','blocks_v.14.mlp.fc1.bias','blocks_v.14.mlp.fc2.weight','blocks_v.14.mlp.fc2.bias',
'blocks_v.15.norm1.weight','blocks_v.15.norm1.bias','blocks_v.15.attn.qkv.weight','blocks_v.15.attn.qkv.bias','blocks_v.15.attn.proj.weight','blocks_v.15.attn.proj.bias','blocks_v.15.norm2.weight','blocks_v.15.norm2.bias','blocks_v.15.mlp.fc1.weight','blocks_v.15.mlp.fc1.bias','blocks_v.15.mlp.fc2.weight','blocks_v.15.mlp.fc2.bias',
'blocks_v.16.norm1.weight','blocks_v.16.norm1.bias','blocks_v.16.attn.qkv.weight','blocks_v.16.attn.qkv.bias','blocks_v.16.attn.proj.weight','blocks_v.16.attn.proj.bias','blocks_v.16.norm2.weight','blocks_v.16.norm2.bias','blocks_v.16.mlp.fc1.weight','blocks_v.16.mlp.fc1.bias','blocks_v.16.mlp.fc2.weight','blocks_v.16.mlp.fc2.bias',
'blocks_v.17.norm1.weight','blocks_v.17.norm1.bias','blocks_v.17.attn.qkv.weight','blocks_v.17.attn.qkv.bias','blocks_v.17.attn.proj.weight','blocks_v.17.attn.proj.bias','blocks_v.17.norm2.weight','blocks_v.17.norm2.bias','blocks_v.17.mlp.fc1.weight','blocks_v.17.mlp.fc1.bias','blocks_v.17.mlp.fc2.weight','blocks_v.17.mlp.fc2.bias',
'blocks_v.18.norm1.weight','blocks_v.18.norm1.bias','blocks_v.18.attn.qkv.weight','blocks_v.18.attn.qkv.bias','blocks_v.18.attn.proj.weight','blocks_v.18.attn.proj.bias','blocks_v.18.norm2.weight','blocks_v.18.norm2.bias','blocks_v.18.mlp.fc1.weight','blocks_v.18.mlp.fc1.bias','blocks_v.18.mlp.fc2.weight','blocks_v.18.mlp.fc2.bias',
'blocks_v.19.norm1.weight','blocks_v.19.norm1.bias','blocks_v.19.attn.qkv.weight','blocks_v.19.attn.qkv.bias','blocks_v.19.attn.proj.weight','blocks_v.19.attn.proj.bias','blocks_v.19.norm2.weight','blocks_v.19.norm2.bias','blocks_v.19.mlp.fc1.weight','blocks_v.19.mlp.fc1.bias','blocks_v.19.mlp.fc2.weight','blocks_v.19.mlp.fc2.bias',
'blocks_v.20.norm1.weight','blocks_v.20.norm1.bias','blocks_v.20.attn.qkv.weight','blocks_v.20.attn.qkv.bias','blocks_v.20.attn.proj.weight','blocks_v.20.attn.proj.bias','blocks_v.20.norm2.weight','blocks_v.20.norm2.bias','blocks_v.20.mlp.fc1.weight','blocks_v.20.mlp.fc1.bias','blocks_v.20.mlp.fc2.weight','blocks_v.20.mlp.fc2.bias',
'blocks_v.21.norm1.weight','blocks_v.21.norm1.bias','blocks_v.21.attn.qkv.weight','blocks_v.21.attn.qkv.bias','blocks_v.21.attn.proj.weight','blocks_v.21.attn.proj.bias','blocks_v.21.norm2.weight','blocks_v.21.norm2.bias','blocks_v.21.mlp.fc1.weight','blocks_v.21.mlp.fc1.bias','blocks_v.21.mlp.fc2.weight','blocks_v.21.mlp.fc2.bias',
'blocks_v.22.norm1.weight','blocks_v.22.norm1.bias','blocks_v.22.attn.qkv.weight','blocks_v.22.attn.qkv.bias','blocks_v.22.attn.proj.weight','blocks_v.22.attn.proj.bias','blocks_v.22.norm2.weight','blocks_v.22.norm2.bias','blocks_v.22.mlp.fc1.weight','blocks_v.22.mlp.fc1.bias','blocks_v.22.mlp.fc2.weight','blocks_v.22.mlp.fc2.bias',
'blocks_v.23.norm1.weight','blocks_v.23.norm1.bias','blocks_v.23.attn.qkv.weight','blocks_v.23.attn.qkv.bias','blocks_v.23.attn.proj.weight','blocks_v.23.attn.proj.bias','blocks_v.23.norm2.weight','blocks_v.23.norm2.bias','blocks_v.23.mlp.fc1.weight','blocks_v.23.mlp.fc1.bias','blocks_v.23.mlp.fc2.weight','blocks_v.23.mlp.fc2.bias',
'patch_embed.proj.weight','patch_embed.proj.bias', 'pos_emb',
'blocks.0.norm1.weight','blocks.0.norm1.bias','blocks.0.attn.qkv.weight','blocks.0.attn.qkv.bias','blocks.0.attn.proj.weight','blocks.0.attn.proj.bias','blocks.0.norm2.weight','blocks.0.norm2.bias','blocks.0.mlp.fc1.weight','blocks.0.mlp.fc1.bias','blocks.0.mlp.fc2.weight','blocks.0.mlp.fc2.bias',
'blocks.1.norm1.weight','blocks.1.norm1.bias','blocks.1.attn.qkv.weight','blocks.1.attn.qkv.bias','blocks.1.attn.proj.weight','blocks.1.attn.proj.bias','blocks.1.norm2.weight','blocks.1.norm2.bias','blocks.1.mlp.fc1.weight','blocks.1.mlp.fc1.bias','blocks.1.mlp.fc2.weight','blocks.1.mlp.fc2.bias',
'blocks.2.norm1.weight','blocks.2.norm1.bias','blocks.2.attn.qkv.weight','blocks.2.attn.qkv.bias','blocks.2.attn.proj.weight','blocks.2.attn.proj.bias','blocks.2.norm2.weight','blocks.2.norm2.bias','blocks.2.mlp.fc1.weight','blocks.2.mlp.fc1.bias','blocks.2.mlp.fc2.weight','blocks.2.mlp.fc2.bias',
'blocks.3.norm1.weight','blocks.3.norm1.bias','blocks.3.attn.qkv.weight','blocks.3.attn.qkv.bias','blocks.3.attn.proj.weight','blocks.3.attn.proj.bias','blocks.3.norm2.weight','blocks.3.norm2.bias','blocks.3.mlp.fc1.weight','blocks.3.mlp.fc1.bias','blocks.3.mlp.fc2.weight','blocks.3.mlp.fc2.bias',
'blocks.4.norm1.weight','blocks.4.norm1.bias','blocks.4.attn.qkv.weight','blocks.4.attn.qkv.bias','blocks.4.attn.proj.weight','blocks.4.attn.proj.bias','blocks.4.norm2.weight','blocks.4.norm2.bias','blocks.4.mlp.fc1.weight','blocks.4.mlp.fc1.bias','blocks.4.mlp.fc2.weight','blocks.4.mlp.fc2.bias',
'blocks.5.norm1.weight','blocks.5.norm1.bias','blocks.5.attn.qkv.weight','blocks.5.attn.qkv.bias','blocks.5.attn.proj.weight','blocks.5.attn.proj.bias','blocks.5.norm2.weight','blocks.5.norm2.bias','blocks.5.mlp.fc1.weight','blocks.5.mlp.fc1.bias','blocks.5.mlp.fc2.weight','blocks.5.mlp.fc2.bias',
'blocks.6.norm1.weight','blocks.6.norm1.bias','blocks.6.attn.qkv.weight','blocks.6.attn.qkv.bias','blocks.6.attn.proj.weight','blocks.6.attn.proj.bias','blocks.6.norm2.weight','blocks.6.norm2.bias','blocks.6.mlp.fc1.weight','blocks.6.mlp.fc1.bias','blocks.6.mlp.fc2.weight','blocks.6.mlp.fc2.bias',
'blocks.7.norm1.weight','blocks.7.norm1.bias','blocks.7.attn.qkv.weight','blocks.7.attn.qkv.bias','blocks.7.attn.proj.weight','blocks.7.attn.proj.bias','blocks.7.norm2.weight','blocks.7.norm2.bias','blocks.7.mlp.fc1.weight','blocks.7.mlp.fc1.bias','blocks.7.mlp.fc2.weight','blocks.7.mlp.fc2.bias',
'blocks.8.norm1.weight','blocks.8.norm1.bias','blocks.8.attn.qkv.weight','blocks.8.attn.qkv.bias','blocks.8.attn.proj.weight','blocks.8.attn.proj.bias','blocks.8.norm2.weight','blocks.8.norm2.bias','blocks.8.mlp.fc1.weight','blocks.8.mlp.fc1.bias','blocks.8.mlp.fc2.weight','blocks.8.mlp.fc2.bias',
'blocks.9.norm1.weight','blocks.9.norm1.bias','blocks.9.attn.qkv.weight','blocks.9.attn.qkv.bias','blocks.9.attn.proj.weight','blocks.9.attn.proj.bias','blocks.9.norm2.weight','blocks.9.norm2.bias','blocks.9.mlp.fc1.weight','blocks.9.mlp.fc1.bias','blocks.9.mlp.fc2.weight','blocks.9.mlp.fc2.bias',
'blocks.10.norm1.weight','blocks.10.norm1.bias','blocks.10.attn.qkv.weight','blocks.10.attn.qkv.bias','blocks.10.attn.proj.weight','blocks.10.attn.proj.bias','blocks.10.norm2.weight','blocks.10.norm2.bias','blocks.10.mlp.fc1.weight','blocks.10.mlp.fc1.bias','blocks.10.mlp.fc2.weight','blocks.10.mlp.fc2.bias',
'blocks.11.norm1.weight','blocks.11.norm1.bias','blocks.11.attn.qkv.weight','blocks.11.attn.qkv.bias','blocks.11.attn.proj.weight','blocks.11.attn.proj.bias','blocks.11.norm2.weight','blocks.11.norm2.bias','blocks.11.mlp.fc1.weight','blocks.11.mlp.fc1.bias','blocks.11.mlp.fc2.weight','blocks.11.mlp.fc2.bias'
'blocks.12.norm1.weight','blocks.12.norm1.bias','blocks.12.attn.qkv.weight','blocks.12.attn.qkv.bias','blocks.12.attn.proj.weight','blocks.12.attn.proj.bias','blocks.12.norm2.weight','blocks.12.norm2.bias','blocks.12.mlp.fc1.weight','blocks.12.mlp.fc1.bias','blocks.12.mlp.fc2.weight','blocks.12.mlp.fc2.bias',
'blocks.13.norm1.weight','blocks.13.norm1.bias','blocks.13.attn.qkv.weight','blocks.13.attn.qkv.bias','blocks.13.attn.proj.weight','blocks.13.attn.proj.bias','blocks.13.norm2.weight','blocks.13.norm2.bias','blocks.13.mlp.fc1.weight','blocks.13.mlp.fc1.bias','blocks.13.mlp.fc2.weight','blocks.13.mlp.fc2.bias',
'blocks.14.norm1.weight','blocks.14.norm1.bias','blocks.14.attn.qkv.weight','blocks.14.attn.qkv.bias','blocks.14.attn.proj.weight','blocks.14.attn.proj.bias','blocks.14.norm2.weight','blocks.14.norm2.bias','blocks.14.mlp.fc1.weight','blocks.14.mlp.fc1.bias','blocks.14.mlp.fc2.weight','blocks.14.mlp.fc2.bias',
'blocks.15.norm1.weight','blocks.15.norm1.bias','blocks.15.attn.qkv.weight','blocks.15.attn.qkv.bias','blocks.15.attn.proj.weight','blocks.15.attn.proj.bias','blocks.15.norm2.weight','blocks.15.norm2.bias','blocks.15.mlp.fc1.weight','blocks.15.mlp.fc1.bias','blocks.15.mlp.fc2.weight','blocks.15.mlp.fc2.bias',
'blocks.16.norm1.weight','blocks.16.norm1.bias','blocks.16.attn.qkv.weight','blocks.16.attn.qkv.bias','blocks.16.attn.proj.weight','blocks.16.attn.proj.bias','blocks.16.norm2.weight','blocks.16.norm2.bias','blocks.16.mlp.fc1.weight','blocks.16.mlp.fc1.bias','blocks.16.mlp.fc2.weight','blocks.16.mlp.fc2.bias',
'blocks.17.norm1.weight','blocks.17.norm1.bias','blocks.17.attn.qkv.weight','blocks.17.attn.qkv.bias','blocks.17.attn.proj.weight','blocks.17.attn.proj.bias','blocks.17.norm2.weight','blocks.17.norm2.bias','blocks.17.mlp.fc1.weight','blocks.17.mlp.fc1.bias','blocks.17.mlp.fc2.weight','blocks.17.mlp.fc2.bias',
'blocks.18.norm1.weight','blocks.18.norm1.bias','blocks.18.attn.qkv.weight','blocks.18.attn.qkv.bias','blocks.18.attn.proj.weight','blocks.18.attn.proj.bias','blocks.18.norm2.weight','blocks.18.norm2.bias','blocks.18.mlp.fc1.weight','blocks.18.mlp.fc1.bias','blocks.18.mlp.fc2.weight','blocks.18.mlp.fc2.bias',
'blocks.19.norm1.weight','blocks.19.norm1.bias','blocks.19.attn.qkv.weight','blocks.19.attn.qkv.bias','blocks.19.attn.proj.weight','blocks.19.attn.proj.bias','blocks.19.norm2.weight','blocks.19.norm2.bias','blocks.19.mlp.fc1.weight','blocks.19.mlp.fc1.bias','blocks.19.mlp.fc2.weight','blocks.19.mlp.fc2.bias',
'blocks.20.norm1.weight','blocks.20.norm1.bias','blocks.20.attn.qkv.weight','blocks.20.attn.qkv.bias','blocks.20.attn.proj.weight','blocks.20.attn.proj.bias','blocks.20.norm2.weight','blocks.20.norm2.bias','blocks.20.mlp.fc1.weight','blocks.20.mlp.fc1.bias','blocks.20.mlp.fc2.weight','blocks.20.mlp.fc2.bias',
'blocks.21.norm1.weight','blocks.21.norm1.bias','blocks.21.attn.qkv.weight','blocks.21.attn.qkv.bias','blocks.21.attn.proj.weight','blocks.21.attn.proj.bias','blocks.21.norm2.weight','blocks.21.norm2.bias','blocks.21.mlp.fc1.weight','blocks.21.mlp.fc1.bias','blocks.21.mlp.fc2.weight','blocks.21.mlp.fc2.bias',
'blocks.22.norm1.weight','blocks.22.norm1.bias','blocks.22.attn.qkv.weight','blocks.22.attn.qkv.bias','blocks.22.attn.proj.weight','blocks.22.attn.proj.bias','blocks.22.norm2.weight','blocks.22.norm2.bias','blocks.22.mlp.fc1.weight','blocks.22.mlp.fc1.bias','blocks.22.mlp.fc2.weight','blocks.22.mlp.fc2.bias',
'blocks.23.norm1.weight','blocks.23.norm1.bias','blocks.23.attn.qkv.weight','blocks.23.attn.qkv.bias','blocks.23.attn.proj.weight','blocks.23.attn.proj.bias','blocks.23.norm2.weight','blocks.23.norm2.bias','blocks.23.mlp.fc1.weight','blocks.23.mlp.fc1.bias','blocks.23.mlp.fc2.weight','blocks.23.mlp.fc2.bias',
]


def param_groups_lrd2(model, weight_decay=0.05, no_weight_decay_list=[], layer_decay=.75):
    """
    Parameter groups for layer-wise lr decay
    Following BEiT: https://github.com/microsoft/unilm/blob/master/beit/optim_factory.py#L58
    """
    param_group_names = {}
    param_groups = {}
    num_layers = len(model.blocks) + 1
    layer_scales = list(layer_decay ** (num_layers - i) for i in range(num_layers + 1))

    param_group_names_v = {}
    param_groups_v = {}

    for n, p in model.named_parameters():
        if not p.requires_grad:
            continue

        # no decay: all 1D parameters and model specific ones
        if p.ndim == 1 or n in no_weight_decay_list:
            g_decay = "no_decay"
            this_decay = 0.
        else:
            g_decay = "decay"
            this_decay = weight_decay
            
        layer_id = get_layer_id_for_vit(n, num_layers)
        group_name = "layer_%d_%s" % (layer_id, g_decay)

        #if '_v' in n:
        #    print(n)
        if n in vision_list:
            #print(n)
            if group_name not in param_group_names_v:
                this_scale = layer_scales[layer_id]

                param_group_names_v[group_name] = {
                    "lr_scale": 0.5*this_scale,
                    "weight_decay": this_decay,
                    "params": [],
                }
                param_groups_v[group_name] = {
                    "lr_scale": 0.5*this_scale,
                    "weight_decay": this_decay,
                    "params": [],
                }            
            param_group_names_v[group_name]["params"].append(n)
            param_groups_v[group_name]["params"].append(p)
        else:
            if group_name not in param_group_names:
                this_scale = layer_scales[layer_id]

                param_group_names[group_name] = {
                    "lr_scale": this_scale,
                    "weight_decay": this_decay,
                    "params": [],               
                }
                param_groups[group_name] = {
                    "lr_scale": this_scale,
                    "weight_decay": this_decay,
                    "params": [],                
                }
            param_group_names[group_name]["params"].append(n)
            param_groups[group_name]["params"].append(p)

    #print("parameter groups: \n%s" % json.dumps(param_group_names, indent=2))
    #print("parameter groups2: \n%s" % json.dumps(param_group_names_v, indent=2))

    #return param_groups.values(), param_groups_v.values()
    return list(param_groups.values()), list(param_groups_v.values())


def param_groups_lrd_av_slow(model, weight_decay=0.05, no_weight_decay_list=[], layer_decay=.75):
    """
    Parameter groups for layer-wise lr decay
    Following BEiT: https://github.com/microsoft/unilm/blob/master/beit/optim_factory.py#L58
    """
    param_group_names = {}
    param_groups = {}
    num_layers = len(model.blocks) + 1
    layer_scales = list(layer_decay ** (num_layers - i) for i in range(num_layers + 1))

    param_group_names_v = {}
    param_groups_v = {}

    for n, p in model.named_parameters():
        if not p.requires_grad:
            continue

        # no decay: all 1D parameters and model specific ones
        if p.ndim == 1 or n in no_weight_decay_list:
            g_decay = "no_decay"
            this_decay = 0.
        else:
            g_decay = "decay"
            this_decay = weight_decay
            
        layer_id = get_layer_id_for_vit(n, num_layers)
        group_name = "layer_%d_%s" % (layer_id, g_decay)

        #if '_v' in n:
        #    print(n)
        if n in audio_video_list:
            #print(n)
            if group_name not in param_group_names_v:
                this_scale = layer_scales[layer_id]

                param_group_names_v[group_name] = {
                    "lr_scale": 0.1*this_scale,
                    "weight_decay": this_decay,
                    "params": [],
                }
                param_groups_v[group_name] = {
                    "lr_scale": 0.1*this_scale,
                    "weight_decay": this_decay,
                    "params": [],
                }            
            param_group_names_v[group_name]["params"].append(n)
            param_groups_v[group_name]["params"].append(p)
        else:
            if group_name not in param_group_names:
                this_scale = layer_scales[layer_id]

                param_group_names[group_name] = {
                    "lr_scale": this_scale,
                    "weight_decay": this_decay,
                    "params": [],               
                }
                param_groups[group_name] = {
                    "lr_scale": this_scale,
                    "weight_decay": this_decay,
                    "params": [],                
                }
            param_group_names[group_name]["params"].append(n)
            param_groups[group_name]["params"].append(p)

    #print("parameter groups: \n%s" % json.dumps(param_group_names, indent=2))
    #print("parameter groups2: \n%s" % json.dumps(param_group_names_v, indent=2))

    #return param_groups.values(), param_groups_v.values()
    return list(param_groups.values()), list(param_groups_v.values())



def get_layer_id_for_vit(name, num_layers):
    """
    Assign a parameter with its layer id
    Following BEiT: https://github.com/microsoft/unilm/blob/master/beit/optim_factory.py#L33
    """
    if name in ['cls_token', 'pos_embed', 'cls_token_v', 'pos_embed_v']:
        return 0
    elif name.startswith('patch_embed') or name.startswith('patch_embed_v'):
        return 0
    elif name.startswith('blocks'):
        return int(name.split('.')[1]) + 1
    else:
        return num_layers

upstream_models/mavil/util/lr_decay.py
roger-tseng/av-superb
# 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.
# --------------------------------------------------------
# References:
# DeiT: https://github.com/facebookresearch/deit
# BEiT: https://github.com/microsoft/unilm/tree/master/beit
# --------------------------------------------------------

import builtins
import datetime
import os
import time
from collections import defaultdict, deque
from pathlib import Path

import torch
import torch.distributed as dist
from torch import inf
from .pos_embed import get_3d_sincos_pos_embed


class SmoothedValue(object):
    """Track a series of values and provide access to smoothed values over a
    window or the global series average.
    """

    def __init__(self, window_size=20, fmt=None):
        if fmt is None:
            fmt = "{median:.4f} ({global_avg:.4f})"
        self.deque = deque(maxlen=window_size)
        self.total = 0.0
        self.count = 0
        self.fmt = fmt

    def update(self, value, n=1):
        self.deque.append(value)
        self.count += n
        self.total += value * n

    def synchronize_between_processes(self):
        """
        Warning: does not synchronize the deque!
        """
        if not is_dist_avail_and_initialized():
            return
        t = torch.tensor([self.count, self.total], dtype=torch.float64, device='cuda')
        dist.barrier()
        dist.all_reduce(t)
        t = t.tolist()
        self.count = int(t[0])
        self.total = t[1]

    @property
    def median(self):
        d = torch.tensor(list(self.deque))
        return d.median().item()

    @property
    def avg(self):
        d = torch.tensor(list(self.deque), dtype=torch.float32)
        return d.mean().item()

    @property
    def global_avg(self):
        return self.total / self.count

    @property
    def max(self):
        return max(self.deque)

    @property
    def value(self):
        return self.deque[-1]

    def __str__(self):
        return self.fmt.format(
            median=self.median,
            avg=self.avg,
            global_avg=self.global_avg,
            max=self.max,
            value=self.value)


class MetricLogger(object):
    def __init__(self, delimiter="\t"):
        self.meters = defaultdict(SmoothedValue)
        self.delimiter = delimiter

    def update(self, **kwargs):
        for k, v in kwargs.items():
            if v is None:
                continue
            if isinstance(v, torch.Tensor):
                v = v.item()
            assert isinstance(v, (float, int))
            self.meters[k].update(v)

    def __getattr__(self, attr):
        if attr in self.meters:
            return self.meters[attr]
        if attr in self.__dict__:
            return self.__dict__[attr]
        raise AttributeError("'{}' object has no attribute '{}'".format(
            type(self).__name__, attr))

    def __str__(self):
        loss_str = []
        for name, meter in self.meters.items():
            loss_str.append(
                "{}: {}".format(name, str(meter))
            )
        return self.delimiter.join(loss_str)

    def synchronize_between_processes(self):
        for meter in self.meters.values():
            meter.synchronize_between_processes()

    def add_meter(self, name, meter):
        self.meters[name] = meter

    def log_every(self, iterable, print_freq, header=None):
        i = 0
        if not header:
            header = ''
        start_time = time.time()
        end = time.time()
        iter_time = SmoothedValue(fmt='{avg:.4f}')
        data_time = SmoothedValue(fmt='{avg:.4f}')
        space_fmt = ':' + str(len(str(len(iterable)))) + 'd'
        log_msg = [
            header,
            '[{0' + space_fmt + '}/{1}]',
            'eta: {eta}',
            '{meters}',
            'time: {time}',
            'data: {data}'
        ]
        if torch.cuda.is_available():
            log_msg.append('max mem: {memory:.0f}')
        log_msg = self.delimiter.join(log_msg)
        MB = 1024.0 * 1024.0
        for obj in iterable:
            data_time.update(time.time() - end)
            yield obj
            iter_time.update(time.time() - end)
            if i % print_freq == 0 or i == len(iterable) - 1:
                eta_seconds = iter_time.global_avg * (len(iterable) - i)
                eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
                if torch.cuda.is_available():
                    print(log_msg.format(
                        i, len(iterable), eta=eta_string,
                        meters=str(self),
                        time=str(iter_time), data=str(data_time),
                        memory=torch.cuda.max_memory_allocated() / MB))
                else:
                    print(log_msg.format(
                        i, len(iterable), eta=eta_string,
                        meters=str(self),
                        time=str(iter_time), data=str(data_time)))
            i += 1
            end = time.time()
        total_time = time.time() - start_time
        total_time_str = str(datetime.timedelta(seconds=int(total_time)))
        print('{} Total time: {} ({:.4f} s / it)'.format(
            header, total_time_str, total_time / len(iterable)))


def setup_for_distributed(is_master):
    """
    This function disables printing when not in master process
    """
    builtin_print = builtins.print

    def print(*args, **kwargs):
        force = kwargs.pop('force', False)
        force = force or (get_world_size() > 8)
        if is_master or force:
            now = datetime.datetime.now().time()
            builtin_print('[{}] '.format(now), end='')  # print with time stamp
            builtin_print(*args, **kwargs)

    builtins.print = print


def is_dist_avail_and_initialized():
    if not dist.is_available():
        return False
    if not dist.is_initialized():
        return False
    return True


def get_world_size():
    if not is_dist_avail_and_initialized():
        return 1
    return dist.get_world_size()


def get_rank():
    if not is_dist_avail_and_initialized():
        return 0
    return dist.get_rank()


def is_main_process():
    return get_rank() == 0


def save_on_master(*args, **kwargs):
    if is_main_process():
        torch.save(*args, **kwargs)


def init_distributed_mode(args):
    if args.dist_on_itp:
        print('use_dist_on_itp')
        args.rank = int(os.environ['OMPI_COMM_WORLD_RANK'])
        args.world_size = int(os.environ['OMPI_COMM_WORLD_SIZE'])
        args.gpu = int(os.environ['OMPI_COMM_WORLD_LOCAL_RANK'])
        args.dist_url = "tcp://%s:%s" % (os.environ['MASTER_ADDR'], os.environ['MASTER_PORT'])
        os.environ['LOCAL_RANK'] = str(args.gpu)
        os.environ['RANK'] = str(args.rank)
        os.environ['WORLD_SIZE'] = str(args.world_size)
        # ["RANK", "WORLD_SIZE", "MASTER_ADDR", "MASTER_PORT", "LOCAL_RANK"]
    elif 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:
        print(f"Environment variable initialization: RANK={os.environ['RANK']} WORLD_SIZE={os.environ['WORLD_SIZE']}, LOCAL_RANK={os.environ['LOCAL_RANK']}")
        args.rank = int(os.environ["RANK"])
        args.world_size = int(os.environ['WORLD_SIZE'])
        args.gpu = int(os.environ['LOCAL_RANK'])
    #elif 'SLURM_PROCID' in os.environ:
    #    print("'SLURM_PROCID' in os.envision")
    #    env = os.environ
    #    print('env:', env)
    #    num_nodes = env["SLURM_JOB_NUM_NODES"]
    #    gpu_per_node = int(env["SLURM_GPUS_PER_NODE"])
    #    node_rank = int(env["SLURM_PROCID"])
    #    args.rank = gpu_per_node * node_rank + args.local_rank
    #    args.gpu = args.rank % torch.cuda.device_count()
    #    #args.world_size = int(env['SLURM_GPUS'])
    #    args.world_size = int(env['SLURM_NPROCS'])
    #    args.dist_url=f"tcp://{env['MASTER_ADDR']}:{env['MASTER_PORT']}"
    #    print(f"RANK and WORLD_SIZE in environ: {args.rank}/{args.world_size}")
    #    #args.rank = int(os.environ['SLURM_PROCID'])
    #    #args.gpu = args.rank % torch.cuda.device_count()
    else:
        print('Not using distributed mode')
        setup_for_distributed(is_master=True)  # hack
        args.distributed = False
        return

    args.distributed = True

    torch.cuda.set_device(args.gpu)
    args.dist_backend = 'nccl'
    print('distributed init (rank {}): {}, gpu {}, world_size {}'.format(
        args.rank, args.dist_url, args.gpu, args.world_size
    ))
    #torch.distributed.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
    #                                     world_size=args.world_size, rank=args.rank)
    #torch.distributed.init_process_group(
    #    backend=args.dist_backend,
    #    init_method=args.dist_url,
    #    world_size=args.world_size,
    #    rank=args.rank,
    #)

    torch.distributed.init_process_group(backend=args.dist_backend, world_size=args.world_size, rank=args.rank)

    torch.distributed.barrier()
    setup_for_distributed(args.rank == 0)


class NativeScalerWithGradNormCount:
    state_dict_key = "amp_scaler"

    def __init__(self):
        self._scaler = torch.cuda.amp.GradScaler()

    def __call__(self, loss, optimizer, clip_grad=None, parameters=None, create_graph=False, update_grad=True):
        self._scaler.scale(loss).backward(create_graph=create_graph)
        if update_grad:
            if clip_grad is not None:
                assert parameters is not None
                self._scaler.unscale_(optimizer)  # unscale the gradients of optimizer's assigned params in-place
                norm = torch.nn.utils.clip_grad_norm_(parameters, clip_grad)
            else:
                self._scaler.unscale_(optimizer)
                norm = get_grad_norm_(parameters)
            self._scaler.step(optimizer)
            self._scaler.update()
        else:
            norm = None
        return norm

    def state_dict(self):
        return self._scaler.state_dict()

    def load_state_dict(self, state_dict):
        self._scaler.load_state_dict(state_dict)


def get_grad_norm_(parameters, norm_type: float = 2.0) -> torch.Tensor:
    if isinstance(parameters, torch.Tensor):
        parameters = [parameters]
    parameters = [p for p in parameters if p.grad is not None]
    norm_type = float(norm_type)
    if len(parameters) == 0:
        return torch.tensor(0.)
    device = parameters[0].grad.device
    if norm_type == inf:
        total_norm = max(p.grad.detach().abs().max().to(device) for p in parameters)
    else:
        total_norm = torch.norm(torch.stack([torch.norm(p.grad.detach(), norm_type).to(device) for p in parameters]), norm_type)
    return total_norm


def save_model(args, epoch, model, model_without_ddp, optimizer, loss_scaler):
    output_dir = Path(args.output_dir)
    epoch_name = str(epoch)
    if loss_scaler is not None:
        checkpoint_paths = [output_dir / ('checkpoint-%s.pth' % epoch_name)]
        for checkpoint_path in checkpoint_paths:
            to_save = {
                'model': model_without_ddp.state_dict(),
                'optimizer': optimizer.state_dict(),
                'epoch': epoch,
                'scaler': loss_scaler.state_dict(),
                'args': args,
            }

            save_on_master(to_save, checkpoint_path)
    else:
        client_state = {'epoch': epoch}
        model.save_checkpoint(save_dir=args.output_dir, tag="checkpoint-%s" % epoch_name, client_state=client_state)


def load_model(args, model_without_ddp, optimizer, loss_scaler):
    if args.resume:
        if args.resume.startswith('https'):
            checkpoint = torch.hub.load_state_dict_from_url(
                args.resume, map_location='cpu', check_hash=True)
        else:
            checkpoint = torch.load(args.resume, map_location='cpu')
        model_without_ddp.load_state_dict(checkpoint['model'])
        print("Resume checkpoint %s" % args.resume)
        if 'optimizer' in checkpoint and 'epoch' in checkpoint and not (hasattr(args, 'eval') and args.eval):
            optimizer.load_state_dict(checkpoint['optimizer'])
            args.start_epoch = checkpoint['epoch'] + 1
            if 'scaler' in checkpoint:
                loss_scaler.load_state_dict(checkpoint['scaler'])
            print("With optim & sched!")


def all_reduce_mean(x):
    world_size = get_world_size()
    if world_size > 1:
        x_reduce = torch.tensor(x).cuda()
        dist.all_reduce(x_reduce)
        x_reduce /= world_size
        return x_reduce.item()
    else:
        return x


# utils
@torch.no_grad()
def concat_all_gather(tensor):
    """
    Performs all_gather operation on the provided tensors.
    *** Warning ***: torch.distributed.all_gather has no gradient.
    """
    tensors_gather = [torch.ones_like(tensor)
        for _ in range(torch.distributed.get_world_size())]
    torch.distributed.all_gather(tensors_gather, tensor, async_op=False)

    output = torch.cat(tensors_gather, dim=0)
    return output
    



def merge_vmae_to_avmae(avmae_state_dict, vmae_ckpt):
    #keys_to_copy=['pos_embed','patch_embed']
    #replaced=0

    vmae_ckpt['cls_token'] = vmae_ckpt['cls_token_v']
    vmae_ckpt['mask_token'] = vmae_ckpt['mask_token_v']

    # pos_emb % not trainable, use default
    pos_embed_v = vmae_ckpt['pos_embed_v'] #1,589,768
    pos_embed = pos_embed_v[:,1:,:] #1,588,768
    cls_embed = pos_embed_v[:,0,:].unsqueeze(1)
    pos_embed = pos_embed.reshape(1, 2, 14, 14,768).sum(dim=1) # 1, 14, 14, 768
    print("Position interpolate from 14,14 to 64,8")
    pos_embed = pos_embed.permute(0, 3, 1, 2) # 1, 14,14,768 -> 1,768,14,14
    pos_embed = torch.nn.functional.interpolate(
        pos_embed, size=(64, 8), mode='bicubic', align_corners=False)
    pos_embed = pos_embed.permute(0, 2, 3, 1).flatten(1, 2) # 1, 14, 14, 768 => 1, 196,768
    pos_embed = torch.cat((cls_embed, pos_embed), dim=1)
    assert(vmae_ckpt['pos_embed'].shape == pos_embed.shape)
    vmae_ckpt['pos_embed'] = pos_embed 
    # patch_emb
    # aggregate 3 channels in video-rgb ckpt to 1 channel for audio
    v_weight = vmae_ckpt['patch_embed_v.proj.weight'] # 768,3,2,16,16
    new_proj_weight = torch.nn.Parameter(v_weight.sum(dim=2).sum(dim=1).unsqueeze(1))
    assert(new_proj_weight.shape == vmae_ckpt['patch_embed.proj.weight'].shape)
    vmae_ckpt['patch_embed.proj.weight'] = new_proj_weight
    vmae_ckpt['patch_embed.proj.bias'] = vmae_ckpt['patch_embed_v.proj.bias']

    # hack 
    vmae_ckpt['norm.weight'] = vmae_ckpt['norm_v.weight']
    vmae_ckpt['norm.bias'] = vmae_ckpt['norm_v.bias']

    # replace transformer encoder
    for k,v in vmae_ckpt.items():
        if k.startswith('blocks.'):
            kk = k.replace('blocks.','blocks_v.')
            vmae_ckpt[k] = vmae_ckpt[kk]
        elif k.startswith('blocks_v.'):
            pass
        else:
            print(k)
            pass
    print(k)

def merge_mae_avmae(avmae_ckpt, mae_ckpt):
    avmae_ckpt_model = avmae_ckpt['model']
    mae_ckpt_model = mae_ckpt['model']

    keys_to_change=['cls_token_v']
    replaced=0
    for k,v in avmae_ckpt_model.items():
        if k in keys_to_change:
            kk = k.replace('_v','')
            avmae_ckpt_model[k].shape
            mae_ckpt_model[kk].shape
            assert(avmae_ckpt_model[k].shape == mae_ckpt_model[kk].shape)
            avmae_ckpt_model[k] = mae_ckpt_model[kk]
            replaced+=1
        elif k.startswith('blocks_v'):
            kk = k.replace('blocks_v','blocks')
            assert(avmae_ckpt_model[k].shape == mae_ckpt_model[kk].shape)
            avmae_ckpt_model[k] = mae_ckpt_model[kk]
            replaced+=1

    #hacking patch_emb
    
    org_weight = avmae_ckpt_model['patch_embed_v.proj.weight']
    new_weight = mae_ckpt_model['patch_embed.proj.weight'].unsqueeze(2).repeat(1,1,2,1,1)
    #new_weight = mae_ckpt_model['patch_embed.proj.weight'].unsqueeze(2)
    #assert(org_weight.shape == new_weight.shape)
    avmae_ckpt_model['patch_embed_v.proj.weight']=new_weight
    avmae_ckpt_model['patch_embed_v.proj.bias']=mae_ckpt_model['patch_embed.proj.bias']

    #hacking pos_emb
    org_pos = avmae_ckpt_model['pos_embed_v']
    #new_pos_final = get_3d_sincos_pos_embed(org_pos.shape[-1], 14, 4, True) # assume 8 frame, tride = 2
    #new_pos_final = torch.from_numpy(new_pos_final).float().unsqueeze(0)
    
    new_pos=mae_ckpt_model['pos_embed']
    new_cls = new_pos[:,:1,:]
    new_pos_ = new_pos[:,1:,:]
    new_pos_final = torch.cat((new_cls,new_pos_,new_pos_,new_pos_,new_pos_), dim=1)
    assert(org_pos.shape == new_pos_final.shape)
    avmae_ckpt_model['pos_embed_v']=new_pos_final


    print(f'overide avmae ckpt with mae for {replaced}+3 modules')
    avmae_ckpt['model']=avmae_ckpt_model
    return avmae_ckpt



block_v_list=[
'blocks_v.11.norm1.weight','blocks_v.11.norm1.bias','blocks_v.11.attn.qkv.weight','blocks_v.11.attn.qkv.bias','blocks_v.11.attn.proj.weight','blocks_v.11.attn.proj.bias','blocks_v.11.norm2.weight','blocks_v.11.norm2.bias','blocks_v.11.mlp.fc1.weight','blocks_v.11.mlp.fc1.bias','blocks_v.11.mlp.fc2.weight','blocks_v.11.mlp.fc2.bias'
]

block_v_list_large=[
'blocks_v.23.norm1.weight','blocks_v.23.norm1.bias','blocks_v.23.attn.qkv.weight','blocks_v.23.attn.qkv.bias','blocks_v.23.attn.proj.weight','blocks_v.23.attn.proj.bias','blocks_v.23.norm2.weight','blocks_v.23.norm2.bias','blocks_v.23.mlp.fc1.weight','blocks_v.23.mlp.fc1.bias','blocks_v.23.mlp.fc2.weight','blocks_v.23.mlp.fc2.bias'
]

def load_mae(mae_ckpt_path, av=False, large=False):
    mae_ckpt = torch.load(mae_ckpt_path, map_location='cpu')
    mae_ckpt_model = mae_ckpt['model']
    from collections import OrderedDict
    new_ckpt_model = OrderedDict()
    for key, value in mae_ckpt_model.items():
        keys = key.split('.')
        new_key = keys[0] + '_v'

        if len (keys) > 1:
            key_other ='.'.join(keys[1:])
            new_key = new_key + '.' + key_other

        # hacking mae to video mae
        if new_key == 'pos_embed_v':
            new_pos = value
            new_cls = new_pos[:,:1,:]
            new_pos_ = new_pos[:,1:,:]
            value = torch.cat((new_cls,new_pos_,new_pos_,new_pos_,new_pos_), dim=1)
        elif new_key == 'patch_embed_v.proj.weight':
            value = value.unsqueeze(2).repeat(1,1,2,1,1)
        elif new_key == 'decoder_pos_embed_v':
            new_pos = value
            new_cls = new_pos[:,:1,:]
            new_pos_ = new_pos[:,1:,:]
            value = torch.cat((new_cls,new_pos_,new_pos_,new_pos_,new_pos_), dim=1)

        if av:
            if not large: # ViT-B
                if new_key in block_v_list:
                    av_key = new_key.replace('_v.11', '_av.0')
                    new_ckpt_model[av_key] = value
                    av_key = new_key.replace('_v.11', '_av.1')
                    new_ckpt_model[av_key] = value
            else: # ViT-L
                if new_key in block_v_list_large:
                    av_key = new_key.replace('_v.23', '_av.0')
                    new_ckpt_model[av_key] = value
                    av_key = new_key.replace('_v.23', '_av.1')
                    new_ckpt_model[av_key] = value


        print(new_key)
        new_ckpt_model[new_key] = value
    return new_ckpt_model

    # mae_ckpt_model = mae_ckpt['model']


    # ### start here 1025
    # state_dict = model.state_dict()

    # if not args.eval:
    #     for k in ['head.weight', 'head.bias']:
    #         if k in checkpoint_model and checkpoint_model[k].shape != state_dict[k].shape:
    #             print(f"Removing key {k} from pretrained checkpoint")
    #             del checkpoint_model[k]


    # keys_to_change=['cls_token_v']
    # replaced=0
    # for k,v in avmae_ckpt_model.items():
    #     if k in keys_to_change:
    #         kk = k.replace('_v','')
    #         avmae_ckpt_model[k].shape
    #         mae_ckpt_model[kk].shape
            
    #         avmae_ckpt_model[k] = mae_ckpt_model[kk]
    #         replaced+=1
    #     elif k.startswith('blocks_v'):
    #         kk = k.replace('blocks_v','blocks')
            
    #         avmae_ckpt_model[k] = mae_ckpt_model[kk]
    #         replaced+=1

    # #hacking patch_emb
    
    # org_weight = avmae_ckpt_model['patch_embed_v.proj.weight']
    # new_weight = mae_ckpt_model['patch_embed.proj.weight'].unsqueeze(2).repeat(1,1,2,1,1)
    # avmae_ckpt_model['patch_embed_v.proj.weight']=new_weight
    # avmae_ckpt_model['patch_embed_v.proj.bias']=mae_ckpt_model['patch_embed.proj.bias']

    # #hacking pos_emb
    # org_pos = avmae_ckpt_model['pos_embed_v']
    # new_pos=mae_ckpt_model['pos_embed']
    # new_cls = new_pos[:,:1,:]
    # new_pos_ = new_pos[:,1:,:]
    # new_pos_final = torch.cat((new_cls,new_pos_,new_pos_,new_pos_,new_pos_), dim=1)
    # avmae_ckpt_model['pos_embed_v']=new_pos_final


    # print(f'overide avmae ckpt with mae for {replaced}+3 modules')
    # new_ckpt['model']=avmae_ckpt_model
    # return new_ckpt
upstream_models/mavil/util/misc.py
roger-tseng/av-superb

import torch
import torch.nn as nn
from timm.models.layers import to_2tuple

class PatchEmbed_org(nn.Module):
    """ Image to Patch Embedding
    """
    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
        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])
        self.patch_hw = (img_size[1] // patch_size[1], img_size[0] // patch_size[0])
        self.img_size = img_size
        self.patch_size = patch_size
        self.num_patches = num_patches

        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)

    def forward(self, x):
        B, C, H, W = x.shape
        # FIXME look at relaxing size constraints
        #assert H == self.img_size[0] and W == self.img_size[1], \
        #    f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
        x = self.proj(x).flatten(2).transpose(1, 2)
        return x


class PatchEmbed_new(nn.Module):
    """ Flexible Image to Patch Embedding
    """
    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, stride=10):
        super().__init__()
        img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)
        stride = to_2tuple(stride)
        
        self.img_size = img_size
        self.patch_size = patch_size
        

        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride) # with overlapped patches
        #self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)

        #self.patch_hw = (img_size[1] // patch_size[1], img_size[0] // patch_size[0])
        #self.num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
        _, _, h, w = self.get_output_shape(img_size) # n, emb_dim, h, w
        self.patch_hw = (h, w)
        self.num_patches = h*w

    def get_output_shape(self, img_size):
        # todo: don't be lazy..
        return self.proj(torch.randn(1,1,img_size[0],img_size[1])).shape 

    def forward(self, x):
        B, C, H, W = x.shape
        # FIXME look at relaxing size constraints
        #assert H == self.img_size[0] and W == self.img_size[1], \
        #    f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
        #x = self.proj(x).flatten(2).transpose(1, 2)
        x = self.proj(x) # 32, 1, 1024, 128 -> 32, 768, 101, 12
        x = x.flatten(2) # 32, 768, 101, 12 -> 32, 768, 1212
        x = x.transpose(1, 2) # 32, 768, 1212 -> 32, 1212, 768
        return x

class PatchEmbed3D_new(nn.Module):
    """ Flexible Image to Patch Embedding
    """
    def __init__(self, video_size=(16,224,224), patch_size=(2,16,16), in_chans=3, embed_dim=768, stride=(2,16,16)):
        super().__init__()
        
        self.video_size = video_size
        self.patch_size = patch_size
        self.in_chans = in_chans
        

        self.proj = nn.Conv3d(in_chans, embed_dim, kernel_size=patch_size, stride=stride)
        _, _, t, h, w = self.get_output_shape(video_size) # n, emb_dim, h, w
        self.patch_thw = (t, h, w)
        self.num_patches = t*h*w

    def get_output_shape(self, video_size):
        # todo: don't be lazy..
        return self.proj(torch.randn(1, self.in_chans, video_size[0], video_size[1], video_size[2])).shape 

    def forward(self, x):
        B, C, T, H, W = x.shape
        x = self.proj(x) # 32, 3, 16, 224, 224 -> 32, 768, 8, 14, 14
        x = x.flatten(2) # 32, 768, 1568
        x = x.transpose(1, 2) # 32, 768, 1568 -> 32, 1568, 768
        return x

if __name__ == '__main__':
    #patch_emb = PatchEmbed_new(img_size=224, patch_size=16, in_chans=1, embed_dim=64, stride=(16,16))
    #input = torch.rand(8,1,1024,128)
    #output = patch_emb(input)
    #print(output.shape) # (8,512,64)

    patch_emb = PatchEmbed3D_new(video_size=(6,224,224), patch_size=(2,16,16), in_chans=3, embed_dim=768, stride=(2,16,16))
    input = torch.rand(8,3,6,224,224)
    output = patch_emb(input)
    print(output.shape) # (8,64)
upstream_models/mavil/util/patch_embed.py
roger-tseng/av-superb
# 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.
# --------------------------------------------------------
# Position embedding utils
# --------------------------------------------------------

import numpy as np

import torch

# --------------------------------------------------------
# 2D sine-cosine position embedding
# References:
# Transformer: https://github.com/tensorflow/models/blob/master/official/nlp/transformer/model_utils.py
# MoCo v3: https://github.com/facebookresearch/moco-v3
# --------------------------------------------------------
def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False):
    """
    grid_size: int of the grid height and width
    return:
    pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
    """
    grid_h = np.arange(grid_size, dtype=np.float32)
    grid_w = np.arange(grid_size, dtype=np.float32)
    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
    grid = np.stack(grid, axis=0)

    grid = grid.reshape([2, 1, grid_size, grid_size])
    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
    if cls_token:
        pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)
    return pos_embed


def get_2d_sincos_pos_embed_flexible(embed_dim, grid_size, cls_token=False):
    """
    grid_size: int of the grid height and width
    return:
    pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
    """
    grid_h = np.arange(grid_size[0], dtype=np.float32)
    grid_w = np.arange(grid_size[1], dtype=np.float32)
    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
    grid = np.stack(grid, axis=0)

    grid = grid.reshape([2, 1, grid_size[0], grid_size[1]])
    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
    if cls_token:
        pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)
    return pos_embed


def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
    assert embed_dim % 2 == 0

    # use half of dimensions to encode grid_h
    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)

    emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D)
    return emb


def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
    """
    embed_dim: output dimension for each position
    pos: a list of positions to be encoded: size (M,)
    out: (M, D)
    """
    assert embed_dim % 2 == 0
    omega = np.arange(embed_dim // 2, dtype=np.float)
    omega /= embed_dim / 2.
    omega = 1. / 10000**omega  # (D/2,)

    pos = pos.reshape(-1)  # (M,)
    out = np.einsum('m,d->md', pos, omega)  # (M, D/2), outer product

    emb_sin = np.sin(out) # (M, D/2)
    emb_cos = np.cos(out) # (M, D/2)

    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
    return emb


# --------------------------------------------------------
# Interpolate position embeddings for high-resolution
# References:
# DeiT: https://github.com/facebookresearch/deit
# --------------------------------------------------------
def interpolate_pos_embed(model, checkpoint_model):
    if 'pos_embed' in checkpoint_model:
        pos_embed_checkpoint = checkpoint_model['pos_embed']
        embedding_size = pos_embed_checkpoint.shape[-1]
        num_patches = model.patch_embed.num_patches
        num_extra_tokens = model.pos_embed.shape[-2] - num_patches
        # height (== width) for the checkpoint position embedding
        orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
        # height (== width) for the new position embedding
        new_size = int(num_patches ** 0.5)
        # class_token and dist_token are kept unchanged
        if orig_size != new_size:
            print("Position interpolate from %dx%d to %dx%d" % (orig_size, orig_size, new_size, new_size))
            extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
            # only the position tokens are interpolated
            pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
            pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)
            pos_tokens = torch.nn.functional.interpolate(
                pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False)
            pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
            new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
            checkpoint_model['pos_embed'] = new_pos_embed

def interpolate_pos_embed_img2audio(model, checkpoint_model, orig_size, new_size):
    if 'pos_embed' in checkpoint_model:
        pos_embed_checkpoint = checkpoint_model['pos_embed']
        embedding_size = pos_embed_checkpoint.shape[-1]
        num_patches = model.patch_embed.num_patches
        num_extra_tokens = model.pos_embed.shape[-2] - num_patches
        # height (== width) for the checkpoint position embedding
        #orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
        # height (== width) for the new position embedding
        #new_size = int(num_patches ** 0.5)
        # class_token and dist_token are kept unchanged
        if orig_size != new_size:
            print("Position interpolate from %dx%d to %dx%d" % (orig_size[0], orig_size[1], new_size[0], new_size[1]))
            extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
            # only the position tokens are interpolated
            pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
            pos_tokens = pos_tokens.reshape(-1, orig_size[0], orig_size[1], embedding_size).permute(0, 3, 1, 2)
            pos_tokens = torch.nn.functional.interpolate(
                pos_tokens, size=(new_size[0], new_size[1]), mode='bicubic', align_corners=False)
            pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
            new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
            checkpoint_model['pos_embed'] = new_pos_embed

def interpolate_pos_embed_audio(model, checkpoint_model, orig_size, new_size):
    if 'pos_embed' in checkpoint_model:
        pos_embed_checkpoint = checkpoint_model['pos_embed']
        embedding_size = pos_embed_checkpoint.shape[-1]
        num_patches = model.patch_embed.num_patches
        num_extra_tokens = model.pos_embed.shape[-2] - num_patches
        if orig_size != new_size:
            print("Position interpolate from %dx%d to %dx%d" % (orig_size[0], orig_size[1], new_size[0], new_size[1]))
            #extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
            # only the position tokens are interpolated
            cls_token = pos_embed_checkpoint[:,0,:].unsqueeze(1)
            pos_tokens = pos_embed_checkpoint[:,1:,:] # remove 
            pos_tokens = pos_tokens.reshape(-1, orig_size[0], orig_size[1], embedding_size) #.permute(0, 3, 1, 2)
            #pos_tokens = torch.nn.functional.interpolate(
            #    pos_tokens, size=(new_size[0], new_size[1]), mode='bicubic', align_corners=False)
            
            #pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
            pos_tokens = pos_tokens[:,:,:new_size[1],:] # assume only time diff
            pos_tokens = pos_tokens.flatten(1, 2)
            new_pos_embed = torch.cat((cls_token, pos_tokens), dim=1)
            checkpoint_model['pos_embed'] = new_pos_embed


def interpolate_patch_embed_audio(model, checkpoint_model, orig_channel, new_channel=1, kernel_size=(16,16), stride=(16,16), padding=(0,0)):
    if orig_channel != new_channel:
        if 'patch_embed.proj.weight' in checkpoint_model:
            # aggregate 3 channels in rgb ckpt to 1 channel for audio
            new_proj_weight = torch.nn.Parameter(torch.sum(checkpoint_model['patch_embed.proj.weight'], dim=1).unsqueeze(1))
            checkpoint_model['patch_embed.proj.weight'] = new_proj_weight



def get_3d_sincos_pos_embed(embed_dim, grid_size, t_size, cls_token=False):
    """
    grid_size: int of the grid height and width
    t_size: int of the temporal size
    return:
    pos_embed: [t_size*grid_size*grid_size, embed_dim] or [1+t_size*grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
    """
    assert embed_dim % 4 == 0
    embed_dim_spatial = embed_dim // 4 * 3
    embed_dim_temporal = embed_dim // 4

    # spatial
    grid_h = np.arange(grid_size, dtype=np.float32)
    grid_w = np.arange(grid_size, dtype=np.float32)
    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
    grid = np.stack(grid, axis=0)

    grid = grid.reshape([2, 1, grid_size, grid_size])
    pos_embed_spatial = get_2d_sincos_pos_embed_from_grid(embed_dim_spatial, grid)

    # temporal
    grid_t = np.arange(t_size, dtype=np.float32)
    pos_embed_temporal = get_1d_sincos_pos_embed_from_grid(embed_dim_temporal, grid_t)

    # concate: [T, H, W] order
    pos_embed_temporal = pos_embed_temporal[:, np.newaxis, :]
    pos_embed_temporal = np.repeat(
        pos_embed_temporal, grid_size ** 2, axis=1
    )  # [T, H*W, D // 4]
    pos_embed_spatial = pos_embed_spatial[np.newaxis, :, :]
    pos_embed_spatial = np.repeat(
        pos_embed_spatial, t_size, axis=0
    )  # [T, H*W, D // 4 * 3]

    pos_embed = np.concatenate([pos_embed_temporal, pos_embed_spatial], axis=-1)
    pos_embed = pos_embed.reshape([-1, embed_dim])  # [T*H*W, D]

    if cls_token:
        pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)
    return pos_embed

if __name__ == "__main__":
   #pos_emb = get_2d_sincos_pos_embed(768, 14)
   #print(pos_emb.shape) #196x768
   pos_emb = get_3d_sincos_pos_embed(768, 14, 8)
   print(pos_emb.shape) #1568x768
upstream_models/mavil/util/pos_embed.py
roger-tseng/av-superb
import numpy as np
from scipy import stats
from sklearn import metrics
import torch

def d_prime(auc):
    standard_normal = stats.norm()
    d_prime = standard_normal.ppf(auc) * np.sqrt(2.0)
    return d_prime

@torch.no_grad()
def concat_all_gather(tensor):
    """
    Performs all_gather operation on the provided tensors.
    *** Warning ***: torch.distributed.all_gather has no gradient.
    """
    tensors_gather = [torch.ones_like(tensor)
        for _ in range(torch.distributed.get_world_size())]
    torch.distributed.all_gather(tensors_gather, tensor, async_op=False)

    output = torch.cat(tensors_gather, dim=0)
    return output

def calculate_stats(output, target):
    """Calculate statistics including mAP, AUC, etc.

    Args:
      output: 2d array, (samples_num, classes_num)
      target: 2d array, (samples_num, classes_num)

    Returns:
      stats: list of statistic of each class.
    """

    classes_num = target.shape[-1]
    stats = []

    # Accuracy, only used for single-label classification such as esc-50, not for multiple label one such as AudioSet
    acc = metrics.accuracy_score(np.argmax(target, 1), np.argmax(output, 1))

    # Class-wise statistics
    for k in range(classes_num):

        # Average precision
        avg_precision = metrics.average_precision_score(
            target[:, k], output[:, k], average=None)

        # AUC
        #auc = metrics.roc_auc_score(target[:, k], output[:, k], average=None)

        # Precisions, recalls
        (precisions, recalls, thresholds) = metrics.precision_recall_curve(
            target[:, k], output[:, k])

        # FPR, TPR
        #(fpr, tpr, thresholds) = metrics.roc_curve(target[:, k], output[:, k])

        save_every_steps = 1000     # Sample statistics to reduce size
        dict = {'precisions': precisions[0::save_every_steps],
                'recalls': recalls[0::save_every_steps],
                'AP': avg_precision,
                #'fpr': fpr[0::save_every_steps],
                #'fnr': 1. - tpr[0::save_every_steps],
                #'auc': auc,
                # note acc is not class-wise, this is just to keep consistent with other metrics
                'acc': acc
                }
        stats.append(dict)

    return stats
upstream_models/mavil/util/stat.py
roger-tseng/av-superb
#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.

import math

# import cv2
import random

import numpy as np
import torch
import torchvision.transforms.functional as F
import PIL
from PIL import Image, ImageEnhance, ImageOps
from torchvision import transforms
import re


#from .rand_augment import rand_augment_transform RandomErasing


_pil_interpolation_to_str = {
    Image.NEAREST: "PIL.Image.NEAREST",
    Image.BILINEAR: "PIL.Image.BILINEAR",
    Image.BICUBIC: "PIL.Image.BICUBIC",
    Image.LANCZOS: "PIL.Image.LANCZOS",
    Image.HAMMING: "PIL.Image.HAMMING",
    Image.BOX: "PIL.Image.BOX",
}


_RANDOM_INTERPOLATION = (Image.BILINEAR, Image.BICUBIC)


def _pil_interp(method):
    if method == "bicubic":
        return Image.BICUBIC
    elif method == "lanczos":
        return Image.LANCZOS
    elif method == "hamming":
        return Image.HAMMING
    else:
        return Image.BILINEAR


def random_short_side_scale_jitter(
    images, min_size, max_size, boxes=None, inverse_uniform_sampling=False
):
    """
    Perform a spatial short scale jittering on the given images and
    corresponding boxes.
    Args:
        images (tensor): images to perform scale jitter. Dimension is
            `num frames` x `channel` x `height` x `width`.
        min_size (int): the minimal size to scale the frames.
        max_size (int): the maximal size to scale the frames.
        boxes (ndarray): optional. Corresponding boxes to images.
            Dimension is `num boxes` x 4.
        inverse_uniform_sampling (bool): if True, sample uniformly in
            [1 / max_scale, 1 / min_scale] and take a reciprocal to get the
            scale. If False, take a uniform sample from [min_scale, max_scale].
    Returns:
        (tensor): the scaled images with dimension of
            `num frames` x `channel` x `new height` x `new width`.
        (ndarray or None): the scaled boxes with dimension of
            `num boxes` x 4.
    """
    if inverse_uniform_sampling:
        size = int(round(1.0 / np.random.uniform(1.0 / max_size, 1.0 / min_size)))
    else:
        size = int(round(np.random.uniform(min_size, max_size)))

    height = images.shape[2]
    width = images.shape[3]
    if (width <= height and width == size) or (height <= width and height == size):
        return images, boxes
    new_width = size
    new_height = size
    if width < height:
        new_height = int(math.floor((float(height) / width) * size))
        if boxes is not None:
            boxes = boxes * float(new_height) / height
    else:
        new_width = int(math.floor((float(width) / height) * size))
        if boxes is not None:
            boxes = boxes * float(new_width) / width

    return (
        torch.nn.functional.interpolate(
            images,
            size=(new_height, new_width),
            mode="bilinear",
            align_corners=False,
        ),
        boxes,
    )


def crop_boxes(boxes, x_offset, y_offset):
    """
    Peform crop on the bounding boxes given the offsets.
    Args:
        boxes (ndarray or None): bounding boxes to peform crop. The dimension
            is `num boxes` x 4.
        x_offset (int): cropping offset in the x axis.
        y_offset (int): cropping offset in the y axis.
    Returns:
        cropped_boxes (ndarray or None): the cropped boxes with dimension of
            `num boxes` x 4.
    """
    cropped_boxes = boxes.copy()
    cropped_boxes[:, [0, 2]] = boxes[:, [0, 2]] - x_offset
    cropped_boxes[:, [1, 3]] = boxes[:, [1, 3]] - y_offset

    return cropped_boxes


def random_crop(images, size, boxes=None):
    """
    Perform random spatial crop on the given images and corresponding boxes.
    Args:
        images (tensor): images to perform random crop. The dimension is
            `num frames` x `channel` x `height` x `width`.
        size (int): the size of height and width to crop on the image.
        boxes (ndarray or None): optional. Corresponding boxes to images.
            Dimension is `num boxes` x 4.
    Returns:
        cropped (tensor): cropped images with dimension of
            `num frames` x `channel` x `size` x `size`.
        cropped_boxes (ndarray or None): the cropped boxes with dimension of
            `num boxes` x 4.
    """
    if images.shape[2] == size and images.shape[3] == size:
        return images
    height = images.shape[2]
    width = images.shape[3]
    y_offset = 0
    if height > size:
        y_offset = int(np.random.randint(0, height - size))
    x_offset = 0
    if width > size:
        x_offset = int(np.random.randint(0, width - size))
    cropped = images[:, :, y_offset : y_offset + size, x_offset : x_offset + size]

    cropped_boxes = crop_boxes(boxes, x_offset, y_offset) if boxes is not None else None

    return cropped, cropped_boxes


def horizontal_flip(prob, images, boxes=None):
    """
    Perform horizontal flip on the given images and corresponding boxes.
    Args:
        prob (float): probility to flip the images.
        images (tensor): images to perform horizontal flip, the dimension is
            `num frames` x `channel` x `height` x `width`.
        boxes (ndarray or None): optional. Corresponding boxes to images.
            Dimension is `num boxes` x 4.
    Returns:
        images (tensor): images with dimension of
            `num frames` x `channel` x `height` x `width`.
        flipped_boxes (ndarray or None): the flipped boxes with dimension of
            `num boxes` x 4.
    """
    if boxes is None:
        flipped_boxes = None
    else:
        flipped_boxes = boxes.copy()

    if np.random.uniform() < prob:
        images = images.flip((-1))

        if len(images.shape) == 3:
            width = images.shape[2]
        elif len(images.shape) == 4:
            width = images.shape[3]
        else:
            raise NotImplementedError("Dimension does not supported")
        if boxes is not None:
            flipped_boxes[:, [0, 2]] = width - boxes[:, [2, 0]] - 1

    return images, flipped_boxes


def uniform_crop(images, size, spatial_idx, boxes=None, scale_size=None):
    """
    Perform uniform spatial sampling on the images and corresponding boxes.
    Args:
        images (tensor): images to perform uniform crop. The dimension is
            `num frames` x `channel` x `height` x `width`.
        size (int): size of height and weight to crop the images.
        spatial_idx (int): 0, 1, or 2 for left, center, and right crop if width
            is larger than height. Or 0, 1, or 2 for top, center, and bottom
            crop if height is larger than width.
        boxes (ndarray or None): optional. Corresponding boxes to images.
            Dimension is `num boxes` x 4.
        scale_size (int): optinal. If not None, resize the images to scale_size before
            performing any crop.
    Returns:
        cropped (tensor): images with dimension of
            `num frames` x `channel` x `size` x `size`.
        cropped_boxes (ndarray or None): the cropped boxes with dimension of
            `num boxes` x 4.
    """
    assert spatial_idx in [0, 1, 2]
    ndim = len(images.shape)
    if ndim == 3:
        images = images.unsqueeze(0)
    height = images.shape[2]
    width = images.shape[3]

    if scale_size is not None:
        if width <= height:
            width, height = scale_size, int(height / width * scale_size)
        else:
            width, height = int(width / height * scale_size), scale_size
        images = torch.nn.functional.interpolate(
            images,
            size=(height, width),
            mode="bilinear",
            align_corners=False,
        )

    y_offset = int(math.ceil((height - size) / 2))
    x_offset = int(math.ceil((width - size) / 2))

    if height > width:
        if spatial_idx == 0:
            y_offset = 0
        elif spatial_idx == 2:
            y_offset = height - size
    else:
        if spatial_idx == 0:
            x_offset = 0
        elif spatial_idx == 2:
            x_offset = width - size
    cropped = images[:, :, y_offset : y_offset + size, x_offset : x_offset + size]
    cropped_boxes = crop_boxes(boxes, x_offset, y_offset) if boxes is not None else None
    if ndim == 3:
        cropped = cropped.squeeze(0)
    return cropped, cropped_boxes


def clip_boxes_to_image(boxes, height, width):
    """
    Clip an array of boxes to an image with the given height and width.
    Args:
        boxes (ndarray): bounding boxes to perform clipping.
            Dimension is `num boxes` x 4.
        height (int): given image height.
        width (int): given image width.
    Returns:
        clipped_boxes (ndarray): the clipped boxes with dimension of
            `num boxes` x 4.
    """
    clipped_boxes = boxes.copy()
    clipped_boxes[:, [0, 2]] = np.minimum(
        width - 1.0, np.maximum(0.0, boxes[:, [0, 2]])
    )
    clipped_boxes[:, [1, 3]] = np.minimum(
        height - 1.0, np.maximum(0.0, boxes[:, [1, 3]])
    )
    return clipped_boxes


def blend(images1, images2, alpha):
    """
    Blend two images with a given weight alpha.
    Args:
        images1 (tensor): the first images to be blended, the dimension is
            `num frames` x `channel` x `height` x `width`.
        images2 (tensor): the second images to be blended, the dimension is
            `num frames` x `channel` x `height` x `width`.
        alpha (float): the blending weight.
    Returns:
        (tensor): blended images, the dimension is
            `num frames` x `channel` x `height` x `width`.
    """
    return images1 * alpha + images2 * (1 - alpha)


def grayscale(images):
    """
    Get the grayscale for the input images. The channels of images should be
    in order BGR.
    Args:
        images (tensor): the input images for getting grayscale. Dimension is
            `num frames` x `channel` x `height` x `width`.
    Returns:
        img_gray (tensor): blended images, the dimension is
            `num frames` x `channel` x `height` x `width`.
    """
    # R -> 0.299, G -> 0.587, B -> 0.114.
    img_gray = torch.tensor(images)
    gray_channel = 0.299 * images[:, 2] + 0.587 * images[:, 1] + 0.114 * images[:, 0]
    img_gray[:, 0] = gray_channel
    img_gray[:, 1] = gray_channel
    img_gray[:, 2] = gray_channel
    return img_gray


def color_jitter(images, img_brightness=0, img_contrast=0, img_saturation=0):
    """
    Perfrom a color jittering on the input images. The channels of images
    should be in order BGR.
    Args:
        images (tensor): images to perform color jitter. Dimension is
            `num frames` x `channel` x `height` x `width`.
        img_brightness (float): jitter ratio for brightness.
        img_contrast (float): jitter ratio for contrast.
        img_saturation (float): jitter ratio for saturation.
    Returns:
        images (tensor): the jittered images, the dimension is
            `num frames` x `channel` x `height` x `width`.
    """

    jitter = []
    if img_brightness != 0:
        jitter.append("brightness")
    if img_contrast != 0:
        jitter.append("contrast")
    if img_saturation != 0:
        jitter.append("saturation")

    if len(jitter) > 0:
        order = np.random.permutation(np.arange(len(jitter)))
        for idx in range(0, len(jitter)):
            if jitter[order[idx]] == "brightness":
                images = brightness_jitter(img_brightness, images)
            elif jitter[order[idx]] == "contrast":
                images = contrast_jitter(img_contrast, images)
            elif jitter[order[idx]] == "saturation":
                images = saturation_jitter(img_saturation, images)
    return images


def brightness_jitter(var, images):
    """
    Perfrom brightness jittering on the input images. The channels of images
    should be in order BGR.
    Args:
        var (float): jitter ratio for brightness.
        images (tensor): images to perform color jitter. Dimension is
            `num frames` x `channel` x `height` x `width`.
    Returns:
        images (tensor): the jittered images, the dimension is
            `num frames` x `channel` x `height` x `width`.
    """
    alpha = 1.0 + np.random.uniform(-var, var)

    img_bright = torch.zeros(images.shape)
    images = blend(images, img_bright, alpha)
    return images


def contrast_jitter(var, images):
    """
    Perfrom contrast jittering on the input images. The channels of images
    should be in order BGR.
    Args:
        var (float): jitter ratio for contrast.
        images (tensor): images to perform color jitter. Dimension is
            `num frames` x `channel` x `height` x `width`.
    Returns:
        images (tensor): the jittered images, the dimension is
            `num frames` x `channel` x `height` x `width`.
    """
    alpha = 1.0 + np.random.uniform(-var, var)

    img_gray = grayscale(images)
    img_gray[:] = torch.mean(img_gray, dim=(1, 2, 3), keepdim=True)
    images = blend(images, img_gray, alpha)
    return images


def saturation_jitter(var, images):
    """
    Perfrom saturation jittering on the input images. The channels of images
    should be in order BGR.
    Args:
        var (float): jitter ratio for saturation.
        images (tensor): images to perform color jitter. Dimension is
            `num frames` x `channel` x `height` x `width`.
    Returns:
        images (tensor): the jittered images, the dimension is
            `num frames` x `channel` x `height` x `width`.
    """
    alpha = 1.0 + np.random.uniform(-var, var)
    img_gray = grayscale(images)
    images = blend(images, img_gray, alpha)

    return images


def lighting_jitter(images, alphastd, eigval, eigvec):
    """
    Perform AlexNet-style PCA jitter on the given images.
    Args:
        images (tensor): images to perform lighting jitter. Dimension is
            `num frames` x `channel` x `height` x `width`.
        alphastd (float): jitter ratio for PCA jitter.
        eigval (list): eigenvalues for PCA jitter.
        eigvec (list[list]): eigenvectors for PCA jitter.
    Returns:
        out_images (tensor): the jittered images, the dimension is
            `num frames` x `channel` x `height` x `width`.
    """
    if alphastd == 0:
        return images
    # generate alpha1, alpha2, alpha3.
    alpha = np.random.normal(0, alphastd, size=(1, 3))
    eig_vec = np.array(eigvec)
    eig_val = np.reshape(eigval, (1, 3))
    rgb = np.sum(
        eig_vec * np.repeat(alpha, 3, axis=0) * np.repeat(eig_val, 3, axis=0),
        axis=1,
    )
    out_images = torch.zeros_like(images)
    if len(images.shape) == 3:
        # C H W
        channel_dim = 0
    elif len(images.shape) == 4:
        # T C H W
        channel_dim = 1
    else:
        raise NotImplementedError(f"Unsupported dimension {len(images.shape)}")

    for idx in range(images.shape[channel_dim]):
        # C H W
        if len(images.shape) == 3:
            out_images[idx] = images[idx] + rgb[2 - idx]
        # T C H W
        elif len(images.shape) == 4:
            out_images[:, idx] = images[:, idx] + rgb[2 - idx]
        else:
            raise NotImplementedError(f"Unsupported dimension {len(images.shape)}")

    return out_images


def color_normalization(images, mean, stddev):
    """
    Perform color nomration on the given images.
    Args:
        images (tensor): images to perform color normalization. Dimension is
            `num frames` x `channel` x `height` x `width`.
        mean (list): mean values for normalization.
        stddev (list): standard deviations for normalization.

    Returns:
        out_images (tensor): the noramlized images, the dimension is
            `num frames` x `channel` x `height` x `width`.
    """
    if len(images.shape) == 3:
        assert len(mean) == images.shape[0], "channel mean not computed properly"
        assert len(stddev) == images.shape[0], "channel stddev not computed properly"
    elif len(images.shape) == 4:
        assert len(mean) == images.shape[1], "channel mean not computed properly"
        assert len(stddev) == images.shape[1], "channel stddev not computed properly"
    else:
        raise NotImplementedError(f"Unsupported dimension {len(images.shape)}")

    out_images = torch.zeros_like(images)
    for idx in range(len(mean)):
        # C H W
        if len(images.shape) == 3:
            out_images[idx] = (images[idx] - mean[idx]) / stddev[idx]
        elif len(images.shape) == 4:
            out_images[:, idx] = (images[:, idx] - mean[idx]) / stddev[idx]
        else:
            raise NotImplementedError(f"Unsupported dimension {len(images.shape)}")
    return out_images


def _get_param_spatial_crop(
    scale, ratio, height, width, num_repeat=10, log_scale=True, switch_hw=False
):
    """
    Given scale, ratio, height and width, return sampled coordinates of the videos.
    """
    for _ in range(num_repeat):
        area = height * width
        target_area = random.uniform(*scale) * area
        if log_scale:
            log_ratio = (math.log(ratio[0]), math.log(ratio[1]))
            aspect_ratio = math.exp(random.uniform(*log_ratio))
        else:
            aspect_ratio = random.uniform(*ratio)

        w = int(round(math.sqrt(target_area * aspect_ratio)))
        h = int(round(math.sqrt(target_area / aspect_ratio)))

        if np.random.uniform() < 0.5 and switch_hw:
            w, h = h, w

        if 0 < w <= width and 0 < h <= height:
            i = random.randint(0, height - h)
            j = random.randint(0, width - w)
            return i, j, h, w

    # Fallback to central crop
    in_ratio = float(width) / float(height)
    if in_ratio < min(ratio):
        w = width
        h = int(round(w / min(ratio)))
    elif in_ratio > max(ratio):
        h = height
        w = int(round(h * max(ratio)))
    else:  # whole image
        w = width
        h = height
    i = (height - h) // 2
    j = (width - w) // 2
    return i, j, h, w


def random_resized_crop(
    images,
    target_height,
    target_width,
    scale=(0.8, 1.0),
    ratio=(3.0 / 4.0, 4.0 / 3.0),
):
    """
    Crop the given images to random size and aspect ratio. A crop of random
    size (default: of 0.08 to 1.0) of the original size and a random aspect
    ratio (default: of 3/4 to 4/3) of the original aspect ratio is made. This
    crop is finally resized to given size. This is popularly used to train the
    Inception networks.

    Args:
        images: Images to perform resizing and cropping.
        target_height: Desired height after cropping.
        target_width: Desired width after cropping.
        scale: Scale range of Inception-style area based random resizing.
        ratio: Aspect ratio range of Inception-style area based random resizing.
    """

    height = images.shape[2]
    width = images.shape[3]

    i, j, h, w = _get_param_spatial_crop(scale, ratio, height, width)
    cropped = images[:, :, i : i + h, j : j + w]
    return torch.nn.functional.interpolate(
        cropped,
        size=(target_height, target_width),
        mode="bilinear",
        align_corners=False,
    )


def random_resized_crop_with_shift(
    images,
    target_height,
    target_width,
    scale=(0.8, 1.0),
    ratio=(3.0 / 4.0, 4.0 / 3.0),
):
    """
    This is similar to random_resized_crop. However, it samples two different
    boxes (for cropping) for the first and last frame. It then linearly
    interpolates the two boxes for other frames.

    Args:
        images: Images to perform resizing and cropping.
        target_height: Desired height after cropping.
        target_width: Desired width after cropping.
        scale: Scale range of Inception-style area based random resizing.
        ratio: Aspect ratio range of Inception-style area based random resizing.
    """
    t = images.shape[1]
    height = images.shape[2]
    width = images.shape[3]

    i, j, h, w = _get_param_spatial_crop(scale, ratio, height, width)
    i_, j_, h_, w_ = _get_param_spatial_crop(scale, ratio, height, width)
    i_s = [int(i) for i in torch.linspace(i, i_, steps=t).tolist()]
    j_s = [int(i) for i in torch.linspace(j, j_, steps=t).tolist()]
    h_s = [int(i) for i in torch.linspace(h, h_, steps=t).tolist()]
    w_s = [int(i) for i in torch.linspace(w, w_, steps=t).tolist()]
    out = torch.zeros((3, t, target_height, target_width))
    for ind in range(t):
        out[:, ind : ind + 1, :, :] = torch.nn.functional.interpolate(
            images[
                :,
                ind : ind + 1,
                i_s[ind] : i_s[ind] + h_s[ind],
                j_s[ind] : j_s[ind] + w_s[ind],
            ],
            size=(target_height, target_width),
            mode="bilinear",
            align_corners=False,
        )
    return out


def create_random_augment(
    input_size,
    auto_augment=None,
    interpolation="bilinear",
):
    """
    Get video randaug transform.

    Args:
        input_size: The size of the input video in tuple.
        auto_augment: Parameters for randaug. An example:
            "rand-m7-n4-mstd0.5-inc1" (m is the magnitude and n is the number
            of operations to apply).
        interpolation: Interpolation method.
    """
    if isinstance(input_size, tuple):
        img_size = input_size[-2:]
    else:
        img_size = input_size

    if auto_augment:
        assert isinstance(auto_augment, str)
        if isinstance(img_size, tuple):
            img_size_min = min(img_size)
        else:
            img_size_min = img_size
        aa_params = {"translate_const": int(img_size_min * 0.45)}
        if interpolation and interpolation != "random":
            aa_params["interpolation"] = _pil_interp(interpolation)
        if auto_augment.startswith("rand"):
            return transforms.Compose([rand_augment_transform(auto_augment, aa_params)])
    raise NotImplementedError


def random_sized_crop_img(
    im,
    size,
    jitter_scale=(0.08, 1.0),
    jitter_aspect=(3.0 / 4.0, 4.0 / 3.0),
    max_iter=10,
):
    """
    Performs Inception-style cropping (used for training).
    """
    assert len(im.shape) == 3, "Currently only support image for random_sized_crop"
    h, w = im.shape[1:3]
    i, j, h, w = _get_param_spatial_crop(
        scale=jitter_scale,
        ratio=jitter_aspect,
        height=h,
        width=w,
        num_repeat=max_iter,
        log_scale=False,
        switch_hw=True,
    )
    cropped = im[:, i : i + h, j : j + w]
    return torch.nn.functional.interpolate(
        cropped.unsqueeze(0),
        size=(size, size),
        mode="bilinear",
        align_corners=False,
    ).squeeze(0)


# The following code are modified based on timm lib, we will replace the following
# contents with dependency from PyTorchVideo.
# https://github.com/facebookresearch/pytorchvideo
class RandomResizedCropAndInterpolation:
    """Crop the given PIL Image to random size and aspect ratio with random interpolation.
    A crop of random size (default: of 0.08 to 1.0) of the original size and a random
    aspect ratio (default: of 3/4 to 4/3) of the original aspect ratio is made. This crop
    is finally resized to given size.
    This is popularly used to train the Inception networks.
    Args:
        size: expected output size of each edge
        scale: range of size of the origin size cropped
        ratio: range of aspect ratio of the origin aspect ratio cropped
        interpolation: Default: PIL.Image.BILINEAR
    """

    def __init__(
        self,
        size,
        scale=(0.08, 1.0),
        ratio=(3.0 / 4.0, 4.0 / 3.0),
        interpolation="bilinear",
    ):
        if isinstance(size, tuple):
            self.size = size
        else:
            self.size = (size, size)
        if (scale[0] > scale[1]) or (ratio[0] > ratio[1]):
            print("range should be of kind (min, max)")

        if interpolation == "random":
            self.interpolation = _RANDOM_INTERPOLATION
        else:
            self.interpolation = _pil_interp(interpolation)
        self.scale = scale
        self.ratio = ratio

    @staticmethod
    def get_params(img, scale, ratio):
        """Get parameters for ``crop`` for a random sized crop.
        Args:
            img (PIL Image): Image to be cropped.
            scale (tuple): range of size of the origin size cropped
            ratio (tuple): range of aspect ratio of the origin aspect ratio cropped
        Returns:
            tuple: params (i, j, h, w) to be passed to ``crop`` for a random
                sized crop.
        """
        area = img.size[0] * img.size[1]

        for _ in range(10):
            target_area = random.uniform(*scale) * area
            log_ratio = (math.log(ratio[0]), math.log(ratio[1]))
            aspect_ratio = math.exp(random.uniform(*log_ratio))

            w = int(round(math.sqrt(target_area * aspect_ratio)))
            h = int(round(math.sqrt(target_area / aspect_ratio)))

            if w <= img.size[0] and h <= img.size[1]:
                i = random.randint(0, img.size[1] - h)
                j = random.randint(0, img.size[0] - w)
                return i, j, h, w

        # Fallback to central crop
        in_ratio = img.size[0] / img.size[1]
        if in_ratio < min(ratio):
            w = img.size[0]
            h = int(round(w / min(ratio)))
        elif in_ratio > max(ratio):
            h = img.size[1]
            w = int(round(h * max(ratio)))
        else:  # whole image
            w = img.size[0]
            h = img.size[1]
        i = (img.size[1] - h) // 2
        j = (img.size[0] - w) // 2
        return i, j, h, w

    def __call__(self, img):
        """
        Args:
            img (PIL Image): Image to be cropped and resized.
        Returns:
            PIL Image: Randomly cropped and resized image.
        """
        i, j, h, w = self.get_params(img, self.scale, self.ratio)
        if isinstance(self.interpolation, (tuple, list)):
            interpolation = random.choice(self.interpolation)
        else:
            interpolation = self.interpolation
        return F.resized_crop(img, i, j, h, w, self.size, interpolation)

    def __repr__(self):
        if isinstance(self.interpolation, (tuple, list)):
            interpolate_str = " ".join(
                [_pil_interpolation_to_str[x] for x in self.interpolation]
            )
        else:
            interpolate_str = _pil_interpolation_to_str[self.interpolation]
        format_string = self.__class__.__name__ + "(size={0}".format(self.size)
        format_string += ", scale={0}".format(tuple(round(s, 4) for s in self.scale))
        format_string += ", ratio={0}".format(tuple(round(r, 4) for r in self.ratio))
        format_string += ", interpolation={0})".format(interpolate_str)
        return format_string


def transforms_imagenet_train(
    img_size=224,
    scale=None,
    ratio=None,
    hflip=0.5,
    vflip=0.0,
    color_jitter=0.4,
    auto_augment=None,
    interpolation="random",
    use_prefetcher=False,
    mean=(0.485, 0.456, 0.406),
    std=(0.229, 0.224, 0.225),
    re_prob=0.0,
    re_mode="const",
    re_count=1,
    re_num_splits=0,
    separate=False,
):
    """
    If separate==True, the transforms are returned as a tuple of 3 separate transforms
    for use in a mixing dataset that passes
     * all data through the first (primary) transform, called the 'clean' data
     * a portion of the data through the secondary transform
     * normalizes and converts the branches above with the third, final transform
    """
    if isinstance(img_size, tuple):
        img_size = img_size[-2:]
    else:
        img_size = img_size

    scale = tuple(scale or (0.08, 1.0))  # default imagenet scale range
    ratio = tuple(ratio or (3.0 / 4.0, 4.0 / 3.0))  # default imagenet ratio range
    primary_tfl = [
        RandomResizedCropAndInterpolation(
            img_size, scale=scale, ratio=ratio, interpolation=interpolation
        )
    ]
    if hflip > 0.0:
        primary_tfl += [transforms.RandomHorizontalFlip(p=hflip)]
    if vflip > 0.0:
        primary_tfl += [transforms.RandomVerticalFlip(p=vflip)]

    secondary_tfl = []
    if auto_augment:
        assert isinstance(auto_augment, str)
        if isinstance(img_size, tuple):
            img_size_min = min(img_size)
        else:
            img_size_min = img_size
        aa_params = dict(
            translate_const=int(img_size_min * 0.45),
            img_mean=tuple([min(255, round(255 * x)) for x in mean]),
        )
        if interpolation and interpolation != "random":
            aa_params["interpolation"] = _pil_interp(interpolation)
        if auto_augment.startswith("rand"):
            secondary_tfl += [rand_augment_transform(auto_augment, aa_params)]
        elif auto_augment.startswith("augmix"):
            raise NotImplementedError("Augmix not implemented")
        else:
            raise NotImplementedError("Auto aug not implemented")
    elif color_jitter is not None:
        # color jitter is enabled when not using AA
        if isinstance(color_jitter, (list, tuple)):
            # color jitter should be a 3-tuple/list if spec brightness/contrast/saturation
            # or 4 if also augmenting hue
            assert len(color_jitter) in (3, 4)
        else:
            # if it's a scalar, duplicate for brightness, contrast, and saturation, no hue
            color_jitter = (float(color_jitter),) * 3
        secondary_tfl += [transforms.ColorJitter(*color_jitter)]

    final_tfl = []
    final_tfl += [
        transforms.ToTensor(),
        transforms.Normalize(mean=torch.tensor(mean), std=torch.tensor(std)),
    ]
    if re_prob > 0.0:
        final_tfl.append(
            RandomErasing(
                re_prob,
                mode=re_mode,
                max_count=re_count,
                num_splits=re_num_splits,
                device="cpu",
                cube=False,
            )
        )

    if separate:
        return (
            transforms.Compose(primary_tfl),
            transforms.Compose(secondary_tfl),
            transforms.Compose(final_tfl),
        )
    else:
        return transforms.Compose(primary_tfl + secondary_tfl + final_tfl)



_PIL_VER = tuple([int(x) for x in PIL.__version__.split(".")[:2]])

_FILL = (128, 128, 128)

# This signifies the max integer that the controller RNN could predict for the
# augmentation scheme.
_MAX_LEVEL = 10.0

_HPARAMS_DEFAULT = {
    "translate_const": 250,
    "img_mean": _FILL,
}

_RANDOM_INTERPOLATION = (Image.BILINEAR, Image.BICUBIC)


def _interpolation(kwargs):
    interpolation = kwargs.pop("resample", Image.BILINEAR)
    if isinstance(interpolation, (list, tuple)):
        return random.choice(interpolation)
    else:
        return interpolation


def _check_args_tf(kwargs):
    if "fillcolor" in kwargs and _PIL_VER < (5, 0):
        kwargs.pop("fillcolor")
    kwargs["resample"] = _interpolation(kwargs)


def shear_x(img, factor, **kwargs):
    _check_args_tf(kwargs)
    return img.transform(img.size, Image.AFFINE, (1, factor, 0, 0, 1, 0), **kwargs)


def shear_y(img, factor, **kwargs):
    _check_args_tf(kwargs)
    return img.transform(img.size, Image.AFFINE, (1, 0, 0, factor, 1, 0), **kwargs)


def translate_x_rel(img, pct, **kwargs):
    pixels = pct * img.size[0]
    _check_args_tf(kwargs)
    return img.transform(img.size, Image.AFFINE, (1, 0, pixels, 0, 1, 0), **kwargs)


def translate_y_rel(img, pct, **kwargs):
    pixels = pct * img.size[1]
    _check_args_tf(kwargs)
    return img.transform(img.size, Image.AFFINE, (1, 0, 0, 0, 1, pixels), **kwargs)


def translate_x_abs(img, pixels, **kwargs):
    _check_args_tf(kwargs)
    return img.transform(img.size, Image.AFFINE, (1, 0, pixels, 0, 1, 0), **kwargs)


def translate_y_abs(img, pixels, **kwargs):
    _check_args_tf(kwargs)
    return img.transform(img.size, Image.AFFINE, (1, 0, 0, 0, 1, pixels), **kwargs)


def rotate(img, degrees, **kwargs):
    _check_args_tf(kwargs)
    if _PIL_VER >= (5, 2):
        return img.rotate(degrees, **kwargs)
    elif _PIL_VER >= (5, 0):
        w, h = img.size
        post_trans = (0, 0)
        rotn_center = (w / 2.0, h / 2.0)
        angle = -math.radians(degrees)
        matrix = [
            round(math.cos(angle), 15),
            round(math.sin(angle), 15),
            0.0,
            round(-math.sin(angle), 15),
            round(math.cos(angle), 15),
            0.0,
        ]

        def transform(x, y, matrix):
            (a, b, c, d, e, f) = matrix
            return a * x + b * y + c, d * x + e * y + f

        matrix[2], matrix[5] = transform(
            -rotn_center[0] - post_trans[0],
            -rotn_center[1] - post_trans[1],
            matrix,
        )
        matrix[2] += rotn_center[0]
        matrix[5] += rotn_center[1]
        return img.transform(img.size, Image.AFFINE, matrix, **kwargs)
    else:
        return img.rotate(degrees, resample=kwargs["resample"])


def auto_contrast(img, **__):
    return ImageOps.autocontrast(img)


def invert(img, **__):
    return ImageOps.invert(img)


def equalize(img, **__):
    return ImageOps.equalize(img)


def solarize(img, thresh, **__):
    return ImageOps.solarize(img, thresh)


def solarize_add(img, add, thresh=128, **__):
    lut = []
    for i in range(256):
        if i < thresh:
            lut.append(min(255, i + add))
        else:
            lut.append(i)
    if img.mode in ("L", "RGB"):
        if img.mode == "RGB" and len(lut) == 256:
            lut = lut + lut + lut
        return img.point(lut)
    else:
        return img


def posterize(img, bits_to_keep, **__):
    if bits_to_keep >= 8:
        return img
    return ImageOps.posterize(img, bits_to_keep)


def contrast(img, factor, **__):
    return ImageEnhance.Contrast(img).enhance(factor)


def color(img, factor, **__):
    return ImageEnhance.Color(img).enhance(factor)


def brightness(img, factor, **__):
    return ImageEnhance.Brightness(img).enhance(factor)


def sharpness(img, factor, **__):
    return ImageEnhance.Sharpness(img).enhance(factor)


def _randomly_negate(v):
    """With 50% prob, negate the value"""
    return -v if random.random() > 0.5 else v


def _rotate_level_to_arg(level, _hparams):
    # range [-30, 30]
    level = (level / _MAX_LEVEL) * 30.0
    level = _randomly_negate(level)
    return (level,)


def _enhance_level_to_arg(level, _hparams):
    # range [0.1, 1.9]
    return ((level / _MAX_LEVEL) * 1.8 + 0.1,)


def _enhance_increasing_level_to_arg(level, _hparams):
    # the 'no change' level is 1.0, moving away from that towards 0. or 2.0 increases the enhancement blend
    # range [0.1, 1.9]
    level = (level / _MAX_LEVEL) * 0.9
    level = 1.0 + _randomly_negate(level)
    return (level,)


def _shear_level_to_arg(level, _hparams):
    # range [-0.3, 0.3]
    level = (level / _MAX_LEVEL) * 0.3
    level = _randomly_negate(level)
    return (level,)


def _translate_abs_level_to_arg(level, hparams):
    translate_const = hparams["translate_const"]
    level = (level / _MAX_LEVEL) * float(translate_const)
    level = _randomly_negate(level)
    return (level,)


def _translate_rel_level_to_arg(level, hparams):
    # default range [-0.45, 0.45]
    translate_pct = hparams.get("translate_pct", 0.45)
    level = (level / _MAX_LEVEL) * translate_pct
    level = _randomly_negate(level)
    return (level,)


def _posterize_level_to_arg(level, _hparams):
    # As per Tensorflow TPU EfficientNet impl
    # range [0, 4], 'keep 0 up to 4 MSB of original image'
    # intensity/severity of augmentation decreases with level
    return (int((level / _MAX_LEVEL) * 4),)


def _posterize_increasing_level_to_arg(level, hparams):
    # As per Tensorflow models research and UDA impl
    # range [4, 0], 'keep 4 down to 0 MSB of original image',
    # intensity/severity of augmentation increases with level
    return (4 - _posterize_level_to_arg(level, hparams)[0],)


def _posterize_original_level_to_arg(level, _hparams):
    # As per original AutoAugment paper description
    # range [4, 8], 'keep 4 up to 8 MSB of image'
    # intensity/severity of augmentation decreases with level
    return (int((level / _MAX_LEVEL) * 4) + 4,)


def _solarize_level_to_arg(level, _hparams):
    # range [0, 256]
    # intensity/severity of augmentation decreases with level
    return (int((level / _MAX_LEVEL) * 256),)


def _solarize_increasing_level_to_arg(level, _hparams):
    # range [0, 256]
    # intensity/severity of augmentation increases with level
    return (256 - _solarize_level_to_arg(level, _hparams)[0],)


def _solarize_add_level_to_arg(level, _hparams):
    # range [0, 110]
    return (int((level / _MAX_LEVEL) * 110),)


LEVEL_TO_ARG = {
    "AutoContrast": None,
    "Equalize": None,
    "Invert": None,
    "Rotate": _rotate_level_to_arg,
    # There are several variations of the posterize level scaling in various Tensorflow/Google repositories/papers
    "Posterize": _posterize_level_to_arg,
    "PosterizeIncreasing": _posterize_increasing_level_to_arg,
    "PosterizeOriginal": _posterize_original_level_to_arg,
    "Solarize": _solarize_level_to_arg,
    "SolarizeIncreasing": _solarize_increasing_level_to_arg,
    "SolarizeAdd": _solarize_add_level_to_arg,
    "Color": _enhance_level_to_arg,
    "ColorIncreasing": _enhance_increasing_level_to_arg,
    "Contrast": _enhance_level_to_arg,
    "ContrastIncreasing": _enhance_increasing_level_to_arg,
    "Brightness": _enhance_level_to_arg,
    "BrightnessIncreasing": _enhance_increasing_level_to_arg,
    "Sharpness": _enhance_level_to_arg,
    "SharpnessIncreasing": _enhance_increasing_level_to_arg,
    "ShearX": _shear_level_to_arg,
    "ShearY": _shear_level_to_arg,
    "TranslateX": _translate_abs_level_to_arg,
    "TranslateY": _translate_abs_level_to_arg,
    "TranslateXRel": _translate_rel_level_to_arg,
    "TranslateYRel": _translate_rel_level_to_arg,
}


NAME_TO_OP = {
    "AutoContrast": auto_contrast,
    "Equalize": equalize,
    "Invert": invert,
    "Rotate": rotate,
    "Posterize": posterize,
    "PosterizeIncreasing": posterize,
    "PosterizeOriginal": posterize,
    "Solarize": solarize,
    "SolarizeIncreasing": solarize,
    "SolarizeAdd": solarize_add,
    "Color": color,
    "ColorIncreasing": color,
    "Contrast": contrast,
    "ContrastIncreasing": contrast,
    "Brightness": brightness,
    "BrightnessIncreasing": brightness,
    "Sharpness": sharpness,
    "SharpnessIncreasing": sharpness,
    "ShearX": shear_x,
    "ShearY": shear_y,
    "TranslateX": translate_x_abs,
    "TranslateY": translate_y_abs,
    "TranslateXRel": translate_x_rel,
    "TranslateYRel": translate_y_rel,
}


class AugmentOp:
    """
    Apply for video.
    """

    def __init__(self, name, prob=0.5, magnitude=10, hparams=None):
        hparams = hparams or _HPARAMS_DEFAULT
        self.aug_fn = NAME_TO_OP[name]
        self.level_fn = LEVEL_TO_ARG[name]
        self.prob = prob
        self.magnitude = magnitude
        self.hparams = hparams.copy()
        self.kwargs = {
            "fillcolor": hparams["img_mean"] if "img_mean" in hparams else _FILL,
            "resample": hparams["interpolation"]
            if "interpolation" in hparams
            else _RANDOM_INTERPOLATION,
        }

        # If magnitude_std is > 0, we introduce some randomness
        # in the usually fixed policy and sample magnitude from a normal distribution
        # with mean `magnitude` and std-dev of `magnitude_std`.
        # NOTE This is my own hack, being tested, not in papers or reference impls.
        self.magnitude_std = self.hparams.get("magnitude_std", 0)

    def __call__(self, img_list):
        if self.prob < 1.0 and random.random() > self.prob:
            return img_list
        magnitude = self.magnitude
        if self.magnitude_std and self.magnitude_std > 0:
            magnitude = random.gauss(magnitude, self.magnitude_std)
        magnitude = min(_MAX_LEVEL, max(0, magnitude))  # clip to valid range
        level_args = (
            self.level_fn(magnitude, self.hparams) if self.level_fn is not None else ()
        )

        if isinstance(img_list, list):
            return [self.aug_fn(img, *level_args, **self.kwargs) for img in img_list]
        else:
            return self.aug_fn(img_list, *level_args, **self.kwargs)


_RAND_TRANSFORMS = [
    "AutoContrast",
    "Equalize",
    "Invert",
    "Rotate",
    "Posterize",
    "Solarize",
    "SolarizeAdd",
    "Color",
    "Contrast",
    "Brightness",
    "Sharpness",
    "ShearX",
    "ShearY",
    "TranslateXRel",
    "TranslateYRel",
]


_RAND_INCREASING_TRANSFORMS = [
    "AutoContrast",
    "Equalize",
    "Invert",
    "Rotate",
    "PosterizeIncreasing",
    "SolarizeIncreasing",
    "SolarizeAdd",
    "ColorIncreasing",
    "ContrastIncreasing",
    "BrightnessIncreasing",
    "SharpnessIncreasing",
    "ShearX",
    "ShearY",
    "TranslateXRel",
    "TranslateYRel",
]


# These experimental weights are based loosely on the relative improvements mentioned in paper.
# They may not result in increased performance, but could likely be tuned to so.
_RAND_CHOICE_WEIGHTS_0 = {
    "Rotate": 0.3,
    "ShearX": 0.2,
    "ShearY": 0.2,
    "TranslateXRel": 0.1,
    "TranslateYRel": 0.1,
    "Color": 0.025,
    "Sharpness": 0.025,
    "AutoContrast": 0.025,
    "Solarize": 0.005,
    "SolarizeAdd": 0.005,
    "Contrast": 0.005,
    "Brightness": 0.005,
    "Equalize": 0.005,
    "Posterize": 0,
    "Invert": 0,
}


def _select_rand_weights(weight_idx=0, transforms=None):
    transforms = transforms or _RAND_TRANSFORMS
    assert weight_idx == 0  # only one set of weights currently
    rand_weights = _RAND_CHOICE_WEIGHTS_0
    probs = [rand_weights[k] for k in transforms]
    probs /= np.sum(probs)
    return probs


def rand_augment_ops(magnitude=10, hparams=None, transforms=None):
    hparams = hparams or _HPARAMS_DEFAULT
    transforms = transforms or _RAND_TRANSFORMS
    return [
        AugmentOp(name, prob=0.5, magnitude=magnitude, hparams=hparams)
        for name in transforms
    ]


class RandAugment:
    def __init__(self, ops, num_layers=2, choice_weights=None):
        self.ops = ops
        self.num_layers = num_layers
        self.choice_weights = choice_weights

    def __call__(self, img):
        # no replacement when using weighted choice
        ops = np.random.choice(
            self.ops,
            self.num_layers,
            replace=self.choice_weights is None,
            p=self.choice_weights,
        )
        for op in ops:
            img = op(img)
        return img


def rand_augment_transform(config_str, hparams):
    """
    RandAugment: Practical automated data augmentation... - https://arxiv.org/abs/1909.13719

    Create a RandAugment transform
    :param config_str: String defining configuration of random augmentation. Consists of multiple sections separated by
    dashes ('-'). The first section defines the specific variant of rand augment (currently only 'rand'). The remaining
    sections, not order sepecific determine
        'm' - integer magnitude of rand augment
        'n' - integer num layers (number of transform ops selected per image)
        'w' - integer probabiliy weight index (index of a set of weights to influence choice of op)
        'mstd' -  float std deviation of magnitude noise applied
        'inc' - integer (bool), use augmentations that increase in severity with magnitude (default: 0)
    Ex 'rand-m9-n3-mstd0.5' results in RandAugment with magnitude 9, num_layers 3, magnitude_std 0.5
    'rand-mstd1-w0' results in magnitude_std 1.0, weights 0, default magnitude of 10 and num_layers 2
    :param hparams: Other hparams (kwargs) for the RandAugmentation scheme
    :return: A PyTorch compatible Transform
    """
    magnitude = _MAX_LEVEL  # default to _MAX_LEVEL for magnitude (currently 10)
    num_layers = 2  # default to 2 ops per image
    weight_idx = None  # default to no probability weights for op choice
    transforms = _RAND_TRANSFORMS
    config = config_str.split("-")
    assert config[0] == "rand"
    config = config[1:]
    for c in config:
        cs = re.split(r"(\d.*)", c)
        if len(cs) < 2:
            continue
        key, val = cs[:2]
        if key == "mstd":
            # noise param injected via hparams for now
            hparams.setdefault("magnitude_std", float(val))
        elif key == "inc":
            if bool(val):
                transforms = _RAND_INCREASING_TRANSFORMS
        elif key == "m":
            magnitude = int(val)
        elif key == "n":
            num_layers = int(val)
        elif key == "w":
            weight_idx = int(val)
        else:
            assert NotImplementedError
    ra_ops = rand_augment_ops(
        magnitude=magnitude, hparams=hparams, transforms=transforms
    )
    choice_weights = None if weight_idx is None else _select_rand_weights(weight_idx)
    return RandAugment(ra_ops, num_layers, choice_weights=choice_weights)




def _get_pixels(per_pixel, rand_color, patch_size, dtype=torch.float32, device="cuda"):
    # NOTE I've seen CUDA illegal memory access errors being caused by the normal_()
    # paths, flip the order so normal is run on CPU if this becomes a problem
    # Issue has been fixed in master https://github.com/pytorch/pytorch/issues/19508
    if per_pixel:
        return torch.empty(patch_size, dtype=dtype, device=device).normal_()
    elif rand_color:
        return torch.empty((patch_size[0], 1, 1), dtype=dtype, device=device).normal_()
    else:
        return torch.zeros((patch_size[0], 1, 1), dtype=dtype, device=device)


class RandomErasing:
    """Randomly selects a rectangle region in an image and erases its pixels.
        'Random Erasing Data Augmentation' by Zhong et al.
        See https://arxiv.org/pdf/1708.04896.pdf
        This variant of RandomErasing is intended to be applied to either a batch
        or single image tensor after it has been normalized by dataset mean and std.
    Args:
         probability: Probability that the Random Erasing operation will be performed.
         min_area: Minimum percentage of erased area wrt input image area.
         max_area: Maximum percentage of erased area wrt input image area.
         min_aspect: Minimum aspect ratio of erased area.
         mode: pixel color mode, one of 'const', 'rand', or 'pixel'
            'const' - erase block is constant color of 0 for all channels
            'rand'  - erase block is same per-channel random (normal) color
            'pixel' - erase block is per-pixel random (normal) color
        max_count: maximum number of erasing blocks per image, area per box is scaled by count.
            per-image count is randomly chosen between 1 and this value.
    """

    def __init__(
        self,
        probability=0.5,
        min_area=0.02,
        max_area=1 / 3,
        min_aspect=0.3,
        max_aspect=None,
        mode="const",
        min_count=1,
        max_count=None,
        num_splits=0,
        device="cuda",
        cube=True,
    ):
        self.probability = probability
        self.min_area = min_area
        self.max_area = max_area
        max_aspect = max_aspect or 1 / min_aspect
        self.log_aspect_ratio = (math.log(min_aspect), math.log(max_aspect))
        self.min_count = min_count
        self.max_count = max_count or min_count
        self.num_splits = num_splits
        mode = mode.lower()
        self.rand_color = False
        self.per_pixel = False
        self.cube = cube
        if mode == "rand":
            self.rand_color = True  # per block random normal
        elif mode == "pixel":
            self.per_pixel = True  # per pixel random normal
        else:
            assert not mode or mode == "const"
        self.device = device

    def _erase(self, img, chan, img_h, img_w, dtype):
        if random.random() > self.probability:
            return
        area = img_h * img_w
        count = (
            self.min_count
            if self.min_count == self.max_count
            else random.randint(self.min_count, self.max_count)
        )
        for _ in range(count):
            for _ in range(10):
                target_area = (
                    random.uniform(self.min_area, self.max_area) * area / count
                )
                aspect_ratio = math.exp(random.uniform(*self.log_aspect_ratio))
                h = int(round(math.sqrt(target_area * aspect_ratio)))
                w = int(round(math.sqrt(target_area / aspect_ratio)))
                if w < img_w and h < img_h:
                    top = random.randint(0, img_h - h)
                    left = random.randint(0, img_w - w)
                    img[:, top : top + h, left : left + w] = _get_pixels(
                        self.per_pixel,
                        self.rand_color,
                        (chan, h, w),
                        dtype=dtype,
                        device=self.device,
                    )
                    break

    def _erase_cube(
        self,
        img,
        batch_start,
        batch_size,
        chan,
        img_h,
        img_w,
        dtype,
    ):
        if random.random() > self.probability:
            return
        area = img_h * img_w
        count = (
            self.min_count
            if self.min_count == self.max_count
            else random.randint(self.min_count, self.max_count)
        )
        for _ in range(count):
            for _ in range(100):
                target_area = (
                    random.uniform(self.min_area, self.max_area) * area / count
                )
                aspect_ratio = math.exp(random.uniform(*self.log_aspect_ratio))
                h = int(round(math.sqrt(target_area * aspect_ratio)))
                w = int(round(math.sqrt(target_area / aspect_ratio)))
                if w < img_w and h < img_h:
                    top = random.randint(0, img_h - h)
                    left = random.randint(0, img_w - w)
                    for i in range(batch_start, batch_size):
                        img_instance = img[i]
                        img_instance[:, top : top + h, left : left + w] = _get_pixels(
                            self.per_pixel,
                            self.rand_color,
                            (chan, h, w),
                            dtype=dtype,
                            device=self.device,
                        )
                    break

    def __call__(self, input):
        if len(input.size()) == 3:
            self._erase(input, *input.size(), input.dtype)
        else:
            batch_size, chan, img_h, img_w = input.size()
            # skip first slice of batch if num_splits is set (for clean portion of samples)
            batch_start = batch_size // self.num_splits if self.num_splits > 1 else 0
            if self.cube:
                self._erase_cube(
                    input,
                    batch_start,
                    batch_size,
                    chan,
                    img_h,
                    img_w,
                    input.dtype,
                )
            else:
                for i in range(batch_start, batch_size):
                    self._erase(input[i], chan, img_h, img_w, input.dtype)
        return input

upstream_models/mavil/util/transform.py
roger-tseng/av-superb
#!/usr/bin/env python3

import os
import random
import time
from collections import defaultdict
import numpy as np
import torch

from . import transform



def get_sequence(center_idx, half_len, sample_rate, num_frames):
    """
    Sample frames among the corresponding clip.

    Args:
        center_idx (int): center frame idx for current clip
        half_len (int): half of the clip length
        sample_rate (int): sampling rate for sampling frames inside of the clip
        num_frames (int): number of expected sampled frames

    Returns:
        seq (list): list of indexes of sampled frames in this clip.
    """
    seq = list(range(center_idx - half_len, center_idx + half_len, sample_rate))

    for seq_idx in range(len(seq)):
        if seq[seq_idx] < 0:
            seq[seq_idx] = 0
        elif seq[seq_idx] >= num_frames:
            seq[seq_idx] = num_frames - 1
    return seq


def pack_pathway_output(
    cfg,
    frames,
):
    """
    Prepare output as a list of tensors. Each tensor corresponding to a
    unique pathway.
    Args:
        frames (tensor): frames of images sampled from the video. The
            dimension is `channel` x `num frames` x `height` x `width`.
    Returns:
        frame_list (list): list of tensors with the dimension of
            `channel` x `num frames` x `height` x `width`.
    """
    if cfg.DATA.REVERSE_INPUT_CHANNEL:
        frames = frames[[2, 1, 0], :, :, :]
    if cfg.MODEL.ARCH in cfg.MODEL.SINGLE_PATHWAY_ARCH:
        frame_list = [frames]
    elif cfg.MODEL.ARCH in cfg.MODEL.MULTI_PATHWAY_ARCH:
        fast_pathway = frames
        # Perform temporal sampling from the fast pathway.
        slow_pathway = torch.index_select(
            frames,
            1,
            torch.linspace(
                0, frames.shape[1] - 1, frames.shape[1] // cfg.SLOWFAST.ALPHA
            ).long(),
        )
        frame_list = [slow_pathway, fast_pathway]
    else:
        raise NotImplementedError(
            "Model arch {} is not in {}".format(
                cfg.MODEL.ARCH,
                cfg.MODEL.SINGLE_PATHWAY_ARCH + cfg.MODEL.MULTI_PATHWAY_ARCH,
            )
        )
    return frame_list


def spatial_sampling(
    frames,
    spatial_idx=-1,
    min_scale=256,
    max_scale=320,
    crop_size=224,
    random_horizontal_flip=True,
    inverse_uniform_sampling=False,
    aspect_ratio=None,
    scale=None,
    motion_shift=False,
):
    """
    Perform spatial sampling on the given video frames. If spatial_idx is
    -1, perform random scale, random crop, and random flip on the given
    frames. If spatial_idx is 0, 1, or 2, perform spatial uniform sampling
    with the given spatial_idx.
    Args:
        frames (tensor): frames of images sampled from the video. The
            dimension is `num frames` x `height` x `width` x `channel`.
        spatial_idx (int): if -1, perform random spatial sampling. If 0, 1,
            or 2, perform left, center, right crop if width is larger than
            height, and perform top, center, buttom crop if height is larger
            than width.
        min_scale (int): the minimal size of scaling.
        max_scale (int): the maximal size of scaling.
        crop_size (int): the size of height and width used to crop the
            frames.
        inverse_uniform_sampling (bool): if True, sample uniformly in
            [1 / max_scale, 1 / min_scale] and take a reciprocal to get the
            scale. If False, take a uniform sample from [min_scale,
            max_scale].
        aspect_ratio (list): Aspect ratio range for resizing.
        scale (list): Scale range for resizing.
        motion_shift (bool): Whether to apply motion shift for resizing.
    Returns:
        frames (tensor): spatially sampled frames.
    """
    assert spatial_idx in [-1, 0, 1, 2]
    if spatial_idx == -1:
        if aspect_ratio is None and scale is None:
            frames, _ = transform.random_short_side_scale_jitter(
                images=frames,
                min_size=min_scale,
                max_size=max_scale,
                inverse_uniform_sampling=inverse_uniform_sampling,
            )
            frames, _ = transform.random_crop(frames, crop_size)
        else:
            transform_func = (
                transform.random_resized_crop_with_shift
                if motion_shift
                else transform.random_resized_crop
            )
            frames = transform_func(
                images=frames,
                target_height=crop_size,
                target_width=crop_size,
                scale=scale,
                ratio=aspect_ratio,
            )
        if random_horizontal_flip:
            frames, _ = transform.horizontal_flip(0.5, frames)
    else:
        # The testing is deterministic and no jitter should be performed.
        # min_scale, max_scale, and crop_size are expect to be the same.
        assert len({min_scale, max_scale}) == 1
        frames, _ = transform.random_short_side_scale_jitter(
            frames, min_scale, max_scale
        )
        frames, _ = transform.uniform_crop(frames, crop_size, spatial_idx)
    return frames


def as_binary_vector(labels, num_classes):
    """
    Construct binary label vector given a list of label indices.
    Args:
        labels (list): The input label list.
        num_classes (int): Number of classes of the label vector.
    Returns:
        labels (numpy array): the resulting binary vector.
    """
    label_arr = np.zeros((num_classes,))

    for lbl in set(labels):
        label_arr[lbl] = 1.0
    return label_arr


def aggregate_labels(label_list):
    """
    Join a list of label list.
    Args:
        labels (list): The input label list.
    Returns:
        labels (list): The joint list of all lists in input.
    """
    all_labels = []
    for labels in label_list:
        for l in labels:
            all_labels.append(l)
    return list(set(all_labels))


def convert_to_video_level_labels(labels):
    """
    Aggregate annotations from all frames of a video to form video-level labels.
    Args:
        labels (list): The input label list.
    Returns:
        labels (list): Same as input, but with each label replaced by
        a video-level one.
    """
    for video_id in range(len(labels)):
        video_level_labels = aggregate_labels(labels[video_id])
        for i in range(len(labels[video_id])):
            labels[video_id][i] = video_level_labels
    return labels


def load_image_lists(frame_list_file, prefix="", return_list=False):
    """
    Load image paths and labels from a "frame list".
    Each line of the frame list contains:
    `original_vido_id video_id frame_id path labels`
    Args:
        frame_list_file (string): path to the frame list.
        prefix (str): the prefix for the path.
        return_list (bool): if True, return a list. If False, return a dict.
    Returns:
        image_paths (list or dict): list of list containing path to each frame.
            If return_list is False, then return in a dict form.
        labels (list or dict): list of list containing label of each frame.
            If return_list is False, then return in a dict form.
    """
    image_paths = defaultdict(list)
    labels = defaultdict(list)
    with pathmgr.open(frame_list_file, "r") as f:
        assert f.readline().startswith("original_vido_id")
        for line in f:
            row = line.split()
            # original_vido_id video_id frame_id path labels
            assert len(row) == 5
            video_name = row[0]
            if prefix == "":
                path = row[3]
            else:
                path = os.path.join(prefix, row[3])
            image_paths[video_name].append(path)
            frame_labels = row[-1].replace('"', "")
            if frame_labels != "":
                labels[video_name].append([int(x) for x in frame_labels.split(",")])
            else:
                labels[video_name].append([])

    if return_list:
        keys = image_paths.keys()
        image_paths = [image_paths[key] for key in keys]
        labels = [labels[key] for key in keys]
        return image_paths, labels
    return dict(image_paths), dict(labels)


def tensor_normalize(tensor, mean, std):
    """
    Normalize a given tensor by subtracting the mean and dividing the std.
    Args:
        tensor (tensor): tensor to normalize.
        mean (tensor or list): mean value to subtract.
        std (tensor or list): std to divide.
    """
    if tensor.dtype == torch.uint8:
        tensor = tensor.float()
        tensor = tensor / 255.0
    if type(mean) == tuple:
        mean = torch.tensor(mean)
    if type(std) == tuple:
        std = torch.tensor(std)
    tensor = tensor - mean
    tensor = tensor / std
    return tensor


def get_random_sampling_rate(long_cycle_sampling_rate, sampling_rate):
    """
    When multigrid training uses a fewer number of frames, we randomly
    increase the sampling rate so that some clips cover the original span.
    """
    if long_cycle_sampling_rate > 0:
        assert long_cycle_sampling_rate >= sampling_rate
        return random.randint(sampling_rate, long_cycle_sampling_rate)
    else:
        return sampling_rate


def revert_tensor_normalize(tensor, mean, std):
    """
    Revert normalization for a given tensor by multiplying by the std and adding the mean.
    Args:
        tensor (tensor): tensor to revert normalization.
        mean (tensor or list): mean value to add.
        std (tensor or list): std to multiply.
    """
    if type(mean) == list:
        mean = torch.tensor(mean)
    if type(std) == list:
        std = torch.tensor(std)
    tensor = tensor * std
    tensor = tensor + mean
    return tensor




upstream_models/mavil/util/utils.py
roger-tseng/av-superb
"""
Custom class for loading audio-visual model and extract features
Modified from https://github.com/s3prl/s3prl/blob/main/s3prl/upstream/example/expert.py
"""

import sys
from collections import OrderedDict
from typing import Dict, List, Tuple, Union

import torch
import torch.nn as nn

from . import replai
from .replai import models
from .replai.data.builder import build_transforms

sys.modules["replai"] = replai  # create alias for unpickling


from munch import DefaultMunch
from torch import Tensor
from torch.nn.utils.rnn import pad_sequence


class UpstreamExpert(nn.Module):
    def __init__(self, ckpt: str = None, model_config: str = None, **kwargs):
        """
        Args:
            ckpt:
                checkpoint path for loading pretrained weights.

            model_config:
                config path for your model.
        """
        super().__init__()

        # load model weights
        _weights = torch.load(ckpt)
        _weights = _weights["model"]
        _weights = OrderedDict(
            {
                k[16:]: _weights[k]
                for k in _weights.keys()
                if k.startswith("module.backbone")
            }
        )

        # hardcode model config
        model_conf = DefaultMunch.fromDict(
            {
                "audio": {
                    "arch": "avid_spec_cnn_9",
                    "args": {"channels": 1, "pretrained": False},
                    "sync_bn": False,
                    "outp_dim": 512,
                },
                "video": {
                    "arch": "avid_r2plus1d_18",
                    "args": {"pretrained": False},
                    "sync_bn": False,
                    "outp_dim": 512,
                },
            }
        )

        # create model and load weights
        self.backbone = models.build_audio_video_model(model_conf, remove_head=True)
        self.backbone.load_state_dict(
            _weights, strict=True
        )

        self.audio_sample_rate = 16000
        self.video_frame_size = (112, 112)
        self.video_frame_rate = 16

    def preprocess_video(self, video, video_frame_rate):
        """
        Replace this function to preprocess videos into your input format
        video: (video_length, video_channels, height, width), where video_channels is usually 3 for RGB or 1 for greyscale
        in RepLAI, the default length is 0.5 secs for video, resulting in 8 frames (16FPS)
        """
        # Resample video
        # (from https://github.com/pytorch/vision/blob/5b07d6c9c6c14cf88fc545415d63021456874744/torchvision/datasets/video_utils.py#L278)
        step = float(video_frame_rate) / self.video_frame_rate
        if step.is_integer():
            # optimization: if step is integer, don't need to perform
            # advanced indexing
            step = int(step)
            idxs = slice(None, None, step)
        else:
            num_frames = max(int(len(video) / step),1)
            idxs = torch.arange(num_frames, dtype=torch.float32) * step
            idxs = idxs.floor().to(torch.int64)
        video = video[idxs]

        _video_transform = build_transforms(
            cfg=DefaultMunch.fromDict(
                {
                    "video": {
                        "name": "ResizeCropFlip",
                        "args": {
                            "min_size": 128,
                            "max_size": 180,
                            "crop_size": self.video_frame_size[0],
                        },
                        "data_shape": [
                            3,
                            len(video),
                            self.video_frame_size[0],
                            self.video_frame_size[1],
                        ],
                    },
                }
            ),
            augment=False,
        )
        # Original uses OpenCV for resizing numpy tensors
        clips = {
            "video": (video.numpy().transpose(0, 2, 3, 1), self.video_frame_rate),
        }

        clips = _video_transform(clips)

        # output video shape (channel, length, w, h)
        return clips["video"]

    def preprocess_audio(self, audio, audio_sample_rate):
        """
        Replace this function to preprocess audio waveforms into your input format
        audio: (audio_channels, audio_length), where audio_channels is usually 1 or 2
        In RepLAI, they use 2.0 sec audio with raw sample rate of 32khz
        It then first downsample to 16kHz and take 128 temporal frames on mel.
        So I follow the same proportion
        """
        if len(audio.shape) == 2:
            audio = audio.mean(dim=0)

        _audio_length_sec = len(audio) / audio_sample_rate
        num_temporal_frames = int(_audio_length_sec / 2.0 * 128)
        _audio_transform = build_transforms(
            cfg=DefaultMunch.fromDict(
                {
                    "audio": {
                        "name": "ResampleLogMelSpectrogram",
                        "args": {
                            "raw_sample_rate": audio_sample_rate,
                            "audio_rate": self.audio_sample_rate,
                            "mel_window_size": 32,
                            "mel_step_size": 16,
                            "num_mels": 80,
                            "num_temporal_frames": num_temporal_frames,
                        },
                        "data_shape": [1, num_temporal_frames, 80],
                    }
                }
            ),
            augment=False,
        )

        clips = {
            "audio": (audio, audio_sample_rate),
        }

        clips = _audio_transform(clips)

        return clips["audio"]

    def forward(
        self, source: List[Tuple[Tensor, Tensor]]
    ) -> Dict[str, Union[Tensor, List[Tensor]]]:
        """
        Replace this function run a forward pass with your model
        source: list of audio-video Tensor tuples
                [(wav1,vid1), (wav2,vid2), ...]
                in your input format
        """
        bsz = len(source)
        audio, video = zip(*source)

        # Collate audio and video into batch
        audio = [a.squeeze() for a in audio]
        wavs = pad_sequence(audio, batch_first=True).unsqueeze(dim=1)
        # Pad video along time axis, video starts with channel x time x height x width
        video = [v.permute(1, 0, 2, 3) for v in video]
        videos = pad_sequence(video, batch_first=True).permute(0, 2, 1, 3, 4)
        # videos = torch.stack(video)

        assert wavs.shape[0] == bsz
        assert videos.shape[0] == bsz
        assert videos.shape[-2] == 112
        assert videos.shape[-1] == 112

        # Run through audio and video encoders
        video_feats = self.backbone["video"](videos, return_embs=True)
        audio_feats = self.backbone["audio"](wavs, return_embs=True)

        # use the output of last CNN layer before pooling
        video_feats = video_feats["conv5x"]
        audio_feats = audio_feats["conv5x"]

        # convert video_feats to shape (bsz, T', hid_dim)
        video_feats = video_feats.flatten(start_dim=2, end_dim=-1)
        video_feats = video_feats.permute(0, 2, 1)

        # convert video_feats to shape (bsz, T', hid_dim)
        audio_feats = audio_feats.flatten(start_dim=2, end_dim=-1)
        audio_feats = audio_feats.permute(0, 2, 1)

        # Return intermediate layer representations for potential layer-wise experiments
        return {
            "video_feats": [video_feats],
            "audio_feats": [audio_feats],
            "fusion_feats": [],
        }
upstream_models/replai/expert.py
roger-tseng/av-superb
# Copyright (c) Facebook, Inc. All Rights Reserved

# -*- coding: utf-8 -*- #
"""*********************************************************************************************"""
#   FileName     [ upstream/avhubert/hubconf.py ]
#   Synopsis     [ the RepLai torch hubconf ]
#   Author       [ S3PRL / Yi-Jen Shih, NTU]
"""*********************************************************************************************"""

# for more information, go to https://github.com/HimangiM/RepLAI

import os

from utils.download import _urls_to_filepaths

from .expert import UpstreamExpert as _UpstreamExpert

# -------------#


def google_large_file_link(file_id):
    return f"https://drive.google.com/u/0/uc?id={file_id}&export=download&confirm=t&uuid=b73dc818-c9b0-49e8-a488-13b829fdbb7e&at=ANzk5s4TfW9_erYs4M9PH7KYuKzy:1681055195563"


def replai_local(ckpt, *args, **kwargs):
    """
    The model from local ckpt
        ckpt (str): PATH
    """
    assert os.path.isfile(ckpt)
    return _UpstreamExpert(ckpt, *args, **kwargs)


def replai_url(ckpt, refresh=False, *args, **kwargs):
    """
    The model from URL
        ckpt (str): URL
    """
    return replai_local(
        _urls_to_filepaths(ckpt, refresh=refresh, agent="gdown"), *args, **kwargs
    )


def replai(refresh=False, *args, **kwargs):
    """
    The default model
        refresh (bool): whether to download ckpt/config again if existed
    """
    return replai_ek100(refresh=refresh, *args, **kwargs)


############### RepLAI trained on EPIC-KITCHENS-100 ###############
def replai_ek100(refresh=False, *args, **kwargs):
    """
    The replai model trained on EPIC-KITCHENS-100
        refresh (bool): whether to download ckpt/config again if existed
    """
    kwargs["ckpt"] = google_large_file_link("1FlRJxKo0gYGZTzxPUB4hloQWA4cNG4ZU")
    return replai_url(refresh=refresh, *args, **kwargs)


def replai_ek100_scratch(refresh=False, *args, **kwargs):
    """
    The replai model trained on EPIC-KITCHENS-100 (scratch)
        refresh (bool): whether to download ckpt/config again if existed
    """
    kwargs["ckpt"] = google_large_file_link("1aVKCLD6DWZYafvrA-rODLyYMw0X_hj9d")
    return replai_url(refresh=refresh, *args, **kwargs)


def replai_ek100_woMoI(refresh=False, *args, **kwargs):
    """
    The replai model (w/o MoI) trained on EPIC-KITCHENS-100
        refresh (bool): whether to download ckpt/config again if existed
    """
    kwargs["ckpt"] = google_large_file_link("1KAtIvgK4RxcgYPK8aHTbnJxawEpY0Q4w")
    return replai_url(refresh=refresh, *args, **kwargs)


def replai_ek100_woAStC(refresh=False, *args, **kwargs):
    """
    The replai model (w/o AStC) trained on EPIC-KITCHENS-100
        refresh (bool): whether to download ckpt/config again if existed
    """
    kwargs["ckpt"] = google_large_file_link("1ls8MFxxaUr_D8KKzqANUfslU90DQYUxU")
    return replai_url(refresh=refresh, *args, **kwargs)


############### RepLAI trained on Ego4D ###############
def replai_ego4d(refresh=False, *args, **kwargs):
    """
    The replai model trained on Ego4D
        refresh (bool): whether to download ckpt/config again if existed
    """
    kwargs["ckpt"] = google_large_file_link("1uw_fAB9N3y9--vGegKYdK_H2SSMbk4d5")
    return replai_url(refresh=refresh, *args, **kwargs)


def replai_ego4d_scratch(refresh=False, *args, **kwargs):
    """
    The replai model trained on Ego4D (scratch)
        refresh (bool): whether to download ckpt/config again if existed
    """
    kwargs["ckpt"] = google_large_file_link("1g5DhV7z5W5kNBJSMQ2cHK0iWU0KAyecX")
    return replai_url(refresh=refresh, *args, **kwargs)


def replai_ego4d_woMoI(refresh=False, *args, **kwargs):
    """
    The replai model (w/o MoI) trained on Ego4D
        refresh (bool): whether to download ckpt/config again if existed
    """
    kwargs["ckpt"] = google_large_file_link("1F6ZOuyS9C92OLrunA7zQkWp7vb6eQkOT")
    return replai_url(refresh=refresh, *args, **kwargs)


def replai_ego4d_woAStC(refresh=False, *args, **kwargs):
    """
    The replai model (w/o AStC) trained on Ego4D
        refresh (bool): whether to download ckpt/config again if existed
    """
    kwargs["ckpt"] = google_large_file_link("1hQH1UWgC_-EztKpmma3TNST7OYGkQyI_")
    return replai_url(refresh=refresh, *args, **kwargs)

upstream_models/replai/hubconf.py
roger-tseng/av-superb
import numpy as np
import torch
from pytorchvideo import transforms as vT
from torchaudio import transforms as aT
from torchvision import transforms as T

from .transforms import audio as aT2


class ResampleLogMelSpectrogram:
    def __init__(
        self,
        raw_sample_rate=44100,
        audio_rate=16000,
        mel_window_size=32,
        mel_step_size=16,
        num_mels=80,
        num_temporal_frames=128,
        augment=False,
    ):
        self.outp_size = (1, num_temporal_frames, num_mels)
        self.raw_sample_rate = raw_sample_rate

        mel_mean = -7.03
        mel_std = 4.66
        n_fft = int(float(audio_rate) / 1000 * mel_window_size)
        hop_length = int(float(audio_rate) / 1000 * mel_step_size)
        eps = 1e-10

        transforms = [
            aT2.ToMono(),
            aT.Resample(
                orig_freq=raw_sample_rate,
                new_freq=audio_rate,
            ),
            aT.MelSpectrogram(
                sample_rate=audio_rate,
                n_fft=n_fft,
                hop_length=hop_length,
                n_mels=num_mels,
                center=False,
            ),
            aT2.Log(eps=eps),
            aT2.Permute(0, 2, 1),  # (1, F, T) -> (1, T, F)
            aT2.UniformTemporalSubsample(
                num_temporal_frames, time_dim=-2
            ),  # temporal and frequency axis, by here different
            # mean and different std dev, mean and
            T.Normalize((mel_mean,), (mel_std,)),  # remove
        ]
        self.t = T.Compose(transforms)

    def __call__(self, x, fps):
        assert fps == self.raw_sample_rate
        if isinstance(x, torch.Tensor):
            return self.t(x)
        elif isinstance(x, np.ndarray):
            return self.t(torch.from_numpy(x))
        else:
            raise TypeError(f"x is not a tensor or ndarray, but {type(x)}")

    def output_shape(self):
        return self.outp_size


class SpectrogramPrep:
    def __init__(self, args, augment=True):
        self.args = args
        self.augment = augment

        n_fft = int(
            float(args.audio_resampled_rate) / 1000 * args.audio_mel_window_size
        )
        hop_length = int(
            float(args.audio_resampled_rate) / 1000 * args.audio_mel_step_size
        )
        transforms = [
            aT2.RandomVolume(volume=args.volume_coeff),
            aT2.ToMono(),
            aT.Resample(
                orig_freq=args.audio_raw_sample_rate,
                new_freq=args.audio_resampled_rate,
            ),
            aT.MelSpectrogram(
                sample_rate=args.audio_resampled_rate,
                n_fft=n_fft,
                hop_length=hop_length,
                n_mels=args.audio_num_mels,
                center=False,
            ),
            aT.AmplitudeToDB(),
            T.Lambda(lambda x: x / 10.0),
            T.Lambda(
                lambda x: x.transpose(2, 1).unsqueeze(2)
            ),  # (C, F, T) -> (C, T, 1, F)
            vT.UniformTemporalSubsample(args.audio_mel_num_subsample),
            T.Lambda(lambda x: x.squeeze(2)),  # (C, T, 1, F) -> (C, T, F)
        ]
        if not augment:
            transforms.pop(0)
        self.t = T.Compose(transforms)

    def __call__(self, x, fps):
        assert fps == self.args.audio_raw_sample_rate
        return self.t(x)

    def output_shape(self):
        return 1, self.args.audio_mel_num_subsample, self.args.audio_num_mels

upstream_models/replai/replai/data/audio_transforms.py
roger-tseng/av-superb
import torch

from . import audio_transforms, video_transforms


class Transform:
    def __init__(self, transforms, data_shapes):
        self.transforms = transforms
        self.data_shapes = data_shapes

    def __call__(self, x):
        for k in self.transforms:
            if k not in x or x[k] is None:
                x[k] = torch.zeros(tuple(self.data_shapes[k])).float()
            else:
                x[k] = self.transforms[k](x[k][0], x[k][1])
        return x


def build_transforms(cfg, augment):
    transforms, data_shapes = {}, {}
    for k in cfg:
        if cfg[k].name in vars(video_transforms):
            transforms[k] = video_transforms.__dict__[cfg[k].name](
                **cfg[k].args, augment=augment
            )
        elif cfg[k].name in vars(audio_transforms):
            transforms[k] = audio_transforms.__dict__[cfg[k].name](
                **cfg[k].args, augment=augment
            )
        else:
            raise NotImplementedError(f"Transform {cfg.name} not found.")
        data_shapes[k] = cfg[k].data_shape
    return Transform(transforms, data_shapes)

upstream_models/replai/replai/data/builder.py
roger-tseng/av-superb
import io
import logging
import math
import pathlib
from typing import BinaryIO, Dict, List, Optional, Tuple

import av
import numpy as np
import torch
from iopath.common.file_io import g_pathmgr
from pytorchvideo.data.encoded_video_pyav import EncodedVideoPyAV, _pyav_decode_stream
from pytorchvideo.data.utils import pts_to_secs, secs_to_pts, thwc_to_cthw

logger = logging.getLogger(__name__)


class PyAVDecoder(EncodedVideoPyAV):
    def __init__(
        self,
        file_path: str,
        decode_audio: bool = True,
    ) -> None:

        with g_pathmgr.open(file_path, "rb") as fh:
            video_file = io.BytesIO(fh.read())

        super().__init__(
            file=video_file,
            video_name=pathlib.Path(file_path).name,
            decode_audio=decode_audio,
        )
        self._video_fps = self._container.streams.video[0].average_rate
        self._video_start = self._video_start_pts * self._video_time_base
        self._video_duration = (
            self._container.streams.video[0].duration * self._video_time_base
        )
        if self._decode_audio and self._has_audio:
            self._audio_fps = self._container.streams.audio[0].rate
            self._audio_start = self._audio_start_pts * self._audio_time_base
            self._audio_duration = (
                self._container.streams.audio[0].duration * self._audio_time_base
            )

    def get_video_clip(self, start_sec: float, end_sec: float):
        """
        Retrieves frames from the encoded video at the specified start and end times
        in seconds (the video always starts at 0 seconds).

        Args:
            start_sec (float): the clip start time in seconds
            end_sec (float): the clip end time in seconds
        Returns:
            A tensor of the clip's RGB frames with shape:
            (channel, time, height, width). The frames are of type torch.float32 and
            in the range [0 - 255].

            Returns None if no video or audio found within time range.

        """
        self._video = self._decode_video_clip(start_sec, end_sec)

        video_frames = None
        if self._video is not None:
            video_start_pts = secs_to_pts(
                start_sec, self._video_time_base, self._video_start_pts
            )
            video_end_pts = secs_to_pts(
                end_sec, self._video_time_base, self._video_start_pts
            )

            video_frames = [
                f
                for f, pts in self._video
                if pts >= video_start_pts and pts <= video_end_pts
            ]

        if video_frames is None or len(video_frames) == 0:
            logger.debug(
                f"No video found within {start_sec} and {end_sec} seconds. "
                f"Video starts at time 0 and ends at {self.duration}."
            )
            video_frames = None

        return video_frames

    def get_audio_clip(self, start_sec: float, end_sec: float) -> torch.Tensor:
        """
        Retrieves frames from the encoded video at the specified start and end times
        in seconds (the video always starts at 0 seconds).

        Args:
            start_sec (float): the clip start time in seconds
            end_sec (float): the clip end time in seconds
        Returns:
            A tensor of the clip's RGB frames with shape:
            (channel, time, height, width). The frames are of type torch.float32 and
            in the range [0 - 255].

            Returns None if no video or audio found within time range.

        """
        self._audio = self._decode_audio_clip(start_sec, end_sec)

        audio_samples = None
        if self._has_audio and self._audio is not None:
            audio_start_pts = secs_to_pts(
                start_sec, self._audio_time_base, self._audio_start_pts
            )
            audio_end_pts = secs_to_pts(
                end_sec, self._audio_time_base, self._audio_start_pts
            )
            audio_samples = [
                f
                for f, pts in self._audio
                if pts >= audio_start_pts and pts <= audio_end_pts
            ]

            audio_samples = np.concatenate(audio_samples, axis=1).astype(np.float32)
            shift_pts = self._audio[0][1] - audio_start_pts
            duration_pts = audio_end_pts - audio_start_pts
            audio_samples = audio_samples[:, shift_pts : shift_pts + duration_pts]

        if audio_samples is None or len(audio_samples) == 0:
            logger.debug(
                f"No audio found within {start_sec} and {end_sec} seconds. "
                f"Video starts at time 0 and ends at {self.duration}."
            )

        return audio_samples

    def _decode_video_clip(self, start_secs: float = 0.0, end_secs: float = math.inf):
        """
        Selectively decodes a video between start_pts and end_pts in time units of the
        self._video's timebase.
        """
        video_and_pts = None
        try:
            pyav_video_frames, _ = _pyav_decode_stream(
                self._container,
                secs_to_pts(start_secs, self._video_time_base, self._video_start_pts),
                secs_to_pts(end_secs, self._video_time_base, self._video_start_pts),
                self._container.streams.video[0],
                {"video": 0},
            )
            if len(pyav_video_frames) > 0:
                video_and_pts = [
                    (frame.to_image().convert("RGB"), frame.pts)
                    for frame in pyav_video_frames
                ]

        except Exception as e:
            logger.debug(f"Failed to decode video: {self._video_name}. {e}")

        return video_and_pts

    def _decode_audio_clip(self, start_secs: float = 0.0, end_secs: float = math.inf):
        """
        Selectively decodes a video between start_pts and end_pts in time units of the
        self._video's timebase.
        """
        audio_and_pts = None
        try:
            if self._has_audio:
                pyav_audio_frames, _ = _pyav_decode_stream(
                    self._container,
                    secs_to_pts(
                        start_secs, self._audio_time_base, self._audio_start_pts
                    ),
                    secs_to_pts(end_secs, self._audio_time_base, self._audio_start_pts),
                    self._container.streams.audio[0],
                    {"audio": 0},
                )

                if len(pyav_audio_frames) > 0:
                    audio_and_pts = [
                        (frame.to_ndarray(), frame.pts) for frame in pyav_audio_frames
                    ]

        except Exception as e:
            logger.debug(f"Failed to decode audio: {self._video_name}. {e}")

        return audio_and_pts

upstream_models/replai/replai/data/pyav_decoder.py
roger-tseng/av-superb
import torch
import torch.nn as nn


class RandomVolume(object):
    def __init__(self, volume):
        assert len(volume) == 2
        assert volume[1] >= volume[0] > 0
        self.volume = volume

    def __call__(self, x):
        if x.ndim == 1:
            x = x[None]
        assert x.ndim == 2
        n_waves, n_frames = x.shape
        max_mult = torch.minimum(1 / x.max(1)[0], torch.tensor(self.volume[1]))
        min_mult = torch.ones_like(max_mult) * self.volume[0]
        multiplier = torch.rand(n_waves) * (max_mult - min_mult) + min_mult
        return x * multiplier[:, None]


class UniformTemporalSubsample:
    def __init__(self, num_frames, time_dim=-2):
        self.num_frames = num_frames
        self.time_dim = time_dim

    def __call__(self, x):
        t = x.shape[self.time_dim]
        assert self.num_frames > 0 and t > 0

        # Sample by nearest neighbor interpolation if num_samples > t.
        indices = torch.linspace(0, t - 1, self.num_frames)
        indices = torch.clamp(indices, 0, t - 1).long()
        return torch.index_select(x, self.time_dim, indices)


class Log:
    def __init__(self, eps=1e-10):
        self.eps = eps

    def __call__(self, x):
        return torch.log(x.clamp(min=self.eps))


class Permute:
    def __init__(self, *dims):
        self.dims = dims

    def __call__(self, x):
        return x.permute(*self.dims)


class ToMono:
    def __call__(self, x):
        if x.ndim == 1:
            return x[None]
        elif x.ndim == 2:
            return x.mean(0, keepdims=True)
        else:
            raise ValueError("Audio tensor should have at most 2 dimensions (c,t)")


class DropStripes(nn.Module):
    def __init__(self, dim, drop_width, stripes_num):
        """Drop stripes.
        Args:
          dim: int, dimension along which to drop
          drop_width: int, maximum width of stripes to drop
          stripes_num: int, how many stripes to drop
        """
        super(DropStripes, self).__init__()

        assert dim in [2, 3]  # dim 2: time; dim 3: frequency

        self.dim = dim
        self.drop_width = drop_width
        self.stripes_num = stripes_num

    def forward(self, input):
        """input: (batch_size, channels, time_steps, freq_bins)"""

        assert input.ndimension() == 4

        if self.training is False:
            return input

        else:
            batch_size = input.shape[0]
            total_width = input.shape[self.dim]

            for n in range(batch_size):
                self.transform_slice(input[n], total_width)

            return input

    def transform_slice(self, e, total_width):
        """e: (channels, time_steps, freq_bins)"""

        for _ in range(self.stripes_num):
            distance = torch.randint(low=0, high=self.drop_width, size=(1,))[0]
            bgn = torch.randint(low=0, high=total_width - distance, size=(1,))[0]

            if self.dim == 2:
                e[:, bgn : bgn + distance, :] = 0
            elif self.dim == 3:
                e[:, :, bgn : bgn + distance] = 0


class SpecAugmentation(nn.Module):
    def __init__(
        self, time_drop_width, time_stripes_num, freq_drop_width, freq_stripes_num
    ):
        """Spec augmetation.
        [ref] Park, D.S., Chan, W., Zhang, Y., Chiu, C.C., Zoph, B., Cubuk, E.D.
        and Le, Q.V., 2019. Specaugment: A simple data augmentation method
        for automatic speech recognition. arXiv preprint arXiv:1904.08779.
        Args:
          time_drop_width: int
          time_stripes_num: int
          freq_drop_width: int
          freq_stripes_num: int
        """

        super(SpecAugmentation, self).__init__()

        self.time_dropper = DropStripes(
            dim=2, drop_width=time_drop_width, stripes_num=time_stripes_num
        )

        self.freq_dropper = DropStripes(
            dim=3, drop_width=freq_drop_width, stripes_num=freq_stripes_num
        )

    def forward(self, input):
        x = self.time_dropper(input)
        x = self.freq_dropper(x)
        return x

upstream_models/replai/replai/data/transforms/audio.py
roger-tseng/av-superb
# Copyright (c) Facebook, Inc. and its 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 math
import numbers
import random
import warnings

import cv2
import numpy as np
import PIL
import torch
import torchvision
from PIL import Image
from torchvision.transforms import functional as vF


def convert_img(img):
    """Converts (H, W, C) numpy.ndarray to (C, W, H) format"""
    if len(img.shape) == 3:
        img = img.transpose(2, 0, 1)
    if len(img.shape) == 2:
        img = np.expand_dims(img, 0)
    return img


def crop_clip(clip, min_h, min_w, h, w):
    if isinstance(clip[0], np.ndarray):
        cropped = [img[min_h : min_h + h, min_w : min_w + w, :] for img in clip]

    elif isinstance(clip[0], PIL.Image.Image):
        cropped = [img.crop((min_w, min_h, min_w + w, min_h + h)) for img in clip]
    else:
        raise TypeError(
            "Expected numpy.ndarray or PIL.Image"
            + "but got list of {0}".format(type(clip[0]))
        )
    return cropped


def resize_clip(clip, size, interpolation="bilinear"):
    if isinstance(clip[0], np.ndarray): # or isinstance(clip[0], torch.Tensor):
        if isinstance(size, numbers.Number):
            im_h, im_w, im_c = clip[0].shape
            # Min spatial dim already matches minimal size
            if (im_w <= im_h and im_w == size) or (im_h <= im_w and im_h == size):
                return clip
            new_h, new_w = get_resize_sizes(im_h, im_w, size)
            size = (new_w, new_h)
        else:
            size = size[1], size[0]
        if interpolation == "bilinear":
            np_inter = cv2.INTER_LINEAR
        else:
            np_inter = cv2.INTER_NEAREST
        scaled = [cv2.resize(img, size, interpolation=np_inter) for img in clip]
    elif isinstance(clip[0], PIL.Image.Image):
        if isinstance(size, numbers.Number):
            im_w, im_h = clip[0].size
            # Min spatial dim already matches minimal size
            if (im_w <= im_h and im_w == size) or (im_h <= im_w and im_h == size):
                return clip
            new_h, new_w = get_resize_sizes(im_h, im_w, size)
            size = (new_w, new_h)
        else:
            size = size[1], size[0]
        if interpolation == "bilinear":
            pil_inter = PIL.Image.BILINEAR
        else:
            pil_inter = PIL.Image.NEAREST

        scaled = [img.resize(size, pil_inter) for img in clip]
    else:
        raise TypeError(
            "Expected numpy.ndarray or PIL.Image"
            + "but got list of {0}".format(type(clip[0]))
        )
    return scaled


def get_resize_sizes(im_h, im_w, size):
    if im_w < im_h:
        ow = size
        oh = int(size * im_h / im_w)
    else:
        oh = size
        ow = int(size * im_w / im_h)
    return oh, ow


class Compose(object):
    """Composes several transforms

    Args:
    transforms (list of ``Transform`` objects): list of transforms
    to compose
    """

    def __init__(self, transforms):
        self.transforms = transforms

    def __call__(self, clip):
        for t in self.transforms:
            clip = t(clip)
        return clip


class RandomDrop(object):
    """Randomly drops frames in order to keep a fixed number of frames"""

    def __init__(self, num_frames):
        self.num_frames = num_frames

    def __call__(self, clip):
        assert len(clip) >= self.num_frames
        if len(clip) == self.num_frames:
            return clip
        idx = sorted(random.sample(range(len(clip)), k=self.num_frames))
        return [clip[i] for i in idx]


class UniformDrop(object):
    """Randomly drops frames in order to keep a fixed number of frames"""

    def __init__(self, num_frames=None, ss_ratio=None):
        self.num_frames = num_frames
        self.ss_ratio = ss_ratio

    def __call__(self, clip):
        if self.num_frames is not None:
            if len(clip) <= self.num_frames:
                return clip
            idx = np.linspace(0, len(clip) - 1, self.num_frames, endpoint=True).astype(
                int
            )
            return [clip[i] for i in idx]
        elif self.ss_ratio is not None:
            if self.ss_ratio == 1:
                return clip
            idx = np.arange(0, len(clip), self.ss_ratio).astype(int)
            return [clip[i] for i in idx]


class RandomHorizontalFlip(object):
    """Horizontally flip the list of given images randomly
    with a probability 0.5
    """

    def __call__(self, clip):
        """
        Args:
        img (PIL.Image or numpy.ndarray): List of images to be cropped
        in format (h, w, c) in numpy.ndarray

        Returns:
        PIL.Image or numpy.ndarray: Randomly flipped clip
        """
        if random.random() < 0.5:
            if isinstance(clip[0], np.ndarray):
                return [np.fliplr(img) for img in clip]
            elif isinstance(clip[0], PIL.Image.Image):
                return [img.transpose(PIL.Image.FLIP_LEFT_RIGHT) for img in clip]
            else:
                raise TypeError(
                    "Expected numpy.ndarray or PIL.Image"
                    + " but got list of {0}".format(type(clip[0]))
                )
        return clip


class RandomGray(object):
    """Horizontally flip the list of given images randomly
    with a probability 0.5
    """

    def __init__(self, p):
        self.p = p

    def __call__(self, clip):
        """
        Args:
        img (PIL.Image or numpy.ndarray): List of images to be cropped
        in format (h, w, c) in numpy.ndarray

        Returns:
        PIL.Image or numpy.ndarray: Randomly flipped clip
        """

        rand_gray_clip = []
        for img in clip:
            if random.random() < self.p:
                img = vF.to_grayscale(img)
            rand_gray_clip.append(img)
        return rand_gray_clip


class RandomResize(object):
    """Resizes a list of (H x W x C) numpy.ndarray to the final size

    The larger the original image is, the more times it takes to
    interpolate

    Args:
    interpolation (str): Can be one of 'nearest', 'bilinear'
    defaults to nearest
    size (tuple): (widht, height)
    """

    def __init__(self, ratio=(3.0 / 4.0, 4.0 / 3.0), interpolation="nearest"):
        self.ratio = ratio
        self.interpolation = interpolation

    def __call__(self, clip):
        scaling_factor = random.uniform(self.ratio[0], self.ratio[1])

        if isinstance(clip[0], np.ndarray):
            im_h, im_w, im_c = clip[0].shape
        elif isinstance(clip[0], PIL.Image.Image):
            im_w, im_h = clip[0].size

        new_w = int(im_w * scaling_factor)
        new_h = int(im_h * scaling_factor)
        new_size = (new_w, new_h)
        resized = resize_clip(clip, new_size, interpolation=self.interpolation)
        return resized


class Resize(object):
    """Resizes a list of (H x W x C) numpy.ndarray to the final size

    The larger the original image is, the more times it takes to
    interpolate

    Args:
    interpolation (str): Can be one of 'nearest', 'bilinear'
    defaults to nearest
    size (tuple): (widht, height)
    """

    def __init__(self, size, interpolation="nearest"):
        self.size = size
        self.interpolation = interpolation

    def __call__(self, clip):
        resized = resize_clip(clip, self.size, interpolation=self.interpolation)
        return resized


class RandomCrop(object):
    """Extract random crop at the same location for a list of images

    Args:
    size (sequence or int): Desired output size for the
    crop in format (h, w)
    """

    def __init__(self, size):
        if isinstance(size, numbers.Number):
            size = (size, size)

        self.size = size

    def __call__(self, clip):
        """
        Args:
        img (PIL.Image or numpy.ndarray): List of images to be cropped
        in format (h, w, c) in numpy.ndarray

        Returns:
        PIL.Image or numpy.ndarray: Cropped list of images
        """
        h, w = self.size
        if isinstance(clip[0], np.ndarray):
            im_h, im_w, im_c = clip[0].shape
        elif isinstance(clip[0], PIL.Image.Image):
            im_w, im_h = clip[0].size
        else:
            raise TypeError(
                "Expected numpy.ndarray or PIL.Image"
                + "but got list of {0}".format(type(clip[0]))
            )
        if w > im_w or h > im_h:
            error_msg = (
                "Initial image size should be larger then "
                "cropped size but got cropped sizes : ({w}, {h}) while "
                "initial image is ({im_w}, {im_h})".format(
                    im_w=im_w, im_h=im_h, w=w, h=h
                )
            )
            raise ValueError(error_msg)

        x1 = random.randint(0, im_w - w)
        y1 = random.randint(0, im_h - h)
        cropped = crop_clip(clip, y1, x1, h, w)

        return cropped


class CenterCrop(object):
    """Extract center crop at the same location for a list of images

    Args:
    size (sequence or int): Desired output size for the
    crop in format (h, w)
    """

    def __init__(self, size):
        if isinstance(size, numbers.Number):
            size = (size, size)

        self.size = size

    def __call__(self, clip):
        """
        Args:
        img (PIL.Image or numpy.ndarray): List of images to be cropped
        in format (h, w, c) in numpy.ndarray

        Returns:
        PIL.Image or numpy.ndarray: Cropped list of images
        """
        h, w = self.size
        if isinstance(clip[0], np.ndarray):
            im_h, im_w, im_c = clip[0].shape
        elif isinstance(clip[0], PIL.Image.Image):
            im_w, im_h = clip[0].size
        else:
            raise TypeError(
                "Expected numpy.ndarray or PIL.Image"
                + "but got list of {0}".format(type(clip[0]))
            )
        if w > im_w or h > im_h:
            error_msg = (
                "Initial image size should be larger then "
                "cropped size but got cropped sizes : ({w}, {h}) while "
                "initial image is ({im_w}, {im_h})".format(
                    im_w=im_w, im_h=im_h, w=w, h=h
                )
            )
            raise ValueError(error_msg)

        x1 = int(round((im_w - w) / 2.0))
        y1 = int(round((im_h - h) / 2.0))
        cropped = crop_clip(clip, y1, x1, h, w)

        return cropped


class TenCrop(object):
    """Extract center crop at the same location for a list of images

    Args:
    size (sequence or int): Desired output size for the
    crop in format (h, w)
    """

    def __init__(self, size):
        if isinstance(size, numbers.Number):
            size = (size, size)

        self.size = size

    def __call__(self, clip):
        """
        Args:
        img (PIL.Image or numpy.ndarray): List of images to be cropped
        in format (h, w, c) in numpy.ndarray

        Returns:
        PIL.Image or numpy.ndarray: Cropped list of images
        """
        cropped = []
        for frame in clip:
            cropped += list(vF.ten_crop(frame, self.size))

        return cropped


class RandomResizedCrop(object):
    """Crops a series of PIL Image to random size and aspect ratio.

    A crop of random size (default: of 0.08 to 1.0) of the original size and a random
    aspect ratio (default: of 3/4 to 4/3) of the original aspect ratio is made. This crop
    is finally resized to given size.
    This is popularly used to train the Inception networks.

    Args:
        size: expected output size of each edge
        scale: range of size of the origin size cropped
        ratio: range of aspect ratio of the origin aspect ratio cropped
        interpolation: Default: PIL.Image.BILINEAR
    """

    def __init__(
        self,
        size,
        scale=(0.08, 1.0),
        ratio=(3.0 / 4.0, 4.0 / 3.0),
        interpolation=PIL.Image.BILINEAR,
    ):
        if isinstance(size, tuple):
            self.size = size
        else:
            self.size = (size, size)
        if (scale[0] > scale[1]) or (ratio[0] > ratio[1]):
            warnings.warn("range should be of kind (min, max)")

        self.interpolation = interpolation
        self.scale = scale
        self.ratio = ratio

    @staticmethod
    def get_params(img, scale, ratio):
        """Get parameters for ``crop`` for a random sized crop.

        Args:
            img (PIL Image): Image to be cropped.
            scale (tuple): range of size of the origin size cropped
            ratio (tuple): range of aspect ratio of the origin aspect ratio cropped

        Returns:
            tuple: params (i, j, h, w) to be passed to ``crop`` for a random
                sized crop.
        """
        area = img.size[0] * img.size[1]

        for attempt in range(10):
            target_area = random.uniform(*scale) * area
            log_ratio = (math.log(ratio[0]), math.log(ratio[1]))
            aspect_ratio = math.exp(random.uniform(*log_ratio))

            w = int(round(math.sqrt(target_area * aspect_ratio)))
            h = int(round(math.sqrt(target_area / aspect_ratio)))

            if w <= img.size[0] and h <= img.size[1]:
                i = random.randint(0, img.size[1] - h)
                j = random.randint(0, img.size[0] - w)
                return i, j, h, w

        # Fallback to central crop
        in_ratio = img.size[0] / img.size[1]
        if in_ratio < min(ratio):
            w = img.size[0]
            h = int(round(w / min(ratio)))
        elif in_ratio > max(ratio):
            h = img.size[1]
            w = int(round(h * max(ratio)))
        else:  # whole image
            w = img.size[0]
            h = img.size[1]
        i = (img.size[1] - h) // 2
        j = (img.size[0] - w) // 2
        return i, j, h, w

    def __call__(self, clip):
        """
        Args:
            img (PIL Image): Image to be cropped and resized.

        Returns:
            PIL Image: Randomly cropped and resized image.
        """
        i, j, h, w = self.get_params(clip[0], self.scale, self.ratio)
        return [
            vF.resized_crop(img, i, j, h, w, self.size, self.interpolation)
            for img in clip
        ]

    def __repr__(self):
        interpolate_str = "?"
        format_string = self.__class__.__name__ + "(size={0}".format(self.size)
        format_string += ", scale={0}".format(tuple(round(s, 4) for s in self.scale))
        format_string += ", ratio={0}".format(tuple(round(r, 4) for r in self.ratio))
        format_string += ", interpolation={0})".format(interpolate_str)
        return format_string


class RandomShortSideScale(object):
    """Randomly resizes a series of PIL Image to random size.

    A crop of random size (default: of 0.08 to 1.0) of the original size and a random
    aspect ratio (default: of 3/4 to 4/3) of the original aspect ratio is made. This crop
    is finally resized to given size.
    This is popularly used to train the Inception networks.

    Args:
        min_size: minimum size for the short edge
        max_size: maximum size for the short edge
    """

    def __init__(self, min_size, max_size=None):
        self.min_size = min_size
        self.max_size = max_size if max_size is not None else min_size

    def get_params(self, clip):
        new_min_size = random.uniform(self.min_size, self.max_size)
        h, w = clip[0].size
        new_h = int(h / min(h, w) * new_min_size)
        new_w = int(w / min(h, w) * new_min_size)
        return new_h, new_w

    def __call__(self, clip):
        """
        Args:
            img (PIL Image): Image to be cropped and resized.

        Returns:
            PIL Image: Randomly cropped and resized image.
        """
        h, w = self.get_params(clip)
        return resize_clip(clip, (h, w))

    def __repr__(self):
        interpolate_str = "?"
        format_string = self.__class__.__name__ + "(size={0}".format(self.size)
        format_string += ", scale={0}".format(tuple(round(s, 4) for s in self.scale))
        format_string += ", ratio={0}".format(tuple(round(r, 4) for r in self.ratio))
        format_string += ", interpolation={0})".format(interpolate_str)
        return format_string


class ColorJitter(object):
    """Randomly change the brightness, contrast and saturation and hue of the clip

    Args:
    brightness (float): How much to jitter brightness. brightness_factor
    is chosen uniformly from [max(0, 1 - brightness), 1 + brightness].
    contrast (float): How much to jitter contrast. contrast_factor
    is chosen uniformly from [max(0, 1 - contrast), 1 + contrast].
    saturation (float): How much to jitter saturation. saturation_factor
    is chosen uniformly from [max(0, 1 - saturation), 1 + saturation].
    hue(float): How much to jitter hue. hue_factor is chosen uniformly from
    [-hue, hue]. Should be >=0 and <= 0.5.
    """

    def __init__(self, brightness=0, contrast=0, saturation=0, hue=0):
        self.brightness = brightness
        self.contrast = contrast
        self.saturation = saturation
        self.hue = hue

    def get_params(self, brightness, contrast, saturation, hue):
        if brightness > 0:
            brightness_factor = random.uniform(max(0, 1 - brightness), 1 + brightness)
        else:
            brightness_factor = None

        if contrast > 0:
            contrast_factor = random.uniform(max(0, 1 - contrast), 1 + contrast)
        else:
            contrast_factor = None

        if saturation > 0:
            saturation_factor = random.uniform(max(0, 1 - saturation), 1 + saturation)
        else:
            saturation_factor = None

        if hue > 0:
            hue_factor = random.uniform(-hue, hue)
        else:
            hue_factor = None
        return brightness_factor, contrast_factor, saturation_factor, hue_factor

    def __call__(self, clip):
        """
        Args:
        clip (list): list of PIL.Image

        Returns:
        list PIL.Image : list of transformed PIL.Image
        """
        if isinstance(clip[0], np.ndarray):
            raise TypeError("Color jitter not yet implemented for numpy arrays")
        elif isinstance(clip[0], PIL.Image.Image):
            brightness, contrast, saturation, hue = self.get_params(
                self.brightness, self.contrast, self.saturation, self.hue
            )

            # Create img transform function sequence
            img_transforms = []
            if brightness is not None:
                img_transforms.append(
                    lambda img: torchvision.transforms.functional.adjust_brightness(
                        img, brightness
                    )
                )
            if saturation is not None:
                img_transforms.append(
                    lambda img: torchvision.transforms.functional.adjust_saturation(
                        img, saturation
                    )
                )
            if hue is not None:
                img_transforms.append(
                    lambda img: torchvision.transforms.functional.adjust_hue(img, hue)
                )
            if contrast is not None:
                img_transforms.append(
                    lambda img: torchvision.transforms.functional.adjust_contrast(
                        img, contrast
                    )
                )
            random.shuffle(img_transforms)

            # Apply to all images
            jittered_clip = []
            for img in clip:
                for func in img_transforms:
                    img = func(img)
                jittered_clip.append(img)

        else:
            raise TypeError(
                "Expected numpy.ndarray or PIL.Image"
                + "but got list of {0}".format(type(clip[0]))
            )
        return jittered_clip


class TemporalJitter(object):
    """Crop video sequence temporally with jitter

    Args:
    n_frames (int): number  of output frames
    min_scale (float): minimum subsample rate (default 1.0)
    max_scale (float): maximum subsample rate (default 1.0)
    """

    def __init__(self, n_frames, time_scale=(1.0, 1.0)):
        self.n_frames = n_frames
        self.time_scale = time_scale

    def __call__(self, clip):
        """
        Args:
        img (PIL.Image or numpy.ndarray): List of images to be cropped
        in format (h, w, c) in numpy.ndarray

        Returns:
        PIL.Image or numpy.ndarray: Cropped list of images
        """
        rate = (
            self.time_scale[0]
            if self.time_scale[0] == self.time_scale[1]
            else random.uniform(
                self.time_scale[0],
                min(self.time_scale[1], float(len(clip)) / float(self.n_frames)),
            )
        )
        clip_ss = [clip[int(t)] for t in np.arange(0, len(clip), rate)]
        if len(clip_ss) == self.n_frames:
            clip_out = clip_ss
        elif len(clip_ss) < self.n_frames:
            # Wrap to fill frames
            clip_out = [clip_ss[t % len(clip_ss)] for t in range(self.n_frames)]
        else:
            # Extract random crop
            max_init_t = len(clip_ss) - self.n_frames
            init_t = random.randint(0, max_init_t)
            clip_out = clip_ss[init_t : init_t + self.n_frames]
        return clip_out


class TemporalCenterCrop(object):
    """Crop video sequence temporally with jitter

    Args:
    n_frames (int): number  of output frames
    min_scale (float): minimum subsample rate (default 1.0)
    max_scale (float): maximum subsample rate (default 1.0)
    """

    def __init__(self, n_frames, time_scale=1.0):
        self.n_frames = n_frames
        self.time_scale = time_scale

    def __call__(self, clip):
        """
        Args:
        img (PIL.Image or numpy.ndarray): List of images to be cropped
        in format (h, w, c) in numpy.ndarray

        Returns:
        PIL.Image or numpy.ndarray: Cropped list of images
        """
        clip_ss = [clip[int(t)] for t in np.arange(0, len(clip), self.time_scale)]
        if len(clip_ss) == self.n_frames:
            clip_out = clip_ss
        elif len(clip_ss) < self.n_frames:
            # Wrap to fill frames
            clip_out = [clip_ss[t % len(clip_ss)] for t in range(self.n_frames)]
        else:
            # Extract random crop
            init_t = (len(clip_ss) - self.n_frames) // 2
            clip_out = clip_ss[init_t : init_t + self.n_frames]
        return clip_out


class ClipToTensor(object):
    """Convert a list of m (H x W x C) numpy.ndarrays in the range [0, 255]
    to a torch.FloatTensor of shape (C x m x H x W) in the range [0, 1.0]
    """

    def __init__(self, channel_nb=3, div_255=True, numpy=False):
        self.channel_nb = channel_nb
        self.div_255 = div_255
        self.numpy = numpy

    def __call__(self, clip):
        """
        Args: clip (list of numpy.ndarray): clip (list of images)
        to be converted to tensor.
        """
        # Retrieve shape
        if isinstance(clip[0], np.ndarray):
            h, w, ch = clip[0].shape
            assert ch == self.channel_nb, "Got {0} instead of 3 channels".format(ch)
        elif isinstance(clip[0], Image.Image):
            w, h = clip[0].size
        else:
            raise TypeError(
                "Expected numpy.ndarray or PIL.Image\
            but got list of {0}".format(
                    type(clip[0])
                )
            )

        np_clip = np.zeros([self.channel_nb, len(clip), int(h), int(w)])

        # Convert
        for img_idx, img in enumerate(clip):
            if isinstance(img, np.ndarray):
                pass
            elif isinstance(img, Image.Image):
                img = np.array(img, copy=False)
            else:
                raise TypeError(
                    "Expected numpy.ndarray or PIL.Image\
                but got list of {0}".format(
                        type(clip[0])
                    )
                )
            img = convert_img(img)
            np_clip[:, img_idx, :, :] = img
        if self.numpy:
            if self.div_255:
                np_clip = np_clip / 255
            return np_clip

        else:
            tensor_clip = torch.from_numpy(np_clip)

            if not isinstance(tensor_clip, torch.FloatTensor):
                tensor_clip = tensor_clip.float()
            if self.div_255:
                tensor_clip = tensor_clip.div(255)
            return tensor_clip


class ToTensor(object):
    """Converts numpy array to tensor"""

    def __call__(self, array):
        tensor = torch.from_numpy(array)
        return tensor

upstream_models/replai/replai/data/transforms/video.py
roger-tseng/av-superb
import numpy as np
import torch
from pytorchvideo import transforms as vT
from torchvision import transforms as T

from .transforms import video as vT2

__all__ = ["ResizeCropFlip", "MultiResizeCropFlip", "MultiScaleCropFlipColorJitter"]


class MultiScaleCropFlipColorJitter:
    def __init__(
        self,
        num_frames=8,
        crop=(224, 224),
        color=(0.4, 0.4, 0.4, 0.1),
        min_area=0.08,
        augment=True,
    ):
        from collections.abc import Iterable

        if isinstance(crop, Iterable):
            crop = tuple(crop)
        self.crop = crop
        self.augment = augment
        self.num_frames = num_frames

        if augment:
            transforms = [
                vT2.RandomResizedCrop(crop, scale=(min_area, 1.0)),
                vT2.RandomHorizontalFlip(),
                vT2.ColorJitter(*color),
            ]
        else:
            transforms = [
                vT2.Resize(int(crop[0] / 0.875)),
                vT2.CenterCrop(crop),
            ]

        transforms += [
            vT2.ClipToTensor(),
            vT.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
            vT.UniformTemporalSubsample(num_frames),
        ]
        self.t = T.Compose(transforms)

    def __call__(self, x, fps):
        return self.t(x)


class ResizeCropFlip:
    def __init__(
        self, num_frames=8, min_size=256, max_size=360, crop_size=224, augment=True
    ):
        if augment:
            transforms = [
                vT2.RandomShortSideScale(min_size=min_size, max_size=max_size),
                vT2.RandomCrop(crop_size),
                vT2.RandomHorizontalFlip(),
            ]
        else:
            transforms = [
                vT2.Resize(min_size),
                vT2.CenterCrop(crop_size),
            ]
        transforms += [
            vT2.ClipToTensor(),
            vT.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
            # vT.UniformTemporalSubsample(num_frames),
        ]
        self.t = T.Compose(transforms)

    def __call__(self, x, fps):
        return self.t(x)


class UniformClips:
    def __init__(
        self, num_clips, clip_duration, transform, transform_args, augment=True
    ):
        self.num_clips = num_clips
        self.clip_duration = clip_duration
        self.transform = eval(transform)(**transform_args, augment=augment)

    def __call__(self, x, fps):
        nt = len(x)
        clip_len = int(
            self.clip_duration * fps
        )  # [0, 1, ..., 96] -> [0, 1, ..., 88] -> [0, 11, 22, 33, .. 88]
        start_times = [int(ss) for ss in np.linspace(0, nt - clip_len, self.num_clips)]
        clips = [self.transform(x[ss : ss + clip_len], fps) for ss in start_times]
        return clips


class MultiResizeCropFlip:
    def __init__(
        self,
        num_augm=2,
        num_frames=8,
        min_size=256,
        max_size=360,
        crop_size=224,
        augment=True,
    ):
        self.t = ResizeCropFlip(num_frames, min_size, max_size, crop_size, augment)
        self.num_augm = num_augm

    def __call__(self, x):
        return [self.t(x) for _ in range(self.num_augm)]

upstream_models/replai/replai/data/video_transforms.py
roger-tseng/av-superb
from .builder import build_audio_model, build_audio_video_model, build_video_model

upstream_models/replai/replai/models/__init__.py
roger-tseng/av-superb
from .audio_models_benchmarking import *
from .avid_audio import *
from .spec_cnn import *

upstream_models/replai/replai/models/audio/__init__.py
roger-tseng/av-superb
import os

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

__all__ = [
    "Cnn6",
    "Cnn10",
    "Cnn14",
    "ResNet22",
    "ResNet38",
    "ResNet54",
    "MobileNetV1",
    "MobileNetV2",
    "Res1dNet31",
    "Res1dNet51",
    "Wavegram_Cnn14",
    "Wavegram_Logmel_Cnn14",
]


def init_layer(layer):
    """Initialize a Linear or Convolutional layer."""
    nn.init.xavier_uniform_(layer.weight)

    if hasattr(layer, "bias"):
        if layer.bias is not None:
            layer.bias.data.fill_(0.0)


def init_bn(bn):
    """Initialize a Batchnorm layer."""
    bn.bias.data.fill_(0.0)
    bn.weight.data.fill_(1.0)


def do_mixup(x, mixup_lambda):
    """Mixup x of even indexes (0, 2, 4, ...) with x of odd indexes
    (1, 3, 5, ...).

    Args:
      x: (batch_size * 2, ...)
      mixup_lambda: (batch_size * 2,)

    Returns:
      out: (batch_size, ...)
    """
    out = (
        x[0::2].transpose(0, -1) * mixup_lambda[0::2]
        + x[1::2].transpose(0, -1) * mixup_lambda[1::2]
    ).transpose(0, -1)
    return out


def interpolate(x, ratio):
    """Interpolate data in time domain. This is used to compensate the
    resolution reduction in downsampling of a CNN.

    Args:
      x: (batch_size, time_steps, classes_num)
      ratio: int, ratio to interpolate

    Returns:
      upsampled: (batch_size, time_steps * ratio, classes_num)
    """
    (batch_size, time_steps, classes_num) = x.shape
    upsampled = x[:, :, None, :].repeat(1, 1, ratio, 1)
    upsampled = upsampled.reshape(batch_size, time_steps * ratio, classes_num)
    return upsampled


def pad_framewise_output(framewise_output, frames_num):
    """Pad framewise_output to the same length as input frames. The pad value
    is the same as the value of the last frame.

    Args:
      framewise_output: (batch_size, frames_num, classes_num)
      frames_num: int, number of frames to pad

    Outputs:
      output: (batch_size, frames_num, classes_num)
    """
    pad = framewise_output[:, -1:, :].repeat(
        1, frames_num - framewise_output.shape[1], 1
    )
    """tensor for padding"""

    output = torch.cat((framewise_output, pad), dim=1)
    """(batch_size, frames_num, classes_num)"""

    return output


class ConvBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(ConvBlock, self).__init__()

        self.conv1 = nn.Conv2d(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=(3, 3),
            stride=(1, 1),
            padding=(1, 1),
            bias=False,
        )

        self.conv2 = nn.Conv2d(
            in_channels=out_channels,
            out_channels=out_channels,
            kernel_size=(3, 3),
            stride=(1, 1),
            padding=(1, 1),
            bias=False,
        )

        self.bn1 = nn.BatchNorm2d(out_channels)
        self.bn2 = nn.BatchNorm2d(out_channels)

        self.init_weight()

    def init_weight(self):
        init_layer(self.conv1)
        init_layer(self.conv2)
        init_bn(self.bn1)
        init_bn(self.bn2)

    def forward(self, input, pool_size=(2, 2), pool_type="avg"):
        x = input
        x = F.relu_(self.bn1(self.conv1(x)))
        x = F.relu_(self.bn2(self.conv2(x)))
        if pool_type == "max":
            x = F.max_pool2d(x, kernel_size=pool_size)
        elif pool_type == "avg":
            x = F.avg_pool2d(x, kernel_size=pool_size)
        elif pool_type == "avg+max":
            x1 = F.avg_pool2d(x, kernel_size=pool_size)
            x2 = F.max_pool2d(x, kernel_size=pool_size)
            x = x1 + x2
        else:
            raise Exception("Incorrect argument!")

        return x


class ConvBlock5x5(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(ConvBlock5x5, self).__init__()

        self.conv1 = nn.Conv2d(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=(5, 5),
            stride=(1, 1),
            padding=(2, 2),
            bias=False,
        )

        self.bn1 = nn.BatchNorm2d(out_channels)

        self.init_weight()

    def init_weight(self):
        init_layer(self.conv1)
        init_bn(self.bn1)

    def forward(self, input, pool_size=(2, 2), pool_type="avg"):
        x = input
        x = F.relu_(self.bn1(self.conv1(x)))
        if pool_type == "max":
            x = F.max_pool2d(x, kernel_size=pool_size)
        elif pool_type == "avg":
            x = F.avg_pool2d(x, kernel_size=pool_size)
        elif pool_type == "avg+max":
            x1 = F.avg_pool2d(x, kernel_size=pool_size)
            x2 = F.max_pool2d(x, kernel_size=pool_size)
            x = x1 + x2
        else:
            raise Exception("Incorrect argument!")

        return x


class AttBlock(nn.Module):
    def __init__(self, n_in, n_out, activation="linear", temperature=1.0):
        super(AttBlock, self).__init__()

        self.activation = activation
        self.temperature = temperature
        self.att = nn.Conv1d(
            in_channels=n_in,
            out_channels=n_out,
            kernel_size=(1,),
            stride=(1,),
            padding=0,
            bias=True,
        )
        self.cla = nn.Conv1d(
            in_channels=n_in,
            out_channels=n_out,
            kernel_size=(1,),
            stride=(1,),
            padding=0,
            bias=True,
        )

        self.bn_att = nn.BatchNorm1d(n_out)
        self.init_weights()

    def init_weights(self):
        init_layer(self.att)
        init_layer(self.cla)
        init_bn(self.bn_att)

    def forward(self, x):
        # x: (n_samples, n_in, n_time)
        norm_att = torch.softmax(torch.clamp(self.att(x), -10, 10), dim=-1)
        cla = self.nonlinear_transform(self.cla(x))
        x = torch.sum(norm_att * cla, dim=2)
        return x, norm_att, cla

    def nonlinear_transform(self, x):
        if self.activation == "linear":
            return x
        elif self.activation == "sigmoid":
            return torch.sigmoid(x)


class Cnn6(nn.Module):
    def __init__(self, classes_num=527, pretrained=False):
        super(Cnn6, self).__init__()

        self.conv_block1 = ConvBlock5x5(in_channels=1, out_channels=64)
        self.conv_block2 = ConvBlock5x5(in_channels=64, out_channels=128)
        self.conv_block3 = ConvBlock5x5(in_channels=128, out_channels=256)
        self.conv_block4 = ConvBlock5x5(in_channels=256, out_channels=512)

        self.fc1 = nn.Linear(512, 512, bias=True)
        self.fc_audioset = nn.Linear(512, classes_num, bias=True)

        self.init_weight()

        if pretrained:
            import urllib.request

            ckpt_fn = "models/pretrained/audio_benchmarking/Cnn6_mAP_0.343.pth"
            url = "https://zenodo.org/record/3987831/files/Cnn6_mAP%3D0.343.pth"
            if not os.path.isfile(ckpt_fn):
                os.makedirs(os.path.dirname(ckpt_fn), exist_ok=True)
                urllib.request.urlretrieve(url, ckpt_fn)
            self.load_state_dict(torch.load(ckpt_fn, map_location="cpu"))

    def init_weight(self):
        init_layer(self.fc1)
        init_layer(self.fc_audioset)

    def load_state_dict(self, state_dict, strict: bool = True):
        for k in state_dict:
            if (
                "spectrogram_extractor" in k
                or "logmel_extractor" in k
                or "spec_augmenter" in k
                or "bn0" in k
            ):
                del state_dict[k]
        super().load_state_dict(state_dict=state_dict, strict=strict)

    def forward(self, x):
        """
        Input: (batch_size, data_length)"""
        x = self.conv_block1(x, pool_size=(2, 2), pool_type="avg")
        x = F.dropout(x, p=0.2, training=self.training)
        x = self.conv_block2(x, pool_size=(2, 2), pool_type="avg")
        x = F.dropout(x, p=0.2, training=self.training)
        x = self.conv_block3(x, pool_size=(2, 2), pool_type="avg")
        x = F.dropout(x, p=0.2, training=self.training)
        x = self.conv_block4(x, pool_size=(2, 2), pool_type="avg")
        x = F.dropout(x, p=0.2, training=self.training)
        x = torch.mean(x, dim=3)

        (x1, _) = torch.max(x, dim=2)
        x2 = torch.mean(x, dim=2)
        x = x1 + x2
        x = F.dropout(x, p=0.5, training=self.training)
        x = F.relu_(self.fc1(x))
        x = self.fc_audioset(x)
        return x


class Cnn10(nn.Module):
    def __init__(self, classes_num=527, pretrained=False):
        super(Cnn10, self).__init__()
        self.conv_block1 = ConvBlock(in_channels=1, out_channels=64)
        self.conv_block2 = ConvBlock(in_channels=64, out_channels=128)
        self.conv_block3 = ConvBlock(in_channels=128, out_channels=256)
        self.conv_block4 = ConvBlock(in_channels=256, out_channels=512)

        self.fc1 = nn.Linear(512, 512, bias=True)
        self.fc_audioset = nn.Linear(512, classes_num, bias=True)

        self.init_weight()

        if pretrained:
            import urllib.request

            ckpt_fn = "models/pretrained/audio_benchmarking/Cnn10_mAP_0.380.pth"
            url = "https://zenodo.org/record/3987831/files/Cnn10_mAP%3D0.380.pth"
            if not os.path.isfile(ckpt_fn):
                os.makedirs(os.path.dirname(ckpt_fn), exist_ok=True)
                urllib.request.urlretrieve(url, ckpt_fn)
            self.load_state_dict(torch.load(ckpt_fn, map_location="cpu"))

    def load_state_dict(self, state_dict, strict: bool = True):
        for k in state_dict:
            if (
                "spectrogram_extractor" in k
                or "logmel_extractor" in k
                or "spec_augmenter" in k
                or "bn0" in k
            ):
                del state_dict[k]
        super().load_state_dict(state_dict=state_dict, strict=strict)

    def init_weight(self):
        init_layer(self.fc1)
        init_layer(self.fc_audioset)

    def forward(self, x):
        """
        Input: (batch_size, data_length)"""
        x = self.conv_block1(x, pool_size=(2, 2), pool_type="avg")
        x = F.dropout(x, p=0.2, training=self.training)
        x = self.conv_block2(x, pool_size=(2, 2), pool_type="avg")
        x = F.dropout(x, p=0.2, training=self.training)
        x = self.conv_block3(x, pool_size=(2, 2), pool_type="avg")
        x = F.dropout(x, p=0.2, training=self.training)
        x = self.conv_block4(x, pool_size=(2, 2), pool_type="avg")
        x = F.dropout(x, p=0.2, training=self.training)
        x = torch.mean(x, dim=3)

        (x1, _) = torch.max(x, dim=2)
        x2 = torch.mean(x, dim=2)
        x = x1 + x2
        x = F.dropout(x, p=0.5, training=self.training)
        x = F.relu_(self.fc1(x))
        x = self.fc_audioset(x)
        return x


class Cnn14(nn.Module):
    def __init__(self, classes_num=527, pretrained=False):
        super(Cnn14, self).__init__()

        self.conv_block1 = ConvBlock(in_channels=1, out_channels=64)
        self.conv_block2 = ConvBlock(in_channels=64, out_channels=128)
        self.conv_block3 = ConvBlock(in_channels=128, out_channels=256)
        self.conv_block4 = ConvBlock(in_channels=256, out_channels=512)
        self.conv_block5 = ConvBlock(in_channels=512, out_channels=1024)
        self.conv_block6 = ConvBlock(in_channels=1024, out_channels=2048)

        self.fc1 = nn.Linear(2048, 2048, bias=True)
        self.fc_audioset = nn.Linear(2048, classes_num, bias=True)

        self.init_weight()

        if pretrained:
            import urllib.request

            ckpt_fn = "models/pretrained/audio_benchmarking/Cnn14_mAP_0.431.pth"
            url = "https://zenodo.org/record/3987831/files/Cnn14_mAP%3D0.431.pth"
            if not os.path.isfile(ckpt_fn):
                os.makedirs(os.path.dirname(ckpt_fn), exist_ok=True)
                urllib.request.urlretrieve(url, ckpt_fn)
            self.load_state_dict(torch.load(ckpt_fn, map_location="cpu"))

    def load_state_dict(self, state_dict, strict: bool = True):
        for k in state_dict:
            if (
                "spectrogram_extractor" in k
                or "logmel_extractor" in k
                or "spec_augmenter" in k
                or "bn0" in k
            ):
                del state_dict[k]
        super().load_state_dict(state_dict=state_dict, strict=strict)

    def init_weight(self):
        init_layer(self.fc1)
        init_layer(self.fc_audioset)

    def forward(self, x):
        """
        Input: (batch_size, data_length)"""
        x = self.conv_block1(x, pool_size=(2, 2), pool_type="avg")
        x = F.dropout(x, p=0.2, training=self.training)
        x = self.conv_block2(x, pool_size=(2, 2), pool_type="avg")
        x = F.dropout(x, p=0.2, training=self.training)
        x = self.conv_block3(x, pool_size=(2, 2), pool_type="avg")
        x = F.dropout(x, p=0.2, training=self.training)
        x = self.conv_block4(x, pool_size=(2, 2), pool_type="avg")
        x = F.dropout(x, p=0.2, training=self.training)
        x = self.conv_block5(x, pool_size=(2, 2), pool_type="avg")
        x = F.dropout(x, p=0.2, training=self.training)
        x = self.conv_block6(x, pool_size=(1, 1), pool_type="avg")
        x = F.dropout(x, p=0.2, training=self.training)
        x = torch.mean(x, dim=3)

        (x1, _) = torch.max(x, dim=2)
        x2 = torch.mean(x, dim=2)
        x = x1 + x2
        x = F.dropout(x, p=0.5, training=self.training)
        x = F.relu_(self.fc1(x))
        x = self.fc_audioset(x)
        return x


def _resnet_conv3x3(in_planes, out_planes):
    # 3x3 convolution with padding
    return nn.Conv2d(
        in_planes,
        out_planes,
        kernel_size=(3,),
        stride=(1,),
        padding=1,
        groups=1,
        bias=False,
        dilation=(1,),
    )


def _resnet_conv1x1(in_planes, out_planes):
    # 1x1 convolution
    return nn.Conv2d(in_planes, out_planes, kernel_size=(1,), stride=(1,), bias=False)


class _ResnetBasicBlock(nn.Module):
    expansion = 1

    def __init__(
        self,
        inplanes,
        planes,
        stride=1,
        downsample=None,
        groups=1,
        base_width=64,
        dilation=1,
        norm_layer=None,
    ):
        super(_ResnetBasicBlock, self).__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        if groups != 1 or base_width != 64:
            raise ValueError(
                "_ResnetBasicBlock only supports groups=1 and base_width=64"
            )
        if dilation > 1:
            raise NotImplementedError("Dilation > 1 not supported in _ResnetBasicBlock")
        # Both self.conv1 and self.downsample layers downsample the input when stride != 1

        self.stride = stride

        self.conv1 = _resnet_conv3x3(inplanes, planes)
        self.bn1 = norm_layer(planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = _resnet_conv3x3(planes, planes)
        self.bn2 = norm_layer(planes)
        self.downsample = downsample
        self.stride = stride

        self.init_weights()

    def init_weights(self):
        init_layer(self.conv1)
        init_bn(self.bn1)
        init_layer(self.conv2)
        init_bn(self.bn2)
        nn.init.constant_(self.bn2.weight, 0)

    def forward(self, x):
        identity = x

        if self.stride == 2:
            out = F.avg_pool2d(x, kernel_size=(2, 2))
        else:
            out = x

        out = self.conv1(out)
        out = self.bn1(out)
        out = self.relu(out)
        out = F.dropout(out, p=0.1, training=self.training)

        out = self.conv2(out)
        out = self.bn2(out)

        if self.downsample is not None:
            identity = self.downsample(identity)

        out += identity
        out = self.relu(out)

        return out


class _ResnetBottleneck(nn.Module):
    expansion = 4

    def __init__(
        self,
        inplanes,
        planes,
        stride=1,
        downsample=None,
        groups=1,
        base_width=64,
        dilation=1,
        norm_layer=None,
    ):
        super(_ResnetBottleneck, self).__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        width = int(planes * (base_width / 64.0)) * groups
        self.stride = stride
        # Both self.conv2 and self.downsample layers downsample the input when stride != 1
        self.conv1 = _resnet_conv1x1(inplanes, width)
        self.bn1 = norm_layer(width)
        self.conv2 = _resnet_conv3x3(width, width)
        self.bn2 = norm_layer(width)
        self.conv3 = _resnet_conv1x1(width, planes * self.expansion)
        self.bn3 = norm_layer(planes * self.expansion)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

        self.init_weights()

    def init_weights(self):
        init_layer(self.conv1)
        init_bn(self.bn1)
        init_layer(self.conv2)
        init_bn(self.bn2)
        init_layer(self.conv3)
        init_bn(self.bn3)
        nn.init.constant_(self.bn3.weight, 0)

    def forward(self, x):
        identity = x

        if self.stride == 2:
            x = F.avg_pool2d(x, kernel_size=(2, 2))

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)
        out = F.dropout(out, p=0.1, training=self.training)

        out = self.conv3(out)
        out = self.bn3(out)

        if self.downsample is not None:
            identity = self.downsample(identity)

        out += identity
        out = self.relu(out)

        return out


class _ResNet(nn.Module):
    def __init__(
        self,
        block,
        layers,
        zero_init_residual=False,
        groups=1,
        width_per_group=64,
        replace_stride_with_dilation=None,
        norm_layer=None,
    ):
        super(_ResNet, self).__init__()

        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        self._norm_layer = norm_layer

        self.inplanes = 64
        self.dilation = 1
        if replace_stride_with_dilation is None:
            # each element in the tuple indicates if we should replace
            # the 2x2 stride with a dilated convolution instead
            replace_stride_with_dilation = [False, False, False]
        if len(replace_stride_with_dilation) != 3:
            raise ValueError(
                "replace_stride_with_dilation should be None "
                "or a 3-element tuple, got {}".format(replace_stride_with_dilation)
            )
        self.groups = groups
        self.base_width = width_per_group

        self.layer1 = self._make_layer(block, 64, layers[0], stride=1)
        self.layer2 = self._make_layer(
            block, 128, layers[1], stride=2, dilate=replace_stride_with_dilation[0]
        )
        self.layer3 = self._make_layer(
            block, 256, layers[2], stride=2, dilate=replace_stride_with_dilation[1]
        )
        self.layer4 = self._make_layer(
            block, 512, layers[3], stride=2, dilate=replace_stride_with_dilation[2]
        )

    def _make_layer(self, block, planes, blocks, stride=1, dilate=False):
        norm_layer = self._norm_layer
        downsample = None
        previous_dilation = self.dilation
        if dilate:
            self.dilation *= stride
            stride = 1
        if stride != 1 or self.inplanes != planes * block.expansion:
            if stride == 1:
                downsample = nn.Sequential(
                    _resnet_conv1x1(self.inplanes, planes * block.expansion),
                    norm_layer(planes * block.expansion),
                )
                init_layer(downsample[0])
                init_bn(downsample[1])
            elif stride == 2:
                downsample = nn.Sequential(
                    nn.AvgPool2d(kernel_size=2),
                    _resnet_conv1x1(self.inplanes, planes * block.expansion),
                    norm_layer(planes * block.expansion),
                )
                init_layer(downsample[1])
                init_bn(downsample[2])

        layers = []
        layers.append(
            block(
                self.inplanes,
                planes,
                stride,
                downsample,
                self.groups,
                self.base_width,
                previous_dilation,
                norm_layer,
            )
        )
        self.inplanes = planes * block.expansion
        for _ in range(1, blocks):
            layers.append(
                block(
                    self.inplanes,
                    planes,
                    groups=self.groups,
                    base_width=self.base_width,
                    dilation=self.dilation,
                    norm_layer=norm_layer,
                )
            )

        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        return x


class ResNet22(nn.Module):
    def __init__(self, classes_num=527, pretrained=False):
        super(ResNet22, self).__init__()

        self.conv_block1 = ConvBlock(in_channels=1, out_channels=64)
        # self.conv_block2 = ConvBlock(in_channels=64, out_channels=64)

        self.resnet = _ResNet(
            block=_ResnetBasicBlock, layers=[2, 2, 2, 2], zero_init_residual=True
        )

        self.conv_block_after1 = ConvBlock(in_channels=512, out_channels=2048)

        self.fc1 = nn.Linear(2048, 2048)
        self.fc_audioset = nn.Linear(2048, classes_num, bias=True)

        self.init_weights()

        if pretrained:
            import urllib.request

            ckpt_fn = "models/pretrained/audio_benchmarking/ResNet22_mAP_0.430.pth"
            url = "https://zenodo.org/record/3987831/files/ResNet22_mAP%3D0.430.pth"
            if not os.path.isfile(ckpt_fn):
                os.makedirs(os.path.dirname(ckpt_fn), exist_ok=True)
                urllib.request.urlretrieve(url, ckpt_fn)
            self.load_state_dict(torch.load(ckpt_fn, map_location="cpu"))

    def load_state_dict(self, state_dict, strict: bool = True):
        for k in state_dict:
            if (
                "spectrogram_extractor" in k
                or "logmel_extractor" in k
                or "spec_augmenter" in k
                or "bn0" in k
            ):
                del state_dict[k]
        super().load_state_dict(state_dict=state_dict, strict=strict)

    def init_weights(self):
        init_layer(self.fc1)
        init_layer(self.fc_audioset)

    def forward(self, x):
        """
        Input: (batch_size, data_length)"""
        x = self.conv_block1(x, pool_size=(2, 2), pool_type="avg")
        x = F.dropout(x, p=0.2, training=self.training, inplace=True)
        x = self.resnet(x)
        x = F.avg_pool2d(x, kernel_size=(2, 2))
        x = F.dropout(x, p=0.2, training=self.training, inplace=True)
        x = self.conv_block_after1(x, pool_size=(1, 1), pool_type="avg")
        x = F.dropout(x, p=0.2, training=self.training, inplace=True)
        x = torch.mean(x, dim=3)

        (x1, _) = torch.max(x, dim=2)
        x2 = torch.mean(x, dim=2)
        x = x1 + x2
        x = F.dropout(x, p=0.5, training=self.training)
        x = F.relu_(self.fc1(x))
        x = self.fc_audioset(x)
        return x


class ResNet38(nn.Module):
    def __init__(self, classes_num=527, pretrained=False):
        super(ResNet38, self).__init__()

        self.conv_block1 = ConvBlock(in_channels=1, out_channels=64)
        # self.conv_block2 = ConvBlock(in_channels=64, out_channels=64)

        self.resnet = _ResNet(
            block=_ResnetBasicBlock, layers=[3, 4, 6, 3], zero_init_residual=True
        )

        self.conv_block_after1 = ConvBlock(in_channels=512, out_channels=2048)

        self.fc1 = nn.Linear(2048, 2048)
        self.fc_audioset = nn.Linear(2048, classes_num, bias=True)

        self.init_weights()

        if pretrained:
            import urllib.request

            ckpt_fn = "models/pretrained/audio_benchmarking/ResNet38_mAP_0.434.pth"
            url = "https://zenodo.org/record/3987831/files/ResNet38_mAP%3D0.434.pth"
            if not os.path.isfile(ckpt_fn):
                os.makedirs(os.path.dirname(ckpt_fn), exist_ok=True)
                urllib.request.urlretrieve(url, ckpt_fn)
            self.load_state_dict(torch.load(ckpt_fn, map_location="cpu"))

    def load_state_dict(self, state_dict, strict: bool = True):
        for k in state_dict:
            if (
                "spectrogram_extractor" in k
                or "logmel_extractor" in k
                or "spec_augmenter" in k
                or "bn0" in k
            ):
                del state_dict[k]
        super().load_state_dict(state_dict=state_dict, strict=strict)

    def init_weights(self):
        init_layer(self.fc1)
        init_layer(self.fc_audioset)

    def forward(self, x):
        """
        Input: (batch_size, data_length)"""

        x = self.conv_block1(x, pool_size=(2, 2), pool_type="avg")
        x = F.dropout(x, p=0.2, training=self.training, inplace=True)
        x = self.resnet(x)
        x = F.avg_pool2d(x, kernel_size=(2, 2))
        x = F.dropout(x, p=0.2, training=self.training, inplace=True)
        x = self.conv_block_after1(x, pool_size=(1, 1), pool_type="avg")
        x = F.dropout(x, p=0.2, training=self.training, inplace=True)
        x = torch.mean(x, dim=3)

        (x1, _) = torch.max(x, dim=2)
        x2 = torch.mean(x, dim=2)
        x = x1 + x2
        x = F.dropout(x, p=0.5, training=self.training)
        x = F.relu_(self.fc1(x))
        x = self.fc_audioset(x)
        return x


class ResNet54(nn.Module):
    def __init__(self, classes_num=527, pretrained=False):
        super(ResNet54, self).__init__()

        self.conv_block1 = ConvBlock(in_channels=1, out_channels=64)
        # self.conv_block2 = ConvBlock(in_channels=64, out_channels=64)

        self.resnet = _ResNet(
            block=_ResnetBottleneck, layers=[3, 4, 6, 3], zero_init_residual=True
        )

        self.conv_block_after1 = ConvBlock(in_channels=2048, out_channels=2048)

        self.fc1 = nn.Linear(2048, 2048)
        self.fc_audioset = nn.Linear(2048, classes_num, bias=True)

        self.init_weights()

        if pretrained:
            import urllib.request

            ckpt_fn = "models/pretrained/audio_benchmarking/ResNet54_mAP_0.429.pth"
            url = "https://zenodo.org/record/3987831/files/ResNet54_mAP%3D0.429.pth"
            if not os.path.isfile(ckpt_fn):
                os.makedirs(os.path.dirname(ckpt_fn), exist_ok=True)
                urllib.request.urlretrieve(url, ckpt_fn)
            self.load_state_dict(torch.load(ckpt_fn, map_location="cpu"))

    def load_state_dict(self, state_dict, strict: bool = True):
        for k in state_dict:
            if (
                "spectrogram_extractor" in k
                or "logmel_extractor" in k
                or "spec_augmenter" in k
                or "bn0" in k
            ):
                del state_dict[k]
        super().load_state_dict(state_dict=state_dict, strict=strict)

    def init_weights(self):
        init_layer(self.fc1)
        init_layer(self.fc_audioset)

    def forward(self, x):
        """
        Input: (batch_size, data_length)"""

        x = self.conv_block1(x, pool_size=(2, 2), pool_type="avg")
        x = F.dropout(x, p=0.2, training=self.training, inplace=True)
        x = self.resnet(x)
        x = F.avg_pool2d(x, kernel_size=(2, 2))
        x = F.dropout(x, p=0.2, training=self.training, inplace=True)
        x = self.conv_block_after1(x, pool_size=(1, 1), pool_type="avg")
        x = F.dropout(x, p=0.2, training=self.training, inplace=True)
        x = torch.mean(x, dim=3)

        (x1, _) = torch.max(x, dim=2)
        x2 = torch.mean(x, dim=2)
        x = x1 + x2
        x = F.dropout(x, p=0.5, training=self.training)
        x = F.relu_(self.fc1(x))
        x = self.fc_audioset(x)
        return x


class MobileNetV1(nn.Module):
    def __init__(self, classes_num=527, pretrained=False):
        super(MobileNetV1, self).__init__()

        def conv_bn(inp, oup, stride):
            _layers = [
                nn.Conv2d(inp, oup, 3, 1, 1, bias=False),
                nn.AvgPool2d(stride),
                nn.BatchNorm2d(oup),
                nn.ReLU(inplace=True),
            ]
            _layers = nn.Sequential(*_layers)
            init_layer(_layers[0])
            init_bn(_layers[2])
            return _layers

        def conv_dw(inp, oup, stride):
            _layers = [
                nn.Conv2d(inp, inp, 3, 1, 1, groups=inp, bias=False),
                nn.AvgPool2d(stride),
                nn.BatchNorm2d(inp),
                nn.ReLU(inplace=True),
                nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
                nn.BatchNorm2d(oup),
                nn.ReLU(inplace=True),
            ]
            _layers = nn.Sequential(*_layers)
            init_layer(_layers[0])
            init_bn(_layers[2])
            init_layer(_layers[4])
            init_bn(_layers[5])
            return _layers

        self.features = nn.Sequential(
            conv_bn(1, 32, 2),
            conv_dw(32, 64, 1),
            conv_dw(64, 128, 2),
            conv_dw(128, 128, 1),
            conv_dw(128, 256, 2),
            conv_dw(256, 256, 1),
            conv_dw(256, 512, 2),
            conv_dw(512, 512, 1),
            conv_dw(512, 512, 1),
            conv_dw(512, 512, 1),
            conv_dw(512, 512, 1),
            conv_dw(512, 512, 1),
            conv_dw(512, 1024, 2),
            conv_dw(1024, 1024, 1),
        )

        self.fc1 = nn.Linear(1024, 1024, bias=True)
        self.fc_audioset = nn.Linear(1024, classes_num, bias=True)

        self.init_weights()

        if pretrained:
            import urllib.request

            ckpt_fn = "models/pretrained/audio_benchmarking/MobileNetV1_mAP_0.389.pth"
            url = "https://zenodo.org/record/3987831/files/MobileNetV1_mAP%3D0.389.pth"
            if not os.path.isfile(ckpt_fn):
                os.makedirs(os.path.dirname(ckpt_fn), exist_ok=True)
                urllib.request.urlretrieve(url, ckpt_fn)
            self.load_state_dict(torch.load(ckpt_fn, map_location="cpu"))

    def load_state_dict(self, state_dict, strict: bool = True):
        for k in state_dict:
            if (
                "spectrogram_extractor" in k
                or "logmel_extractor" in k
                or "spec_augmenter" in k
                or "bn0" in k
            ):
                del state_dict[k]
        super().load_state_dict(state_dict=state_dict, strict=strict)

    def init_weights(self):
        init_layer(self.fc1)
        init_layer(self.fc_audioset)

    def forward(self, x):
        """
        Input: (batch_size, data_length)"""

        x = self.features(x)
        x = torch.mean(x, dim=3)

        (x1, _) = torch.max(x, dim=2)
        x2 = torch.mean(x, dim=2)
        x = x1 + x2
        x = F.dropout(x, p=0.5, training=self.training)
        x = F.relu_(self.fc1(x))
        x = self.fc_audioset(x)
        return x


class InvertedResidual(nn.Module):
    def __init__(self, inp, oup, stride, expand_ratio):
        super(InvertedResidual, self).__init__()
        self.stride = stride
        assert stride in [1, 2]

        hidden_dim = round(inp * expand_ratio)
        self.use_res_connect = self.stride == 1 and inp == oup

        if expand_ratio == 1:
            _layers = [
                nn.Conv2d(
                    hidden_dim, hidden_dim, 3, 1, 1, groups=hidden_dim, bias=False
                ),
                nn.AvgPool2d(stride),
                nn.BatchNorm2d(hidden_dim),
                nn.ReLU6(inplace=True),
                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
                nn.BatchNorm2d(oup),
            ]
            _layers = nn.Sequential(*_layers)
            init_layer(_layers[0])
            init_bn(_layers[2])
            init_layer(_layers[4])
            init_bn(_layers[5])
            self.conv = _layers
        else:
            _layers = [
                nn.Conv2d(inp, hidden_dim, 1, 1, 0, bias=False),
                nn.BatchNorm2d(hidden_dim),
                nn.ReLU6(inplace=True),
                nn.Conv2d(
                    hidden_dim, hidden_dim, 3, 1, 1, groups=hidden_dim, bias=False
                ),
                nn.AvgPool2d(stride),
                nn.BatchNorm2d(hidden_dim),
                nn.ReLU6(inplace=True),
                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
                nn.BatchNorm2d(oup),
            ]
            _layers = nn.Sequential(*_layers)
            init_layer(_layers[0])
            init_bn(_layers[1])
            init_layer(_layers[3])
            init_bn(_layers[5])
            init_layer(_layers[7])
            init_bn(_layers[8])
            self.conv = _layers

    def forward(self, x):
        if self.use_res_connect:
            return x + self.conv(x)
        else:
            return self.conv(x)


class MobileNetV2(nn.Module):
    def __init__(self, classes_num=527, pretrained=False):
        super(MobileNetV2, self).__init__()

        width_mult = 1.0
        block = InvertedResidual
        input_channel = 32
        last_channel = 1280
        interverted_residual_setting = [
            # t, c, n, s
            [1, 16, 1, 1],
            [6, 24, 2, 2],
            [6, 32, 3, 2],
            [6, 64, 4, 2],
            [6, 96, 3, 2],
            [6, 160, 3, 1],
            [6, 320, 1, 1],
        ]

        def conv_bn(inp, oup, stride):
            _layers = [
                nn.Conv2d(inp, oup, 3, 1, 1, bias=False),
                nn.AvgPool2d(stride),
                nn.BatchNorm2d(oup),
                nn.ReLU6(inplace=True),
            ]
            _layers = nn.Sequential(*_layers)
            init_layer(_layers[0])
            init_bn(_layers[2])
            return _layers

        def conv_1x1_bn(inp, oup):
            _layers = nn.Sequential(
                nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
                nn.BatchNorm2d(oup),
                nn.ReLU6(inplace=True),
            )
            init_layer(_layers[0])
            init_bn(_layers[1])
            return _layers

        # building first layer
        input_channel = int(input_channel * width_mult)
        self.last_channel = (
            int(last_channel * width_mult) if width_mult > 1.0 else last_channel
        )
        self.features = [conv_bn(1, input_channel, 2)]
        # building inverted residual blocks
        for t, c, n, s in interverted_residual_setting:
            output_channel = int(c * width_mult)
            for i in range(n):
                if i == 0:
                    self.features.append(
                        block(input_channel, output_channel, s, expand_ratio=t)
                    )
                else:
                    self.features.append(
                        block(input_channel, output_channel, 1, expand_ratio=t)
                    )
                input_channel = output_channel
        # building last several layers
        self.features.append(conv_1x1_bn(input_channel, self.last_channel))
        # make it nn.Sequential
        self.features = nn.Sequential(*self.features)

        self.fc1 = nn.Linear(1280, 1024, bias=True)
        self.fc_audioset = nn.Linear(1024, classes_num, bias=True)

        self.init_weight()

        if pretrained:
            import urllib.request

            ckpt_fn = "models/pretrained/audio_benchmarking/MobileNetV2_mAP_0.383.pth"
            url = "https://zenodo.org/record/3987831/files/MobileNetV2_mAP%3D0.383.pth"
            if not os.path.isfile(ckpt_fn):
                os.makedirs(os.path.dirname(ckpt_fn), exist_ok=True)
                urllib.request.urlretrieve(url, ckpt_fn)
            self.load_state_dict(torch.load(ckpt_fn, map_location="cpu"))

    def load_state_dict(self, state_dict, strict: bool = True):
        for k in state_dict:
            if (
                "spectrogram_extractor" in k
                or "logmel_extractor" in k
                or "spec_augmenter" in k
                or "bn0" in k
            ):
                del state_dict[k]
        super().load_state_dict(state_dict=state_dict, strict=strict)

    def init_weight(self):
        init_layer(self.fc1)
        init_layer(self.fc_audioset)

    def forward(self, x):
        """
        Input: (batch_size, data_length)"""

        x = self.features(x)

        x = torch.mean(x, dim=3)

        (x1, _) = torch.max(x, dim=2)
        x2 = torch.mean(x, dim=2)
        x = x1 + x2
        # x = F.dropout(x, p=0.5, training=self.training)
        x = F.relu_(self.fc1(x))
        x = self.fc_audioset(x)
        return x


def _resnet_conv3x1_wav1d(in_planes, out_planes, dilation):
    # 3x3 convolution with padding
    return nn.Conv1d(
        in_planes,
        out_planes,
        kernel_size=(3,),
        stride=(1,),
        padding=dilation,
        groups=1,
        bias=False,
        dilation=dilation,
    )


def _resnet_conv1x1_wav1d(in_planes, out_planes):
    # 1x1 convolution
    return nn.Conv1d(in_planes, out_planes, kernel_size=(1,), stride=(1,), bias=False)


class _ResnetBasicBlockWav1d(nn.Module):
    expansion = 1

    def __init__(
        self,
        inplanes,
        planes,
        stride=1,
        downsample=None,
        groups=1,
        base_width=64,
        dilation=1,
        norm_layer=None,
    ):
        super(_ResnetBasicBlockWav1d, self).__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm1d
        if groups != 1 or base_width != 64:
            raise ValueError(
                "_ResnetBasicBlock only supports groups=1 and base_width=64"
            )
        if dilation > 1:
            raise NotImplementedError("Dilation > 1 not supported in _ResnetBasicBlock")
        # Both self.conv1 and self.downsample layers downsample the input when stride != 1

        self.stride = stride

        self.conv1 = _resnet_conv3x1_wav1d(inplanes, planes, dilation=1)
        self.bn1 = norm_layer(planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = _resnet_conv3x1_wav1d(planes, planes, dilation=2)
        self.bn2 = norm_layer(planes)
        self.downsample = downsample
        self.stride = stride

        self.init_weights()

    def init_weights(self):
        init_layer(self.conv1)
        init_bn(self.bn1)
        init_layer(self.conv2)
        init_bn(self.bn2)
        nn.init.constant_(self.bn2.weight, 0)

    def forward(self, x):
        identity = x

        if self.stride != 1:
            out = F.max_pool1d(x, kernel_size=self.stride)
        else:
            out = x

        out = self.conv1(out)
        out = self.bn1(out)
        out = self.relu(out)
        out = F.dropout(out, p=0.1, training=self.training)

        out = self.conv2(out)
        out = self.bn2(out)

        if self.downsample is not None:
            identity = self.downsample(identity)

        out += identity
        out = self.relu(out)

        return out


class _ResNetWav1d(nn.Module):
    def __init__(
        self,
        block,
        layers,
        zero_init_residual=False,
        groups=1,
        width_per_group=64,
        replace_stride_with_dilation=None,
        norm_layer=None,
    ):
        super(_ResNetWav1d, self).__init__()

        if norm_layer is None:
            norm_layer = nn.BatchNorm1d
        self._norm_layer = norm_layer

        self.inplanes = 64
        self.dilation = 1
        if replace_stride_with_dilation is None:
            # each element in the tuple indicates if we should replace
            # the 2x2 stride with a dilated convolution instead
            replace_stride_with_dilation = [False, False, False]
        if len(replace_stride_with_dilation) != 3:
            raise ValueError(
                "replace_stride_with_dilation should be None "
                "or a 3-element tuple, got {}".format(replace_stride_with_dilation)
            )
        self.groups = groups
        self.base_width = width_per_group

        self.layer1 = self._make_layer(block, 64, layers[0], stride=1)
        self.layer2 = self._make_layer(block, 128, layers[1], stride=4)
        self.layer3 = self._make_layer(block, 256, layers[2], stride=4)
        self.layer4 = self._make_layer(block, 512, layers[3], stride=4)
        self.layer5 = self._make_layer(block, 1024, layers[4], stride=4)
        self.layer6 = self._make_layer(block, 1024, layers[5], stride=4)
        self.layer7 = self._make_layer(block, 2048, layers[6], stride=4)

    def _make_layer(self, block, planes, blocks, stride=1, dilate=False):
        norm_layer = self._norm_layer
        downsample = None
        previous_dilation = self.dilation
        if dilate:
            self.dilation *= stride
            stride = 1
        if stride != 1 or self.inplanes != planes * block.expansion:
            if stride == 1:
                downsample = nn.Sequential(
                    _resnet_conv1x1_wav1d(self.inplanes, planes * block.expansion),
                    norm_layer(planes * block.expansion),
                )
                init_layer(downsample[0])
                init_bn(downsample[1])
            else:
                downsample = nn.Sequential(
                    nn.AvgPool1d(kernel_size=stride),
                    _resnet_conv1x1_wav1d(self.inplanes, planes * block.expansion),
                    norm_layer(planes * block.expansion),
                )
                init_layer(downsample[1])
                init_bn(downsample[2])

        layers = []
        layers.append(
            block(
                self.inplanes,
                planes,
                stride,
                downsample,
                self.groups,
                self.base_width,
                previous_dilation,
                norm_layer,
            )
        )
        self.inplanes = planes * block.expansion
        for _ in range(1, blocks):
            layers.append(
                block(
                    self.inplanes,
                    planes,
                    groups=self.groups,
                    base_width=self.base_width,
                    dilation=self.dilation,
                    norm_layer=norm_layer,
                )
            )

        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.layer5(x)
        x = self.layer6(x)
        x = self.layer7(x)

        return x


class Res1dNet31(nn.Module):
    def __init__(self, classes_num=527, pretrained=False):
        super(Res1dNet31, self).__init__()

        self.conv0 = nn.Conv1d(
            in_channels=1,
            out_channels=64,
            kernel_size=(11,),
            stride=(5,),
            padding=5,
            bias=False,
        )
        self.bn0 = nn.BatchNorm1d(64)

        self.resnet = _ResNetWav1d(_ResnetBasicBlockWav1d, [2, 2, 2, 2, 2, 2, 2])

        self.fc1 = nn.Linear(2048, 2048, bias=True)
        self.fc_audioset = nn.Linear(2048, classes_num, bias=True)

        self.init_weight()

        if pretrained:
            import urllib.request

            ckpt_fn = "models/pretrained/audio_benchmarking/Res1dNet31_mAP_0.365.pth"
            url = "https://zenodo.org/record/3987831/files/Res1dNet31_mAP%3D0.365.pth"
            if not os.path.isfile(ckpt_fn):
                os.makedirs(os.path.dirname(ckpt_fn), exist_ok=True)
                urllib.request.urlretrieve(url, ckpt_fn)
            self.load_state_dict(torch.load(ckpt_fn, map_location="cpu"))

    def init_weight(self):
        init_layer(self.conv0)
        init_bn(self.bn0)
        init_layer(self.fc1)
        init_layer(self.fc_audioset)

    def forward(self, input):
        """
        Input: (batch_size, data_length)"""

        x = input[:, None, :]

        x = self.bn0(self.conv0(x))
        x = self.resnet(x)

        (x1, _) = torch.max(x, dim=2)
        x2 = torch.mean(x, dim=2)
        x = x1 + x2
        x = F.dropout(x, p=0.5, training=self.training)
        x = F.relu_(self.fc1(x))
        x = self.fc_audioset(x)
        return x


class Res1dNet51(nn.Module):
    def __init__(self, classes_num=527, pretrained=False):
        super(Res1dNet51, self).__init__()

        self.conv0 = nn.Conv1d(
            in_channels=1,
            out_channels=64,
            kernel_size=(11,),
            stride=(5,),
            padding=(5,),
            bias=False,
        )
        self.bn0 = nn.BatchNorm1d(64)

        self.resnet = _ResNetWav1d(_ResnetBasicBlockWav1d, [2, 3, 4, 6, 4, 3, 2])

        self.fc1 = nn.Linear(2048, 2048, bias=True)
        self.fc_audioset = nn.Linear(2048, classes_num, bias=True)

        self.init_weight()

        if pretrained:
            import urllib.request

            ckpt_fn = "models/pretrained/audio_benchmarking/Res1dNet51_mAP_0.355.pth"
            url = "https://zenodo.org/record/3987831/files/Res1dNet51_mAP%3D0.355.pth"
            if not os.path.isfile(ckpt_fn):
                os.makedirs(os.path.dirname(ckpt_fn), exist_ok=True)
                urllib.request.urlretrieve(url, ckpt_fn)
            self.load_state_dict(torch.load(ckpt_fn, map_location="cpu"))

    def init_weight(self):
        init_layer(self.conv0)
        init_bn(self.bn0)
        init_layer(self.fc1)
        init_layer(self.fc_audioset)

    def forward(self, input):
        """
        Input: (batch_size, data_length)"""

        x = input[:, None, :]

        x = self.bn0(self.conv0(x))
        x = self.resnet(x)

        (x1, _) = torch.max(x, dim=2)
        x2 = torch.mean(x, dim=2)
        x = x1 + x2
        x = F.dropout(x, p=0.5, training=self.training)
        x = F.relu_(self.fc1(x))
        x = self.fc_audioset(x)
        return x


class ConvPreWavBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(ConvPreWavBlock, self).__init__()

        self.conv1 = nn.Conv1d(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=(3,),
            stride=(1,),
            padding=1,
            bias=False,
        )

        self.conv2 = nn.Conv1d(
            in_channels=out_channels,
            out_channels=out_channels,
            kernel_size=(3,),
            stride=(1,),
            dilation=(2,),
            padding=2,
            bias=False,
        )

        self.bn1 = nn.BatchNorm1d(out_channels)
        self.bn2 = nn.BatchNorm1d(out_channels)

        self.init_weight()

    def init_weight(self):
        init_layer(self.conv1)
        init_layer(self.conv2)
        init_bn(self.bn1)
        init_bn(self.bn2)

    def forward(self, input, pool_size):
        x = input
        x = F.relu_(self.bn1(self.conv1(x)))
        x = F.relu_(self.bn2(self.conv2(x)))
        x = F.max_pool1d(x, kernel_size=pool_size)

        return x


class Wavegram_Cnn14(nn.Module):
    def __init__(self, classes_num=527, pretrained=False):
        super(Wavegram_Cnn14, self).__init__()

        self.pre_conv0 = nn.Conv1d(
            in_channels=1,
            out_channels=64,
            kernel_size=(11,),
            stride=(5,),
            padding=5,
            bias=False,
        )
        self.pre_bn0 = nn.BatchNorm1d(64)
        self.pre_block1 = ConvPreWavBlock(64, 64)
        self.pre_block2 = ConvPreWavBlock(64, 128)
        self.pre_block3 = ConvPreWavBlock(128, 128)
        self.pre_block4 = ConvBlock(in_channels=4, out_channels=64)

        self.bn0 = nn.BatchNorm2d(64)
        self.conv_block1 = ConvBlock(in_channels=1, out_channels=64)
        self.conv_block2 = ConvBlock(in_channels=64, out_channels=128)
        self.conv_block3 = ConvBlock(in_channels=128, out_channels=256)
        self.conv_block4 = ConvBlock(in_channels=256, out_channels=512)
        self.conv_block5 = ConvBlock(in_channels=512, out_channels=1024)
        self.conv_block6 = ConvBlock(in_channels=1024, out_channels=2048)

        self.fc1 = nn.Linear(2048, 2048, bias=True)
        self.fc_audioset = nn.Linear(2048, classes_num, bias=True)

        self.init_weight()

        if pretrained:
            import urllib.request

            ckpt_fn = (
                "models/pretrained/audio_benchmarking/Wavegram_Cnn14_mAP_0.389.pth"
            )
            url = (
                "https://zenodo.org/record/3987831/files/Wavegram_Cnn14_mAP%3D0.389.pth"
            )
            if not os.path.isfile(ckpt_fn):
                os.makedirs(os.path.dirname(ckpt_fn), exist_ok=True)
                urllib.request.urlretrieve(url, ckpt_fn)
            self.load_state_dict(torch.load(ckpt_fn, map_location="cpu"))

    def init_weight(self):
        init_layer(self.pre_conv0)
        init_bn(self.pre_bn0)
        init_bn(self.bn0)
        init_layer(self.fc1)
        init_layer(self.fc_audioset)

    def forward(self, input):
        """
        Input: (batch_size, data_length)"""

        # Wavegram
        a1 = F.relu_(self.pre_bn0(self.pre_conv0(input[:, None, :])))
        a1 = self.pre_block1(a1, pool_size=4)
        a1 = self.pre_block2(a1, pool_size=4)
        a1 = self.pre_block3(a1, pool_size=4)
        a1 = a1.reshape((a1.shape[0], -1, 32, a1.shape[-1])).transpose(2, 3)
        a1 = self.pre_block4(a1, pool_size=(2, 1))

        x = a1
        x = F.dropout(x, p=0.2, training=self.training)
        x = self.conv_block2(x, pool_size=(2, 2), pool_type="avg")
        x = F.dropout(x, p=0.2, training=self.training)
        x = self.conv_block3(x, pool_size=(2, 2), pool_type="avg")
        x = F.dropout(x, p=0.2, training=self.training)
        x = self.conv_block4(x, pool_size=(2, 2), pool_type="avg")
        x = F.dropout(x, p=0.2, training=self.training)
        x = self.conv_block5(x, pool_size=(2, 2), pool_type="avg")
        x = F.dropout(x, p=0.2, training=self.training)
        x = self.conv_block6(x, pool_size=(1, 1), pool_type="avg")
        x = F.dropout(x, p=0.2, training=self.training)
        x = torch.mean(x, dim=3)

        (x1, _) = torch.max(x, dim=2)
        x2 = torch.mean(x, dim=2)
        x = x1 + x2
        x = F.dropout(x, p=0.5, training=self.training)
        x = F.relu_(self.fc1(x))
        x = self.fc_audioset(x)

        return x


class Wavegram_Logmel_Cnn14(nn.Module):
    def __init__(self, classes_num=527, pretrained=True):
        super(Wavegram_Logmel_Cnn14, self).__init__()

        self.pre_conv0 = nn.Conv1d(
            in_channels=1,
            out_channels=64,
            kernel_size=(11,),
            stride=(5,),
            padding=5,
            bias=False,
        )
        self.pre_bn0 = nn.BatchNorm1d(64)
        self.pre_block1 = ConvPreWavBlock(64, 64)
        self.pre_block2 = ConvPreWavBlock(64, 128)
        self.pre_block3 = ConvPreWavBlock(128, 128)
        self.pre_block4 = ConvBlock(in_channels=4, out_channels=64)

        self.bn0 = nn.BatchNorm2d(64)

        self.conv_block1 = ConvBlock(in_channels=1, out_channels=64)
        self.conv_block2 = ConvBlock(in_channels=128, out_channels=128)
        self.conv_block3 = ConvBlock(in_channels=128, out_channels=256)
        self.conv_block4 = ConvBlock(in_channels=256, out_channels=512)
        self.conv_block5 = ConvBlock(in_channels=512, out_channels=1024)
        self.conv_block6 = ConvBlock(in_channels=1024, out_channels=2048)

        self.fc1 = nn.Linear(2048, 2048, bias=True)
        self.fc_audioset = nn.Linear(2048, classes_num, bias=True)

        self.init_weight()

        if pretrained:
            import urllib.request

            ckpt_fn = "models/pretrained/audio_benchmarking/Wavegram_Logmel_Cnn14_mAP_0.439.pth"
            url = "https://zenodo.org/record/3987831/files/Wavegram_Logmel_Cnn14_mAP%3D0.439.pth"
            if not os.path.isfile(ckpt_fn):
                os.makedirs(os.path.dirname(ckpt_fn), exist_ok=True)
                urllib.request.urlretrieve(url, ckpt_fn)
            self.load_state_dict(torch.load(ckpt_fn, map_location="cpu"))

    def init_weight(self):
        init_layer(self.pre_conv0)
        init_bn(self.pre_bn0)
        init_bn(self.bn0)
        init_layer(self.fc1)
        init_layer(self.fc_audioset)

    def forward(self, wav, spec):
        """
        Input: (batch_size, data_length)"""

        # Wavegram
        a1 = F.relu_(self.pre_bn0(self.pre_conv0(wav)))
        a1 = self.pre_block1(a1, pool_size=4)
        a1 = self.pre_block2(a1, pool_size=4)
        a1 = self.pre_block3(a1, pool_size=4)
        a1 = a1.reshape((a1.shape[0], -1, 32, a1.shape[-1])).transpose(2, 3)
        a1 = self.pre_block4(a1, pool_size=(2, 1))

        # Spectrogram
        x = self.bn0(spec)
        x = self.conv_block1(x, pool_size=(2, 2), pool_type="avg")

        # Concatenate Wavegram and Log mel spectrogram along the channel dimension
        x = torch.cat((x, a1), dim=1)
        x = F.dropout(x, p=0.2, training=self.training)
        x = self.conv_block2(x, pool_size=(2, 2), pool_type="avg")
        x = F.dropout(x, p=0.2, training=self.training)
        x = self.conv_block3(x, pool_size=(2, 2), pool_type="avg")
        x = F.dropout(x, p=0.2, training=self.training)
        x = self.conv_block4(x, pool_size=(2, 2), pool_type="avg")
        x = F.dropout(x, p=0.2, training=self.training)
        x = self.conv_block5(x, pool_size=(2, 2), pool_type="avg")
        x = F.dropout(x, p=0.2, training=self.training)
        x = self.conv_block6(x, pool_size=(1, 1), pool_type="avg")
        x = F.dropout(x, p=0.2, training=self.training)
        x = torch.mean(x, dim=3)

        (x1, _) = torch.max(x, dim=2)
        x2 = torch.mean(x, dim=2)
        x = x1 + x2
        x = F.dropout(x, p=0.5, training=self.training)
        x = F.relu_(self.fc1(x))
        x = self.fc_audioset(x)
        return x

upstream_models/replai/replai/models/audio/audio_models_benchmarking.py
roger-tseng/av-superb
# Copyright (c) Facebook, Inc. and its 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 torch.nn import functional as F

__all__ = [
    "avid_spec_cnn_9",
]


class Basic2DBlock(nn.Module):
    def __init__(self, in_planes, out_planes, stride=(1, 1)):
        self.__dict__.update(locals())
        super(Basic2DBlock, self).__init__()

        self.conv1 = nn.Conv2d(
            in_planes,
            out_planes,
            kernel_size=(3, 3),
            padding=(1, 1),
            stride=stride,
            bias=False,
        )
        self.bn1 = nn.BatchNorm2d(out_planes)
        self.conv2 = nn.Conv2d(
            out_planes, out_planes, kernel_size=(3, 3), padding=(1, 1), bias=False
        )
        self.bn2 = nn.BatchNorm2d(out_planes)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        x = self.relu(self.bn1(self.conv1(x)))
        x = self.relu(self.bn2(self.conv2(x)))
        return x


class Conv2D(nn.Module):
    def __init__(self, channels=1, depth=10):
        super(Conv2D, self).__init__()
        assert depth == 10

        self.conv1 = nn.Sequential(
            nn.Conv2d(channels, 64, kernel_size=7, padding=3, stride=2, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
        )

        self.block1 = Basic2DBlock(64, 64, stride=(2, 2))
        self.block2 = Basic2DBlock(64, 128, stride=(2, 2))
        self.block3 = Basic2DBlock(128, 256, stride=(2, 2))
        self.block4 = Basic2DBlock(256, 512)
        # original
        # self.pool = nn.AdaptiveMaxPool2d((1, 1))
        self.output_dim = 512

    def forward(self, x, return_embs=False):
        seq_len = x.shape[-2]
        x_c1 = self.conv1(x)
        x_b1 = self.block1(x_c1)
        x_b2 = self.block2(x_b1)
        x_b3 = self.block3(x_b2)
        x_b4 = self.block4(x_b3)

        # original
        # x_pool = self.pool(x_b4)
        # revised: ( generalize temporal pooling to variabble sequence length)
        x_pool = F.adaptive_max_pool2d(x_b4, (seq_len // 128, 1), False)

        if return_embs:
            return {
                "conv2x": x_b1,
                "conv3x": x_b2,
                "conv4x": x_b3,
                "conv5x": x_b4,
                "pool": x_pool,
            }
        else:
            return x_b4


def avid_spec_cnn_9(channels=1, pretrained=False):
    model = Conv2D(channels=channels)

    if pretrained:
        import os

        base_dir = os.environ["PYTHONPATH"].split(":")[0]
        try:
            ckpt_fn = "/glusterfs/hmittal/ssrl/pretrained_av_weights/AVID_Audioset_Cross-N1024_checkpoint.pth.tar"
            ckp = torch.load(ckpt_fn, map_location="cpu")
        except FileNotFoundError:
            ckpt_fn = "/Users/morgado/Projects/ssrl/models/ckpt/AVID_Audioset_Cross-N1024_checkpoint.pth.tar"
            ckp = torch.load(ckpt_fn, map_location="cpu")
        prefix = "module.audio_model."
        model.load_state_dict(
            {
                k.replace(prefix, ""): ckp["model"][k]
                for k in ckp["model"]
                if k.startswith(prefix)
            }
        )
        print(f"Audio model initialized from {ckpt_fn}")

    return model

upstream_models/replai/replai/models/audio/avid_audio.py
roger-tseng/av-superb
import torch.nn as nn

__all__ = ["SpecCNN", "SpectCNN_AVID"]


class BasicBlock(nn.Module):
    def __init__(self, in_planes, out_planes, stride=(1, 1)):
        super().__init__()
        self.conv1 = nn.Conv2d(
            in_planes,
            out_planes,
            kernel_size=(3, 3),
            padding=(1, 1),
            stride=stride,
            bias=False,
        )
        self.bn1 = nn.BatchNorm2d(out_planes)
        self.conv2 = nn.Conv2d(
            out_planes, out_planes, kernel_size=(3, 3), padding=(1, 1), bias=False
        )
        self.bn2 = nn.BatchNorm2d(out_planes)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        x = self.relu(self.bn1(self.conv1(x)))
        x = self.relu(self.bn2(self.conv2(x)))
        return x


class SpecCNN(nn.Sequential):
    def __init__(self, d=64):
        super().__init__(
            nn.Conv2d(1, d, kernel_size=7, padding=3, stride=2, bias=False),
            nn.BatchNorm2d(d),
            nn.ReLU(inplace=True),
            BasicBlock(d, d),
            BasicBlock(d, d * 2, stride=2),
            BasicBlock(d * 2, d * 4, stride=2),
            BasicBlock(d * 4, d * 8),
            nn.AdaptiveMaxPool2d((1, 1)),
        )
        self.output_dim = d * 8

        # Initialize network
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)


class SpectCNN_AVID(nn.Sequential):
    def __init__(self):
        super().__init__(
            nn.Conv2d(1, 64, kernel_size=7, padding=3, stride=2, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            BasicBlock(64, 64, stride=2),
            BasicBlock(64, 128, stride=2),
            BasicBlock(128, 256, stride=2),
            BasicBlock(256, 512),
            nn.AdaptiveMaxPool2d((1, 1)),
        )
        self.output_dim = 512

        # Initialize network
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

upstream_models/replai/replai/models/audio/spec_cnn.py
roger-tseng/av-superb
from torch import nn
from torchvision import models

from ..utils import distributed_ops, network_ops
from . import audio as snd_models
from . import video as vid_models


def build_video_model(cfg, distributed=False, device=None, remove_head=False):
    if cfg.arch in vars(vid_models):
        model_builder = vid_models.__dict__[cfg.arch]
    elif cfg.arch in vars(models):
        model_builder = models.__dict__[cfg.arch]
    else:
        raise ValueError(f"Model {cfg.arch} not found.")

    model = model_builder(**cfg.args)
    if cfg.sync_bn:
        model = nn.SyncBatchNorm.convert_sync_batchnorm(model)

    if remove_head:
        model = network_ops.remove_classification_head(model)

    # Send model to gpus
    return distributed_ops.send_to_device(model, distributed=distributed, device=device)


def build_audio_model(cfg, distributed=False, device=None, remove_head=False):
    if cfg.arch in vars(snd_models):
        model_builder = snd_models.__dict__[cfg.arch]
    elif cfg.arch in vars(models):
        model_builder = models.__dict__[cfg.arch]
    else:
        raise ValueError(f"Model {cfg.arch} not found.")

    model = model_builder(**cfg.args)
    if cfg.sync_bn:
        model = nn.SyncBatchNorm.convert_sync_batchnorm(model)

    if remove_head:
        model = network_ops.remove_classification_head(model)

    # Send model to gpus
    return distributed_ops.send_to_device(model, distributed=distributed, device=device)


def build_audio_video_model(cfg, distributed=False, device=None, remove_head=False):
    video_model = build_video_model(
        cfg.video, distributed=distributed, device=device, remove_head=remove_head
    )
    audio_model = build_audio_model(
        cfg.audio, distributed=distributed, device=device, remove_head=remove_head
    )
    return nn.ModuleDict({"video": video_model, "audio": audio_model})

upstream_models/replai/replai/models/builder.py
roger-tseng/av-superb
from torchvision.models.video import mc3_18, r2plus1d_18, r3d_18

from .avid_video import *

upstream_models/replai/replai/models/video/__init__.py
roger-tseng/av-superb
# Copyright (c) Facebook, Inc. and its 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 torch.nn import functional as F

__all__ = ["avid_r2plus1d_18"]


class BasicR2P1DBlock(nn.Module):
    def __init__(self, in_planes, out_planes, stride=(1, 1, 1)):
        super(BasicR2P1DBlock, self).__init__()
        spt_stride = (1, stride[1], stride[2])
        tmp_stride = (stride[0], 1, 1)
        self.spt_conv1 = nn.Conv3d(
            in_planes,
            out_planes,
            kernel_size=(1, 3, 3),
            stride=spt_stride,
            padding=(0, 1, 1),
            bias=False,
        )
        self.spt_bn1 = nn.BatchNorm3d(out_planes)
        self.tmp_conv1 = nn.Conv3d(
            out_planes,
            out_planes,
            kernel_size=(3, 1, 1),
            stride=tmp_stride,
            padding=(1, 0, 0),
            bias=False,
        )
        self.tmp_bn1 = nn.BatchNorm3d(out_planes)

        self.spt_conv2 = nn.Conv3d(
            out_planes,
            out_planes,
            kernel_size=(1, 3, 3),
            stride=(1, 1, 1),
            padding=(0, 1, 1),
            bias=False,
        )
        self.spt_bn2 = nn.BatchNorm3d(out_planes)
        self.tmp_conv2 = nn.Conv3d(
            out_planes,
            out_planes,
            kernel_size=(3, 1, 1),
            stride=(1, 1, 1),
            padding=(1, 0, 0),
            bias=False,
        )
        self.out_bn = nn.BatchNorm3d(out_planes)

        self.relu = nn.ReLU(inplace=True)

        if in_planes != out_planes or any([s != 1 for s in stride]):
            self.res = True
            self.res_conv = nn.Conv3d(
                in_planes,
                out_planes,
                kernel_size=(1, 1, 1),
                stride=stride,
                padding=(0, 0, 0),
                bias=False,
            )
        else:
            self.res = False

    def forward(self, x):
        x_main = self.tmp_conv1(self.relu(self.spt_bn1(self.spt_conv1(x))))
        x_main = self.relu(self.tmp_bn1(x_main))
        x_main = self.tmp_conv2(self.relu(self.spt_bn2(self.spt_conv2(x_main))))

        x_res = self.res_conv(x) if self.res else x
        x_out = self.relu(self.out_bn(x_main + x_res))
        return x_out


class R2Plus1D(nn.Module):
    """
    Adapted from https://github.com/facebookresearch/VMZ/blob/4c14ee6f8eae8e2ac97fc4c05713b8a112eb1f28/lib/models/video_model.py
    Adaptation has a full Conv3D stem, and does not adjust for the number of dimensions between the spatial and temporal convolution.
    """

    def __init__(self, depth=18):
        super(R2Plus1D, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv3d(
                3,
                64,
                kernel_size=(3, 7, 7),
                padding=(1, 3, 3),
                stride=(1, 2, 2),
                bias=False,
            ),
            nn.BatchNorm3d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool3d(kernel_size=(1, 3, 3), stride=(1, 2, 2), padding=(0, 1, 1)),
        )

        if depth == 10:
            self.conv2x = BasicR2P1DBlock(64, 64)
            self.conv3x = BasicR2P1DBlock(64, 128, stride=(2, 2, 2))
            self.conv4x = BasicR2P1DBlock(128, 256, stride=(2, 2, 2))
            self.conv5x = BasicR2P1DBlock(256, 512, stride=(2, 2, 2))
        elif depth == 18:
            self.conv2x = nn.Sequential(
                BasicR2P1DBlock(64, 64), BasicR2P1DBlock(64, 64)
            )
            self.conv3x = nn.Sequential(
                BasicR2P1DBlock(64, 128, stride=(2, 2, 2)), BasicR2P1DBlock(128, 128)
            )
            self.conv4x = nn.Sequential(
                BasicR2P1DBlock(128, 256, stride=(2, 2, 2)), BasicR2P1DBlock(256, 256)
            )
            self.conv5x = nn.Sequential(
                BasicR2P1DBlock(256, 512, stride=(2, 2, 2)), BasicR2P1DBlock(512, 512)
            )
        elif depth == 34:
            self.conv2x = nn.Sequential(
                BasicR2P1DBlock(64, 64),
                BasicR2P1DBlock(64, 64),
                BasicR2P1DBlock(64, 64),
            )
            self.conv3x = nn.Sequential(
                BasicR2P1DBlock(64, 128, stride=(2, 2, 2)),
                BasicR2P1DBlock(128, 128),
                BasicR2P1DBlock(128, 128),
                BasicR2P1DBlock(128, 128),
            )
            self.conv4x = nn.Sequential(
                BasicR2P1DBlock(128, 256, stride=(2, 2, 2)),
                BasicR2P1DBlock(256, 256),
                BasicR2P1DBlock(256, 256),
                BasicR2P1DBlock(256, 256),
                BasicR2P1DBlock(256, 256),
                BasicR2P1DBlock(256, 256),
            )
            self.conv5x = nn.Sequential(
                BasicR2P1DBlock(256, 512, stride=(2, 2, 2)),
                BasicR2P1DBlock(512, 512),
                BasicR2P1DBlock(512, 512),
            )
        # original:
        # self.pool = nn.AdaptiveMaxPool3d((1, 1, 1))

        self.output_dim = 512

    def forward(self, x, return_embs=False):
        seq_len = x.shape[-3]
        x_c1 = self.conv1(x)
        x_b1 = self.conv2x(x_c1)
        x_b2 = self.conv3x(x_b1)
        x_b3 = self.conv4x(x_b2)
        x_b4 = self.conv5x(x_b3)

        # original
        # x_pool = self.pool(x_b4)
        # revised: ( generalize temporal pooling to variabble sequence length)
        x_pool = F.adaptive_max_pool3d(x_b4, (seq_len // 8, 1, 1), False)

        # x_pool = self.pool(x_b4)
        if return_embs:
            return {
                "conv1": x_c1,
                "conv2x": x_b1,
                "conv3x": x_b2,
                "conv4x": x_b3,
                "conv5x": x_b4,
                "pool": x_pool,
            }
        else:
            return x_pool


def avid_r2plus1d_18(pretrained=False):
    model = R2Plus1D(depth=18)

    if pretrained:
        import os

        base_dir = os.environ["PYTHONPATH"].split(":")[0]
        try:
            ckpt_fn = "/glusterfs/hmittal/ssrl/pretrained_av_weights/AVID_Audioset_Cross-N1024_checkpoint.pth.tar"
            ckp = torch.load(ckpt_fn, map_location="cpu")
        except FileNotFoundError:
            ckpt_fn = "/Users/morgado/Projects/ssrl/models/ckpt/AVID_Audioset_Cross-N1024_checkpoint.pth.tar"
            ckp = torch.load(ckpt_fn, map_location="cpu")
        prefix = "module.video_model."
        model.load_state_dict(
            {
                k.replace(prefix, ""): ckp["model"][k]
                for k in ckp["model"]
                if k.startswith(prefix)
            }
        )
        print(f"Video model initialized from {ckpt_fn}")

    return model

upstream_models/replai/replai/models/video/avid_video.py
roger-tseng/av-superb
import os

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel


def init_distributed_environment(gpu, ngpus_per_node, args):
    if args.environment.dist_url == "env://" and args.environment.rank == -1:
        args.environment.rank = int(os.environ["RANK"])
    if args.environment.multiprocessing_distributed:
        # For multiprocessing distributed training, rank needs to be the
        # global rank among all the processes
        args.environment.rank = args.environment.rank * ngpus_per_node + gpu
    dist.init_process_group(
        backend=args.environment.dist_backend,
        init_method=args.environment.dist_url,
        world_size=args.environment.world_size,
        rank=args.environment.rank,
    )

    if args.environment.gpu is not None:
        # For multiprocessing distributed, DistributedDataParallel constructor
        # should always set the single device scope, otherwise,
        # DistributedDataParallel will use all available devices.
        torch.cuda.set_device(args.environment.gpu)

        # When using a single GPU per process and per
        # DistributedDataParallel, we need to divide the batch size
        # ourselves based on the total number of GPUs we have
        args.optim.batch_size = int(args.optim.batch_size / ngpus_per_node)
        args.environment.workers = int(
            (args.environment.workers + ngpus_per_node - 1) / ngpus_per_node
        )
    return args


def send_to_device(module, distributed, device=None):
    if distributed:
        module.cuda(device)
        module = DistributedDataParallel(
            module, device_ids=[device] if device is not None else None
        )
    elif device is not None:
        raise NotImplementedError("Only DistributedDataParallel is supported.")
    # else:
    #     module.cuda()
    return module


def shuffle_batch(x1, x2):
    """
    Batch shuffle, for making use of BatchNorm.
    *** Only support DistributedDataParallel (DDP) model. ***
    """
    if not dist.is_initialized():
        batch_size_all = x1.shape[0]
        idx_shuffle = torch.randperm(batch_size_all).to(x1.device)
        return x1[idx_shuffle], x2[idx_shuffle]

    else:
        # gather from all gpus
        batch_size_this = x1.shape[0]
        x1_gather = concat_all_gather(x1)
        x2_gather = concat_all_gather(x2)
        batch_size_all = x1_gather.shape[0]

        # random shuffle index
        idx_shuffle = torch.randperm(batch_size_all).to(x1.device)

        # broadcast to all gpus
        dist.broadcast(idx_shuffle, src=0)

        # shuffled index for this gpu
        gpu_idx = dist.get_rank()
        num_gpus = batch_size_all // batch_size_this
        idx_this = idx_shuffle.view(num_gpus, -1)[gpu_idx]

        # shuffle
        return x1_gather[idx_this], x2_gather[idx_this]


@torch.no_grad()
def concat_all_gather(tensor):
    """
    Performs all_gather operation on the provided tensors.
    *** Warning ***: dist.all_gather has no gradient.
    """
    if not dist.is_initialized():
        return tensor
    tensors_gather = [torch.ones_like(tensor) for _ in range(dist.get_world_size())]
    dist.all_gather(tensors_gather, tensor, async_op=False)
    output = torch.cat(tensors_gather, dim=0)
    return output

upstream_models/replai/replai/utils/distributed_ops.py
roger-tseng/av-superb
import torch.nn as nn


def remove_classification_head(model):
    if hasattr(model, "module"):
        model = model.module

    if hasattr(model, "fc"):
        model.fc = nn.Identity()
        return model

    if hasattr(model, "fc_audioset"):
        model.fc_audioset = nn.Identity()
        return model

    return model

upstream_models/replai/replai/utils/network_ops.py
roger-tseng/av-superb
# -*- coding: utf-8 -*- #
"""*********************************************************************************************"""
#   FileName     [ upstream/avhubert/expert.py ]
#   Synopsis     [ the avhubert wrapper ]
#   Author       [ S3PRL ]
#   Copyright    [ Copyleft(c), Puyuan Peng, UT Austin ]
"""*********************************************************************************************"""

import argparse
from concurrent.futures import process

import fairseq
import numpy as np
import torch
import torch.nn.functional as F
import torchaudio
from packaging import version
from torch.nn.utils.rnn import pad_sequence
from torchvision.transforms.functional import rgb_to_grayscale

from interfaces import UpstreamBase

from . import utils as custom_utils
from .hubert import AVHubertConfig, AVHubertModel


class UpstreamExpert(UpstreamBase):
    def __init__(self, ckpt, **kwargs):
        super().__init__(**kwargs)
        assert version.parse(fairseq.__version__) > version.parse(
            "0.10.2"
        ), "Please install the fairseq master branch."
        bundle = torch.load(ckpt)
        cfg = vars(AVHubertConfig())
        for key in bundle["cfg"]["model"]:
            cfg[key] = bundle["cfg"]["model"][key]
        cfg = argparse.Namespace(**cfg)

        # transfer some stuff from cfg.task
        cfg.sample_rate = bundle["cfg"]["task"]["sample_rate"]
        cfg.stack_order_audio = bundle["cfg"]["task"]["stack_order_audio"]
        cfg.normalize = bundle["cfg"]["task"]["normalize"]

        # hard code some stuff
        dictionaries = [[1] * bundle["model"]["label_embs_concat"].shape[0]]
        cfg.required_seq_len_multiple = 1  # newer fairseq wav2vec2 TranformerEncoder requires this parameter, which optimize the processing efficiency for FP16
        cfg.layer_type = "transformer"  # newer fairseq wav2vec2 TranformerEncoder requires this parameter, hard coded to fit avhubert
        cfg.checkpoint_activations = False  # newer fairseq wav2vec2 TranformerEncoder requires this parameter, hard coded to fit avhubert
        self.audio_sample_rate = 16000
        self.video_frame_rate = 25

        model = AVHubertModel(cfg, dictionaries)
        model.load_state_dict(bundle["model"])
        self.model = model.eval()
        self.cfg = cfg
        assert (
            self.video_frame_rate == cfg.sample_rate
        ), f"video sample rate should equal to the task sample rate, but it's not, video_sample_rate: {self.video_frame_rate}, task_cfg.sample_rate: {self.cfg.sample_rate}"
        self.transform = custom_utils.Compose(
            [
                custom_utils.Normalize(0.0, 255.0),
                custom_utils.CenterCrop((88, 88)),
                custom_utils.Normalize(0.421, 0.165),
            ]
        )

        self.model.feature_grad_mult = 0.0
        self.model.encoder.layerdrop = 0.0

        if len(self.hooks) == 0:
            module_name = "self.model.encoder.layers"
            for module_id in range(len(eval(module_name))):
                self.add_hook(
                    f"{module_name}[{module_id}]",
                    lambda input, output: input[0].transpose(0, 1),
                )
            self.add_hook("self.model.encoder", lambda input, output: output[0])

            def postprocess(xs):
                names, hiddens = zip(*xs)
                names = [name + "_fusion" for name in names]
                unpad_len = min([hidden.size(1) for hidden in hiddens])
                hiddens = [hidden[:, :unpad_len, :] for hidden in hiddens]
                return list(zip(names, hiddens))

            self.hook_postprocess = postprocess

    def stacker(self, feats, stack_order):
        """
        TODO: need to make this a batch processing method
        Concatenating consecutive audio frames
        Args:
        feats - numpy.ndarray of shape [T, F]
        stack_order - int (number of neighboring frames to concatenate
        Returns:
        feats - numpy.ndarray of shape [T', F']
        """
        feat_dim = feats.shape[1]
        if len(feats) % stack_order != 0:
            res = stack_order - len(feats) % stack_order
            res = np.zeros([res, feat_dim]).astype(feats.dtype)
            feats = np.concatenate([feats, res], axis=0)
        feats = feats.reshape((-1, stack_order, feat_dim)).reshape(
            -1, stack_order * feat_dim
        )
        return feats

    def preprocess_audio(self, audio, audio_sample_rate):
        # audio: (audio_channels, audio_length), where audio_channels is usually 1 or 2
        # since using av-hubert native implementation, needs to work with numpy objects
        orig_device = audio.device
        audio = audio.cpu()
        if len(audio.shape) >= 3:
            raise NotImplementedError(
                f"input should be single sample, not a batch! shape of audio input to preprocess_audio: {audio.shape}"
            )
        elif len(audio.shape) == 2:
            assert (
                audio.shape[0] == 1 or audio.shape[0] == 2
            ), f"wrong audio shape to the preprocess_audio method: {audio.shape}"
            audio = audio.mean(0)
        # it can indeed do batch processing
        if audio_sample_rate != self.audio_sample_rate:
            audio = torchaudio.functional.resample(
                audio, audio_sample_rate, self.audio_sample_rate
            )
        audio_feats = custom_utils.logfbank(
            audio, samplerate=self.audio_sample_rate
        ).astype(
            np.float32
        )  # [T, F]
        in_data = self.stacker(audio_feats, self.cfg.stack_order_audio)
        # [T/stack_order_audio, stack_order_audio*F]
        in_data = torch.from_numpy(in_data.astype(np.float32))
        if self.cfg.normalize:
            with torch.no_grad():
                in_data = F.layer_norm(in_data, in_data.shape[1:])

        return in_data.to(orig_device)  # TxF

    def preprocess_video(self, video, video_frame_rate):
        # video: (video_length, video_channels, height, width), where video_channels is usually 3 for RGB or 1 for greyscale
        # avhubert will make image and audio have the same framerate so we can add them or concat them in feature dimension. image sample rate if 25Hz, audio sample rate is 100Hz (originally 16kHz, but after fbank it's 100Hz), four neighboring audio sample is stacked to get
        # since using av-hubert native implementation, needs to work with numpy objects
        orig_device = video.device
        # Resample video
        # (from https://github.com/pytorch/vision/blob/5b07d6c9c6c14cf88fc545415d63021456874744/torchvision/datasets/video_utils.py#L278)
        step = float(video_frame_rate) / self.video_frame_rate
        if step.is_integer():
            # optimization: if step is integer, don't need to perform
            # advanced indexing
            step = int(step)
            idxs = slice(None, None, step)
        else:
            num_frames = int(len(video) / step)
            idxs = torch.arange(num_frames, dtype=torch.float32) * step
            idxs = idxs.floor().to(torch.int64)
        video = video[idxs]

        # Transform to greyscale
        if video.shape[1] == 3:
            video = 0.2989 * video[:, 0] + 0.587 * video[:, 1] + 0.114 * video[:, 2]
        else:
            video = video.mean(dim=1)
        feats = self.transform(video)

        # T, H, W
        if isinstance(feats, np.ndarray):
            return torch.from_numpy(feats).to(orig_device)
        elif isinstance(feats, torch.Tensor):
            return feats.to(orig_device)

    def forward(self, processed_data):
        device = processed_data[0][0].device

        audio, video = [], []
        for audio_feats, video_feats in processed_data:
            diff = len(audio_feats) - len(video_feats)
            if diff > 0:
                audio_feats = audio_feats[:-diff]
            elif diff < 0:
                audio_feats = F.pad(audio_feats, (0, 0, 0, -diff), "constant", 0)
            audio.append(audio_feats)
            video.append(video_feats)

        audio_length = torch.LongTensor([len(item) for item in audio])
        video_length = torch.LongTensor([len(item) for item in video])
        assert sum([a == v for a, v in zip(audio_length, video_length)]) == len(
            audio_length
        ), f"each audio should have the same temporal length with the corresponding video, but they are not: audio_length: {audio_length}, video_length: {video_length}"
        padding_mask = ~torch.lt(
            torch.arange(max(audio_length)).unsqueeze(0),
            (audio_length).unsqueeze(1),
        ).to(device)
        padded_audio = pad_sequence(audio, batch_first=True)
        padded_video = pad_sequence(video, batch_first=True)
        source = {
            "audio": padded_audio.transpose(1, 2),
            "video": padded_video.unsqueeze(dim=1),
        }
        result = self.model(
            source, padding_mask=padding_mask, mask=False, features_only=True
        )
        return {
            "video_feats": result["features_video"].transpose(1, 2),
            "audio_feats": result["features_audio"].transpose(1, 2),
            # fusion_feats are handled by UpstreamBase's hooks
        }

class FinetunedUpstreamExpert(UpstreamBase):
    def __init__(self, ckpt, **kwargs):
        super().__init__(**kwargs)
        assert version.parse(fairseq.__version__) > version.parse(
            "0.10.2"
        ), "Please install the fairseq master branch."
        bundle = torch.load(ckpt)
        cfg = vars(AVHubertConfig())
        for key in bundle["cfg"]["model"]["w2v_args"]["model"]:
            cfg[key] = bundle["cfg"]["model"]["w2v_args"]["model"][key]
        cfg = argparse.Namespace(**cfg)

        # transfer some stuff from cfg.task
        cfg.sample_rate = bundle["cfg"]["model"]["w2v_args"]["task"]["sample_rate"]
        cfg.stack_order_audio = bundle["cfg"]["model"]["w2v_args"]["task"]["stack_order_audio"]
        cfg.normalize = bundle["cfg"]["model"]["w2v_args"]["task"]["normalize"]

        # hard code some stuff
        bundle["model"] = {
            k[len("encoder.w2v_model."):]: bundle["model"][k]
            for k in bundle["model"].keys()
            if k.startswith("encoder.w2v_model.")
        }
        dictionaries = [[1] * bundle["model"]["label_embs_concat"].shape[0]]
        cfg.required_seq_len_multiple = 1  # newer fairseq wav2vec2 TranformerEncoder requires this parameter, which optimize the processing efficiency for FP16
        cfg.layer_type = "transformer"  # newer fairseq wav2vec2 TranformerEncoder requires this parameter, hard coded to fit avhubert
        cfg.checkpoint_activations = False  # newer fairseq wav2vec2 TranformerEncoder requires this parameter, hard coded to fit avhubert
        self.audio_sample_rate = 16000
        self.video_frame_rate = 25

        model = AVHubertModel(cfg, dictionaries)
        model.load_state_dict(bundle["model"], strict=False) # final_proj removed for finetuning
        self.model = model.eval()
        self.cfg = cfg
        assert (
            self.video_frame_rate == cfg.sample_rate
        ), f"video sample rate should equal to the task sample rate, but it's not, video_sample_rate: {self.video_frame_rate}, task_cfg.sample_rate: {self.cfg.sample_rate}"
        self.transform = custom_utils.Compose(
            [
                custom_utils.Normalize(0.0, 255.0),
                custom_utils.CenterCrop((88, 88)),
                custom_utils.Normalize(0.421, 0.165),
            ]
        )

        self.model.feature_grad_mult = 0.0
        self.model.encoder.layerdrop = 0.0

        if len(self.hooks) == 0:
            module_name = "self.model.encoder.layers"
            for module_id in range(len(eval(module_name))):
                self.add_hook(
                    f"{module_name}[{module_id}]",
                    lambda input, output: input[0].transpose(0, 1),
                )
            self.add_hook("self.model.encoder", lambda input, output: output[0])

            def postprocess(xs):
                names, hiddens = zip(*xs)
                names = [name + "_fusion" for name in names]
                unpad_len = min([hidden.size(1) for hidden in hiddens])
                hiddens = [hidden[:, :unpad_len, :] for hidden in hiddens]
                return list(zip(names, hiddens))

            self.hook_postprocess = postprocess

    def stacker(self, feats, stack_order):
        """
        TODO: need to make this a batch processing method
        Concatenating consecutive audio frames
        Args:
        feats - numpy.ndarray of shape [T, F]
        stack_order - int (number of neighboring frames to concatenate
        Returns:
        feats - numpy.ndarray of shape [T', F']
        """
        feat_dim = feats.shape[1]
        if len(feats) % stack_order != 0:
            res = stack_order - len(feats) % stack_order
            res = np.zeros([res, feat_dim]).astype(feats.dtype)
            feats = np.concatenate([feats, res], axis=0)
        feats = feats.reshape((-1, stack_order, feat_dim)).reshape(
            -1, stack_order * feat_dim
        )
        return feats

    def preprocess_audio(self, audio, audio_sample_rate):
        # audio: (audio_channels, audio_length), where audio_channels is usually 1 or 2
        # since using av-hubert native implementation, needs to work with numpy objects
        orig_device = audio.device
        audio = audio.cpu()
        if len(audio.shape) >= 3:
            raise NotImplementedError(
                f"input should be single sample, not a batch! shape of audio input to preprocess_audio: {audio.shape}"
            )
        elif len(audio.shape) == 2:
            assert (
                audio.shape[0] == 1 or audio.shape[0] == 2
            ), f"wrong audio shape to the preprocess_audio method: {audio.shape}"
            audio = audio.mean(0)
        # it can indeed do batch processing
        if audio_sample_rate != self.audio_sample_rate:
            audio = torchaudio.functional.resample(
                audio, audio_sample_rate, self.audio_sample_rate
            )
        audio_feats = custom_utils.logfbank(
            audio, samplerate=self.audio_sample_rate
        ).astype(
            np.float32
        )  # [T, F]
        in_data = self.stacker(audio_feats, self.cfg.stack_order_audio)
        # [T/stack_order_audio, stack_order_audio*F]
        in_data = torch.from_numpy(in_data.astype(np.float32))
        if self.cfg.normalize:
            with torch.no_grad():
                in_data = F.layer_norm(in_data, in_data.shape[1:])

        return in_data.to(orig_device)  # TxF

    def preprocess_video(self, video, video_frame_rate):
        # video: (video_length, video_channels, height, width), where video_channels is usually 3 for RGB or 1 for greyscale
        # avhubert will make image and audio have the same framerate so we can add them or concat them in feature dimension. image sample rate if 25Hz, audio sample rate is 100Hz (originally 16kHz, but after fbank it's 100Hz), four neighboring audio sample is stacked to get
        # since using av-hubert native implementation, needs to work with numpy objects
        orig_device = video.device
        # Resample video
        # (from https://github.com/pytorch/vision/blob/5b07d6c9c6c14cf88fc545415d63021456874744/torchvision/datasets/video_utils.py#L278)
        step = float(video_frame_rate) / self.video_frame_rate
        if step.is_integer():
            # optimization: if step is integer, don't need to perform
            # advanced indexing
            step = int(step)
            idxs = slice(None, None, step)
        else:
            num_frames = int(len(video) / step)
            idxs = torch.arange(num_frames, dtype=torch.float32) * step
            idxs = idxs.floor().to(torch.int64)
        video = video[idxs]

        # Transform to greyscale
        if video.shape[1] == 3:
            video = 0.2989 * video[:, 0] + 0.587 * video[:, 1] + 0.114 * video[:, 2]
        else:
            video = video.mean(dim=1)
        feats = self.transform(video)

        # T, H, W
        if isinstance(feats, np.ndarray):
            return torch.from_numpy(feats).to(orig_device)
        elif isinstance(feats, torch.Tensor):
            return feats.to(orig_device)

    def forward(self, processed_data):
        device = processed_data[0][0].device

        audio, video = [], []
        for audio_feats, video_feats in processed_data:
            diff = len(audio_feats) - len(video_feats)
            if diff > 0:
                audio_feats = audio_feats[:-diff]
            elif diff < 0:
                audio_feats = F.pad(audio_feats, (0, 0, 0, -diff), "constant", 0)
            audio.append(audio_feats)
            video.append(video_feats)

        audio_length = torch.LongTensor([len(item) for item in audio])
        video_length = torch.LongTensor([len(item) for item in video])
        assert sum([a == v for a, v in zip(audio_length, video_length)]) == len(
            audio_length
        ), f"each audio should have the same temporal length with the corresponding video, but they are not: audio_length: {audio_length}, video_length: {video_length}"
        padding_mask = ~torch.lt(
            torch.arange(max(audio_length)).unsqueeze(0),
            (audio_length).unsqueeze(1),
        ).to(device)
        padded_audio = pad_sequence(audio, batch_first=True)
        padded_video = pad_sequence(video, batch_first=True)
        source = {
            "audio": padded_audio.transpose(1, 2),
            "video": padded_video.unsqueeze(dim=1),
        }
        result = self.model(
            source, padding_mask=padding_mask, mask=False, features_only=True
        )
        return {
            "video_feats": result["features_video"].transpose(1, 2),
            "audio_feats": result["features_audio"].transpose(1, 2),
            # fusion_feats are handled by UpstreamBase's hooks
        }

upstream_models/vhubert/expert.py
roger-tseng/av-superb
# One needs to agree with the AV-HuBERT license (see COPYRIGHT under root 
# directory of source tree) to use pretrained models listed in this file.
# Copyright (c) Meta Platforms, Inc. All Rights Reserved.

# -*- coding: utf-8 -*- #
"""*********************************************************************************************"""
#   FileName     [ upstream/avhubert/hubconf.py ]
#   Synopsis     [ the avhubert torch hubconf ]
#   Author       [ S3PRL / Puyuan Peng, UT Austin]
"""*********************************************************************************************"""

###############
# IMPORTATION #
###############
import os

from utils.download import _urls_to_filepaths

from .expert import UpstreamExpert as _UpstreamExpert
from .expert import FinetunedUpstreamExpert as _FinetunedUpstreamExpert

# -------------#


def avhubert_local(ckpt, *args, **kwargs):
    """
    The model from local ckpt
        ckpt (str): PATH
    """
    assert os.path.isfile(ckpt)
    if kwargs.get("finetune"):
        return _FinetunedUpstreamExpert(ckpt, *args, **kwargs)
    else:
        return _UpstreamExpert(ckpt, *args, **kwargs)


def avhubert_url(ckpt, refresh=False, *args, **kwargs):
    """
    The model from URL
        ckpt (str): URL
    """
    return avhubert_local(_urls_to_filepaths(ckpt, refresh=refresh), *args, **kwargs)


def avhubert(refresh=False, *args, **kwargs):
    """
    The default model
        refresh (bool): whether to download ckpt/config again if existed
    """
    return avhubert_base_lrs3(refresh=refresh, *args, **kwargs)

def avhubert_audio(refresh=False, *args, **kwargs):
    """
    The default model
        refresh (bool): whether to download ckpt/config again if existed
    """
    model = avhubert(refresh=refresh, *args, **kwargs)
    model.model.modality_dropout = 1
    model.model.audio_dropout = 0
    return model

def avhubert_video(refresh=False, *args, **kwargs):
    """
    The default model
        refresh (bool): whether to download ckpt/config again if existed
    """
    model = avhubert(refresh=refresh, *args, **kwargs)
    model.model.modality_dropout = 1
    model.model.audio_dropout = 1
    return model

def avhubert_fusion(refresh=False, *args, **kwargs):
    """
    The default model
        refresh (bool): whether to download ckpt/config again if existed
    """
    model = avhubert(refresh=refresh, *args, **kwargs)
    model.model.modality_dropout = 0
    return model

############### AV-HuBERT trained on LRS3 ###############
def avhubert_base_lrs3(refresh=False, *args, **kwargs):
    """
    The avhubert base model trained on LRS3
        refresh (bool): whether to download ckpt/config again if existed
    """
    kwargs[
        "ckpt"
    ] = "https://dl.fbaipublicfiles.com/avhubert/model/lrs3/clean-pretrain/base_lrs3_iter5.pt"
    return avhubert_url(refresh=refresh, *args, **kwargs)

def avhubert_ft_lrs3_433(refresh=False, *args, **kwargs):
    """
    The avhubert base model trained on LRS3 then fine-tuned on LRS3-433h
        refresh (bool): whether to download ckpt/config again if existed
    """
    kwargs[
        "ckpt"
    ] = "https://dl.fbaipublicfiles.com/avhubert/model/lrs3/vsr/base_lrs3_433h.pt"
    kwargs["finetune"] = True
    return avhubert_url(refresh=refresh, *args, **kwargs)

def avhubert_ft_lrs3_433_audio(refresh=False, *args, **kwargs):
    """
    The avhubert base model trained on LRS3 then fine-tuned on LRS3-433h
        refresh (bool): whether to download ckpt/config again if existed
    """
    model = avhubert_ft_lrs3_433(refresh=refresh, *args, **kwargs)
    model.model.modality_dropout = 1
    model.model.audio_dropout = 0
    return model

def avhubert_ft_lrs3_433_video(refresh=False, *args, **kwargs):
    """
    The avhubert base model trained on LRS3 then fine-tuned on LRS3-433h
        refresh (bool): whether to download ckpt/config again if existed
    """
    model = avhubert_ft_lrs3_433(refresh=refresh, *args, **kwargs)
    model.model.modality_dropout = 1
    model.model.audio_dropout = 1
    return model

def avhubert_ft_lrs3_433_fusion(refresh=False, *args, **kwargs):
    """
    The avhubert base model trained on LRS3 then fine-tuned on LRS3-433h
        refresh (bool): whether to download ckpt/config again if existed
    """
    model = avhubert_ft_lrs3_433(refresh=refresh, *args, **kwargs)
    model.model.modality_dropout = 0
    return model

def avhubert_large_lrs3(refresh=False, *args, **kwargs):
    """
    The avhubert large model trained on LRS3
        refresh (bool): whether to download ckpt/config again if existed
    """
    kwargs[
        "ckpt"
    ] = "https://dl.fbaipublicfiles.com/avhubert/model/lrs3/clean-pretrain/large_lrs3_iter5.pt"
    return avhubert_url(refresh=refresh, *args, **kwargs)


############### AV-HuBERT trained on both LRS3 and VoxCeleb2 ###############
def avhubert_base_lrs3vox(refresh=False, *args, **kwargs):
    """
    The avhubert base model trained on LRS3 and VoxCeleb2
        refresh (bool): whether to download ckpt/config again if existed
    """
    kwargs[
        "ckpt"
    ] = "https://dl.fbaipublicfiles.com/avhubert/model/lrs3_vox/clean-pretrain/base_vox_iter5.pt"
    return avhubert_url(refresh=refresh, *args, **kwargs)


def avhubert_large_lrs3vox(refresh=False, *args, **kwargs):
    """
    The avhubert large model trained on LRS3 and VoxCeleb2
        refresh (bool): whether to download ckpt/config again if existed
    """
    kwargs[
        "ckpt"
    ] = "https://dl.fbaipublicfiles.com/avhubert/model/lrs3_vox/clean-pretrain/large_vox_iter5.pt"
    return avhubert_url(refresh=refresh, *args, **kwargs)


############### noise-augmented AV-HuBERT ###############
def avhubert_base_lrs3vox_na(refresh=False, *args, **kwargs):
    """
    The avhubert base model trained on LRS3 and VoxCeleb2, trained with noise augmentation
        refresh (bool): whether to download ckpt/config again if existed
    """
    kwargs[
        "ckpt"
    ] = "https://dl.fbaipublicfiles.com/avhubert/model/lrs3_vox/noise-pretrain/base_vox_iter5.pt"
    return avhubert_url(refresh=refresh, *args, **kwargs)


def avhubert_large_lrs3vox_na(refresh=False, *args, **kwargs):
    """
    The avhubert large model trained on LRS3 and VoxCeleb2, trained with noise augmentation
        refresh (bool): whether to download ckpt/config again if existed
    """
    kwargs[
        "ckpt"
    ] = "https://dl.fbaipublicfiles.com/avhubert/model/lrs3_vox/noise-pretrain/large_vox_iter5.pt"
    return avhubert_url(refresh=refresh, *args, **kwargs)

upstream_models/vhubert/hubconf.py
roger-tseng/av-superb
# This file is modified from code licensed under the AV-HuBERT license (see COPYRIGHT)
# Copyright (c) Meta Platforms, Inc. All Rights Reserved.

# Copyright (c) Facebook, Inc. and its 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 os
import sys
from copy import deepcopy
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple

import numpy as np
import torch
import torch.nn as nn
from fairseq import utils
from fairseq.data.data_utils import compute_mask_indices
from fairseq.data.dictionary import Dictionary
from fairseq.dataclass import ChoiceEnum, FairseqDataclass
from fairseq.models import BaseFairseqModel, register_model
from fairseq.models.wav2vec.wav2vec2 import (
    ConvFeatureExtractionModel,
    TransformerEncoder,
)
from fairseq.modules import GradMultiply, LayerNorm
from omegaconf import II

from .resnet import ResEncoder
from .utils import compute_mask_indices

EXTRACTOR_MODE_CHOICES = ChoiceEnum(["default", "layer_norm"])
MASKING_DISTRIBUTION_CHOICES = ChoiceEnum(["static", "uniform", "normal", "poisson"])


@dataclass
class AVHubertConfig(FairseqDataclass):
    label_rate: int = II("task.label_rate")
    input_modality: str = II("task.input_modality")
    extractor_mode: EXTRACTOR_MODE_CHOICES = field(
        default="default",
        metadata={
            "help": "mode for feature extractor. default has a single group "
            "norm with d groups in the first conv block, whereas layer_norm "
            "has layer norms in every block (meant to use with normalize=True)"
        },
    )
    encoder_layers: int = field(
        default=12, metadata={"help": "num encoder layers in the transformer"}
    )
    encoder_embed_dim: int = field(
        default=768, metadata={"help": "encoder embedding dimension"}
    )
    encoder_ffn_embed_dim: int = field(
        default=3072, metadata={"help": "encoder embedding dimension for FFN"}
    )
    encoder_attention_heads: int = field(
        default=12, metadata={"help": "num encoder attention heads"}
    )
    activation_fn: ChoiceEnum(utils.get_available_activation_fns()) = field(
        default="gelu", metadata={"help": "activation function to use"}
    )

    # dropouts
    dropout: float = field(
        default=0.1,
        metadata={"help": "dropout probability for the transformer"},
    )
    attention_dropout: float = field(
        default=0.1,
        metadata={"help": "dropout probability for attention weights"},
    )
    activation_dropout: float = field(
        default=0.0,
        metadata={"help": "dropout probability after activation in FFN"},
    )
    encoder_layerdrop: float = field(
        default=0.0,
        metadata={"help": "probability of dropping a tarnsformer layer"},
    )
    dropout_input: float = field(
        default=0.0,
        metadata={"help": "dropout to apply to the input (after feat extr)"},
    )
    dropout_features: float = field(
        default=0.0,
        metadata={"help": "dropout to apply to the features (after feat extr)"},
    )

    final_dim: int = field(
        default=0,
        metadata={
            "help": "project final representations and targets to this many "
            "dimensions. set to encoder_embed_dim is <= 0"
        },
    )
    untie_final_proj: bool = field(
        default=False,
        metadata={"help": "use separate projection for each target"},
    )
    layer_norm_first: bool = field(
        default=False,
        metadata={"help": "apply layernorm first in the transformer"},
    )
    conv_feature_layers: str = field(
        default="[(512,10,5)] + [(512,3,2)] * 4 + [(512,2,2)] * 2",
        metadata={
            "help": "string describing convolutional feature extraction "
            "layers in form of a python list that contains "
            "[(dim, kernel_size, stride), ...]"
        },
    )
    conv_bias: bool = field(
        default=False, metadata={"help": "include bias in conv encoder"}
    )
    logit_temp: float = field(
        default=0.1, metadata={"help": "temperature to divide logits by"}
    )
    target_glu: bool = field(
        default=False, metadata={"help": "adds projection + glu to targets"}
    )
    feature_grad_mult: float = field(
        default=1.0,
        metadata={"help": "multiply feature extractor var grads by this"},
    )

    # masking
    mask_length_audio: int = field(default=10, metadata={"help": "mask length"})
    mask_prob_audio: float = field(
        default=0.65,
        metadata={"help": "probability of replacing a token with mask"},
    )
    mask_length_image: int = field(default=10, metadata={"help": "mask length"})
    mask_prob_image: float = field(
        default=0.65,
        metadata={"help": "probability of replacing a token with mask"},
    )
    mask_selection: MASKING_DISTRIBUTION_CHOICES = field(
        default="static", metadata={"help": "how to choose mask length"}
    )
    mask_other: float = field(
        default=0,
        metadata={
            "help": "secondary mask argument "
            "(used for more complex distributions), "
            "see help in compute_mask_indicesh"
        },
    )
    no_mask_overlap: bool = field(
        default=False, metadata={"help": "whether to allow masks to overlap"}
    )
    mask_min_space: int = field(
        default=1,
        metadata={"help": "min space between spans (if no overlap is enabled)"},
    )

    # channel masking
    mask_channel_length: int = field(
        default=10,
        metadata={"help": "length of the mask for features (channels)"},
    )
    mask_channel_prob: float = field(
        default=0.0,
        metadata={"help": "probability of replacing a feature with 0"},
    )
    mask_channel_selection: MASKING_DISTRIBUTION_CHOICES = field(
        default="static",
        metadata={"help": "how to choose mask length for channel masking"},
    )
    mask_channel_other: float = field(
        default=0,
        metadata={
            "help": "secondary mask argument "
            "(used for more complex distributions), "
            "see help in compute_mask_indicesh"
        },
    )
    no_mask_channel_overlap: bool = field(
        default=False,
        metadata={"help": "whether to allow channel masks to overlap"},
    )
    mask_channel_min_space: int = field(
        default=1,
        metadata={"help": "min space between spans (if no overlap is enabled)"},
    )

    # positional embeddings
    conv_pos: int = field(
        default=128,
        metadata={"help": "number of filters for convolutional positional embeddings"},
    )
    conv_pos_groups: int = field(
        default=16,
        metadata={"help": "number of groups for convolutional positional embedding"},
    )

    latent_temp: Tuple[float, float, float] = field(
        default=(2, 0.5, 0.999995),
        metadata={"help": "legacy (to be removed)"},
    )

    # loss computation
    skip_masked: bool = field(
        default=False,
        metadata={"help": "skip computing losses over masked frames"},
    )
    skip_nomask: bool = field(
        default=False,
        metadata={"help": "skip computing losses over unmasked frames"},
    )
    resnet_relu_type: str = field(
        default="prelu", metadata={"help": "relu type for resnet"}
    )
    resnet_weights: Optional[str] = field(
        default=None, metadata={"help": "resnet weights"}
    )
    sim_type: str = field(default="cosine", metadata={"help": "similarity type"})

    sub_encoder_layers: int = field(
        default=0, metadata={"help": "number of transformer layers for single modality"}
    )
    audio_feat_dim: int = field(
        default=-1, metadata={"help": "audio feature dimension"}
    )
    modality_dropout: float = field(default=0, metadata={"help": "drop one modality"})
    audio_dropout: float = field(default=0, metadata={"help": "drop audio feature"})
    modality_fuse: str = field(
        default="concat", metadata={"help": "fusing two modalities: add,concat"}
    )
    selection_type: str = field(
        default="same_other_seq",
        metadata={
            "help": "type of selectig images, same_other_seq: replace masked span with span from another sequence, same_seq: repace masked span with span of the same sequence"
        },
    )
    masking_type: str = field(
        default="input", metadata={"help": "input or feature masking"}
    )

    decoder_embed_dim: int = field(
        default=768, metadata={"help": "decoder embedding dimension"}
    )
    decoder_ffn_embed_dim: int = field(
        default=3072, metadata={"help": "decoder embedding dimension for FFN"}
    )
    decoder_layers: int = field(default=6, metadata={"help": "num of decoder layers"})
    decoder_layerdrop: float = field(
        default=0.0, metadata={"help": "decoder layerdrop chance"}
    )
    decoder_attention_heads: int = field(
        default=4, metadata={"help": "num decoder attention heads"}
    )
    decoder_learned_pos: bool = field(
        default=False,
        metadata={"help": "use learned positional embeddings in the decoder"},
    )
    decoder_normalize_before: bool = field(
        default=False,
        metadata={"help": "apply layernorm before each decoder block"},
    )
    no_token_positional_embeddings: bool = field(
        default=False,
        metadata={
            "help": "if set, disables positional embeddings " "(outside self attention)"
        },
    )
    decoder_dropout: float = field(
        default=0.1, metadata={"help": "dropout probability in the decoder"}
    )
    decoder_attention_dropout: float = field(
        default=0.1,
        metadata={
            "help": "dropout probability for attention weights " "inside the decoder"
        },
    )
    decoder_activation_dropout: float = field(
        default=0.0,
        metadata={
            "help": "dropout probability after activation in FFN " "inside the decoder"
        },
    )
    max_target_positions: int = field(
        default=2048, metadata={"help": "max target positions"}
    )
    share_decoder_input_output_embed: bool = field(
        default=False,
        metadata={"help": "share decoder input and output embeddings"},
    )
    no_scale_embedding: bool = field(default=True, metadata={"help": "scale embedding"})


class SubModel(nn.Module):
    def __init__(self, resnet=None, input_dim=None, cfg=None):
        super().__init__()
        self.resnet = resnet
        self.proj = nn.Linear(input_dim, cfg.encoder_embed_dim)
        self.encoder = TransformerEncoder(cfg) if cfg.encoder_layers > 0 else None

    def forward(self, x):
        if self.resnet is not None:
            x = self.resnet(x)
        x = self.proj(x.transpose(1, 2))
        if self.encoder is not None:
            x = self.encoder(x)[0].transpose(1, 2)
        else:
            x = x.transpose(1, 2)
        return x


class AVHubertModel(BaseFairseqModel):
    def __init__(self, cfg, dictionaries) -> None:
        super().__init__()
        feature_ds_rate = 1
        self.feat2tar_ratio = cfg.label_rate * feature_ds_rate / cfg.sample_rate
        sub_cfg = deepcopy(cfg)
        sub_cfg.encoder_layers = sub_cfg.sub_encoder_layers
        resnet = ResEncoder(relu_type=cfg.resnet_relu_type, weights=cfg.resnet_weights)
        self.feature_extractor_audio = SubModel(
            resnet=None, input_dim=cfg.audio_feat_dim, cfg=sub_cfg
        )
        self.feature_extractor_video = SubModel(
            resnet=resnet, input_dim=resnet.backend_out, cfg=sub_cfg
        )
        self.modality_dropout, self.audio_dropout = (
            cfg.modality_dropout,
            cfg.audio_dropout,
        )
        self.modality_fuse = cfg.modality_fuse
        self.encoder_embed_dim = cfg.encoder_embed_dim
        if self.modality_fuse == "concat":
            self.embed = cfg.encoder_embed_dim * 2
        elif self.modality_fuse == "add":
            self.embed = cfg.encoder_embed_dim
        self.post_extract_proj = (
            nn.Linear(self.embed, cfg.encoder_embed_dim)
            if self.embed != cfg.encoder_embed_dim
            else None
        )

        self.mask_prob_image, self.mask_prob_audio = (
            cfg.mask_prob_image,
            cfg.mask_prob_audio,
        )
        self.mask_selection = cfg.mask_selection
        self.mask_other = cfg.mask_other
        self.mask_length_image, self.mask_length_audio = (
            cfg.mask_length_image,
            cfg.mask_length_audio,
        )
        self.no_mask_overlap = cfg.no_mask_overlap
        self.mask_min_space = cfg.mask_min_space

        self.mask_channel_prob = cfg.mask_channel_prob
        self.mask_channel_selection = cfg.mask_channel_selection
        self.mask_channel_other = cfg.mask_channel_other
        self.mask_channel_length = cfg.mask_channel_length
        self.no_mask_channel_overlap = cfg.no_mask_channel_overlap
        self.mask_channel_min_space = cfg.mask_channel_min_space

        self.dropout_input = nn.Dropout(cfg.dropout_input)
        self.dropout_features = nn.Dropout(cfg.dropout_features)

        self.feature_grad_mult = cfg.feature_grad_mult
        self.logit_temp = cfg.logit_temp
        self.skip_masked = cfg.skip_masked
        self.skip_nomask = cfg.skip_nomask
        self.sim_type = cfg.sim_type
        self.selection_type = cfg.selection_type
        self.masking_type = cfg.masking_type

        final_dim = cfg.final_dim if cfg.final_dim > 0 else cfg.encoder_embed_dim

        self.mask_emb = nn.Parameter(
            torch.FloatTensor(cfg.audio_feat_dim).uniform_()
            if self.masking_type == "input"
            else torch.FloatTensor(cfg.encoder_embed_dim).uniform_()
        )

        self.encoder = TransformerEncoder(cfg)
        self.layer_norm = LayerNorm(self.embed)

        self.target_glu = None
        if cfg.target_glu:
            self.target_glu = nn.Sequential(
                nn.Linear(final_dim, final_dim * 2), nn.GLU()
            )

        self.untie_final_proj = cfg.untie_final_proj
        if self.untie_final_proj:
            self.final_proj = nn.Linear(
                cfg.encoder_embed_dim, final_dim * len(dictionaries)
            )
        else:
            self.final_proj = nn.Linear(cfg.encoder_embed_dim, final_dim)

        # modules below are not needed during fine-tuning
        if any([d is None for d in dictionaries]):
            print("cannot find dictionary. assume will be used for fine-tuning")
        else:
            self.num_classes = [len(d) for d in dictionaries]
            self.label_embs_concat = nn.Parameter(
                torch.FloatTensor(sum(self.num_classes), final_dim)
            )
            nn.init.uniform_(self.label_embs_concat)

    def upgrade_state_dict_named(self, state_dict, name):
        """Upgrade a (possibly old) state dict for new versions of fairseq."""

        super().upgrade_state_dict_named(state_dict, name)
        return state_dict

    # @classmethod
    # def build_model(cls, cfg: AVHubertConfig, task: AVHubertPretrainingTask):
    #     """Build a new model instance."""

    #     kwargs = {}
    #     model = AVHubertModel(cfg, task.cfg, task.dictionaries, **kwargs)
    #     return model

    def apply_input_mask(self, x, padding_mask, target_list):
        B, C, T = x.shape[:3]
        is_audio = True if len(x.shape) == 3 else False
        if is_audio:
            mask_prob, mask_length = self.mask_prob_audio, self.mask_length_audio
        else:
            mask_prob, mask_length = self.mask_prob_image, self.mask_length_image
        if mask_prob > 0:
            mask_indices, starts, ends, batch_indexes = compute_mask_indices(
                (B, T),
                padding_mask,
                mask_prob,
                mask_length,
                self.mask_selection,
                self.mask_other,
                min_masks=2,
                no_overlap=self.no_mask_overlap,
                min_space=self.mask_min_space,
            )
            mask_indices_np = mask_indices
            mask_indices = torch.from_numpy(mask_indices).to(x.device)
            x = x.transpose(1, 2).contiguous()  # [B, T, C, H, W]
            if B == 1:
                x[mask_indices] = 0
            elif is_audio:
                x[mask_indices] = self.mask_emb
            elif self.selection_type == "same_other_seq":
                perm = (torch.arange(B) + torch.randint(low=1, high=B, size=(1,))) % B
                x_perm = x[perm]
                x[mask_indices] = x_perm[mask_indices]
            elif self.selection_type == "same_seq":
                batch_indexes_, other_indexes = [], []
                for batch_index, start, end in zip(batch_indexes, starts, ends):
                    length = end - start
                    other_start = np.setdiff1d(
                        np.arange(T), np.arange(max(0, start - length), end)
                    )
                    if len(other_start) > 0:
                        other_start = np.random.choice(other_start, size=1)
                    else:
                        other_start = 0
                    other_end = other_start + length
                    other_indexes.append(
                        np.arange(other_start, other_end).clip(max=T - 1)
                    )
                    batch_indexes_.append(
                        np.zeros([length], dtype=np.int64) + batch_index
                    )
                batch_indexes, other_indexes = np.concatenate(
                    batch_indexes_
                ), np.concatenate(other_indexes)
                x[mask_indices] = x[batch_indexes, other_indexes]

            x = x.transpose(1, 2).contiguous()
        else:
            mask_indices = None

        if self.mask_channel_prob > 0:
            print(f"No mask channel prob for input masking")
        return x, mask_indices

    def apply_feature_mask(self, x, padding_mask, target_list):
        B, T, C = x.shape
        assert (
            self.mask_prob_audio == self.mask_prob_image
            and self.mask_length_audio == self.mask_length_image
        ), f"masking prob/length for image/audio be same for feature masking"
        mask_prob, mask_length = self.mask_prob_audio, self.mask_length_image
        if mask_prob > 0:
            mask_indices, _, _, _ = compute_mask_indices(
                (B, T),
                padding_mask,
                mask_prob,
                mask_length,
                self.mask_selection,
                self.mask_other,
                min_masks=2,
                no_overlap=self.no_mask_overlap,
                min_space=self.mask_min_space,
            )
            mask_indices = torch.from_numpy(mask_indices).to(x.device)
            x[mask_indices] = self.mask_emb
        else:
            mask_indices = None

        if self.mask_channel_prob > 0:
            mask_channel_indices, _, _, _ = compute_mask_indices(
                (B, C),
                None,
                self.mask_channel_prob,
                self.mask_channel_length,
                self.mask_channel_selection,
                self.mask_channel_other,
                no_overlap=self.no_mask_channel_overlap,
                min_space=self.mask_channel_min_space,
            )
            mask_channel_indices = (
                torch.from_numpy(mask_channel_indices)
                .to(x.device)
                .unsqueeze(1)
                .expand(-1, T, -1)
            )
            x[mask_channel_indices] = 0

        return x, mask_indices

    def forward_features(self, source: torch.Tensor, modality: str) -> torch.Tensor:
        extractor = eval(f"self.feature_extractor_{modality}")
        if self.feature_grad_mult > 0:
            features = extractor(source)
            if self.feature_grad_mult != 1.0:
                features = GradMultiply.apply(features, self.feature_grad_mult)
        else:
            with torch.no_grad():
                features = extractor(source)
        return features

    def forward_targets(
        self,
        features: torch.Tensor,
        mask_indices: torch.Tensor,
        target_list: List[torch.Tensor],
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        # Trim features to ensure labels exist and then get aligned labels
        feat_tsz = features.size(2)
        targ_tsz = min([t.size(1) for t in target_list])
        if self.feat2tar_ratio * feat_tsz > targ_tsz:
            feat_tsz = int(targ_tsz / self.feat2tar_ratio)
            features = features[..., :feat_tsz]
            if mask_indices is not None:
                mask_indices = mask_indices[..., :feat_tsz]
        target_inds = torch.arange(feat_tsz).float() * self.feat2tar_ratio
        target_list = [t[:, target_inds.long()] for t in target_list]
        return features, mask_indices, target_list

    def forward_padding_mask(
        self,
        features: torch.Tensor,
        padding_mask: torch.Tensor,
    ) -> torch.Tensor:
        extra = padding_mask.size(1) % features.size(1)
        if extra > 0:
            padding_mask = padding_mask[:, :-extra]
        padding_mask = padding_mask.view(padding_mask.size(0), features.size(1), -1)
        padding_mask = padding_mask.all(-1)
        return padding_mask

    def compute_logits(self, feats, emb_mat):
        # feats: [B, T, F], emb_mat: [V, F]
        if self.sim_type == "dot":
            logits = torch.matmul(feats, emb_mat.transpose(0, 1))
        elif self.sim_type == "cosine":
            batch_size, timesteps, emb_dim = feats.size()
            feats_ = feats.view(-1, emb_dim)
            nom = (feats_.unsqueeze(dim=1) * emb_mat.unsqueeze(dim=0)).sum(
                dim=-1
            )  # [B*T, V]
            denom = (feats_**2).sum(dim=-1).sqrt().unsqueeze(dim=1) * (
                emb_mat**2
            ).sum(dim=-1).sqrt().unsqueeze(
                dim=0
            )  # [B*T, V]
            logits = (nom / denom.clamp(min=1e-6)).view(batch_size, timesteps, -1)
        else:
            raise NotImplementedError
        logits = logits / self.logit_temp
        return logits

    def forward(
        self,
        source: torch.Tensor,
        target_list: Optional[List[torch.Tensor]] = None,
        padding_mask: Optional[torch.Tensor] = None,
        mask: bool = True,
        features_only: bool = False,
        output_layer: Optional[int] = None,
    ) -> Dict[str, torch.Tensor]:
        """output layer is 1-based"""
        src_audio, src_video = source["audio"], source["video"]
        if mask and self.masking_type == "input":
            src_video, mask_indices_video = self.apply_input_mask(
                src_video, padding_mask, target_list
            )
            src_audio, mask_indices_audio = self.apply_input_mask(
                src_audio, padding_mask, target_list
            )
            mask_indices = torch.logical_or(mask_indices_audio, mask_indices_video)
        else:
            src_audio, src_video, mask_indices = src_audio, src_video, None

        if src_audio is not None and src_video is None:
            features_audio = self.forward_features(
                src_audio, modality="audio"
            )  # features: [B, F, T]
            features_video = features_audio.new_zeros(
                features_audio.size(0), self.encoder_embed_dim, features_audio.size(-1)
            )
        elif src_audio is None and src_video is not None:
            features_video = self.forward_features(src_video, modality="video")
            features_audio = features_video.new_zeros(
                features_video.size(0), self.encoder_embed_dim, features_video.size(-1)
            )
        elif src_audio is not None and src_video is not None:
            features_video = self.forward_features(src_video, modality="video")
            features_audio = self.forward_features(
                src_audio, modality="audio"
            )  # features: [B, F, T]
        # features_audio = self.forward_features(src_audio, modality='audio') # features: [B, F, T]
        # features_video = self.forward_features(src_video, modality='video')
        modality_drop_prob, audio_drop_prob = np.random.random(), np.random.random()

        if modality_drop_prob < self.modality_dropout:
            if audio_drop_prob < self.audio_dropout:
                features_audio = 0 * features_audio
            else:
                features_video = 0 * features_video

        if self.modality_fuse == "concat":
            features = torch.cat([features_audio, features_video], dim=1)
        elif self.modality_fuse == "add":
            features = features_audio + features_video
        if target_list is not None:
            features, mask_indices, target_list = self.forward_targets(
                features, mask_indices, target_list
            )

        # features_pen = features.float().pow(2).mean()

        features = features.transpose(1, 2)
        features = self.layer_norm(features)

        if padding_mask is not None:
            padding_mask = self.forward_padding_mask(features, padding_mask)

        if self.post_extract_proj is not None:
            features = self.post_extract_proj(features)

        features = self.dropout_input(features)
        if self.masking_type == "feature" and mask:
            x, mask_indices = self.apply_feature_mask(
                features, padding_mask, target_list
            )
        else:
            x = features

        # feature: (B, T, D), float
        # target: (B, T), long
        # x: (B, T, D), float
        # padding_mask: (B, T), bool
        # mask_indices: (B, T), bool
        x, _ = self.encoder(
            x,
            padding_mask=padding_mask,
            layer=None if output_layer is None else output_layer - 1,
        )

        if features_only:
            return {
                "x": x,
                "padding_mask": padding_mask,
                "features_audio": features_audio,
                "features_video": features_video,
                "features": features,  # input to transformer layers (concat/sum of features_audio & features_video, through layer norm + post_extract_proj)
            }
        # should never go down below!

        label_embs_list = self.label_embs_concat.split(self.num_classes, 0)
        proj_x = self.final_proj(x)
        if self.untie_final_proj:
            proj_x_list = proj_x.chunk(len(self.num_classes), dim=-1)
        else:
            proj_x_list = [proj_x for _ in self.num_classes]
        logit_list = [
            self.compute_logits(proj, emb).view(-1, num_class)
            for proj, emb, num_class in zip(
                proj_x_list, label_embs_list, self.num_classes
            )
        ]  # [[B*T, V]]
        mask, unmask = torch.logical_and(mask_indices, ~padding_mask).view(
            -1
        ), torch.logical_and(~mask_indices, ~padding_mask).view(
            -1
        )  # [B*T]
        logit_m_list, logit_u_list = [logit[mask] for logit in logit_list], [
            logit[unmask] for logit in logit_list
        ]
        target_m_list, target_u_list = [
            target.view(-1)[mask].long() for target in target_list
        ], [target.view(-1)[unmask].long() for target in target_list]
        result = {
            "logit_m_list": logit_m_list,
            "logit_u_list": logit_u_list,
            "target_m_list": target_m_list,
            "target_u_list": target_u_list,
            "padding_mask": padding_mask,
            "features_pen": features_pen,
        }
        return result

    def extract_features(
        self,
        source: torch.Tensor,
        padding_mask: Optional[torch.Tensor] = None,
        mask: bool = False,
        ret_conv: bool = False,
        output_layer: Optional[int] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        res = self.forward(
            source,
            padding_mask=padding_mask,
            mask=mask,
            features_only=True,
            output_layer=output_layer,
        )
        feature = res["features"] if ret_conv else res["x"]
        return feature, res["padding_mask"]

    def extract_finetune(
        self, source, padding_mask=None, mask=False, ret_conv=False, output_layer=None
    ):
        src_audio, src_video = source["audio"], source["video"]
        if mask and self.masking_type == "input":
            src_video, mask_indices_video = self.apply_input_mask(
                src_video, padding_mask, target_list=None
            )
            src_audio, mask_indices_audio = self.apply_input_mask(
                src_audio, padding_mask, target_list=None
            )
            mask_indices = torch.logical_or(
                mask_indices_audio, mask_indices_video
            )  # mask_indices not used in fine-tuning
        else:
            src_audio, src_video, mask_indices = src_audio, src_video, None

        if src_audio is not None and src_video is None:
            features_audio = self.forward_features(
                src_audio, modality="audio"
            )  # features: [B, F, T]
            features_video = features_audio.new_zeros(
                features_audio.size(0), self.encoder_embed_dim, features_audio.size(-1)
            )
        elif src_audio is None and src_video is not None:
            features_video = self.forward_features(src_video, modality="video")
            features_audio = features_video.new_zeros(
                features_video.size(0), self.encoder_embed_dim, features_video.size(-1)
            )
        elif src_audio is not None and src_video is not None:
            features_video = self.forward_features(src_video, modality="video")
            features_audio = self.forward_features(
                src_audio, modality="audio"
            )  # features: [B, F, T]

        if self.modality_fuse == "concat":
            features = torch.cat([features_audio, features_video], dim=1)
        elif self.modality_fuse == "add":
            features = features_audio + features_video
        features_pen = features.float().pow(2).mean()

        features = features.transpose(1, 2)
        features = self.layer_norm(features)
        unmasked_features = features.clone()

        if padding_mask is not None:
            padding_mask = self.forward_padding_mask(features, padding_mask)

        if self.post_extract_proj is not None:
            features = self.post_extract_proj(features)

        features = self.dropout_input(features)
        unmasked_features = self.dropout_features(unmasked_features)
        x = features
        mask_indices = None

        # feature: (B, T, D), float
        # target: (B, T), long
        # x: (B, T, D), float
        # padding_mask: (B, T), bool
        # mask_indices: (B, T), bool
        x, _ = self.encoder(
            x,
            padding_mask=padding_mask,
            layer=None if output_layer is None else output_layer - 1,
        )

        return x, padding_mask

    def get_extra_losses(self, net_output):
        extra_losses = []
        names = []
        if "features_pen" in net_output:
            extra_losses.append(net_output["features_pen"])
            names.append("features_pen")

        return extra_losses, names

    def remove_pretraining_modules(self):
        self.target_glu = None
        self.final_proj = None

    def get_logits(self, net_output, is_masked=True):
        raise NotImplementedError

    def get_targets(self, net_output, is_masked=True):
        raise NotImplementedError

    def compute_nce(self, x, pos, negs):
        neg_is_pos = (pos == negs).all(-1)
        pos = pos.unsqueeze(0)
        targets = torch.cat([pos, negs], dim=0)

        logits = torch.cosine_similarity(x.float(), targets.float(), dim=-1).type_as(x)
        logits /= self.logit_temp
        if neg_is_pos.any():
            logits[1:][neg_is_pos] = float("-inf")
        logits = logits.transpose(0, 1)  # (num_x, num_cls+1)
        return logits

upstream_models/vhubert/hubert.py
roger-tseng/av-superb
# This file is modified from code licensed under the AV-HuBERT license (see COPYRIGHT)
# Copyright (c) Meta Platforms, Inc. All Rights Reserved.

# Copyright (c) Facebook, Inc. and its 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 math
import pdb
from collections import OrderedDict

import torch
import torch.nn as nn

logger = logging.getLogger(__name__)


def conv3x3(in_planes, out_planes, stride=1):
    return nn.Conv2d(
        in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False
    )


def downsample_basic_block(inplanes, outplanes, stride):
    return nn.Sequential(
        nn.Conv2d(inplanes, outplanes, kernel_size=1, stride=stride, bias=False),
        nn.BatchNorm2d(outplanes),
    )


def downsample_basic_block_v2(inplanes, outplanes, stride):
    return nn.Sequential(
        nn.AvgPool2d(
            kernel_size=stride, stride=stride, ceil_mode=True, count_include_pad=False
        ),
        nn.Conv2d(inplanes, outplanes, kernel_size=1, stride=1, bias=False),
        nn.BatchNorm2d(outplanes),
    )


class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, inplanes, planes, stride=1, downsample=None, relu_type="relu"):
        super(BasicBlock, self).__init__()

        assert relu_type in ["relu", "prelu"]

        self.conv1 = conv3x3(inplanes, planes, stride)
        self.bn1 = nn.BatchNorm2d(planes)

        if relu_type == "relu":
            self.relu1 = nn.ReLU(inplace=True)
            self.relu2 = nn.ReLU(inplace=True)
        elif relu_type == "prelu":
            self.relu1 = nn.PReLU(num_parameters=planes)
            self.relu2 = nn.PReLU(num_parameters=planes)
        else:
            raise Exception("relu type not implemented")

        self.conv2 = conv3x3(planes, planes)
        self.bn2 = nn.BatchNorm2d(planes)

        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu1(out)
        out = self.conv2(out)
        out = self.bn2(out)
        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.relu2(out)

        return out


class ResNet(nn.Module):
    def __init__(
        self,
        block,
        layers,
        num_classes=1000,
        relu_type="relu",
        gamma_zero=False,
        avg_pool_downsample=False,
    ):
        self.inplanes = 64
        self.relu_type = relu_type
        self.gamma_zero = gamma_zero
        self.downsample_block = (
            downsample_basic_block_v2 if avg_pool_downsample else downsample_basic_block
        )

        super(ResNet, self).__init__()
        self.layer1 = self._make_layer(block, 64, layers[0])
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
        self.avgpool = nn.AdaptiveAvgPool2d(1)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2.0 / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()

        if self.gamma_zero:
            for m in self.modules():
                if isinstance(m, BasicBlock):
                    m.bn2.weight.data.zero_()

    def _make_layer(self, block, planes, blocks, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = self.downsample_block(
                inplanes=self.inplanes,
                outplanes=planes * block.expansion,
                stride=stride,
            )

        layers = []
        layers.append(
            block(self.inplanes, planes, stride, downsample, relu_type=self.relu_type)
        )
        self.inplanes = planes * block.expansion
        for i in range(1, blocks):
            layers.append(block(self.inplanes, planes, relu_type=self.relu_type))

        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.avgpool(x)
        x = x.view(x.size(0), -1)
        return x


class ResEncoder(nn.Module):
    def __init__(self, relu_type, weights):
        super(ResEncoder, self).__init__()
        self.frontend_nout = 64
        self.backend_out = 512
        frontend_relu = (
            nn.PReLU(num_parameters=self.frontend_nout)
            if relu_type == "prelu"
            else nn.ReLU()
        )
        self.frontend3D = nn.Sequential(
            nn.Conv3d(
                1,
                self.frontend_nout,
                kernel_size=(5, 7, 7),
                stride=(1, 2, 2),
                padding=(2, 3, 3),
                bias=False,
            ),
            nn.BatchNorm3d(self.frontend_nout),
            frontend_relu,
            nn.MaxPool3d(kernel_size=(1, 3, 3), stride=(1, 2, 2), padding=(0, 1, 1)),
        )
        self.trunk = ResNet(BasicBlock, [2, 2, 2, 2], relu_type=relu_type)
        if weights is not None:
            logger.info(f"Load {weights} for resnet")
            std = torch.load(weights, map_location=torch.device("cpu"))[
                "model_state_dict"
            ]
            frontend_std, trunk_std = OrderedDict(), OrderedDict()
            for key, val in std.items():
                new_key = ".".join(key.split(".")[1:])
                if "frontend3D" in key:
                    frontend_std[new_key] = val
                if "trunk" in key:
                    trunk_std[new_key] = val
            self.frontend3D.load_state_dict(frontend_std)
            self.trunk.load_state_dict(trunk_std)

    def forward(self, x):
        B, C, T, H, W = x.size()
        x = self.frontend3D(x)
        Tnew = x.shape[2]
        x = self.threeD_to_2D_tensor(x)
        x = self.trunk(x)
        x = x.view(B, Tnew, x.size(1))
        x = x.transpose(1, 2).contiguous()
        return x

    def threeD_to_2D_tensor(self, x):
        n_batch, n_channels, s_time, sx, sy = x.shape
        x = x.transpose(1, 2).contiguous()
        return x.reshape(n_batch * s_time, n_channels, sx, sy)

upstream_models/vhubert/resnet.py
roger-tseng/av-superb
# This file is modified from code licensed under the AV-HuBERT license (see COPYRIGHT)
# Copyright (c) Meta Platforms, Inc. All Rights Reserved.

# Copyright (c) Facebook, Inc. and its 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
from typing import Dict, List, Optional, Tuple

import cv2
import numpy as np
import torch


#### CV utils from av-hubert
class Compose(object):
    """Compose several preprocess together.
    Args:
        preprocess (list of ``Preprocess`` objects): list of preprocess to compose.
    """

    def __init__(self, preprocess):
        self.preprocess = preprocess

    def __call__(self, sample):
        for t in self.preprocess:
            sample = t(sample)
        return sample

    def __repr__(self):
        format_string = self.__class__.__name__ + "("
        for t in self.preprocess:
            format_string += "\n"
            format_string += "    {0}".format(t)
        format_string += "\n)"
        return format_string


class Normalize(object):
    """Normalize a ndarray image with mean and standard deviation."""

    def __init__(self, mean, std):
        self.mean = mean
        self.std = std

    def __call__(self, frames):
        """
        Args:
            tensor (Tensor): Tensor image of size (C, H, W) to be normalized.
        Returns:
            Tensor: Normalized Tensor image.
        """
        frames = (frames - self.mean) / self.std
        return frames

    def __repr__(self):
        return self.__class__.__name__ + "(mean={0}, std={1})".format(
            self.mean, self.std
        )


class CenterCrop(object):
    """Crop the given image at the center"""

    def __init__(self, size):
        self.size = size

    def __call__(self, frames):
        """
        Args:
            img (np.ndarray): Images to be cropped.
        Returns:
            np.ndarray: Cropped image.
        """
        t, h, w = frames.shape
        th, tw = self.size
        # Pad surrounding with zeros if smaller that target size
        if tw > w or th > h:
            padding = [
                (0, 0),
                ((th - h) // 2, (th - h + 1) // 2) if th > h else (0, 0),
                ((tw - w) // 2, (tw - w + 1) // 2) if tw > w else (0, 0),
            ]
            frames = np.pad(
                frames, padding, mode="constant", constant_values=0
            )  # PIL uses fill value 0
            t, h, w = frames.shape
            if tw == w and th == h:
                return frames

        delta_w = int(round((w - tw)) / 2.0)
        delta_h = int(round((h - th)) / 2.0)
        frames = frames[:, delta_h : delta_h + th, delta_w : delta_w + tw]
        return frames


class RandomCrop(object):
    """Crop the given image at the center"""

    def __init__(self, size):
        self.size = size

    def __call__(self, frames):
        """
        Args:
            img (np.ndarray): Images to be cropped.
        Returns:
            np.ndarray: Cropped image.
        """
        t, h, w = frames.shape
        th, tw = self.size
        # Pad surrounding with zeros if smaller that target size
        if tw > w or th > h:
            padding = [
                (0, 0),
                ((th - h) // 2, (th - h + 1) // 2) if th > h else (0, 0),
                ((tw - w) // 2, (tw - w + 1) // 2) if tw > w else (0, 0),
            ]
            frames = np.pad(
                frames, padding, mode="constant", constant_values=0
            )  # PIL uses fill value 0
            t, h, w = frames.shape
            if tw == w and th == h:
                return frames

        delta_w = random.randint(0, w - tw)
        delta_h = random.randint(0, h - th)
        frames = frames[:, delta_h : delta_h + th, delta_w : delta_w + tw]
        return frames

    def __repr__(self):
        return self.__class__.__name__ + "(size={0})".format(self.size)


class HorizontalFlip(object):
    """Flip image horizontally."""

    def __init__(self, flip_ratio):
        self.flip_ratio = flip_ratio

    def __call__(self, frames):
        """
        Args:
            img (np.ndarray): Images to be flipped with a probability flip_ratio
        Returns:
            np.ndarray: Cropped image.
        """
        t, h, w = frames.shape
        if random.random() < self.flip_ratio:
            for index in range(t):
                frames[index] = cv2.flip(frames[index], 1)
        return frames


def compute_mask_indices(
    shape: Tuple[int, int],
    padding_mask: Optional[torch.Tensor],
    mask_prob: float,
    mask_length: int,
    mask_type: str = "static",
    mask_other: float = 0.0,
    min_masks: int = 0,
    no_overlap: bool = False,
    min_space: int = 0,
) -> np.ndarray:
    """
    Computes random mask spans for a given shape
    Args:
        shape: the the shape for which to compute masks.
            should be of size 2 where first element is batch size and 2nd is timesteps
        padding_mask: optional padding mask of the same size as shape, which will prevent masking padded elements
        mask_prob: probability for each token to be chosen as start of the span to be masked. this will be multiplied by
            number of timesteps divided by length of mask span to mask approximately this percentage of all elements.
            however due to overlaps, the actual number will be smaller (unless no_overlap is True)
        mask_type: how to compute mask lengths
            static = fixed size
            uniform = sample from uniform distribution [mask_other, mask_length*2]
            normal = sample from normal distribution with mean mask_length and stdev mask_other. mask is min 1 element
            poisson = sample from possion distribution with lambda = mask length
        min_masks: minimum number of masked spans
        no_overlap: if false, will switch to an alternative recursive algorithm that prevents spans from overlapping
        min_space: only used if no_overlap is True, this is how many elements to keep unmasked between spans
    """

    bsz, all_sz = shape
    mask = np.full((bsz, all_sz), False)

    all_num_mask = int(
        # add a random number for probabilistic rounding
        mask_prob * all_sz / float(mask_length)
        + np.random.rand()
    )

    all_num_mask = max(min_masks, all_num_mask)

    mask_idcs = []
    for i in range(bsz):
        if padding_mask is not None:
            sz = all_sz - padding_mask[i].long().sum().item()
            num_mask = int(
                # add a random number for probabilistic rounding
                mask_prob * sz / float(mask_length)
                + np.random.rand()
            )
            num_mask = max(min_masks, num_mask)
        else:
            sz = all_sz
            num_mask = all_num_mask

        if mask_type == "static":
            lengths = np.full(num_mask, mask_length)
        elif mask_type == "uniform":
            lengths = np.random.randint(mask_other, mask_length * 2 + 1, size=num_mask)
        elif mask_type == "normal":
            lengths = np.random.normal(mask_length, mask_other, size=num_mask)
            lengths = [max(1, int(round(x))) for x in lengths]
        elif mask_type == "poisson":
            lengths = np.random.poisson(mask_length, size=num_mask)
            lengths = [int(round(x)) for x in lengths]
        else:
            raise Exception("unknown mask selection " + mask_type)

        if sum(lengths) == 0:
            lengths[0] = min(mask_length, sz - 1)

        if no_overlap:
            mask_idc = []

            def arrange(s, e, length, keep_length):
                span_start = np.random.randint(s, e - length)
                mask_idc.extend(span_start + i for i in range(length))

                new_parts = []
                if span_start - s - min_space >= keep_length:
                    new_parts.append((s, span_start - min_space + 1))
                if e - span_start - keep_length - min_space > keep_length:
                    new_parts.append((span_start + length + min_space, e))
                return new_parts

            parts = [(0, sz)]
            min_length = min(lengths)
            for length in sorted(lengths, reverse=True):
                lens = np.fromiter(
                    (e - s if e - s >= length + min_space else 0 for s, e in parts),
                    np.int,
                )
                l_sum = np.sum(lens)
                if l_sum == 0:
                    break
                probs = lens / np.sum(lens)
                c = np.random.choice(len(parts), p=probs)
                s, e = parts.pop(c)
                parts.extend(arrange(s, e, length, min_length))
            mask_idc = np.asarray(mask_idc)
        else:
            min_len = min(lengths)
            if sz - min_len <= num_mask:
                min_len = sz - num_mask - 1

            mask_idc = np.random.choice(sz - min_len, num_mask, replace=False)

            mask_idc = np.asarray(
                [
                    mask_idc[j] + offset
                    for j in range(len(mask_idc))
                    for offset in range(lengths[j])
                ]
            )

        mask_idcs.append(np.unique(mask_idc[mask_idc < sz]))

    min_len = min([len(m) for m in mask_idcs])
    batch_indexes, starts, ends = [], [], []
    for i, mask_idc in enumerate(mask_idcs):
        if len(mask_idc) > min_len:
            mask_idc = np.random.choice(mask_idc, min_len, replace=False)
        mask[i, mask_idc] = True
        vals, run_starts, run_lengths = find_runs(mask[i])
        start_indices, lengths = run_starts[vals == True], run_lengths[vals == True]
        starts.append(start_indices)
        ends.append(start_indices + lengths)
        batch_indexes.append(np.zeros([len(start_indices)]) + i)
    return (
        mask,
        np.concatenate(starts).astype(np.int64),
        np.concatenate(ends).astype(np.int64),
        np.concatenate(batch_indexes).astype(np.int64),
    )


def find_runs(x):
    """Find runs of consecutive items in an array."""

    # ensure array
    x = np.asanyarray(x)
    if x.ndim != 1:
        raise ValueError("only 1D array supported")
    n = x.shape[0]

    # handle empty array
    if n == 0:
        return np.array([]), np.array([]), np.array([])

    else:
        # find run starts
        loc_run_start = np.empty(n, dtype=bool)
        loc_run_start[0] = True
        np.not_equal(x[:-1], x[1:], out=loc_run_start[1:])
        run_starts = np.nonzero(loc_run_start)[0]

        # find run values
        run_values = x[loc_run_start]

        # find run lengths
        run_lengths = np.diff(np.append(run_starts, n))

        return run_values, run_starts, run_lengths


#### audio utils from https://github.com/jameslyons/python_speech_features/
import decimal
import logging
import math

import numpy


def round_half_up(number):
    return int(
        decimal.Decimal(number).quantize(
            decimal.Decimal("1"), rounding=decimal.ROUND_HALF_UP
        )
    )


def rolling_window(a, window, step=1):
    # http://ellisvalentiner.com/post/2017-03-21-np-strides-trick
    shape = a.shape[:-1] + (a.shape[-1] - window + 1, window)
    strides = a.strides + (a.strides[-1],)
    return np.lib.stride_tricks.as_strided(a, shape=shape, strides=strides)[::step]


def framesig(
    sig, frame_len, frame_step, winfunc=lambda x: np.ones((x,)), stride_trick=True
):
    """Frame a signal into overlapping frames.
    :param sig: the audio signal to frame.
    :param frame_len: length of each frame measured in samples.
    :param frame_step: number of samples after the start of the previous frame that the next frame should begin.
    :param winfunc: the analysis window to apply to each frame. By default no window is applied.
    :param stride_trick: use stride trick to compute the rolling window and window multiplication faster
    :returns: an array of frames. Size is NUMFRAMES by frame_len.
    """
    slen = len(sig)
    frame_len = int(round_half_up(frame_len))
    frame_step = int(round_half_up(frame_step))
    if slen <= frame_len:
        numframes = 1
    else:
        numframes = 1 + int(math.ceil((1.0 * slen - frame_len) / frame_step))

    padlen = int((numframes - 1) * frame_step + frame_len)

    zeros = np.zeros((padlen - slen,))
    padsignal = np.concatenate((sig, zeros))
    if stride_trick:
        win = winfunc(frame_len)
        frames = rolling_window(padsignal, window=frame_len, step=frame_step)
    else:
        indices = (
            np.tile(np.arange(0, frame_len), (numframes, 1))
            + np.tile(
                np.arange(0, numframes * frame_step, frame_step), (frame_len, 1)
            ).T
        )
        indices = np.array(indices, dtype=np.int32)
        frames = padsignal[indices]
        win = np.tile(winfunc(frame_len), (numframes, 1))

    return frames * win


def deframesig(frames, siglen, frame_len, frame_step, winfunc=lambda x: np.ones((x,))):
    """Does overlap-add procedure to undo the action of framesig.
    :param frames: the array of frames.
    :param siglen: the length of the desired signal, use 0 if unknown. Output will be truncated to siglen samples.
    :param frame_len: length of each frame measured in samples.
    :param frame_step: number of samples after the start of the previous frame that the next frame should begin.
    :param winfunc: the analysis window to apply to each frame. By default no window is applied.
    :returns: a 1-D signal.
    """
    frame_len = round_half_up(frame_len)
    frame_step = round_half_up(frame_step)
    numframes = np.shape(frames)[0]
    assert (
        np.shape(frames)[1] == frame_len
    ), '"frames" matrix is wrong size, 2nd dim is not equal to frame_len'

    indices = (
        np.tile(np.arange(0, frame_len), (numframes, 1))
        + np.tile(np.arange(0, numframes * frame_step, frame_step), (frame_len, 1)).T
    )
    indices = np.array(indices, dtype=np.int32)
    padlen = (numframes - 1) * frame_step + frame_len

    if siglen <= 0:
        siglen = padlen

    rec_signal = np.zeros((padlen,))
    window_correction = np.zeros((padlen,))
    win = winfunc(frame_len)

    for i in range(0, numframes):
        window_correction[indices[i, :]] = (
            window_correction[indices[i, :]] + win + 1e-15
        )  # add a little bit so it is never zero
        rec_signal[indices[i, :]] = rec_signal[indices[i, :]] + frames[i, :]

    rec_signal = rec_signal / window_correction
    return rec_signal[0:siglen]


def magspec(frames, NFFT):
    """Compute the magnitude spectrum of each frame in frames. If frames is an NxD matrix, output will be Nx(NFFT/2+1).
    :param frames: the array of frames. Each row is a frame.
    :param NFFT: the FFT length to use. If NFFT > frame_len, the frames are zero-padded.
    :returns: If frames is an NxD matrix, output will be Nx(NFFT/2+1). Each row will be the magnitude spectrum of the corresponding frame.
    """
    if np.shape(frames)[1] > NFFT:
        logging.warn(
            "frame length (%d) is greater than FFT size (%d), frame will be truncated. Increase NFFT to avoid.",
            np.shape(frames)[1],
            NFFT,
        )
    complex_spec = np.fft.rfft(frames, NFFT)
    return np.absolute(complex_spec)


def powspec(frames, NFFT):
    """Compute the power spectrum of each frame in frames. If frames is an NxD matrix, output will be Nx(NFFT/2+1).
    :param frames: the array of frames. Each row is a frame.
    :param NFFT: the FFT length to use. If NFFT > frame_len, the frames are zero-padded.
    :returns: If frames is an NxD matrix, output will be Nx(NFFT/2+1). Each row will be the power spectrum of the corresponding frame.
    """
    return 1.0 / NFFT * np.square(magspec(frames, NFFT))


def logpowspec(frames, NFFT, norm=1):
    """Compute the log power spectrum of each frame in frames. If frames is an NxD matrix, output will be Nx(NFFT/2+1).
    :param frames: the array of frames. Each row is a frame.
    :param NFFT: the FFT length to use. If NFFT > frame_len, the frames are zero-padded.
    :param norm: If norm=1, the log power spectrum is normalised so that the max value (across all frames) is 0.
    :returns: If frames is an NxD matrix, output will be Nx(NFFT/2+1). Each row will be the log power spectrum of the corresponding frame.
    """
    ps = powspec(frames, NFFT)
    ps[ps <= 1e-30] = 1e-30
    lps = 10 * np.log10(ps)
    if norm:
        return lps - np.max(lps)
    else:
        return lps


def preemphasis(signal, coeff=0.95):
    """perform preemphasis on the input signal.
    :param signal: The signal to filter.
    :param coeff: The preemphasis coefficient. 0 is no filter, default is 0.95.
    :returns: the filtered signal.
    """
    return np.append(signal[0], signal[1:] - coeff * signal[:-1])


def fbank(
    signal,
    samplerate=16000,
    winlen=0.025,
    winstep=0.01,
    nfilt=26,
    nfft=512,
    lowfreq=0,
    highfreq=None,
    preemph=0.97,
    winfunc=lambda x: np.ones((x,)),
):
    """Compute Mel-filterbank energy features from an audio signal.
    :param signal: the audio signal from which to compute features. Should be an N*1 array
    :param samplerate: the sample rate of the signal we are working with, in Hz.
    :param winlen: the length of the analysis window in seconds. Default is 0.025s (25 milliseconds)
    :param winstep: the step between successive windows in seconds. Default is 0.01s (10 milliseconds)
    :param nfilt: the number of filters in the filterbank, default 26.
    :param nfft: the FFT size. Default is 512.
    :param lowfreq: lowest band edge of mel filters. In Hz, default is 0.
    :param highfreq: highest band edge of mel filters. In Hz, default is samplerate/2
    :param preemph: apply preemphasis filter with preemph as coefficient. 0 is no filter. Default is 0.97.
    :param winfunc: the analysis window to apply to each frame. By default no window is applied. You can use numpy window functions here e.g. winfunc=np.hamming
    :returns: 2 values. The first is a numpy array of size (NUMFRAMES by nfilt) containing features. Each row holds 1 feature vector. The
        second return value is the energy in each frame (total energy, unwindowed)
    """
    highfreq = highfreq or samplerate / 2
    signal = preemphasis(signal, preemph)
    frames = framesig(signal, winlen * samplerate, winstep * samplerate, winfunc)
    pspec = powspec(frames, nfft)
    energy = np.sum(pspec, 1)  # this stores the total energy in each frame
    energy = np.where(
        energy == 0, np.finfo(float).eps, energy
    )  # if energy is zero, we get problems with log

    fb = get_filterbanks(nfilt, nfft, samplerate, lowfreq, highfreq)
    feat = np.dot(pspec, fb.T)  # compute the filterbank energies
    feat = np.where(
        feat == 0, np.finfo(float).eps, feat
    )  # if feat is zero, we get problems with log

    return feat, energy


def logfbank(
    signal,
    samplerate=16000,
    winlen=0.025,
    winstep=0.01,
    nfilt=26,
    nfft=512,
    lowfreq=0,
    highfreq=None,
    preemph=0.97,
    winfunc=lambda x: np.ones((x,)),
):
    """Compute log Mel-filterbank energy features from an audio signal.
    :param signal: the audio signal from which to compute features. Should be an N*1 array
    :param samplerate: the sample rate of the signal we are working with, in Hz.
    :param winlen: the length of the analysis window in seconds. Default is 0.025s (25 milliseconds)
    :param winstep: the step between successive windows in seconds. Default is 0.01s (10 milliseconds)
    :param nfilt: the number of filters in the filterbank, default 26.
    :param nfft: the FFT size. Default is 512.
    :param lowfreq: lowest band edge of mel filters. In Hz, default is 0.
    :param highfreq: highest band edge of mel filters. In Hz, default is samplerate/2
    :param preemph: apply preemphasis filter with preemph as coefficient. 0 is no filter. Default is 0.97.
    :param winfunc: the analysis window to apply to each frame. By default no window is applied. You can use numpy window functions here e.g. winfunc=np.hamming
    :returns: A numpy array of size (NUMFRAMES by nfilt) containing features. Each row holds 1 feature vector.
    """
    feat, energy = fbank(
        signal,
        samplerate,
        winlen,
        winstep,
        nfilt,
        nfft,
        lowfreq,
        highfreq,
        preemph,
        winfunc,
    )
    return np.log(feat)


def ssc(
    signal,
    samplerate=16000,
    winlen=0.025,
    winstep=0.01,
    nfilt=26,
    nfft=512,
    lowfreq=0,
    highfreq=None,
    preemph=0.97,
    winfunc=lambda x: np.ones((x,)),
):
    """Compute Spectral Subband Centroid features from an audio signal.
    :param signal: the audio signal from which to compute features. Should be an N*1 array
    :param samplerate: the sample rate of the signal we are working with, in Hz.
    :param winlen: the length of the analysis window in seconds. Default is 0.025s (25 milliseconds)
    :param winstep: the step between successive windows in seconds. Default is 0.01s (10 milliseconds)
    :param nfilt: the number of filters in the filterbank, default 26.
    :param nfft: the FFT size. Default is 512.
    :param lowfreq: lowest band edge of mel filters. In Hz, default is 0.
    :param highfreq: highest band edge of mel filters. In Hz, default is samplerate/2
    :param preemph: apply preemphasis filter with preemph as coefficient. 0 is no filter. Default is 0.97.
    :param winfunc: the analysis window to apply to each frame. By default no window is applied. You can use numpy window functions here e.g. winfunc=np.hamming
    :returns: A numpy array of size (NUMFRAMES by nfilt) containing features. Each row holds 1 feature vector.
    """
    highfreq = highfreq or samplerate / 2
    signal = preemphasis(signal, preemph)
    frames = framesig(signal, winlen * samplerate, winstep * samplerate, winfunc)
    pspec = powspec(frames, nfft)
    pspec = np.where(
        pspec == 0, np.finfo(float).eps, pspec
    )  # if things are all zeros we get problems

    fb = get_filterbanks(nfilt, nfft, samplerate, lowfreq, highfreq)
    feat = np.dot(pspec, fb.T)  # compute the filterbank energies
    R = np.tile(
        np.linspace(1, samplerate / 2, np.size(pspec, 1)), (np.size(pspec, 0), 1)
    )

    return np.dot(pspec * R, fb.T) / feat


def hz2mel(hz):
    """Convert a value in Hertz to Mels
    :param hz: a value in Hz. This can also be a numpy array, conversion proceeds element-wise.
    :returns: a value in Mels. If an array was passed in, an identical sized array is returned.
    """
    return 2595 * np.log10(1 + hz / 700.0)


def mel2hz(mel):
    """Convert a value in Mels to Hertz
    :param mel: a value in Mels. This can also be a numpy array, conversion proceeds element-wise.
    :returns: a value in Hertz. If an array was passed in, an identical sized array is returned.
    """
    return 700 * (10 ** (mel / 2595.0) - 1)


def get_filterbanks(nfilt=20, nfft=512, samplerate=16000, lowfreq=0, highfreq=None):
    """Compute a Mel-filterbank. The filters are stored in the rows, the columns correspond
    to fft bins. The filters are returned as an array of size nfilt * (nfft/2 + 1)
    :param nfilt: the number of filters in the filterbank, default 20.
    :param nfft: the FFT size. Default is 512.
    :param samplerate: the sample rate of the signal we are working with, in Hz. Affects mel spacing.
    :param lowfreq: lowest band edge of mel filters, default 0 Hz
    :param highfreq: highest band edge of mel filters, default samplerate/2
    :returns: A numpy array of size nfilt * (nfft/2 + 1) containing filterbank. Each row holds 1 filter.
    """
    highfreq = highfreq or samplerate / 2
    assert highfreq <= samplerate / 2, "highfreq is greater than samplerate/2"

    # compute points evenly spaced in mels
    lowmel = hz2mel(lowfreq)
    highmel = hz2mel(highfreq)
    melpoints = np.linspace(lowmel, highmel, nfilt + 2)
    # our points are in Hz, but we use fft bins, so we have to convert
    #  from Hz to fft bin number
    bin = np.floor((nfft + 1) * mel2hz(melpoints) / samplerate)

    fbank = np.zeros([nfilt, nfft // 2 + 1])
    for j in range(0, nfilt):
        for i in range(int(bin[j]), int(bin[j + 1])):
            fbank[j, i] = (i - bin[j]) / (bin[j + 1] - bin[j])
        for i in range(int(bin[j + 1]), int(bin[j + 2])):
            fbank[j, i] = (bin[j + 2] - i) / (bin[j + 2] - bin[j + 1])
    return fbank


def logfbank(
    signal,
    samplerate=16000,
    winlen=0.025,
    winstep=0.01,
    nfilt=26,
    nfft=512,
    lowfreq=0,
    highfreq=None,
    preemph=0.97,
    winfunc=lambda x: np.ones((x,)),
):
    """Compute log Mel-filterbank energy features from an audio signal.
    :param signal: the audio signal from which to compute features. Should be an N*1 array
    :param samplerate: the sample rate of the signal we are working with, in Hz.
    :param winlen: the length of the analysis window in seconds. Default is 0.025s (25 milliseconds)
    :param winstep: the step between successive windows in seconds. Default is 0.01s (10 milliseconds)
    :param nfilt: the number of filters in the filterbank, default 26.
    :param nfft: the FFT size. Default is 512.
    :param lowfreq: lowest band edge of mel filters. In Hz, default is 0.
    :param highfreq: highest band edge of mel filters. In Hz, default is samplerate/2
    :param preemph: apply preemphasis filter with preemph as coefficient. 0 is no filter. Default is 0.97.
    :param winfunc: the analysis window to apply to each frame. By default no window is applied. You can use numpy window functions here e.g. winfunc=np.hamming
    :returns: A numpy array of size (NUMFRAMES by nfilt) containing features. Each row holds 1 feature vector.
    """
    feat, energy = fbank(
        signal,
        samplerate,
        winlen,
        winstep,
        nfilt,
        nfft,
        lowfreq,
        highfreq,
        preemph,
        winfunc,
    )
    return np.log(feat)

upstream_models/vhubert/utils.py
roger-tseng/av-superb
# Copied from S3PRL 
# (Authors: Leo Yang, Andy T. Liu and S3PRL team, https://github.com/s3prl/s3prl/blob/main/s3prl/utility/download.py)

import hashlib
import os

import gdown
import torch
from filelock import FileLock


def _download(filename, url, refresh, agent):
    dirpath = f"{torch.hub.get_dir()}/s3prl_cache"
    os.makedirs(dirpath, exist_ok=True)
    filepath = f"{dirpath}/{filename}"
    with FileLock(filepath + ".lock"):
        if not os.path.isfile(filepath) or refresh:
            if agent == "wget":
                os.system(f"wget {url} -O {filepath}")
            elif agent == "gdown":
                gdown.download(url, filepath, use_cookies=False)
            else:
                print(
                    "[Download] - Unknown download agent. Only 'wget' and 'gdown' are supported."
                )
                raise NotImplementedError
        else:
            print(f"Using cache found in {filepath}\nfor {url}")
    return filepath


def _urls_to_filepaths(*args, refresh=False, agent="wget"):
    """
    Preprocess the URL specified in *args into local file paths after downloading

    Args:
        Any number of URLs (1 ~ any)

    Return:
        Same number of downloaded file paths
    """

    def url_to_filename(url):
        assert type(url) is str
        m = hashlib.sha256()
        m.update(str.encode(url))
        return str(m.hexdigest())

    def url_to_path(url, refresh):
        if type(url) is str and len(url) > 0:
            return _download(url_to_filename(url), url, refresh, agent=agent)
        else:
            return None

    paths = [url_to_path(url, refresh) for url in args]
    return paths if len(paths) > 1 else paths[0]


def _gdriveids_to_filepaths(*args, refresh=False):
    """
    Preprocess the Google Drive id specified in *args into local file paths after downloading

    Args:
        Any number of Google Drive ids (1 ~ any)

    Return:
        Same number of downloaded file paths
    """

    def gdriveid_to_url(gdriveid):
        if type(gdriveid) is str and len(gdriveid) > 0:
            return f"https://drive.google.com/uc?id={gdriveid}"
        else:
            return None

    return _urls_to_filepaths(
        *[gdriveid_to_url(gid) for gid in args], refresh=refresh, agent="gdown"
    )

utils/download.py
roger-tseng/av-superb
# Copied from S3PRL 
# (https://github.com/s3prl/s3prl/blob/main/s3prl/utility/helper.py)

# -*- coding: utf-8 -*- #
"""*********************************************************************************************"""
#   FileName     [ utility/helper.py ]
#   Synopsis     [ helper functions ]
#   Author       [ Andy T. Liu (Andi611) ]
#   Copyright    [ Copyleft(c), Speech Lab, NTU, Taiwan ]
"""*********************************************************************************************"""


import builtins
import math

###############
# IMPORTATION #
###############
import os
import shutil
import sys
from collections import defaultdict
from datetime import datetime
from pathlib import Path
from time import time
from typing import List

import numpy as np
import torch
from torch.distributed import get_rank, get_world_size, is_initialized
from torch.nn.parallel import DistributedDataParallel as DDP


def is_leader_process():
    return not is_initialized() or get_rank() == 0


def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)


def count_used_parameters(model):
    # The model should be at least backward once
    return sum(p.numel() for p in model.parameters() if p.grad is not None)


def get_time_tag():
    return datetime.fromtimestamp(time()).strftime("%Y-%m-%d-%H-%M-%S")


def backup(src_path, tgt_dir):
    stem = Path(src_path).stem
    suffix = Path(src_path).suffix
    shutil.copyfile(src_path, os.path.join(tgt_dir, f"{stem}_{get_time_tag()}{suffix}"))


def get_model_state(model):
    if isinstance(model, DDP):
        return model.module.state_dict()
    return model.state_dict()


def show(*args, **kwargs):
    if is_leader_process():
        print(*args, **kwargs)


def hack_isinstance():
    # Pytorch do not support passing a defaultdict into DDP module
    # https://github.com/pytorch/pytorch/blob/v1.7.1/torch/nn/parallel/scatter_gather.py#L19

    # This hack can be removed after torch 1.8.0, where when each DDP process use single GPU
    # (which is the best practice) DDP will not pass args, kwargs into scatter function
    # https://github.com/pytorch/pytorch/blob/v1.7.1/torch/nn/parallel/distributed.py#L617
    # https://github.com/pytorch/pytorch/blob/v1.8.0-rc1/torch/nn/parallel/distributed.py#L700

    _isinstance = builtins.isinstance

    def isinstance(obj, cls):
        if _isinstance(obj, defaultdict):
            return _isinstance(obj, cls) and issubclass(cls, defaultdict)
        return _isinstance(obj, cls)

    builtins.isinstance = isinstance


def override(string, args, config):
    """
    Example usgae:
        -o "config.optimizer.lr=1.0e-3,,config.optimizer.name='AdamW',,config.runner.eval_dataloaders=['dev', 'test']"
    """
    options = string.split(",,")
    for option in options:
        option = option.strip()
        key, value_str = option.split("=")
        key, value_str = key.strip(), value_str.strip()
        first_field, *remaining = key.split(".")

        try:
            value = eval(value_str)
        except:
            value = value_str

        print(f"[Override] - {key} = {value}", file=sys.stderr)

        if first_field == "args":
            assert len(remaining) == 1
            setattr(args, remaining[0], value)
        elif first_field == "config":
            target_config = config
            for i, field_name in enumerate(remaining):
                if i == len(remaining) - 1:
                    target_config[field_name] = value
                else:
                    target_config.setdefault(field_name, {})
                    target_config = target_config[field_name]


def zero_mean_unit_var_norm(input_values: List[np.ndarray]) -> List[np.ndarray]:
    """
    Every array in the list is normalized to have zero mean and unit variance
    Taken from huggingface to ensure the same behavior across s3prl and huggingface
    Reference: https://github.com/huggingface/transformers/blob/a26f4d620874b32d898a5b712006a4c856d07de1/src/transformers/models/wav2vec2/feature_extraction_wav2vec2.py#L81-L86
    """
    return [(x - np.mean(x)) / np.sqrt(np.var(x) + 1e-5) for x in input_values]


#####################
# PARSE PRUNE HEADS #
#####################
def parse_prune_heads(config):
    if (
        "prune_headids" in config["transformer"]
        and config["transformer"]["prune_headids"] != "None"
    ):
        heads_int = []
        spans = config["transformer"]["prune_headids"].split(",")
        for span in spans:
            endpoints = span.split("-")
            if len(endpoints) == 1:
                heads_int.append(int(endpoints[0]))
            elif len(endpoints) == 2:
                heads_int += torch.arange(int(endpoints[0]), int(endpoints[1])).tolist()
            else:
                raise ValueError
        print(f"[PRUNING] - heads {heads_int} will be pruned")
        config["transformer"]["prune_headids"] = heads_int
    else:
        config["transformer"]["prune_headids"] = None


##########################
# GET TRANSFORMER TESTER #
##########################
def get_transformer_tester(
    from_path="result/result_transformer/libri_sd1337_fmllrBase960-F-N-K-RA/model-1000000.ckpt",
    display_settings=False,
):
    """Wrapper that loads the transformer model from checkpoint path"""

    # load config and paras
    all_states = torch.load(from_path, map_location="cpu")
    config = all_states["Settings"]["Config"]
    paras = all_states["Settings"]["Paras"]

    # handling older checkpoints
    if not hasattr(paras, "multi_gpu"):
        setattr(paras, "multi_gpu", False)
    if "prune_headids" not in config["transformer"]:
        config["transformer"]["prune_headids"] = None

    # display checkpoint settings
    if display_settings:
        for cluster in config:
            print(cluster + ":")
            for item in config[cluster]:
                print("\t" + str(item) + ": ", config[cluster][item])
        print("paras:")
        v_paras = vars(paras)
        for item in v_paras:
            print("\t" + str(item) + ": ", v_paras[item])

    # load model with Tester
    from transformer.solver import Tester

    tester = Tester(config, paras)
    tester.set_model(inference=True, with_head=False, from_path=from_path)
    return tester

utils/helper.py
roger-tseng/av-superb
# Copied from S3PRL 
# (Authors: Leo Yang, Andy T. Liu and S3PRL team, https://github.com/s3prl/s3prl/blob/main/s3prl/optimizers.py)

import abc
import copy
import logging
import math
import sys
from typing import Callable, Iterable, Tuple

import torch
import torch.nn as nn
from torch.nn.utils import clip_grad_norm_
from torch.optim import Adam, Optimizer
from torch.optim.lr_scheduler import LambdaLR
from torch.optim.optimizer import required

logger = logging.getLogger(__name__)


def get_optimizer(model_params, total_steps, optimizer_config):
    optimizer_config = copy.deepcopy(optimizer_config)
    optimizer_name = optimizer_config.pop("name")
    optimizer = eval(f"get_{optimizer_name}")(
        model_params, total_steps=total_steps, **optimizer_config
    )
    return optimizer


def get_grouped_parameters(model_params):
    named_params = []
    for m in model_params:
        named_params += list(m.named_parameters())

    no_decay = ["bias", "LayerNorm.bias", "LayerNorm.weight"]
    grouped_parameters = [
        {
            "params": [
                p for n, p in named_params if not any(nd in n for nd in no_decay)
            ],
            "weight_decay": 0.01,
        },
        {
            "params": [p for n, p in named_params if any(nd in n for nd in no_decay)],
            "weight_decay": 0.0,
        },
    ]
    return grouped_parameters


def get_BertAdam_with_schedule(
    model_params, lr=2e-4, total_steps=20000, warmup_proportion=0.07, **kwargs
):
    grouped_parameters = get_grouped_parameters(model_params)
    optimizer = BertAdam(
        grouped_parameters, lr=lr, warmup=warmup_proportion, t_total=total_steps
    )
    return optimizer


def get_AdamW_with_schedule(
    model_params, lr=2e-4, total_steps=20000, warmup_proportion=0.07, **kwargs
):
    grouped_parameters = get_grouped_parameters(model_params)
    optimizer = Lamb(
        grouped_parameters,
        lr=lr,
        warmup=warmup_proportion,
        t_total=total_steps,
        adam=True,
        correct_bias=True,
        **kwargs,
    )
    return optimizer


def get_Lamb_with_schedule(
    model_params, lr=2e-4, total_steps=20000, warmup_proportion=0.07, **kwargs
):
    grouped_parameters = get_grouped_parameters(model_params)
    optimizer = Lamb(
        grouped_parameters,
        lr=lr,
        warmup=warmup_proportion,
        t_total=total_steps,
        adam=False,
        correct_bias=False,
        **kwargs,
    )
    return optimizer


def get_Adam(model_params, lr=2e-4, **kwargs):
    params = []
    for m in model_params:
        params += list(m.parameters())
    return Adam(params, lr=lr, betas=(0.9, 0.999))


def get_AdamW(model_params, lr=2e-4, **kwargs):
    params = []
    for m in model_params:
        params += list(m.parameters())
    optimizer = AdamW(params, lr=lr)
    return optimizer


def get_TorchOptim(model_params, torch_optim_name, **kwargs):
    params = []
    for m in model_params:
        params += list(m.parameters())
    Opt_class = getattr(torch.optim, torch_optim_name)

    kwargs.pop("total_steps")
    optim = Opt_class(params, **kwargs)
    return optim


class AdamW(Optimizer):
    """
    Implements Adam algorithm with weight decay fix as introduced in
    `Decoupled Weight Decay Regularization <https://arxiv.org/abs/1711.05101>`__.
    Parameters:
        params (:obj:`Iterable[torch.nn.parameter.Parameter]`):
            Iterable of parameters to optimize or dictionaries defining parameter groups.
        lr (:obj:`float`, `optional`, defaults to 1e-3):
            The learning rate to use.
        betas (:obj:`Tuple[float,float]`, `optional`, defaults to (0.9, 0.999)):
            Adam's betas parameters (b1, b2).
        eps (:obj:`float`, `optional`, defaults to 1e-6):
            Adam's epsilon for numerical stability.
        weight_decay (:obj:`float`, `optional`, defaults to 0):
            Decoupled weight decay to apply.
        correct_bias (:obj:`bool`, `optional`, defaults to `True`):
            Whether ot not to correct bias in Adam (for instance, in Bert TF repository they use :obj:`False`).
    """

    def __init__(
        self,
        params: Iterable[torch.nn.parameter.Parameter],
        lr: float = 1e-3,
        betas: Tuple[float, float] = (0.9, 0.999),
        eps: float = 1e-7,
        weight_decay: float = 0.0,
        correct_bias: bool = True,
    ):
        if lr < 0.0:
            raise ValueError("Invalid learning rate: {} - should be >= 0.0".format(lr))
        if not 0.0 <= betas[0] < 1.0:
            raise ValueError(
                "Invalid beta parameter: {} - should be in [0.0, 1.0[".format(betas[0])
            )
        if not 0.0 <= betas[1] < 1.0:
            raise ValueError(
                "Invalid beta parameter: {} - should be in [0.0, 1.0[".format(betas[1])
            )
        if not 0.0 <= eps:
            raise ValueError("Invalid epsilon value: {} - should be >= 0.0".format(eps))
        defaults = dict(
            lr=lr,
            betas=betas,
            eps=eps,
            weight_decay=weight_decay,
            correct_bias=correct_bias,
        )
        super().__init__(params, defaults)

    def step(self, closure: Callable = None):
        """
        Performs a single optimization step.
        Arguments:
            closure (:obj:`Callable`, `optional`): A closure that reevaluates the model and returns the loss.
        """
        loss = None
        if closure is not None:
            loss = closure()

        for group in self.param_groups:
            for p in group["params"]:
                if p.grad is None:
                    continue
                grad = p.grad.data
                if grad.is_sparse:
                    raise RuntimeError(
                        "Adam does not support sparse gradients, please consider SparseAdam instead"
                    )

                state = self.state[p]

                # State initialization
                if len(state) == 0:
                    state["step"] = 0
                    # Exponential moving average of gradient values
                    state["exp_avg"] = torch.zeros_like(p.data)
                    # Exponential moving average of squared gradient values
                    state["exp_avg_sq"] = torch.zeros_like(p.data)

                exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
                beta1, beta2 = group["betas"]

                state["step"] += 1

                # Decay the first and second moment running average coefficient
                # In-place operations to update the averages at the same time
                exp_avg.mul_(beta1).add_(grad, alpha=1.0 - beta1)
                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)
                denom = exp_avg_sq.sqrt().add_(group["eps"])

                step_size = group["lr"]
                if group["correct_bias"]:  # No bias correction for Bert
                    bias_correction1 = 1.0 - beta1 ** state["step"]
                    bias_correction2 = 1.0 - beta2 ** state["step"]
                    step_size = (
                        step_size * math.sqrt(bias_correction2) / bias_correction1
                    )

                p.data.addcdiv_(exp_avg, denom, value=-step_size)

                # Just adding the square of the weights to the loss function is *not*
                # the correct way of using L2 regularization/weight decay with Adam,
                # since that will interact with the m and v parameters in strange ways.
                #
                # Instead we want to decay the weights in a manner that doesn't interact
                # with the m/v parameters. This is equivalent to adding the square
                # of the weights to the loss with plain (non-momentum) SGD.
                # Add weight decay at the end (fixed version)
                if group["weight_decay"] > 0.0:
                    p.data.add_(p.data, alpha=-group["lr"] * group["weight_decay"])

        return loss

    def get_lr(self):
        lr = []
        for group in self.param_groups:
            for p in group["params"]:
                state = self.state[p]
                if len(state) == 0:
                    pass
                else:
                    lr.append(group["lr"])
        return lr


# For the following codes:
"""PyTorch optimization for BERT model."""
# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
#
# 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.


if sys.version_info >= (3, 4):
    ABC = abc.ABC
else:
    ABC = abc.ABCMeta("ABC", (), {})


class _LRSchedule(ABC):
    """Parent of all LRSchedules here."""

    warn_t_total = False  # is set to True for schedules where progressing beyond t_total steps doesn't make sense

    def __init__(self, warmup=0.002, t_total=-1, **kw):
        """
        :param warmup:  what fraction of t_total steps will be used for linear warmup
        :param t_total: how many training steps (updates) are planned
        :param kw:
        """
        super(_LRSchedule, self).__init__(**kw)
        if t_total < 0:
            logger.warning(
                "t_total value of {} results in schedule not being applied".format(
                    t_total
                )
            )
        if not 0.0 <= warmup < 1.0 and not warmup == -1:
            raise ValueError(
                "Invalid warmup: {} - should be in [0.0, 1.0[ or -1".format(warmup)
            )
        warmup = max(warmup, 0.0)
        self.warmup, self.t_total = float(warmup), float(t_total)
        self.warned_for_t_total_at_progress = -1

    def get_lr(self, step, nowarn=False):
        """
        :param step:    which of t_total steps we're on
        :param nowarn:  set to True to suppress warning regarding training beyond specified 't_total' steps
        :return:        learning rate multiplier for current update
        """
        if self.t_total < 0:
            return 1.0
        progress = float(step) / self.t_total
        ret = self.get_lr_(progress)
        # warning for exceeding t_total (only active with warmup_linear
        if (
            not nowarn
            and self.warn_t_total
            and progress > 1.0
            and progress > self.warned_for_t_total_at_progress
        ):
            logger.warning(
                "Training beyond specified 't_total'. Learning rate multiplier set to {}. Please set 't_total' of {} correctly.".format(
                    ret, self.__class__.__name__
                )
            )
            self.warned_for_t_total_at_progress = progress
        # end warning
        return ret

    @abc.abstractmethod
    def get_lr_(self, progress):
        """
        :param progress:    value between 0 and 1 (unless going beyond t_total steps) specifying training progress
        :return:            learning rate multiplier for current update
        """
        return 1.0


class ConstantLR(_LRSchedule):
    def get_lr_(self, progress):
        return 1.0


class WarmupCosineSchedule(_LRSchedule):
    """
    Linearly increases learning rate from 0 to 1 over `warmup` fraction of training steps.
    Decreases learning rate from 1. to 0. over remaining `1 - warmup` steps following a cosine curve.
    If `cycles` (default=0.5) is different from default, learning rate follows cosine function after warmup.
    """

    warn_t_total = True

    def __init__(self, warmup=0.002, t_total=-1, cycles=0.5, **kw):
        """
        :param warmup:      see LRSchedule
        :param t_total:     see LRSchedule
        :param cycles:      number of cycles. Default: 0.5, corresponding to cosine decay from 1. at progress==warmup and 0 at progress==1.
        :param kw:
        """
        super(WarmupCosineSchedule, self).__init__(warmup=warmup, t_total=t_total, **kw)
        self.cycles = cycles

    def get_lr_(self, progress):
        if progress < self.warmup:
            return progress / self.warmup
        else:
            progress = (progress - self.warmup) / (
                1 - self.warmup
            )  # progress after warmup
            return 0.5 * (1.0 + math.cos(math.pi * self.cycles * 2 * progress))


class WarmupCosineWithHardRestartsSchedule(WarmupCosineSchedule):
    """
    Linearly increases learning rate from 0 to 1 over `warmup` fraction of training steps.
    If `cycles` (default=1.) is different from default, learning rate follows `cycles` times a cosine decaying
    learning rate (with hard restarts).
    """

    def __init__(self, warmup=0.002, t_total=-1, cycles=1.0, **kw):
        super(WarmupCosineWithHardRestartsSchedule, self).__init__(
            warmup=warmup, t_total=t_total, cycles=cycles, **kw
        )
        assert cycles >= 1.0

    def get_lr_(self, progress):
        if progress < self.warmup:
            return progress / self.warmup
        else:
            progress = (progress - self.warmup) / (
                1 - self.warmup
            )  # progress after warmup
            ret = 0.5 * (1.0 + math.cos(math.pi * ((self.cycles * progress) % 1)))
            return ret


class WarmupCosineWithWarmupRestartsSchedule(WarmupCosineWithHardRestartsSchedule):
    """
    All training progress is divided in `cycles` (default=1.) parts of equal length.
    Every part follows a schedule with the first `warmup` fraction of the training steps linearly increasing from 0. to 1.,
    followed by a learning rate decreasing from 1. to 0. following a cosine curve.
    """

    def __init__(self, warmup=0.002, t_total=-1, cycles=1.0, **kw):
        assert warmup * cycles < 1.0
        warmup = warmup * cycles if warmup >= 0 else warmup
        super(WarmupCosineWithWarmupRestartsSchedule, self).__init__(
            warmup=warmup, t_total=t_total, cycles=cycles, **kw
        )

    def get_lr_(self, progress):
        progress = progress * self.cycles % 1.0
        if progress < self.warmup:
            return progress / self.warmup
        else:
            progress = (progress - self.warmup) / (
                1 - self.warmup
            )  # progress after warmup
            ret = 0.5 * (1.0 + math.cos(math.pi * progress))
            return ret


class WarmupConstantSchedule(_LRSchedule):
    """
    Linearly increases learning rate from 0 to 1 over `warmup` fraction of training steps.
    Keeps learning rate equal to 1. after warmup.
    """

    def get_lr_(self, progress):
        if progress < self.warmup:
            return progress / self.warmup
        return 1.0


class WarmupLinearSchedule(_LRSchedule):
    """
    Linearly increases learning rate from 0 to 1 over `warmup` fraction of training steps.
    Linearly decreases learning rate from 1. to 0. over remaining `1 - warmup` steps.
    """

    warn_t_total = True

    def get_lr_(self, progress):
        if progress < self.warmup:
            return progress / self.warmup
        return max((progress - 1.0) / (self.warmup - 1.0), 0.0)


SCHEDULES = {
    None: ConstantLR,
    "none": ConstantLR,
    "warmup_cosine": WarmupCosineSchedule,
    "warmup_constant": WarmupConstantSchedule,
    "warmup_linear": WarmupLinearSchedule,
}


class BertAdam(Optimizer):
    """Implements BERT version of Adam algorithm with weight decay fix.
    Params:
        lr: learning rate
        warmup: portion of t_total for the warmup, -1  means no warmup. Default: -1
        t_total: total number of training steps for the learning
            rate schedule, -1  means constant learning rate of 1. (no warmup regardless of warmup setting). Default: -1
        schedule: schedule to use for the warmup (see above).
            Can be `'warmup_linear'`, `'warmup_constant'`, `'warmup_cosine'`, `'none'`, `None` or a `_LRSchedule` object (see below).
            If `None` or `'none'`, learning rate is always kept constant.
            Default : `'warmup_linear'`
        betas: Adams betas. Default: (0.9, 0.999)
        e: Adams epsilon. Default: 1e-6
        weight_decay: Weight decay. Default: 0.01
        max_grad_norm: Maximum norm for the gradients (-1 means no clipping). Default: 1.0
    """

    def __init__(
        self,
        params=None,
        lr="required",
        warmup=-1,
        t_total=-1,
        schedule="warmup_linear",
        betas=(0.9, 0.999),
        e=1e-6,
        weight_decay=0.01,
        max_grad_norm=1.0,
        **kwargs,
    ):
        if lr == "required" or lr < 0.0:
            raise ValueError("Invalid learning rate: {} - should be >= 0.0".format(lr))
        if not isinstance(schedule, _LRSchedule) and schedule not in SCHEDULES:
            raise ValueError("Invalid schedule parameter: {}".format(schedule))
        if not 0.0 <= betas[0] < 1.0:
            raise ValueError(
                "Invalid beta parameter at index 0: {} - should be in [0.0, 1.0[".format(
                    betas[0]
                )
            )
        if not 0.0 <= betas[1] < 1.0:
            raise ValueError(
                "Invalid beta parameter at index 1: {} - should be in [0.0, 1.0[".format(
                    betas[1]
                )
            )
        if not e >= 0.0:
            raise ValueError("Invalid epsilon value: {} - should be >= 0.0".format(e))
        # initialize schedule object
        if not isinstance(schedule, _LRSchedule):
            schedule_type = SCHEDULES[schedule]
            schedule = schedule_type(warmup=warmup, t_total=t_total)
        else:
            if warmup != -1 or t_total != -1:
                logger.warning(
                    "warmup and t_total on the optimizer are ineffective when _LRSchedule object is provided as schedule. "
                    "Please specify custom warmup and t_total in _LRSchedule object."
                )
        defaults = dict(
            lr=lr,
            schedule=schedule,
            betas=betas,
            e=e,
            weight_decay=weight_decay,
            max_grad_norm=max_grad_norm,
        )
        super(BertAdam, self).__init__(params, defaults)

    def get_lr(self):
        lr = []
        for group in self.param_groups:
            for p in group["params"]:
                state = self.state[p]
                if len(state) == 0:
                    pass
                else:
                    lr_scheduled = group["lr"]
                    lr_scheduled *= group["schedule"].get_lr(state["step"])
                    lr.append(lr_scheduled)
        return lr

    def step(self, closure=None):
        """Performs a single optimization step.

        Arguments:
            closure (callable, optional): A closure that reevaluates the model
                and returns the loss.
        """
        loss = None
        if closure is not None:
            loss = closure()

        for group in self.param_groups:
            for p in group["params"]:
                if p.grad is None:
                    continue
                grad = p.grad.data
                if grad.is_sparse:
                    raise RuntimeError(
                        "Adam does not support sparse gradients, please consider SparseAdam instead"
                    )

                state = self.state[p]

                # State initialization
                if len(state) == 0:
                    state["step"] = 0
                    # Exponential moving average of gradient values
                    state["next_m"] = torch.zeros_like(p.data)
                    # Exponential moving average of squared gradient values
                    state["next_v"] = torch.zeros_like(p.data)

                next_m, next_v = state["next_m"], state["next_v"]
                beta1, beta2 = group["betas"]

                # Add grad clipping
                if group["max_grad_norm"] > 0:
                    clip_grad_norm_(p, group["max_grad_norm"])

                # Decay the first and second moment running average coefficient
                # In-place operations to update the averages at the same time
                next_m.mul_(beta1).add_(1 - beta1, grad)
                next_v.mul_(beta2).addcmul_(1 - beta2, grad, grad)
                update = next_m / (next_v.sqrt() + group["e"])

                # Just adding the square of the weights to the loss function is *not*
                # the correct way of using L2 regularization/weight decay with Adam,
                # since that will interact with the m and v parameters in strange ways.
                #
                # Instead we want to decay the weights in a manner that doesn't interact
                # with the m/v parameters. This is equivalent to adding the square
                # of the weights to the loss with plain (non-momentum) SGD.
                if group["weight_decay"] > 0.0:
                    update += group["weight_decay"] * p.data

                lr_scheduled = group["lr"]
                lr_scheduled *= group["schedule"].get_lr(state["step"])

                update_with_lr = lr_scheduled * update
                p.data.add_(-update_with_lr)

                state["step"] += 1

                # step_size = lr_scheduled * math.sqrt(bias_correction2) / bias_correction1
                # No bias correction
                # bias_correction1 = 1 - beta1 ** state['step']
                # bias_correction2 = 1 - beta2 ** state['step']

        return loss


class Lamb(Optimizer):
    r"""Implements Lamb algorithm.
    It has been proposed in `Large Batch Optimization for Deep Learning: Training BERT in 76 minutes`_.
    Arguments:
        params (iterable): iterable of parameters to optimize or dicts defining
            parameter groups
        lr (float, optional): learning rate (default: 1e-3)
        warmup (float, optional): portion of t_total for the warmup, -1  means no warmup. Default: -1
        t_total (int, optional): total number of training steps for the learning
            rate schedule, -1  means constant learning rate of 1. (no warmup regardless of warmup setting). Default: -1
        schedule (string, optional): schedule to use for the warmup (see above).
            Can be `'warmup_linear'`, `'warmup_constant'`, `'warmup_cosine'`, `'none'`, `None` or a `_LRSchedule` object (see below).
            If `None` or `'none'`, learning rate is always kept constant.
            Default : `'warmup_linear'`
        betas (Tuple[float, float], optional): coefficients used for computing
            running averages of gradient and its square (default: (0.9, 0.999))
        eps (float, optional): term added to the denominator to improve
            numerical stability (default: 1e-8)
        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
        adam (bool, optional): always use trust ratio = 1, which turns this into
            Adam. Useful for comparison purposes. Set to True for AdamW.
        correct_bias (bool, optional): adam-correction, no bias correction for Bert. Set to True for AdamW.
    .. _Large Batch Optimization for Deep Learning: Training BERT in 76 minutes:
        https://arxiv.org/abs/1904.00962
    """

    def __init__(
        self,
        params,
        lr=1e-3,
        warmup=-1,
        t_total=-1,
        schedule="warmup_linear",
        betas=(0.9, 0.999),
        eps=1e-8,
        weight_decay=0.0,
        adam=False,
        correct_bias=False,
    ):
        if not 0.0 <= lr:
            raise ValueError("Invalid learning rate: {}".format(lr))
        if not 0.0 <= eps:
            raise ValueError("Invalid epsilon value: {}".format(eps))
        if not 0.0 <= betas[0] < 1.0:
            raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
        if not 0.0 <= betas[1] < 1.0:
            raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
        # initialize schedule object
        if not isinstance(schedule, _LRSchedule):
            schedule_type = SCHEDULES[schedule]
            schedule = schedule_type(warmup=warmup, t_total=t_total)
        else:
            if warmup != -1 or t_total != -1:
                logger.warning(
                    "warmup and t_total on the optimizer are ineffective when _LRSchedule object is provided as schedule. "
                    "Please specify custom warmup and t_total in _LRSchedule object."
                )
        defaults = dict(
            lr=lr,
            betas=betas,
            eps=eps,
            schedule=schedule,
            weight_decay=weight_decay,
            correct_bias=correct_bias,
        )
        self.adam = adam

        super(Lamb, self).__init__(params, defaults)

    def get_lr(self):
        lr = []
        for group in self.param_groups:
            for p in group["params"]:
                state = self.state[p]
                if len(state) == 0:
                    pass
                else:
                    lr_scheduled = group["lr"]
                    lr_scheduled *= group["schedule"].get_lr(state["step"])
                    lr.append(lr_scheduled)
        return lr

    def step(self, closure=None):
        """Performs a single optimization step.
        Arguments:
            closure (callable, optional): A closure that reevaluates the model
                and returns the loss.
        """
        loss = None
        if closure is not None:
            loss = closure()

        for group in self.param_groups:
            for p in group["params"]:
                if p.grad is None:
                    continue
                grad = p.grad.data
                if grad.is_sparse:
                    raise RuntimeError(
                        "Lamb does not support sparse gradients, consider SparseAdam instad."
                    )

                state = self.state[p]

                # State initialization
                if len(state) == 0:
                    state["step"] = 0
                    # Exponential moving average of gradient values
                    state["exp_avg"] = torch.zeros_like(p.data)
                    # Exponential moving average of squared gradient values
                    state["exp_avg_sq"] = torch.zeros_like(p.data)

                exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
                beta1, beta2 = group["betas"]

                state["step"] += 1

                # Decay the first and second moment running average coefficient
                # m_t
                exp_avg.mul_(beta1).add_(1 - beta1, grad)
                # v_t
                exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)

                # Paper v3 does not use debiasing.
                # bias_correction1 = 1 - beta1 ** state['step']
                # bias_correction2 = 1 - beta2 ** state['step']
                # Apply bias to lr to avoid broadcast.
                step_size = group[
                    "lr"
                ]  # * math.sqrt(bias_correction2) / bias_correction1
                if group["correct_bias"]:  # No bias correction for Bert
                    bias_correction1 = 1.0 - beta1 ** state["step"]
                    bias_correction2 = 1.0 - beta2 ** state["step"]
                    step_size = (
                        step_size * math.sqrt(bias_correction2) / bias_correction1
                    )
                lr_scheduled = step_size * group["schedule"].get_lr(state["step"])

                weight_norm = p.data.pow(2).sum().sqrt()

                adam_step = exp_avg / exp_avg_sq.sqrt().add(group["eps"])
                if group["weight_decay"] != 0:
                    adam_step.add_(group["weight_decay"], p.data)

                adam_norm = adam_step.pow(2).sum().sqrt()
                if weight_norm == 0 or adam_norm == 0:
                    trust_ratio = 1
                else:
                    trust_ratio = weight_norm / adam_norm
                state["weight_norm"] = weight_norm
                state["adam_norm"] = adam_norm
                state["trust_ratio"] = trust_ratio
                if self.adam:
                    trust_ratio = 1

                p.data.add_(-lr_scheduled * trust_ratio, adam_step)

        return loss

utils/optimizers.py
roger-tseng/av-superb
# Copied from S3PRL 
# (Authors: Leo Yang, Andy T. Liu and S3PRL team, https://github.com/s3prl/s3prl/blob/main/s3prl/schedulers.py)

import copy
import math
import random
from typing import Callable, Iterable, Tuple

import numpy as np
import torch
from torch.optim import Optimizer
from torch.optim.lr_scheduler import LambdaLR


def get_scheduler(optimizer, total_steps, scheduler_config):
    scheduler_config = copy.deepcopy(scheduler_config)
    scheduler_name = scheduler_config.pop("name")
    scheduler = eval(f"get_{scheduler_name}")(
        optimizer, num_training_steps=total_steps, **scheduler_config
    )
    return scheduler


def get_cosine_with_hard_restarts_schedule_with_warmup(
    optimizer: Optimizer,
    num_warmup_steps: int,
    num_training_steps: int,
    num_cycles: int = 1,
    last_epoch: int = -1,
):
    """
    Create a schedule with a learning rate that decreases following the values of the cosine function between the
    initial lr set in the optimizer to 0, with several hard restarts, after a warmup period during which it increases
    linearly between 0 and the initial lr set in the optimizer.

    Args:
        optimizer (:class:`~torch.optim.Optimizer`):
            The optimizer for which to schedule the learning rate.
        num_warmup_steps (:obj:`int`):
            The number of steps for the warmup phase.
        num_training_steps (:obj:`int`):
            The total number of training steps.
        num_cycles (:obj:`int`, `optional`, defaults to 1):
            The number of hard restarts to use.
        last_epoch (:obj:`int`, `optional`, defaults to -1):
            The index of the last epoch when resuming training.

    Return:
        :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
    """

    def lr_lambda(current_step):
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        progress = float(current_step - num_warmup_steps) / float(
            max(1, num_training_steps - num_warmup_steps)
        )
        if progress >= 1.0:
            return 0.0
        return max(
            0.0,
            0.5 * (1.0 + math.cos(math.pi * ((float(num_cycles) * progress) % 1.0))),
        )

    return LambdaLR(optimizer, lr_lambda, last_epoch)


def get_cosine_schedule_with_warmup(
    optimizer: Optimizer,
    num_warmup_steps: int,
    num_training_steps: int,
    num_cycles: float = 0.5,
    last_epoch: int = -1,
):
    """
    Create a schedule with a learning rate that decreases following the values of the cosine function between the
    initial lr set in the optimizer to 0, after a warmup period during which it increases linearly between 0 and the
    initial lr set in the optimizer.

    Args:
        optimizer (:class:`~torch.optim.Optimizer`):
            The optimizer for which to schedule the learning rate.
        num_warmup_steps (:obj:`int`):
            The number of steps for the warmup phase.
        num_training_steps (:obj:`int`):
            The total number of training steps.
        num_cycles (:obj:`float`, `optional`, defaults to 0.5):
            The number of waves in the cosine schedule (the defaults is to just decrease from the max value to 0
            following a half-cosine).
        last_epoch (:obj:`int`, `optional`, defaults to -1):
            The index of the last epoch when resuming training.

    Return:
        :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
    """

    def lr_lambda(current_step):
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        progress = float(current_step - num_warmup_steps) / float(
            max(1, num_training_steps - num_warmup_steps)
        )
        return max(
            0.0, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress))
        )

    return LambdaLR(optimizer, lr_lambda, last_epoch)


def get_linear_schedule_with_warmup(
    optimizer, num_warmup_steps, num_training_steps, last_epoch=-1
):
    """
    Create a schedule with a learning rate that decreases linearly from the initial lr set in the optimizer to 0,
    after a warmup period during which it increases linearly from 0 to the initial lr set in the optimizer.
    Args:
        optimizer (:class:`~torch.optim.Optimizer`):
            The optimizer for which to schedule the learning rate.
        num_warmup_steps (:obj:`int`):
            The number of steps for the warmup phase.
        num_training_steps (:obj:`int`):
            The total number of training steps.
        last_epoch (:obj:`int`, `optional`, defaults to -1):
            The index of the last epoch when resuming training.
    Return:
        :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
    """

    def lr_lambda(current_step: int):
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        return max(
            0.0,
            float(num_training_steps - current_step)
            / float(max(1, num_training_steps - num_warmup_steps)),
        )

    return LambdaLR(optimizer, lr_lambda, last_epoch)


def get_sqrt_decay_schedule_with_warmup(
    optimizer, num_warmup_steps, num_training_steps, last_epoch=-1
):
    """
    Create a schedule with a learning rate that decreases linearly from the initial lr set in the optimizer to 0,
    after a warmup period during which it increases linearly from 0 to the initial lr set in the optimizer.
    Args:
        optimizer (:class:`~torch.optim.Optimizer`):
            The optimizer for which to schedule the learning rate.
        num_warmup_steps (:obj:`int`):
            The number of steps for the warmup phase.
        num_training_steps (:obj:`int`):
            The total number of training steps.
        last_epoch (:obj:`int`, `optional`, defaults to -1):
            The index of the last epoch when resuming training.
    Return:
        :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
    """

    def lr_lambda(current_step: int):
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        return 1.0 / math.sqrt(max(current_step, num_warmup_steps))

    return LambdaLR(optimizer, lr_lambda, last_epoch)


def get_constant_decay_schedule_with_warmup(
    optimizer, num_warmup_steps, num_training_steps, last_epoch=-1
):
    """
    Create a schedule with a learning rate that decreases linearly from the initial lr set in the optimizer to 0,
    after a warmup period during which it increases linearly from 0 to the initial lr set in the optimizer.
    Args:
        optimizer (:class:`~torch.optim.Optimizer`):
            The optimizer for which to schedule the learning rate.
        num_warmup_steps (:obj:`int`):
            The number of steps for the warmup phase.
        num_training_steps (:obj:`int`):
            The total number of training steps.
        last_epoch (:obj:`int`, `optional`, defaults to -1):
            The index of the last epoch when resuming training.
    Return:
        :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
    """

    def lr_lambda(current_step: int):
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        return 1.0

    return LambdaLR(optimizer, lr_lambda, last_epoch)


def get_noam_decay_schedule_with_warmup(
    optimizer, num_warmup_steps, num_training_steps, last_epoch=-1
):
    """
    Create a schedule with a learning rate that decreases linearly from the initial lr set in the optimizer to 0,
    after a warmup period during which it increases linearly from 0 to the initial lr set in the optimizer.
    Args:
        optimizer (:class:`~torch.optim.Optimizer`):
            The optimizer for which to schedule the learning rate.
        num_warmup_steps (:obj:`int`):
            The number of steps for the warmup phase.
        num_training_steps (:obj:`int`):
            The total number of training steps.
        last_epoch (:obj:`int`, `optional`, defaults to -1):
            The index of the last epoch when resuming training.
    Return:
        :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
    """

    def lr_lambda(current_step: int):
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        return 768 ** (-0.5) * min(
            current_step ** (-0.5), current_step * num_warmup_steps ** (-1.5)
        )

    return LambdaLR(optimizer, lr_lambda, last_epoch)


def get_polynomial_decay_schedule_with_warmup(
    optimizer,
    num_warmup_steps,
    num_training_steps,
    lr_end=1e-7,
    power=1.0,
    last_epoch=-1,
):
    """
    Create a schedule with a learning rate that decreases as a polynomial decay from the initial lr set in the
    optimizer to end lr defined by `lr_end`, after a warmup period during which it increases linearly from 0 to the
    initial lr set in the optimizer.
    Args:
        optimizer (:class:`~torch.optim.Optimizer`):
            The optimizer for which to schedule the learning rate.
        num_warmup_steps (:obj:`int`):
            The number of steps for the warmup phase.
        num_training_steps (:obj:`int`):
            The total number of training steps.
        lr_end (:obj:`float`, `optional`, defaults to 1e-7):
            The end LR.
        power (:obj:`float`, `optional`, defaults to 1.0):
            Power factor.
        last_epoch (:obj:`int`, `optional`, defaults to -1):
            The index of the last epoch when resuming training.
    Note: `power` defaults to 1.0 as in the fairseq implementation, which in turn is based on the original BERT
    implementation at
    https://github.com/google-research/bert/blob/f39e881b169b9d53bea03d2d341b31707a6c052b/optimization.py#L37
    Return:
        :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
    """

    lr_init = optimizer.defaults["lr"]
    assert (
        lr_init > lr_end
    ), f"lr_end ({lr_end}) must be be smaller than initial lr ({lr_init})"

    def lr_lambda(current_step: int):
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        elif current_step > num_training_steps:
            return lr_end / lr_init  # as LambdaLR multiplies by lr_init
        else:
            lr_range = lr_init - lr_end
            decay_steps = num_training_steps - num_warmup_steps
            pct_remaining = 1 - (current_step - num_warmup_steps) / decay_steps
            decay = lr_range * pct_remaining**power + lr_end
            return decay / lr_init  # as LambdaLR multiplies by lr_init

    return LambdaLR(optimizer, lr_lambda, last_epoch)

utils/schedulers.py
roger-tseng/av-superb