The official pytorch implementation of paper "Differentiable Resolution Compression and alignment for Efficient Video Classification and Retrieval".
# checkpoint_config = dict(interval=1)
find_unused_parameters=True
log_config = dict(
    interval=20,
    hooks=[
        dict(type='TextLoggerHook'),
        dict(type='TensorboardLoggerHook'),
    ])
# runtime settings
dist_params = dict(backend='nccl')
log_level = 'INFO'
load_from = None
resume_from = None
workflow = [('train', 1)]

# disable opencv multithreading to avoid system being overloaded
opencv_num_threads = 0
# set multi-process start method as `fork` to speed up the training
mp_start_method = 'fork'

configs/_base_/default_runtime.py
dun-research/DRCA
model = dict(
    type='Recognizer3D',
    backbone=dict(
        type='DRCA',
        pretrained=None,
        num_frames=8,
        img_size=224,
        patch_size=16,
        embed_dims=384,
        num_heads=6,
        in_channels=3,
        comp_insert_layer=-1,
        comp_k=8,
        comp_strategy="diff_rank",
        comp_module="sal_ref",
        sigma=0.05,
        num_samples=500,
        dropout_ratio=0.0,
        transformer_layers=None,
        norm_cfg=dict(type='LN', eps=1e-06)),
    cls_head=dict(type='ClsHead', num_classes=200, in_channels=384),
    train_cfg=None,
    test_cfg=dict(average_clips='prob'))
configs/video_classification/DRCA-S-224.py
dun-research/DRCA
_base_ = ['../_base_/default_runtime.py', './DRCA-S-224.py',
                  './data_pipeline_minik.py', './solver.py']

model = dict(
    backbone=dict(
        comp_insert_layer=3,
        comp_k=4,
    ))

load_from =  'weights/DRCA-S-K8-full-224-minik.pth'
work_dir = 'work_dirs/video_classification/DRCA-S-K4-224'

configs/video_classification/DRCA-S-K4-224.py
dun-research/DRCA
data = dict(
    videos_per_gpu=26,
    workers_per_gpu=16,
    test_dataloader=dict(videos_per_gpu=1),
    train=dict(
        type='RawframeDataset',
        ann_file=
        'data/mini-kinetics/train_list_rawframes.txt',
        data_prefix='data/kinetics400/rawframes_train',
        pipeline=[
            dict(
                type='SampleFrames',
                clip_len=8,
                frame_interval=32,
                num_clips=1),
            dict(type='RawFrameDecode'),
            dict(type='Resize', scale=(-1, 256)),
            dict(type='RandomResizedCrop'),
            dict(type='Resize', scale=(224, 224), keep_ratio=False),
            dict(type='Flip', flip_ratio=0.5),
            dict(
                type='Normalize',
                mean=[127.5, 127.5, 127.5],
                std=[127.5, 127.5, 127.5],
                to_bgr=False),
            dict(type='FormatShape', input_format='NCTHW'),
            dict(type='Collect', keys=['imgs', 'label'], meta_keys=[]),
            dict(type='ToTensor', keys=['imgs', 'label'])
        ]),
    val=dict(
        type='RawframeDataset',
        ann_file=
        'data/mini-kinetics/val_list_rawframes.txt',
        data_prefix='data/kinetics400/rawframes_val',
        pipeline=[
            dict(
                type='SampleFrames',
                clip_len=8,
                frame_interval=32,
                num_clips=1,
                test_mode=True),
            dict(type='RawFrameDecode'),
            dict(type='Resize', scale=(-1, 256)),
            dict(type='CenterCrop', crop_size=224),
            dict(
                type='Normalize',
                mean=[127.5, 127.5, 127.5],
                std=[127.5, 127.5, 127.5],
                to_bgr=False),
            dict(type='FormatShape', input_format='NCTHW'),
            dict(type='Collect', keys=['imgs', 'label'], meta_keys=[]),
            dict(type='ToTensor', keys=['imgs', 'label'])
        ]),
    test=dict(
        type='RawframeDataset',
        ann_file=
        'data/mini-kinetics/val_list_rawframes.txt',
        data_prefix='data/kinetics400/rawframes_val',
        pipeline=[
            dict(
                type='SampleFrames',
                clip_len=8,
                frame_interval=32,
                num_clips=1,
                test_mode=True),
            dict(type='RawFrameDecode'),
            dict(type='Resize', scale=(-1, 256)),
            dict(type='CenterCrop', crop_size=224),
            dict(
                type='Normalize',
                mean=[127.5, 127.5, 127.5],
                std=[127.5, 127.5, 127.5],
                to_bgr=False),
            dict(type='FormatShape', input_format='NCTHW'),
            dict(type='Collect', keys=['imgs', 'label'], meta_keys=[]),
            dict(type='ToTensor', keys=['imgs', 'label'])
        ]))
evaluation = dict(
    interval=1, metrics=['top_k_accuracy', 'mean_class_accuracy'])
eval_config=evaluation

configs/video_classification/data_pipeline_minik.py
dun-research/DRCA
checkpoint_config = dict(interval=5)
optimizer = dict(
    type='SGD',
    lr=0.01,
    momentum=0.9,
    paramwise_cfg=dict(
        custom_keys=dict({
            '.backbone.cls_token': dict(decay_mult=0.0),
            '.backbone.pos_embed': dict(decay_mult=0.0),
            '.backbone.time_embed': dict(decay_mult=0.0)
        })),
    weight_decay=0.0001,
    nesterov=True)
optimizer_config = dict(grad_clip=dict(max_norm=40, norm_type=2))
lr_config = dict(policy='step', step=[10, 20])
total_epochs = 30

configs/video_classification/solver.py
dun-research/DRCA
model = dict(
    type='SimilarityRecognizer3D',
    backbone=dict(
        type='DRCA',
        num_frames=8,
        img_size=224,
        patch_size=16,
        embed_dims=768,
        num_heads=12,
        in_channels=3,
        comp_insert_layer=-1,
        comp_k=8,
        comp_strategy="diff_rank",
        comp_module="sal_ref",
        sigma=0.05,
        num_samples=500,
        dropout_ratio=0.0,
        transformer_layers=None,
        norm_cfg=dict(type='LN', eps=1e-06)),
    sim_head=dict(
        type='SimilarityHead',
        encoder=dict(type='None', input_channels=768, out_channels=768),
        metric_loss=dict(
            type='MultiSimilarityLoss',
            distance='Cosine',
            loss_weight=1.0,
            margin=0.3,
            cross_batch_memory=True,
            memory_size=1024,
            feat_dim=768,
            miner_name='MultiSimilarityMiner',
            distributed=True)),
    # model training and testing settings
    train_cfg=dict(aux_info=['video_id']),
    test_cfg=dict(average_clips='prob'))
configs/video_retrieval/DRCA-B-224.py
dun-research/DRCA
_base_ = ['../_base_/default_runtime.py', './DRCA-B-224.py',
                  './data_pipeline_fivr.py', './solver.py']

model = dict(
    backbone=dict(
        comp_insert_layer=3,
        comp_k=4,
    ))

work_dir = 'work_dirs/DRCA/DRCA-B-K4-224-fivr'

configs/video_retrieval/DRCA-B-K4-224.py
dun-research/DRCA
model = dict(
    type='SimilarityRecognizer3D',
    backbone=dict(
        type='DRCA',
        pretrained=None,
        num_frames=8,
        img_size=224,
        patch_size=16,
        embed_dims=384,
        num_heads=6,
        in_channels=3,
        comp_insert_layer=-1,
        comp_k=8,
        comp_strategy="diff_rank",
        comp_module="sal_ref",
        sigma=0.05,
        num_samples=500,
        dropout_ratio=0.0,
        transformer_layers=None,
        norm_cfg=dict(type='LN', eps=1e-06)),
    sim_head=dict(
        type='SimilarityHead',
        encoder=dict(type='None', input_channels=384, out_channels=384),
        metric_loss=dict(
            type='MultiSimilarityLoss',
            distance='Cosine',
            loss_weight=1.0,
            margin=0.3,
            cross_batch_memory=True,
            memory_size=1024,
            feat_dim=768,
            miner_name='MultiSimilarityMiner',
            distributed=True)),
    # model training and testing settings
    train_cfg=dict(aux_info=['video_id']),
    test_cfg=dict(average_clips='prob'))

configs/video_retrieval/DRCA-S-224.py
dun-research/DRCA
_base_ = ['../_base_/default_runtime.py', './DRCA-S-224.py',
                  './data_pipeline_fivr.py', './solver.py']

model = dict(
    backbone=dict(
        comp_insert_layer=3,
        comp_k=3,
    ))

work_dir = 'work_dirs/DRCA/DRCA-S-K3-224-fivr'

configs/video_retrieval/DRCA-S-K3-224.py
dun-research/DRCA
img_norm_cfg = dict(
    mean=[127.5, 127.5, 127.5], std=[127.5, 127.5, 127.5], to_bgr=False)

# ======================================== path settings ========================================
data_root = 'data/kinetics400/rawframes_train'
ann_file_train = 'data/kinetics400/kinetics400_train_list_rawframes.txt'

fivr_version = "5k" 
if fivr_version == "5k":
    query_file="data/FIVR-200K/annotation/fivr-5k-queries.txt"
    database_file="data/FIVR-200K//annotation/fivr-5k-database.txt"
    anno_file="data/FIVR-200K/annotation/fivr-200k-visil.pkl"
    data_root="data/FIVR-200K/"
elif fivr_version == "200k":
    dataset_name = "fivr_200k"
    query_file="data/FIVR-200K/annotation/fivr-200k-queries.txt"
    database_file="data/FIVR-200K//annotation/fivr-200k-database.txt"
    anno_file="data/FIVR-200K//annotation/fivr-200k-visil.pkl"
    data_prefix="data/FIVR-200K/"
else:
    raise ValueError("fivr_version must be 5k or 200k")
# ======================================== path settings ========================================

train_pipeline = [
    dict(type='Flip', flip_ratio=0.5),
    dict(type='ImageAug', ),
    dict(type='PadAndResize', size=(224,224)),
    dict(type='Normalize', **img_norm_cfg),
    dict(type='FormatShape', input_format='NCTHW'),
    dict(type='Collect', keys=['imgs', 'label', 'video_id', ], meta_keys=[]),
    dict(type='ToTensor', keys=['imgs', 'label', 'video_id', ])
]


train_dataset = dict(
            type="RetrievalTrainRawframeDataset",
            ann_file=ann_file_train,
            data_prefix=data_root,
            clip_len=8, 
            scale_range=(256, 320),
            crop_size=224,
            long_frame_interval=32, 
            short_frame_interval=16,
            pipeline=train_pipeline)


fivr_pipeline = [
            dict(type="GeneralVideoDecoder", clip_len=8,  backend="pyav", fps_interval=1),
            dict(type='Resize', scale=(-1, 224)),
            dict(type='CenterCrop', crop_size=224),
            dict(type='Normalize', **img_norm_cfg),
            dict(type='FormatShape', input_format='NCTHW'),
            dict(type='Collect', keys=['imgs', 'label', ], meta_keys=[]),
            dict(type='ToTensor', keys=['imgs', 'label', ]),

]
fivr_dataset = dict(
    type="FIVREvalDataset",
    query_file=query_file,
    database_file=database_file,
    anno_file=anno_file,
    data_prefix=data_root,
    pipeline=fivr_pipeline,
)

data = dict(
    videos_per_gpu=16,
    workers_per_gpu=8,
    test_dataloader=dict(videos_per_gpu=1),
    val_dataloader=dict(videos_per_gpu=1),
    train=train_dataset,
    val=fivr_dataset,
    test=fivr_dataset,
    )
evaluation = dict()
eval_config = dict()

configs/video_retrieval/data_pipeline_fivr.py
dun-research/DRCA
total_epochs = 100
warmup_epochs=5

base_lr = 1.5e-4
batch_size_per_gpu = 16
n_gpus = 8
batch_size = batch_size_per_gpu*n_gpus
lr = batch_size * base_lr / 256

optimizer = dict(
    type='AdamW',
    lr=lr,
    paramwise_cfg=dict(
       custom_keys={
           '.backbone.cls_token': dict(decay_mult=0.0),
           '.backbone.pos_embed': dict(decay_mult=0.0),
           '.backbone.time_embed': dict(decay_mult=0.0)
       }),
    weight_decay=1e-4,
    ) 

lr_config = dict(
    policy='CosineAnnealing',
    by_epoch=False,
    min_lr_ratio=1e-6,
    warmup='linear',
    warmup_ratio=1e-3,
    warmup_iters=20,
    warmup_by_epoch=True)


optimizer_config = dict(grad_clip=dict(max_norm=40, norm_type=2))


configs/video_retrieval/solver.py
dun-research/DRCA
# register modules
from models.modules.reso_align_transformer import (ResolutionAlignDividedTemporalAttentionWithNorm,
                        ResolutionAlignDividedSpatialAttentionWithNorm, 
                        ResolutionAlignFFN)
from models.modules.reso_align_layers import ResolutionAlignTransformerLayerSequence
from models.backbones.drca import DRCA

from models.heads.cls_head import ClsHead
from models.heads.similarity_head import SimilarityHead
from models.recognizer.similarity_recognizer import SimilarityRecognizer3D
models/__init__.py
dun-research/DRCA

from collections import OrderedDict
import numpy as np
import torch
from torch import Tensor
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange
from mmcv.cnn import  build_norm_layer
from mmcv.runner import _load_checkpoint, load_state_dict
from mmcv.cnn.utils.weight_init import trunc_normal_
from mmaction.utils import get_root_logger
from mmaction.models.builder import BACKBONES
from mmaction.models.backbones.timesformer import PatchEmbed
from models.modules.reso_align_layers import ResolutionAlignTransformerLayerSequence


@BACKBONES.register_module()
class DRCA(nn.Module):
    def __init__(self,
                 num_frames,
                 img_size,
                 patch_size,
                 pretrained=None,
                 embed_dims=768,
                 num_heads=12,
                 num_transformer_layers=12,
                 in_channels=3,
                 dropout_ratio=0.,
                 transformer_layers=None,
                 norm_cfg=dict(type='LN', eps=1e-6),
                 # compression args
                 comp_insert_layer=-1, # if comp_insert_layer is -1, don't use compression
                 comp_k=4,
                 comp_strategy='diff_rank',
                 comp_module="sal_ref",
                 sigma=0.05,
                 num_samples=500,
                 **kwargs):
        super().__init__(**kwargs)

        assert transformer_layers is None or isinstance(
            transformer_layers, (dict, list))
        
        self.num_frames = num_frames
        self.pretrained = pretrained
        self.embed_dims = embed_dims
        self.num_transformer_layers = num_transformer_layers



        self.patch_embed = PatchEmbed(
            img_size=img_size,
            patch_size=patch_size,
            in_channels=in_channels,
            embed_dims=embed_dims)
        num_patches = self.patch_embed.num_patches
        self.img_size = img_size
        self.patch_size = patch_size

        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dims))
        self.pos_embed = nn.Parameter(
            torch.zeros(1, num_patches + 1, embed_dims))
        self.drop_after_pos = nn.Dropout(p=dropout_ratio)

        self.time_embed = nn.Parameter(
                torch.zeros(1, num_frames, embed_dims))
        self.drop_after_time = nn.Dropout(p=dropout_ratio)

        self.norm = build_norm_layer(norm_cfg, embed_dims)[1]

        single_resolutions = [True for i in range(num_transformer_layers)] 
        if comp_insert_layer >=0:
            n = comp_insert_layer
            for i in range(n, num_transformer_layers):
                single_resolutions[i] = False
            
        if transformer_layers is None:
            # stochastic depth decay rule
            dpr = np.linspace(0, 0.1, num_transformer_layers)
            _transformerlayers_cfg = [
                dict(
                    type='BaseTransformerLayer',
                    attn_cfgs=[
                        dict(
                            type='ResolutionAlignDividedTemporalAttentionWithNorm',
                            embed_dims=embed_dims,
                            num_heads=num_heads,
                            num_frames=num_frames,
                            dropout_layer=dict(
                                type='DropPath', drop_prob=dpr[i]),
                            norm_cfg=dict(type='LN', eps=1e-6)),
                        dict(
                            type='ResolutionAlignDividedSpatialAttentionWithNorm',
                            embed_dims=embed_dims,
                            num_heads=num_heads,
                            num_frames=num_frames,
                            single_resolution=single_resolutions[i],
                            dropout_layer=dict(
                                type='DropPath', drop_prob=dpr[i]),
                            norm_cfg=dict(type='LN', eps=1e-6))
                    ],
                    ffn_cfgs=dict(
                            type='ResolutionAlignFFN',
                        single_resolution=single_resolutions[i],
                        embed_dims=embed_dims,
                        feedforward_channels=embed_dims * 4,
                        num_fcs=2,
                        act_cfg=dict(type='GELU'),
                        dropout_layer=dict(
                            type='DropPath', drop_prob=dpr[i]),
                        norm_cfg=dict(type='LN', eps=1e-6)),
                    operation_order=('self_attn', 'self_attn', 'ffn'))
                for i in range(num_transformer_layers)
            ]

        compression_args = dict(
                            comp_insert_layer=comp_insert_layer,
                            comp_k=comp_k,
                            comp_strategy=comp_strategy,
                            comp_module=comp_module,
                            sigma=sigma,
                            num_samples=num_samples,
                            embed_dim=embed_dims,
                            temporal_size=num_frames,
                            )

        self.transformer_layers = ResolutionAlignTransformerLayerSequence(_transformerlayers_cfg, num_transformer_layers, compression_args)

    def init_weights(self, pretrained=None):
        """Initiate the parameters either from existing checkpoint or from
        scratch."""
        
        trunc_normal_(self.pos_embed, std=.02)
        trunc_normal_(self.cls_token, std=.02)
        trunc_normal_(self.time_embed, std=.02)

        if pretrained:
            self.pretrained = pretrained
        
            
        if isinstance(self.pretrained, str):
            logger = get_root_logger()
            logger.info(f'load model from: {self.pretrained}')

            state_dict = _load_checkpoint(self.pretrained)
            if 'state_dict' in state_dict:
                state_dict = state_dict['state_dict']

            if not self.pretrained.startswith('imagenet'):
                self.load_state_dict(state_dict)

            self.attention_type = 'divided_space_time'
            if self.attention_type == 'divided_space_time':
                # modify the key names of norm layers
                old_state_dict_keys = list(state_dict.keys())
                for old_key in old_state_dict_keys:
                    if 'norms' in old_key:
                        new_key = old_key.replace('norms.0',
                                                  'attentions.0.norm')
                        new_key = new_key.replace('norms.1', 'ffns.0.norm')
                        state_dict[new_key] = state_dict.pop(old_key)

                # copy the parameters of space attention to time attention
                old_state_dict_keys = list(state_dict.keys())
                for old_key in old_state_dict_keys:
                    if 'attentions.0' in old_key:
                        new_key = old_key.replace('attentions.0',
                                                  'attentions.1')
                        state_dict[new_key] = state_dict[old_key].clone()

                # init pos_embedding
                ## Resizing the positional embeddings in case they don't match
                num_patches = (self.img_size // self.patch_size) * (self.img_size // self.patch_size)
                if num_patches + 1 != state_dict['pos_embed'].size(1):
                    pos_embed = state_dict['pos_embed']
                    cls_pos_embed = pos_embed[0,0,:].unsqueeze(0).unsqueeze(1)
                    other_pos_embed = pos_embed[0,1:,:].unsqueeze(0).transpose(1, 2)
                    new_pos_embed = F.interpolate(other_pos_embed, size=(num_patches), mode='nearest')
                    new_pos_embed = new_pos_embed.transpose(1, 2)
                    new_pos_embed = torch.cat((cls_pos_embed, new_pos_embed), 1)
                    state_dict['pos_embed'] = new_pos_embed
                    logger.info(f"resize pos_embed from {pos_embed.size()} to {new_pos_embed.size()}")

                load_state_dict(self, state_dict, strict=False, logger=logger)

        
    def forward(self, x):
        """Defines the computation performed at every call."""
        # x [batch_size * num_frames, num_patches, embed_dims]
        batches = x.shape[0]
        x = self.patch_embed(x)

        # x [batch_size * num_frames, num_patches + 1, embed_dims]
        cls_tokens = self.cls_token.expand(x.size(0), -1, -1)
        x = torch.cat((cls_tokens, x), dim=1)
        x = x + self.pos_embed
        x = self.drop_after_pos(x)

        # Add Time Embedding
        # x [batch_size, num_patches * num_frames + 1, embed_dims]
        cls_tokens = x[:batches, 0, :].unsqueeze(1)
        x = rearrange(x[:, 1:, :], '(b t) p m -> (b p) t m', b=batches)
        x = x + self.time_embed
        x = self.drop_after_time(x)
        x = rearrange(x, '(b p) t m -> b (p t) m', b=batches)
        x = torch.cat((cls_tokens, x), dim=1)

        x = self.transformer_layers(x, None, None)
        x = self.norm(x)

        # Return Class Token
        return x[:, 0]

    def get_rank_index(self, x):
        batches = x.shape[0]
        x = self.patch_embed(x)

        # x [batch_size * num_frames, num_patches + 1, embed_dims]
        cls_tokens = self.cls_token.expand(x.size(0), -1, -1)
        x = torch.cat((cls_tokens, x), dim=1)
        x = x + self.pos_embed
        x = self.drop_after_pos(x)

        # Add Time Embedding
        # x [batch_size, num_patches * num_frames + 1, embed_dims]
        cls_tokens = x[:batches, 0, :].unsqueeze(1)
        x = rearrange(x[:, 1:, :], '(b t) p m -> (b p) t m', b=batches)
        x = x + self.time_embed
        x = self.drop_after_time(x)
        x = rearrange(x, '(b p) t m -> b (p t) m', b=batches)
        x = torch.cat((cls_tokens, x), dim=1)

        index = self.transformer_layers.get_rank_index(x)
        return index

models/backbones/drca.py
dun-research/DRCA
# Copyright (c) OpenMMLab. All rights reserved.
import torch.nn as nn
from mmcv.cnn import trunc_normal_init
from mmaction.models import BaseHead, HEADS


@HEADS.register_module()
class ClsHead(BaseHead):
    """Classification head for TimeSformer.

    Args:
        num_classes (int): Number of classes to be classified.
        in_channels (int): Number of channels in input feature.
        loss_cls (dict): Config for building loss.
            Defaults to `dict(type='CrossEntropyLoss')`.
        init_std (float): Std value for Initiation. Defaults to 0.02.
        kwargs (dict, optional): Any keyword argument to be used to initialize
            the head.
    """

    def __init__(self,
                 num_classes,
                 in_channels,
                 loss_cls=dict(type='CrossEntropyLoss'),
                 init_std=0.02,
                 **kwargs):
        super().__init__(num_classes, in_channels, loss_cls, **kwargs)
        self.init_std = init_std
        self.fc_cls = nn.Linear(self.in_channels, self.num_classes)

    def init_weights(self):
        """Initiate the parameters from scratch."""
        trunc_normal_init(self.fc_cls, std=self.init_std)

    def forward(self, x):
        # [N, in_channels]
        cls_score = self.fc_cls(x)
        # [N, num_classes]
        return cls_score
models/heads/cls_head.py
dun-research/DRCA


import torch.nn as nn

from models.losses.metric_loss import MetricLoss
from mmaction.models.heads import BaseHead
from mmaction.models.builder import HEADS
from collections import OrderedDict


@HEADS.register_module()
class SimilarityHead(BaseHead):
    def __init__(self, encoder=None, 
                       metric_loss=dict(type="MultiSimilarityLoss",
                                               distance="Cosine",
                                               miner_name="MultiSimilarityMiner"),
                        ):
        super().__init__(num_classes=1, in_channels=encoder['input_channels'])
        
        encoder_type = encoder.pop("type")
        out_channel = encoder['out_channels']
        if encoder_type == "FC":
            model = nn.Linear(encoder['input_channels'], encoder['out_channels'])
        elif encoder_type == "FCBN":
            model = nn.Sequential(OrderedDict([
                    ("conv1",nn.Linear(encoder['input_channels'], encoder['out_channels'])),
                    ("bn1",nn.BatchNorm1d(encoder['input_channels']))])
                )
        elif encoder_type == "None":
            model = nn.Identity()
            out_channel = encoder['input_channels']
        else:
             model = None
        
        self.encoder = model
        print(f"sim_head.encoder = {model}" )
        if metric_loss.get('feat_dim', -1) <=0 :
            metric_loss['feat_dim'] = out_channel
        self.loss_sim = MetricLoss(**metric_loss)
        self.pool = nn.AdaptiveAvgPool3d([1,None, None])
        self.fp16_enabled = False
    
    def init_weights(self):
        self.encoder.init_weights()

    def forward(self, features, **kwargs):
        n_dim = len(features.size())
        if self.encoder is not None:
            assert n_dim in [2, 5]
            if n_dim == 5:
                features = self.pool(features).squeeze()
            
            features = self.encoder(features)
            return features
        else:
            raise NotImplemented

    def loss(self, embeddings, labels, **kwargs):
        #loss = self.loss_sim(embeddings, labels, **kwargs)
        n_dim = len(embeddings.size())
        # print("embeddings :{}".format(embeddings.size()))
        if n_dim > 2 :
            embeddings = embeddings.flatten(2).mean(dim=2)
        loss = self.loss_sim(embeddings, labels, **kwargs)

        return {"loss_sim": loss}




models/heads/similarity_head.py
dun-research/DRCA


from numpy import empty
import numpy as np
import torch
import torch.nn as nn
from mmaction.models.builder import LOSSES
from mmaction.models.losses import BaseWeightedLoss
from pytorch_metric_learning import miners, losses, distances
from pytorch_metric_learning.utils import distributed as pml_dist

from pytorch_metric_learning.utils import loss_and_miner_utils as lmu
import torch.distributed as dist
def get_rank():
    if not dist.is_available():
        return 0
    if not dist.is_initialized():
        return 0
    return dist.get_rank()


@LOSSES.register_module()
class MetricLoss(BaseWeightedLoss):
    def __init__(self,  loss_weight=1, 
                        type="TripletMarginLoss",
                        miner_name="Empty",
                        margin=0.3, 
                        cross_batch_memory=False,
                        memory_size=512,
                        feat_dim=-1,
                        copy_ref_label=False,
                        distributed=False,
                        distance="Cosine"):
        super().__init__(loss_weight)
        if distance == "Cosine":
            dist_func = distances.CosineSimilarity()
        elif distance == "L2":
            dist_func = distances.LpDistance(p=2, power=2)
        elif distance == "Empty":
            dist_func = None
        else:
            raise NotImplemented("unsupported metric distance:{}".format(distance)) 

        if miner_name == "Empty":
            miner = None
        elif miner_name == "TripletMarginMiner":
            miner = miners.TripletMarginMiner(margin=margin, type_of_triplets="semihard", distance=dist_func)
        elif miner_name == "MultiSimilarityMiner":
            miner = miners.MultiSimilarityMiner(epsilon=margin, distance=dist_func)
        else:
            raise NotImplementedError("unsupported miner: {}".format(miner_name))               

        if type == "TripletMarginLoss":
            metric_loss = losses.TripletMarginLoss(margin=margin, distance=dist_func)
        elif type == "MultiSimilarityLoss":
            metric_loss = losses.MultiSimilarityLoss(distance=dist_func)
        else:
            raise NotImplementedError("unsupported metric loss:{}".format(type)) 
        
        efficient=False
        if cross_batch_memory:
            metric_loss = losses.CrossBatchMemory(metric_loss, embedding_size=feat_dim, memory_size=memory_size)
        self.distributed = distributed
        
        if distributed:
            metric_loss = pml_dist.DistributedLossWrapper(metric_loss,efficient=efficient)
            if miner is not None:
                miner = pml_dist.DistributedMinerWrapper(miner,efficient=efficient)
        self.distance = dist_func
        self.metric_loss = metric_loss
        self.miner = miner
        self.fp16_enabled = True

    def _forward(self, embeddings, targets, enqueue_idx=None):
        embeddings = embeddings / embeddings.norm(dim=-1, keepdim=True)
        labels = targets.reshape(-1,)
        loss = self.metric_loss(embeddings, labels.reshape(-1), indices_tuple=None,)
        return loss




models/losses/metric_loss.py
dun-research/DRCA
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange
from models.modules.weight_init import weight_init
from models.modules.reference_compress import ReferenceBlock
from typing import Any, Callable, List, Optional, Type, Union
from torch import Tensor


def conv3x3_2d(in_planes: int, out_planes: int, stride: int = 1, groups: int = 1, dilation: int = 1) -> nn.Conv2d:
    """3x3 convolution with padding"""
    return nn.Conv2d(
        in_planes,
        out_planes,
        kernel_size=3,
        stride=stride,
        padding=dilation,
        groups=groups,
        bias=False,
        dilation=dilation,
    )


def conv1x1_2d(in_planes: int, out_planes: int, stride: int = 1) -> nn.Conv2d:
    """1x1 convolution"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)

def conv1x1x1_3d(in_planes: int, out_planes: int, stride: int = 1) -> nn.Conv3d:
    """1x1 convolution"""
    return nn.Conv3d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)


class Bottleneck2D(nn.Module):
    # Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2)
    # while original implementation places the stride at the first 1x1 convolution(self.conv1)
    # according to "Deep residual learning for image recognition" https://arxiv.org/abs/1512.03385.
    # This variant is also known as ResNet V1.5 and improves accuracy according to
    # https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch.

    expansion: int = 4

    def __init__(
        self,
        inplanes: int,
        planes: int,
        width: int,
        stride: int = 1,
        downsample: Optional[nn.Module] = None,
        dilation: int = 1,
        norm_layer: Optional[Callable[..., nn.Module]] = None,
    ) -> None:
        super().__init__()

        norm_layer = nn.BatchNorm2d
        self.conv1 = conv1x1_2d(inplanes, width)
        self.bn1 = norm_layer(width)
        self.conv2 = conv3x3_2d(width, width, stride, 1, dilation)
        self.bn2 = norm_layer(width)
        self.conv3 = conv1x1_2d(width, inplanes)
        self.bn3 = norm_layer(inplanes)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x: Tensor) -> Tensor:
        identity = x

        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 = self.conv3(out)
        out = self.bn3(out)
        # print(f"out-bn3: {out.shape}")
        if self.downsample is not None:
            identity = self.downsample(x)
        # print(f"identity: {identity.shape}")
        out += identity
        out = self.relu(out)

        return out


class Bottleneck3D(nn.Module):
    # Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2)
    # while original implementation places the stride at the first 1x1 convolution(self.conv1)
    # according to "Deep residual learning for image recognition" https://arxiv.org/abs/1512.03385.
    # This variant is also known as ResNet V1.5 and improves accuracy according to
    # https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch.

    expansion: int = 4

    def __init__(
        self,
        inplanes: int,
        planes: int,
        width: int,
        stride: int = 1,
        downsample: Optional[nn.Module] = None,
        dilation: int = 1,
        norm_layer: Optional[Callable[..., nn.Module]] = None,
    ) -> None:
        super().__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm3d
        self.conv1 = conv1x1x1_3d(inplanes, width, )
        self.bn1 = norm_layer(width)
        self.conv2 = nn.Conv3d(width, width, kernel_size=(3, 3, 3), stride=(1, 2, 2), padding=(1, 1, 1), bias=False)
        self.bn2 = norm_layer(width)
        self.conv3 = conv1x1x1_3d(width, inplanes)
        self.bn3 = norm_layer(inplanes)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x: Tensor) -> Tensor:
        identity = x

        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 = self.conv3(out)
        out = self.bn3(out)
        if self.downsample is not None:
            identity = self.downsample(x)
        out += identity
        out = self.relu(out)

        return out

class CompressionConv2D(nn.Module):
    def __init__(self, in_channels) -> None:
        super().__init__()
        self.downsample = nn.Sequential(
                nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=1),
                nn.BatchNorm2d(in_channels),
                nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1),
                nn.BatchNorm2d(in_channels),
        ) 
        self.apply(weight_init)

    def forward(self, x, x_ref=None):
        B, C, T, H, W = x.shape
        x = rearrange(x, 'b c t h w -> (b t) c h w')
        out = self.downsample(x)
        W = x.size(-1)
        out = rearrange(out, '(b t) c h w -> b c t h w', b=B,)
        return out
    

class ResBlock2D(nn.Module):
    def __init__(self, in_channels) -> None:
        super().__init__()
        downsample = nn.Sequential(
                conv1x1_2d(in_channels, in_channels, 2),
                norm_layer(in_channels),
            )
        self.bottle_neck = Bottleneck2D(in_channels, in_channels, width=in_channels*2, stride=2, downsample=downsample)
    
    def forward(self, x, x_ref):
        B = x.size(0)
        x = rearrange(x, 'b c t h w -> (b t) c h w')
        out = self.bottle_neck(x)
        W = x.size(-1)
        out = rearrange(out, '(b t) c h w -> b c t h w', b=B,)
        return out



class ResBlock3D(nn.Module):
    def __init__(self, in_channels) -> None:
        super().__init__()
        downsample = nn.Sequential(
                nn.Conv3d(in_channels, in_channels, kernel_size=(3, 3, 3), stride=(1, 2, 2), padding=(1, 1, 1), bias=False ),
                nn.BatchNorm3d(in_channels),
            )
        self.bottle_neck = Bottleneck3D(in_channels, in_channels, width=in_channels*2, stride=2, downsample=downsample)
    
    def forward(self, x, x_ref):
        # x.shape=(b c t h w)
        B = x.size(0)
        out = self.bottle_neck(x)
        return out


class CompressionConv3D(nn.Module):
    def __init__(self, in_channels) -> None:
        super().__init__()
        self.reduce = nn.Sequential(
                nn.Conv3d(in_channels, in_channels*2, kernel_size=(3, 3, 3), stride=(1, 2, 2), padding=(1, 1, 1), bias=False),
                nn.BatchNorm3d(in_channels*2),
                nn.Conv3d(in_channels*2, in_channels, kernel_size=1, stride=1, padding=0, bias=False),
                nn.BatchNorm3d(in_channels),
        ) 
        self.apply(weight_init)

    def forward(self, x, x_ref=None):
        B, C, T, H, W = x.shape
        out = self.reduce(x)
        # print(f"out: {out.shape}" )
        return out


class PatchMerge3D(nn.Module):
    """ Patch Merging Layer
    Args:
        dim (int): Number of input channels.
        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
    """
    def __init__(self, in_channels, norm_layer=nn.LayerNorm):
        super().__init__()
        dim = in_channels
        # swap temporal information
        self.in_conv = nn.Sequential(
            nn.Conv3d(in_channels, in_channels, kernel_size=(1, 1, 1), stride=1, padding=0, bias=False),
            nn.BatchNorm3d(in_channels),
            nn.Conv3d(in_channels, in_channels, kernel_size=(3, 1, 1), stride=1, padding=(1, 0, 0), bias=False),
            nn.BatchNorm3d(in_channels),
            nn.Conv3d(in_channels, in_channels, kernel_size=(1, 1, 1), stride=1, padding=0, bias=False),
            nn.BatchNorm3d(in_channels),
        )
        self.reduction = nn.Linear(4 * dim, dim, bias=False)
        self.norm = norm_layer(4 * dim)
        self.apply(weight_init)

    def forward(self, x, x_ref=None):
        B, D, H, W, C = x.shape
        x = self.in_conv(x)
        x = rearrange(x, 'b c t h w -> b t  h w c')
        # padding
        pad_input = (H % 2 == 1) or (W % 2 == 1)
        if pad_input:
            x = F.pad(x, (0, 0, 0, W % 2, 0, H % 2))

        x0 = x[:, :, 0::2, 0::2, :]  # B D H/2 W/2 C
        x1 = x[:, :, 1::2, 0::2, :]  # B D H/2 W/2 C
        x2 = x[:, :, 0::2, 1::2, :]  # B D H/2 W/2 C
        x3 = x[:, :, 1::2, 1::2, :]  # B D H/2 W/2 C
        x = torch.cat([x0, x1, x2, x3], -1)  # B D H/2 W/2 4*C

        x = self.norm(x)
        x = self.reduction(x)
        x = rearrange(x, 'b t  h w c -> b c t h w ')
 
        return x

 
def get_saliency_frame_reference_compression(in_channels):
    model =  ReferenceBlock(dim=in_channels,
                dim_out=in_channels,
                num_heads=4,
                has_cls_embed=False,
                kernel_q=(3, 3),
                kernel_kv=(1, 1),
                stride_q=(2, 2),
                stride_kv=(1, 1),
            ) 
    return model


def get_compression_module(name, in_channels):
    if name == "conv2d":
        return CompressionConv2D(in_channels)
    elif name == "conv3d":
        return CompressionConv3D(in_channels)
    elif name == "patch_merge":
        return PatchMerge3D(in_channels)
    elif name == "resblock_2d":
        return ResBlock2D(in_channels)
    elif name == "resblock_3d":
        return ResBlock3D(in_channels)
    elif name == "sal_ref":
        return get_saliency_frame_reference_compression(in_channels)
    else:
        raise NotImplementedError(f"Unknown compression module: {name}")

models/modules/compression_modules.py
dun-research/DRCA
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange
from math import sqrt
from models.modules.compression_modules import get_compression_module
from models.modules.diff_rank import PerturbedRank, UniformSample
from models.modules.weight_init import weight_init


class DiffContextAwareCompressionModule(nn.Module):
    def __init__(self, in_channels, T, strategy="diff_rank", comp_module="krc", k=4, sigma=0.05, n_sample=500) -> None:
        super().__init__()
        self.strategy = strategy
        self.comp_module = comp_module
        self.k = k
        self.sigma = sigma
        self.n_sample = n_sample
        self.in_channels = in_channels

        self.compression = get_compression_module(comp_module, in_channels)
        
        if strategy == "diff_rank":
            self.rank_net = PerturbedRank(in_channels, k, n_sample, sigma)
        elif strategy == "wo_diff_rank":
            self.rank_net = PerturbedRank(in_channels, k, n_sample, sigma, do_train=False)
        elif strategy == "uniform":
            self.rank_net = UniformSample(k, T)
        else:
            raise ValueError("Unknown strategy: {}".format(strategy))
        self.apply(weight_init)

    def forward(self, x, x_cls, return_index=False):
        B, C, T, H, W = x.shape
        origin_H, origin_W = H, W
        
        frames_topk, frames_back, sorted_inds = self.rank_net(x)
        frames_back_lowres = self.compression(frames_back, frames_topk)
        
        frames_topk = rearrange(frames_topk, "b c k h w -> b h w k c", )
        frames_back_lowres = rearrange(frames_back_lowres, "b c nk h w -> b h w nk c", )
        
        if return_index:
            return frames_topk, frames_back_lowres, sorted_inds
        else:
            return frames_topk, frames_back


models/modules/dccm.py
dun-research/DRCA
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import einops
from einops import rearrange
from math import sqrt
import torch.distributed as dist
from models.modules.weight_init import weight_init

class UniformSample(nn.Module):
    def __init__(self, k: int, T: int) -> None:
        super().__init__()
        self.k = k
        self.T = T

    def forward(self, frames):
        B,C,T,H,W = frames.shape
        frames = rearrange(frames, "b c t h w -> b t (h w c)")
        indices_topk, indices_back = self.generate_uniform_indices(B)
        frames_topk = extract_from_indices(frames, indices_topk)
        frames_back = extract_from_indices(frames, indices_back)
        sorted_inds = torch.cat([indices_topk, indices_back], dim=-1)
        frames_topk = rearrange(frames_topk, "b k (h w c) -> b c k h w", b=B, k=self.k, h=H, w=W)
        frames_back = rearrange(frames_back, "b nk (h w c) -> b c nk h w", b=B, nk=T-self.k, h=H, w=W)
        return frames_topk, frames_back, sorted_inds

    def generate_uniform_indices(self, b):
        indices_top = np.linspace(0, self.T, self.k+1).astype(np.int32)[:-1]
        indices_back = np.array(list(set(range(self.T)) - set(indices_top.tolist()))).astype(np.int32)
        indices_back.sort()

        indices_top = torch.from_numpy(indices_top).long().unsqueeze(0).repeat(b, 1).cuda()
        indices_back = torch.from_numpy(indices_back).long().unsqueeze(0).repeat(b, 1).cuda()
        return indices_top, indices_back
    

class PerturbedRankFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, k: int, num_samples: int = 1000, sigma: float = 0.05):
        b, d = x.shape
        # for Gaussian: noise and gradient are the same.
        noise = torch.normal(mean=0.0, std=1.0, size=(b, num_samples, d)).to(x.device)
        # noise = torch.zeros_like(noise)
        perturbed_x = x[:, None, :] + noise * sigma # b, nS, d
        
        topk_results = torch.topk(perturbed_x, k=k, dim=-1, sorted=False)
        indices_topk = topk_results.indices # b, nS, d
        #indices_topk, indices_back = torch.split(indices, [k, d-k], dim=2) #indices[:, :, :k], indices[:, :, k:]
        indices_topk = torch.sort(indices_topk, dim=-1).values
        
        back_results = torch.topk(perturbed_x, k=d-k, dim=-1, largest=False, sorted=False)
        indices_back = back_results.indices # b, nS, d
        #indices_topk, indices_back = torch.split(indices, [k, d-k], dim=2) #indices[:, :, :k], indices[:, :, k:]
        indices_back = torch.sort(indices_back, dim=-1).values

        # b, nS, k, d
        perturbed_output_topk = torch.nn.functional.one_hot(indices_topk, num_classes=d).float()
        indicators_topk = perturbed_output_topk.mean(dim=1) # b, k, d

        perturbed_output_back = torch.nn.functional.one_hot(indices_back, num_classes=d).float()
        indicators_back = perturbed_output_back.mean(dim=1) # b, k, d

        # constants for backward
        ctx.k = k
        ctx.num_samples = num_samples
        ctx.sigma = sigma

        ctx.save_for_backward(perturbed_output_topk, perturbed_output_back, noise)
        return (indicators_topk, indicators_back)

    @staticmethod
    def backward(ctx, grad_output_topk, grad_output_back):
        if grad_output_topk is None:
            print("grad_output_topk is None")
            return tuple([None] * 5)

        perturbed_output_topk, perturbed_output_back, noise_gradient = ctx.saved_tensors

        # import pdb; pdb.Pdb(nosigint=True).set_trace()
        if ctx.sigma <= 1e-20:
            b, _, k, d = perturbed_output_topk.size()
            expected_gradient_topk = torch.zeros(b, k, d).to(grad_output_topk.device)

            b, _, k, d = perturbed_output_back.size()
            expected_gradient_back = torch.zeros(b, k, d).to(grad_output_back.device)
        else:
            expected_gradient_topk = (
                torch.einsum("bnkd,bnd->bkd", perturbed_output_topk, noise_gradient)
                / ctx.num_samples
                / (ctx.sigma)
            )

            expected_gradient_back = (
                torch.einsum("bnkd,bnd->bkd", perturbed_output_back, noise_gradient)
                / ctx.num_samples
                / (ctx.sigma)
            )
        # (grad_output_topk * expected_gradient_topk).sum(1)
        grad_input_topk = torch.einsum("bkd,bkd->bd", grad_output_topk, expected_gradient_topk)

        grad_input_back = torch.einsum("bkd,bkd->bd", grad_output_back, expected_gradient_back)
        #print(f"topk: {grad_input_topk.sum().item()}, back: {grad_input_back.sum().item()}")

        # get final grad_input
        grad_input = grad_input_topk + grad_input_back
        #print(f"grad_input: {grad_input.sum()}, {grad_output_topk.shape}, {grad_output_back.shape}")
        return (grad_input/10., None, None, None, None)
    
    
class PerturbedRank(nn.Module):
    def __init__(self, embed_dim:int, k: int, num_samples: int = 1000, sigma: float = 0.05, do_train: bool = True):
        super(PerturbedRank, self).__init__()
        self.num_samples = num_samples
        self.sigma = sigma
        self.k = k
        self.score_network = ScoreNet(embed_dim)
        self.do_train = do_train
        self.apply(weight_init)

    def forward(self, frames):
        B,C,T,H,W = frames.shape
        scores = self.score_network(frames)

        frames = rearrange(frames, "b c t h w -> b t (h w c)")
        if self.training and self.do_train:
            indicator_topk, indicator_back = self.get_indicator(scores)
            frames_topk = extract_from_indicators(frames, indicator_topk)
            frames_back = extract_from_indicators(frames, indicator_back)
        else:
            indices_topk, indices_back = self.get_indices(scores)
            frames_topk = extract_from_indices(frames, indices_topk)
            frames_back = extract_from_indices(frames, indices_back)
        _, sorted_inds = torch.sort(scores, dim=1,)  
        
        frames_topk = rearrange(frames_topk, "b k (h w c) -> b c k h w", b=B, k=self.k, h=H, w=W)
        frames_back = rearrange(frames_back, "b nk (h w c) -> b c nk h w", b=B, nk=T-self.k, h=H, w=W)
        return frames_topk, frames_back, sorted_inds
    
    def get_indicator(self, score):
        indicator_topk, indicator_back = PerturbedRankFunction.apply(score, self.k, self.num_samples, self.sigma)
        indicator_topk = einops.rearrange(indicator_topk, "b k d -> b d k")
        indicator_back = einops.rearrange(indicator_back, "b k d -> b d k")
        return indicator_topk, indicator_back

    def get_indices(self, score):
        k = self.k
        topk_results = torch.topk(score, k=score.shape[1], dim=-1, sorted=True)
        indices = topk_results.indices # b, k
        indices_topk, indices_back = indices[..., :k], indices[..., k:]
        indices_topk = torch.sort(indices_topk, dim=-1).values
        indices_back = torch.sort(indices_back, dim=-1).values
        return indices_topk, indices_back



class ScoreNet(nn.Module):
    """ Image to Patch Embedding
    """
    def __init__(self, embed_dim: int=384):
        super().__init__()
        embed_dim *= 2
        self.downsample = torch.nn.Conv3d(embed_dim//2, embed_dim//2, (1,3,3), (1,2,2))
        self.bn = torch.nn.BatchNorm3d(embed_dim//2)

        self.in_conv = nn.Sequential(
            nn.LayerNorm(embed_dim),
            nn.Linear(embed_dim, embed_dim),
            nn.GELU()
        )
        
        self.out_conv = nn.Sequential(
            nn.Linear(embed_dim, embed_dim // 2),
            nn.GELU(),
            nn.Linear(embed_dim // 2, embed_dim // 4),
            nn.GELU(),
            nn.Linear(embed_dim // 4, 1)
        )
        self.sigmoid = torch.nn.Sigmoid()

    def forward(self, x):
        x = x.detach()  # stop gradient here
        x = self.bn(self.downsample(x))
        
        x = rearrange(x, 'b c t h w -> b t (h w) c')
        avg = torch.mean(x, dim=2, keepdim=False)
        max = torch.max(x, dim=2).values
        x = torch.cat((avg, max), dim=2)
        
        x = self.in_conv(x)
        B, N, C = x.size()
        local_x = x[:,:, :C//2]
        
        global_x = torch.mean(x[:,:, C//2:], dim=1, keepdim=True)
        x = torch.cat([local_x, global_x.expand(B, N, C//2)], dim=-1)

        score =  self.out_conv(x)
        score = score.squeeze(-1)
        score = self.min_max_norm(score)
        return score

    def min_max_norm(self, x):
        flatten_score_min = x.min(axis=-1, keepdim=True).values
        flatten_score_max = x.max(axis=-1, keepdim=True).values
        norm_flatten_score = (x - flatten_score_min) / (flatten_score_max - flatten_score_min + 1e-5)
        return norm_flatten_score


def batched_index_select(input, dim, index):
    for i in range(1, len(input.shape)):
        if i != dim:
            index = index.unsqueeze(i)
    expanse = list(input.shape)
    expanse[0] = -1
    expanse[dim] = -1
    index = index.expand(expanse)
    return torch.gather(input, dim, index)


def extract_from_indices(x, indices):
    batch_size, _, channels = x.shape
    k = indices.shape[-1]
    patches = x
    patches = batched_index_select(patches, 1, indices)
    patches = patches.contiguous().view(batch_size, k, channels)
    return patches


def extract_from_indicators(x, indicators):
    indicators = rearrange(indicators, "b d k -> b k d")
    patches = torch.bmm(indicators, x)
    return patches

models/modules/diff_rank.py
dun-research/DRCA
import numpy
import torch
import torch.nn as nn
from torch.nn.init import trunc_normal_
from einops import rearrange
from torch.nn import init
from math import sqrt

class Mlp(nn.Module):
    def __init__(
        self,
        in_features,
        hidden_features=None,
        out_features=None,
        act_layer=nn.GELU,
        drop_rate=0.0,
    ):
        super().__init__()
        self.drop_rate = drop_rate
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        if self.drop_rate > 0.0:
            self.drop = nn.Dropout(drop_rate)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        if self.drop_rate > 0.0:
            x = self.drop(x)
        x = self.fc2(x)
        if self.drop_rate > 0.0:
            x = self.drop(x)
        return x


class Permute(nn.Module):
    def __init__(self, dims):
        super().__init__()
        self.dims = dims

    def forward(self, x):
        return x.permute(*self.dims)


def drop_path(x, drop_prob: float = 0.0, training: bool = False):
    """
    Stochastic Depth per sample.
    """
    if drop_prob == 0.0 or not training:
        return x
    keep_prob = 1 - drop_prob
    shape = (x.shape[0],) + (1,) * (
        x.ndim - 1
    )  # work with diff dim tensors, not just 2D ConvNets
    mask = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
    mask.floor_()  # binarize
    output = x.div(keep_prob) * mask
    return output


class DropPath(nn.Module):
    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks)."""

    def __init__(self, drop_prob=None):
        super(DropPath, self).__init__()
        self.drop_prob = drop_prob

    def forward(self, x):
        return drop_path(x, self.drop_prob, self.training)


def attention_downsample(tensor, pool, hw_shape, has_cls_embed=True, norm=None):
    if pool is None:
        return tensor, hw_shape
    tensor_dim = tensor.ndim
    tensor_shape = tensor.shape
    if tensor_dim == 4:
        pass
    elif tensor_dim == 5:
        tensor = tensor.reshape((-1,)+tensor_shape[2:])
    else:
        raise NotImplementedError(f"Unsupported input dimension {tensor.shape}")

    if has_cls_embed:
        cls_tok, tensor = tensor[:, :, :1, :], tensor[:, :, 1:, :]

    B, N, L, C = tensor.shape
    #H, W = hw_shape
    H = W = int(sqrt(L))
    tensor = tensor.reshape(B * N, H, W, C).permute(0, 3, 1, 2).contiguous()

    tensor = pool(tensor)
    #print(f"poolout: {tensor.shape}")
    hw_shape = [tensor.shape[2], tensor.shape[3]]
    L_pooled = tensor.shape[2] * tensor.shape[3]
    tensor = tensor.reshape(B, N, C, L_pooled).transpose(2, 3)
    if has_cls_embed:
        tensor = torch.cat((cls_tok, tensor), dim=2)
    if norm is not None:
        tensor = norm(tensor)

    if tensor_dim == 3:
        tensor = tensor.squeeze(1)
    elif tensor_dim == 5:
        tensor = tensor.reshape((tensor_shape[0], tensor_shape[1]) + tensor.shape[1:])
    return tensor, hw_shape


class ReferenceAttention(nn.Module):
    def __init__(
        self,
        dim,
        dim_out,
        num_heads=8,
        qkv_bias=False,
        kernel_q=(1, 1),
        kernel_kv=(1, 1),
        stride_q=(1, 1),
        stride_kv=(1, 1),
        norm_layer=nn.LayerNorm,
        has_cls_embed=True,
    ):
        super().__init__()

        self.num_heads = num_heads
        self.dim_out = dim_out
        head_dim = dim_out // num_heads
        self.scale = head_dim**-0.5
        self.has_cls_embed = has_cls_embed
        padding_q = [int(q // 2) for q in kernel_q]
        padding_kv =[int(kv // 2) for kv in kernel_kv]

        self.q = nn.Linear(dim, dim_out, bias=qkv_bias)
        self.k = nn.Linear(dim, dim_out, bias=qkv_bias)
        self.v = nn.Linear(dim, dim_out, bias=qkv_bias)
    
        self.proj = nn.Linear(dim_out, dim_out)

        # Skip pooling with kernel and stride size of (1, 1, 1).
        if numpy.prod(kernel_q) == 1 and numpy.prod(stride_q) == 1:
            kernel_q = ()
        if numpy.prod(kernel_kv) == 1 and numpy.prod(stride_kv) == 1:
            kernel_kv = ()
        dim_conv = dim // num_heads
        self.pool_q = (
            nn.Conv2d(
                dim_conv,
                dim_conv,
                kernel_q,
                stride=stride_q,
                padding=padding_q,
                groups=dim_conv,
                bias=False,
            )
            if len(kernel_q) > 0
            else None
        )
        self.norm_q = norm_layer(dim_conv) if len(kernel_q) > 0 else None
        self.pool_k = (
            nn.Conv2d(
                dim_conv,
                dim_conv,
                kernel_kv,
                stride=stride_kv,
                padding=padding_kv,
                groups=dim_conv,
                bias=False,
            )
            if len(kernel_kv) > 0
            else None
        )
        self.norm_k = norm_layer(dim_conv) if len(kernel_kv) > 0 else None
        self.pool_v = (
            nn.Conv2d(
                dim_conv,
                dim_conv,
                kernel_kv,
                stride=stride_kv,
                padding=padding_kv,
                groups=dim_conv,
                bias=False,
            )
            if len(kernel_kv) > 0
            else None
        )
        self.norm_v = norm_layer(dim_conv) if len(kernel_kv) > 0 else None


    def forward(self, x, x_ref, hw_shape):
        """
            x: (B, T1, N, C)
            x_ref: (B, T2, N, C)
        """
        B, T1, N, _ = x.shape
        B, T2, N, _ = x_ref.shape

    
        fold_dim = self.num_heads
        
        x = x.reshape(B, T1, N, fold_dim, -1) #.permute(0,  2, 1, 3)
        x_ref = x_ref.reshape(B, T2, N, fold_dim, -1) #.permute(0, 2, 1, 3)
        x = rearrange(x, "b t n f c -> b t f n c")
        x_ref = rearrange(x_ref, "b t n f c -> b t f n c")
        # q = k = v = x
        q = x
        k = v = x_ref
        # print(f"step1: q({q.shape}), k({k.shape}), v({v.shape})")

        q, q_shape = attention_downsample(
            q,
            self.pool_q,
            hw_shape,
            has_cls_embed=self.has_cls_embed,
            norm=self.norm_q if hasattr(self, "norm_q") else None,
        )
        k, k_shape = attention_downsample(
            k,
            self.pool_k,
            hw_shape,
            has_cls_embed=self.has_cls_embed,
            norm=self.norm_k if hasattr(self, "norm_k") else None,
        )
        v, v_shape = attention_downsample(
            v,
            self.pool_v,
            hw_shape,
            has_cls_embed=self.has_cls_embed,
            norm=self.norm_v if hasattr(self, "norm_v") else None,
        )

        # print(f"step2: q({q.shape}), k({k.shape}), v({v.shape})")

        q = rearrange(q, "b t nh qn nd ->  (b t) qn (nh nd)")
        q = self.q(q)
        q = rearrange(q, "(b t) qn (nh nd) -> b t nh qn nd", nh=self.num_heads, b=B)

        v = rearrange(v, "b t nh qn nd ->  (b t) qn (nh nd)")
        v = self.v(v)
        v = rearrange(v, "(b t) qn (nh nd) -> b t nh qn nd", nh=self.num_heads, b=B)

        k = rearrange(k, "b t nh qn nd ->  (b t) qn (nh nd)")
        k = self.k(k)
        k = rearrange(k, "(b t) qn (nh nd) -> b t nh qn nd", nh=self.num_heads, b=B)

        # step3: q(torch.Size([3, 6, 4, 49, 8])), k(torch.Size([3, 2, 4, 196, 8])), v(torch.Size([3, 2, 4, 196, 8]))
        # print(f"step3: q({q.shape}), k({k.shape}), v({v.shape})")
        N = q.shape[2]

        # k.transpose(-2, -1)=torch.Size([2, 4, 8, 196]), attn: torch.Size([2, 4, 49, 196])
        q = rearrange(q, "b t nh qn nd -> b nh (t qn) nd")
        k = rearrange(k, "b t nh kn nd -> b nh (t kn) nd")
        # print(f"q: {q.shape}, k: {k.shape}")
        attn = (q * self.scale) @ k.transpose(-2, -1)

        # print(f"k.transpose(-2, -1)={k.transpose(-2, -1).shape}, attn: {attn.shape}")

        attn = attn.softmax(dim=-1)
        v = rearrange(v, "b t nh vn nd -> b nh (t vn) nd")
        x = attn @ v

        # print(f"after-attn-x: {x.shape}")
        if self.has_cls_embed:
            x[:, :, 1:, :] += q[:, :, 1:, :]
        else:
            x = x + q

        x = x.transpose(1, 2).reshape(B, -1, self.dim_out)
        x = self.proj(x)
        # print(f"final-attn-x: {x.shape}")
        x = rearrange(x, "b (t xn) c -> b t xn c", t=T1)

        return x, q_shape


class ReferenceBlock(nn.Module):
    def __init__(
        self,
        dim,
        dim_out,
        num_heads,
        mlp_ratio=4.0,
        qkv_bias=False,
        drop_path=0.0,
        norm_layer=nn.LayerNorm,
        kernel_q=(1, 1),
        kernel_kv=(1, 1),
        stride_q=(1, 1),
        stride_kv=(1, 1),
        has_cls_embed=True,
        dim_mul_in_att=False,
    ):
        super().__init__()
        self.dim = dim
        self.dim_out = dim_out
        self.norm1 = norm_layer(dim)
        self.dim_mul_in_att = dim_mul_in_att

        att_dim = dim
        self.attn = ReferenceAttention(
            dim,
            att_dim,
            num_heads=num_heads,
            qkv_bias=qkv_bias,
            kernel_q=kernel_q,
            kernel_kv=kernel_kv,
            stride_q=stride_q,
            stride_kv=stride_kv,
            norm_layer=norm_layer,
            has_cls_embed=has_cls_embed,
        )
        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
        self.norm2 = norm_layer(att_dim)
        mlp_hidden_dim = int(att_dim * mlp_ratio)
        self.has_cls_embed = has_cls_embed

        mlp_dim_out = dim_out
        self.mlp = Mlp(
            in_features=att_dim,
            hidden_features=mlp_hidden_dim,
            out_features=mlp_dim_out,
        )
        if dim != dim_out:
            self.proj = nn.Linear(dim, dim_out)

        if len(stride_q) > 0 and numpy.prod(stride_q) > 1:
            kernel_skip = [s + 1 if s > 1 else s for s in stride_q]
            stride_skip = stride_q
            padding_skip = [int(skip // 2) for skip in kernel_skip]
            self.pool_skip = nn.MaxPool2d(
                kernel_skip, stride_skip, padding_skip, ceil_mode=False
            )
        else:
            self.pool_skip = None

    def forward(self, x, x_ref):
        
        hw_shape = (x.shape[3], x.shape[4])
        x = rearrange(x, 'b c t h w -> b t (h w) c')
        x_ref = rearrange(x_ref, 'b c t h w -> b t (h w) c')

        x_norm = self.norm1(x)
        x_block, hw_shape_new = self.attn(x_norm, x_ref, hw_shape)
        if self.dim_mul_in_att and self.dim != self.dim_out:
            x = self.proj(x_norm)

        x_res, _ = attention_downsample(
            x, self.pool_skip, hw_shape, has_cls_embed=self.has_cls_embed
        )
        
        x = x_res + self.drop_path(x_block)
        x_norm = self.norm2(x)
        x_mlp = self.mlp(x_norm)

        if not self.dim_mul_in_att and self.dim != self.dim_out:
            x = self.proj(x_norm)
        x = x + self.drop_path(x_mlp)
        x = rearrange(x, "b t (h w) c -> b c t h w", h=hw_shape_new[0], w=hw_shape_new[1])
        return x


def initialize_weights(module):
    if isinstance(module, nn.Conv2d):
        init.kaiming_normal_(module.weight, mode='fan_out', nonlinearity='relu')
        if module.bias is not None:
            init.constant_(module.bias, 0)
    elif isinstance(module, nn.BatchNorm2d):
        init.constant_(module.weight, 1)
        if module.bias is not None:
            init.constant_(module.bias, 0)
    elif isinstance(module, nn.Linear):
        init.normal_(module.weight, 0, 0.01)
        if module.bias is not None:
            init.constant_(module.bias, 0)


if __name__ == "__main__":
    print("\n======================================\n")
    kr_block = ReferenceBlock(dim=32,
                dim_out=32,
                num_heads=4,
                has_cls_embed=False,
                kernel_q=(1, 1),
                kernel_kv=(3, 3),
                stride_q=(1, 1),
                stride_kv=(2, 2),
            )

    x = torch.rand(3, 32, 6, 7, 7)
    x_ref = torch.rand(3, 32, 2, 7, 7,)
    out = kr_block(x, x_ref)
    print(f"input :{x.shape} ->  output: {out.shape}")

models/modules/reference_compress.py
dun-research/DRCA
import torch
import torch.nn as nn
from mmcv.cnn.bricks.transformer import build_transformer_layer
import copy
import numpy as np
from einops import rearrange
from models.modules.dccm import DiffContextAwareCompressionModule


class ResolutionAlignTransformerLayerSequence(nn.Module):
    def __init__(self, transformerlayers=None, num_layers=None, compression_args=None, init_cfg=None):
        super().__init__()
        if isinstance(transformerlayers, dict):
            transformerlayers = [
                copy.deepcopy(transformerlayers) for _ in range(num_layers)
            ]
        else:
            assert isinstance(transformerlayers, list) and \
                   len(transformerlayers) == num_layers
        self.num_layers = num_layers
        self.layers = nn.ModuleList()
        for i in range(num_layers):
            self.layers.append(build_transformer_layer(transformerlayers[i]))
        self.embed_dims = self.layers[0].embed_dims
        self.pre_norm = self.layers[0].pre_norm

        compression_args = copy.deepcopy(compression_args)
        comp_module = compression_args['comp_module']
        comp_insert_layer = compression_args['comp_insert_layer']
        comp_k = compression_args['comp_k']
        comp_strategy = compression_args['comp_strategy']
        embed_dim = compression_args['embed_dim']
        embedding_temporal_size = compression_args['temporal_size']
        num_samples=compression_args['num_samples']
        sigma = compression_args['sigma']

        self.num_frames = embedding_temporal_size
        self.comp_insert_layer = comp_insert_layer
        self.sigma = sigma


        if comp_insert_layer > 0:
            self.dccm = DiffContextAwareCompressionModule(
                                                        in_channels=int(embed_dim),
                                                        T=embedding_temporal_size,
                                                        strategy=comp_strategy,
                                                        comp_module=comp_module,
                                                        k=comp_k,
                                                        sigma=sigma,
                                                        n_sample=num_samples)
        else:
            self.dccm = None

    def forward(self, query, key=None, value=None,):
        # query: (bs, 1+T*P, c)
 
        B, N, c = query.shape 
        T = self.num_frames
        P = (N - 1) // T
        H = W = int( np.sqrt(P))

        multi_resolution = False
        sort_index =None
        for i, layer in enumerate(self.layers):
           
            if self.dccm is not None and i == self.comp_insert_layer:
                cls_token, ts_tokens = query[:, 0].unsqueeze(1), query[:, 1:]
                ts_tokens = rearrange(ts_tokens, "b (h w t) c -> b c t h w", b=B, t=T, c=c, h=H, w=W)
                # split token to different resolutions
                # ts_tokens_high: (B, H1, W1, T1, C), ts_tokens_low:(B, H2, W2, T2, C)
                ts_tokens_high, ts_tokens_low, sort_index = self.dccm(ts_tokens, cls_token.squeeze(1),return_index=True)
                multi_resolution = True
                query = (cls_token, ts_tokens_high, ts_tokens_low)

            query = layer(query)
            #print(f"layers[{i}]: {query[0].sum()}")

        if multi_resolution:
            x_cls, x_high, x_low = query
            x_high = rearrange(x_high, "b h w t c -> b (h w t) c")
            x_low = rearrange(x_low, "b h w t c -> b (h w t) c")
            x = torch.cat([x_cls, x_high, x_low], dim=1)
            query = x
        return query
    
    def get_rank_index(self, query):
        # query: (bs, 1+T*P, c)
        B, N, c = query.shape 
        T = self.num_frames
        P = (N - 1) // T
        H = W = int( np.sqrt(P))
        if self.dccm is None:
            return torch.arange(T).unsqueeze(0).repeat(B, 1).to(query.device)

        for i, layer in enumerate(self.layers):
            if i == self.comp_insert_layer:
                cls_token, ts_tokens = query[:, 0].unsqueeze(1), query[:, 1:]
                ts_tokens = rearrange(ts_tokens, "b (h w t) c -> b c t h w", b=B, t=T, c=c, h=H, w=W)
                # split token to different resolutions
                # ts_tokens_high: (B, H1, W1, T1, C), ts_tokens_low:(B, H2, W2, T2, C)
                ts_tokens_high, ts_tokens_low, score_index = self.dccm(ts_tokens, cls_token.squeeze(1), return_index=True)
                query = (cls_token, ts_tokens_high, ts_tokens_low)
                break
            query = layer(query)

        return score_index


models/modules/reso_align_layers.py
dun-research/DRCA
import torch
import torch.nn as nn
from mmcv.cnn import build_norm_layer
from mmcv.cnn.bricks.registry import ATTENTION, FEEDFORWARD_NETWORK
from einops import rearrange
from mmaction.models.common import DividedTemporalAttentionWithNorm

import torch.functional as F

from mmcv.cnn import (build_norm_layer)
from mmcv.runner.base_module import BaseModule
from mmcv.cnn import build_norm_layer
from mmcv.cnn.bricks.registry import ATTENTION, FEEDFORWARD_NETWORK
from mmcv.cnn.bricks.transformer import FFN, build_dropout
from mmcv.utils import digit_version


@ATTENTION.register_module()
class ResolutionAlignDividedTemporalAttentionWithNorm(DividedTemporalAttentionWithNorm):
    def __init__(self, embed_dims, num_heads, num_frames, attn_drop=0, proj_drop=0, dropout_layer=dict(type='DropPath', drop_prob=0.1), norm_cfg=dict(type='LN'), init_cfg=None, **kwargs):
        super().__init__(embed_dims, num_heads, num_frames, attn_drop, proj_drop, dropout_layer, norm_cfg, init_cfg, **kwargs)
       
        self.temporal_downsample = nn.AvgPool2d(kernel_size=2)
        self.temporal_upsample = torch.nn.UpsamplingBilinear2d(scale_factor=2)

    def forward(self, query, key=None, value=None, residual=None, **kwargs):
        assert residual is None, (
            'Always adding the shortcut in the forward function')

        if isinstance(query, torch.Tensor):
            # single resolution
            return super().forward(query, key, value, residual, **kwargs)
        else:
            # multiple resolutions
            (cls_token, x_high, x_low) = query[:3]
            if len(query) == 4: 
                sort_index = query[-1]
            else:
                sort_index = None
            B, h_high, w_high, t_high, C = x_high.shape
            B, h_low, w_low, t_low, _ = x_low.shape

            xt_high_identity = x_high
            xt_low_identity = x_low

            # downsample
            x_high = rearrange(x_high, "b h w t c -> (b t) c h w",)
            x_high_ds = self.temporal_downsample(x_high)
            #print(f"x_high_ds:{x_high_ds.size()}, x_high:{x_high.size()}, x_low:{x_low.size()}")
            new_h, new_w = x_high_ds.size(2), x_high_ds.size(3) 

            # concat x_high and x_low
            x_high_ds = rearrange(x_high_ds, "(b t) c h w -> b (h w) t c", t=t_high)
            x_low = rearrange(x_low, "b h w t c -> b (h w) t c", t=t_low)
            query_t = torch.cat([x_high_ds, x_low,], dim=2)
            query_t = rearrange(query_t, "b p t c -> b (p t) c")

            # compute residual
            res_temporal = self.compute_residual(query_t)

            # split x_high and x_low
            res_temporal = rearrange(res_temporal, "(b h w) t c -> b h w t c", 
                                                    b=B, h=h_low, w=w_low, t=t_high+t_low, c=C,)
            res_temporal_high, res_temporal_low = torch.split(res_temporal, [t_high, t_low], dim=3)

            # upsample res_temporal_high
            res_temporal_high = rearrange(res_temporal_high, "b h w t c -> (b t) c h w", h=new_h, w=new_w)
            res_temporal_high = self.temporal_upsample(res_temporal_high)
            res_temporal_high = rearrange(res_temporal_high, "(b t) c h w -> b h w t c", t=t_high)
            
            new_xt_high = xt_high_identity + res_temporal_high
            new_xt_low = xt_low_identity + res_temporal_low
            if sort_index is None:
                return cls_token, new_xt_high, new_xt_low           
            else:
                return cls_token, new_xt_high, new_xt_low,sort_index       

    def compute_residual(self, query_t, ):
        # query_t [batch_size, num_patches * num_frames, embed_dims]
        b, pt, m = query_t.size()
        p, t = pt // self.num_frames, self.num_frames

        # res_temporal [batch_size * num_patches, num_frames, embed_dims]
        query_t = self.norm(query_t.reshape(b * p, t, m)).permute(1, 0, 2)
        res_temporal = self.attn(query_t, query_t, query_t)[0].permute(1, 0, 2)
        res_temporal = self.dropout_layer(
            self.proj_drop(res_temporal.contiguous()))
        res_temporal = self.temporal_fc(res_temporal)
        return res_temporal



@ATTENTION.register_module()
class ResolutionAlignDividedSpatialAttentionWithNorm(BaseModule):

    def __init__(self,
                 embed_dims,
                 num_heads,
                 num_frames,
                 single_resolution=True,
                 attn_drop=0.,
                 proj_drop=0.,
                 dropout_layer=dict(type='DropPath', drop_prob=0.1),
                 norm_cfg=dict(type='LN'),
                 init_cfg=None,
                 **kwargs):
        super().__init__(init_cfg)
        self.embed_dims = embed_dims
        self.num_heads = num_heads
        self.num_frames = num_frames
        self.norm = build_norm_layer(norm_cfg, self.embed_dims)[1]
        
        if digit_version(torch.__version__) < digit_version('1.9.0'):
            kwargs.pop('batch_first', None)

        if not single_resolution:
            self.norm_low = build_norm_layer(norm_cfg, self.embed_dims)[1]
            self.attn_low = nn.MultiheadAttention(embed_dims, num_heads, attn_drop,
                                             **kwargs)
        self.attn = nn.MultiheadAttention(embed_dims, num_heads, attn_drop,
                                          **kwargs)
        self.proj_drop = nn.Dropout(proj_drop)
        self.dropout_layer = build_dropout(
            dropout_layer) if dropout_layer else nn.Identity()

    def forward_one(self, query, key=None, value=None, residual=None, **kwargs):
        assert residual is None, (
            'Always adding the shortcut in the forward function')

        identity = query
        init_cls_token = query[:, 0, :].unsqueeze(1)
        query_s = query[:, 1:, :]

        # query_s [batch_size, num_patches * num_frames, embed_dims]
        b, pt, m = query_s.size()
        p, t = pt // self.num_frames, self.num_frames

        # cls_token [batch_size * num_frames, 1, embed_dims]
        cls_token = init_cls_token.repeat(1, t, 1).reshape(b * t,
                                                           m).unsqueeze(1)

        # query_s [batch_size * num_frames, num_patches + 1, embed_dims]
        query_s = rearrange(query_s, 'b (p t) m -> (b t) p m', p=p, t=t)
        query_s = torch.cat((cls_token, query_s), 1)

        # res_spatial [batch_size * num_frames, num_patches + 1, embed_dims]
        query_s = self.norm(query_s).permute(1, 0, 2)
        res_spatial = self.attn(query_s, query_s, query_s)[0].permute(1, 0, 2)
        res_spatial = self.dropout_layer(
            self.proj_drop(res_spatial.contiguous()))

        # cls_token [batch_size, 1, embed_dims]
        cls_token = res_spatial[:, 0, :].reshape(b, t, m)
        cls_token = torch.mean(cls_token, 1, True)

        # res_spatial [batch_size * num_frames, num_patches + 1, embed_dims]
        res_spatial = rearrange(
            res_spatial[:, 1:, :], '(b t) p m -> b (p t) m', p=p, t=t)
        res_spatial = torch.cat((cls_token, res_spatial), 1)

        new_query = identity + res_spatial
        return new_query

    def forward(self, query, key=None, value=None, residual=None, **kwargs):
        assert residual is None, (
            'Always adding the shortcut in the forward function')

        if isinstance(query, torch.Tensor):
            # single resolution
            return self.forward_one(query, key, value, residual, **kwargs)
        else:
            # multiple resolutions
            (x_cls, xs_high, xs_low) = query

            x_high_shape = xs_high.shape
            x_low_shape = xs_low.shape
            b = xs_high.shape[0]

            xs_high = rearrange(xs_high, "b h w t c -> b (h w t) c")
            xs_low = rearrange(xs_low, "b h w t c -> b (h w t) c")

            res_spatial_high, cls_token_high =  self.compute_residual(xs_high, x_cls, t=x_high_shape[3],mode="high")
            res_spatial_low, cls_token_low = self.compute_residual(xs_low, x_cls,t=x_low_shape[3], mode="low")

            cls_token = torch.cat([cls_token_high, cls_token_low], dim=1)
            cls_token = torch.mean(cls_token, 1, True) # averaging for every frame

            res_spatial = torch.cat([cls_token, res_spatial_high, res_spatial_low], dim=1)
            identity = torch.cat([x_cls, xs_high, xs_low], dim=1)

            new_query = identity + res_spatial     
            new_cls, new_high, new_low = torch.split(new_query, [1, xs_high.shape[1], xs_low.shape[1]], dim=1)
            new_high = new_high.reshape(*x_high_shape)
            new_low = new_low.reshape(*x_low_shape)
            return new_cls, new_high, new_low

    def compute_residual(self, query_s, init_cls_token, t, mode):
        # query_s [batch_size, num_patches * num_frames, embed_dims]
        b, pt, m = query_s.size()
        p = pt // t

        # cls_token [batch_size * num_frames, 1, embed_dims]
        cls_token = init_cls_token.repeat(1, t, 1).reshape(b * t, m).unsqueeze(1)

        # query_s [batch_size * num_frames, num_patches + 1, embed_dims]
        query_s = rearrange(query_s, 'b (p t) m -> (b t) p m', p=p, t=t)
        query_s = torch.cat((cls_token, query_s), 1)

        if mode == "high":
            query_s = self.norm(query_s).permute(1, 0, 2)
            res_spatial = self.attn(query_s, query_s, query_s)[0].permute(1, 0, 2)
        else:
            query_s = self.norm_low(query_s).permute(1, 0, 2)
            res_spatial = self.attn_low(query_s, query_s, query_s)[0].permute(1, 0, 2)

        res_spatial = self.dropout_layer(
            self.proj_drop(res_spatial.contiguous()))
        
        cls_token = res_spatial[:, 0, :].reshape(b, t, m)
        res_spatial = rearrange(
            res_spatial[:, 1:, :], '(b t) p m -> b (p t) m', p=p, t=t)
        return res_spatial, cls_token


@FEEDFORWARD_NETWORK.register_module()
class ResolutionAlignFFN(FFN):
    def __init__(self, *args, norm_cfg=dict(type='LN'), **kwargs):
        super().__init__(*args, **kwargs)
        self.norm = build_norm_layer(norm_cfg, self.embed_dims)[1]

    def forward(self, x, residual=None):
        assert residual is None, ('Cannot apply pre-norm with FFNWithNorm')
        if isinstance(x, torch.Tensor):
            return super().forward(self.norm(x), x)
        else:
            (cls_token, x_high, x_low) = x[:3]
            if len(x) ==4:
                sort_index = x[-1]
            else:
                sort_index = None

            x_high_shape = x_high.shape
            x_low_shape = x_low.shape

            x_high = rearrange(x_high, "b h w t c -> b (h w t) c")
            x_low = rearrange(x_low, "b h w t c -> b (h w t) c")
            len_high, len_low = x_high.size(1), x_low.size(1)
            
            x = torch.cat([cls_token, x_high, x_low], dim=1)
            x = super().forward(self.norm(x), x)
            
            x_cls, x_high, x_low = torch.split(x, [1, len_high, len_low], dim=1)
            x_high = x_high.reshape(*x_high_shape)
            x_low = x_low.reshape(*x_low_shape)
            if sort_index == None:
                return x_cls, x_high, x_low
            else:
                return x_cls, x_high, x_low, sort_index



models/modules/reso_align_transformer.py
dun-research/DRCA

import torch.nn.functional as F
from torch import nn
from torch.nn import Dropout
from torch.nn import init

def weight_init(m):
    '''
    Usage:
        model = Model()
        model.apply(weight_init)
    '''

    if isinstance(m, (nn.Conv2d, nn.Conv3d)):
        init.xavier_normal_(m.weight.data)
        if m.bias is not None:
            init.normal_(m.bias.data)
    elif isinstance(m, (nn.Linear, nn.Embedding, nn.Conv1d)):
        #m.weight.data.normal_(mean=0.0, std=0.1)
        init.xavier_normal_(m.weight.data)
        if m.bias is not None:
            init.zeros_(m.bias.data)
    elif isinstance(m, nn.ConvTranspose2d):
        init.xavier_normal_(m.weight.data)
        if m.bias is not None:
            init.normal_(m.bias.data)
    elif isinstance(m, nn.BatchNorm2d):
        init.normal_(m.weight.data, mean=1, std=0.02)
        init.constant_(m.bias.data, 0)
    elif isinstance(m, nn.LSTM):
        for param in m.parameters():
            if len(param.shape) >= 2:
                init.orthogonal_(param.data)
            else:
                init.normal_(param.data)
    elif isinstance(m, nn.LayerNorm):
        m.bias.data.zero_()
        m.weight.data.fill_(1.0)
models/modules/weight_init.py
dun-research/DRCA

import torch
from torch import  nn

import numpy as np
from mmaction.models import RECOGNIZERS
from mmaction.models.recognizers import BaseRecognizer, Recognizer3D
from mmaction.models.builder import build_head
import torch.distributed as dist

def get_rank():
    if not dist.is_available():
        return 0
    if not dist.is_initialized():
        return 0
    return dist.get_rank()

def get_world_size():
    if not dist.is_available():
        return 0
    if not dist.is_initialized():
        return 0
    return dist.get_world_size()


@RECOGNIZERS.register_module()
class SimilarityRecognizer3D(BaseRecognizer):
    def __init__(self, backbone, sim_head=None,max_batch_size=6, cls_head=None, neck=None,   train_cfg=None, test_cfg=None):
        super().__init__(backbone, cls_head, neck, train_cfg, test_cfg)
        self.sim_head = build_head(sim_head) if sim_head is not None else nn.Identity()
        self.max_batch_size = max_batch_size
        self.fp16_enabled = False
        self.cos_similar = torch.nn.CosineSimilarity()
        self.relu = torch.nn.ReLU()

    def extract_feat(self, imgs):
        max_bs = self.max_batch_size
        n_imgs = imgs.size(0)
        st = 0
        all_feats = []

        if n_imgs < max_bs:
            feats = self.backbone(imgs)
            feats = self.sim_head(feats)
            return feats

        while st < n_imgs:
            cur_imgs = imgs[st : st+max_bs]
            feats = self.backbone(cur_imgs)
            feats = self.sim_head(feats)
            all_feats.append(feats)
            st += max_bs
        all_feats = torch.cat(all_feats, dim=0)
        return all_feats

    def forward_train(self, imgs, labels, video_id=None, frame_inds=None,**kwargs):
        """Defines the computation performed at every call when training."""
        assert video_id is not None
        c = imgs.size(1)
        #print("imgs :{}".format(imgs.size()))
        assert c  == 2 or c == 3
        bs = imgs.size(0)
        
        if video_id is None:
            sim_labels = torch.arange(0, bs).to(imgs.device)  + (get_rank() + 1) *  get_world_size()
        else:
            sim_labels = video_id
        sim_labels = sim_labels.reshape(-1, 1).repeat(1, c).reshape(-1,)
        
        imgs = imgs.reshape((-1, ) + imgs.shape[2:])
        
        losses = dict()
        x = self.extract_feat(imgs)
    
        embeddings = self.sim_head(x)
        loss_sim = self.sim_head.loss(embeddings, sim_labels, **kwargs)
        losses.update(loss_sim)
        return losses

    def _do_test(self, imgs):
        """Defines the computation performed at every call when evaluation,
        testing and gradcam."""
        batches = imgs.shape[0]
        num_segs = imgs.shape[1]
        imgs = imgs.reshape((-1, ) + imgs.shape[2:])

        if self.max_testing_views is not None:
            total_views = imgs.shape[0]
            assert num_segs == total_views, (
                'max_testing_views is only compatible '
                'with batch_size == 1')
            view_ptr = 0
            feats = []
            while view_ptr < total_views:
                batch_imgs = imgs[view_ptr:view_ptr + self.max_testing_views]
                emb = self.extract_feat(batch_imgs)
                feats.append(emb)
                view_ptr += self.max_testing_views
            # should consider the case that feat is a tuple
            if isinstance(feats[0], tuple):
                len_tuple = len(feats[0])
                feat = [
                    torch.cat([x[i] for x in feats]) for i in range(len_tuple)
                ]
                embs = tuple(feat)
            else:
                embs = torch.cat(feats)

        else:
            embs = self.extract_feat(imgs)

        feat = embs
        feat_dim = len(feat[0].size()) if isinstance(feat, tuple) else len(
            feat.size())
        if feat_dim != 2:  # 3D-CNN architecture
            feat = feat.flatten(2).mean(2)
            if num_segs == 1:
                feat = feat.reshape((batches, -1))
            else:
                feat = feat.reshape((batches, num_segs, -1))
        else:
            feat = feat.reshape((batches, num_segs, -1))
            if num_segs == 1:
                feat = feat[:, 0, :]
        return feat

    def forward_test(self, imgs):
        """Defines the computation performed at every call when evaluation and
        testing."""
        return self._do_test(imgs).detach().cpu().numpy()

    def forward_gradcam(self, imgs):
        return self.forward_test(imgs)
    
    def forward_dummy(self, imgs, softmax=False):
        """Used for computing network FLOPs.

        See ``tools/analysis/get_flops.py``.

        Args:
            imgs (torch.Tensor): Input images.

        Returns:
            Tensor: Class score.
        """
        # assert self.with_cls_head
        imgs = imgs.reshape((-1, ) + imgs.shape[2:])
        feats = self.extract_feat(imgs)
        return (feats, )
models/recognizer/similarity_recognizer.py
dun-research/DRCA
# Copyright (c) OpenMMLab. All rights reserved.
import argparse
from register import register_custom_modules
from mmaction.models import build_recognizer
from mmcv import Config

try:
    from mmengine.analysis import get_model_complexity_info
except ImportError:
    raise ImportError('Please upgrade mmcv to >0.6.2')
import torch


def parse_args():
    parser = argparse.ArgumentParser(description='Get model flops and params')
    parser.add_argument('config', help='config file path')
    parser.add_argument("-q", "--quiet", action="store_true", help="silent mode")
    parser.add_argument(
        '--shape',
        type=int,
        nargs='+',
        default=[224, 224],
        help='input image size')
    args = parser.parse_args()
    return args


def main():

    args = parse_args()

    if len(args.shape) == 1:
        input_shape = (1, 3, args.shape[0], args.shape[0])
    elif len(args.shape) == 2:
        input_shape = (1, 3) + tuple(args.shape)
    elif len(args.shape) == 4:
        # n, c, h, w = args.shape for 2D recognizer
        input_shape = tuple(args.shape)
    elif len(args.shape) == 5:
        # n, c, t, h, w = args.shape for 3D recognizer or
        # n, m, t, v, c = args.shape for GCN-based recognizer
        input_shape = tuple(args.shape)
    else:
        raise ValueError('invalid input shape')

    cfg = Config.fromfile(args.config)
    model = build_recognizer(
        cfg.model,
        train_cfg=cfg.get('train_cfg'),
        test_cfg=cfg.get('test_cfg'))

    model = model.cuda()
    model.eval()

    if hasattr(model, 'forward_dummy'):
        model.forward = model.forward_dummy
    else:
        raise NotImplementedError(
            'FLOPs counter is currently not currently supported with {}'.
            format(model.__class__.__name__))

    if args.quiet:
        show_arch, show_table = False, False
    else:
        show_arch, show_table = True, True 
    analysis_results = get_model_complexity_info(model, input_shape=input_shape, inputs=None,
                                                 show_arch=show_arch, show_table=show_table, )

    flops = analysis_results['flops_str']
    params = analysis_results['params_str']
    table = analysis_results['out_table']
    print(table)
    split_line = '=' * 30
    print(f'\n{split_line}\nInput shape: {input_shape}\n'
          f'Flops: {flops}\nParams: {params}\n{split_line}')
    print('!!!Please be cautious if you use the results in papers. '
          'You may need to check if all ops are supported and verify that the '
          'flops computation is correct.')


if __name__ == '__main__':
    main()

tools/get_flops.py
dun-research/DRCA
import sys
sys.path.append(".")

# Important! Register custom modules here
import models
import datasets
import warnings
warnings.filterwarnings("ignore")

def register_custom_modules():
    pass    

tools/register.py
dun-research/DRCA
# Copyright (c) OpenMMLab. All rights reserved.
import argparse


from numpy import isin


import os
import os.path as osp
import warnings
from time import time
import mmcv
import torch
from mmcv import Config, DictAction
from mmcv.cnn import fuse_conv_bn
from mmcv.fileio.io import file_handlers
from mmcv.runner import get_dist_info, init_dist, load_checkpoint
from mmcv.runner.fp16_utils import wrap_fp16_model

from mmaction.datasets import build_dataloader, build_dataset
from mmaction.models import build_model
from mmaction.utils import (build_ddp, build_dp, default_device,
                            register_module_hooks, setup_multi_processes)
import numpy as np
from register import register_custom_modules

# TODO import test functions from mmcv and delete them from mmaction2
try:
    from mmcv.engine import multi_gpu_test, single_gpu_test
except (ImportError, ModuleNotFoundError):
    warnings.warn(
        'DeprecationWarning: single_gpu_test, multi_gpu_test, '
        'collect_results_cpu, collect_results_gpu from mmaction2 will be '
        'deprecated. Please install mmcv through master branch.')
    from mmaction.apis import multi_gpu_test, single_gpu_test

import logging
from mmaction.utils import get_root_logger
def parse_args():
    parser = argparse.ArgumentParser(
        description='MMAction2 test (and eval) a model')
    parser.add_argument('config', help='test config file path')
    parser.add_argument('checkpoint', help='checkpoint file')
    # parser.add_argument(
    #     '--anno_file',
    #     default=None,
    #     help='output result file in pkl/yaml/json format')
    parser.add_argument(
        '--out',
        default=None,
        help='output result file in pkl/yaml/json format')
    parser.add_argument(
        '--fuse-conv-bn',
        action='store_true',
        help='Whether to fuse conv and bn, this will slightly increase'
        'the inference speed')
    parser.add_argument(
        '--eval',
        type=str,
        nargs='+',
        help='evaluation metrics, which depends on the dataset, e.g.,'
        ' "top_k_accuracy", "mean_class_accuracy" for video dataset')
    parser.add_argument(
        '--gpu-collect',
        action='store_true',
        help='whether to use gpu to collect results')
    parser.add_argument(
        '--tmpdir',
        help='tmp directory used for collecting results from multiple '
        'workers, available when gpu-collect is not specified')
    parser.add_argument(
        '--options',
        nargs='+',
        action=DictAction,
        default={},
        help='custom options for evaluation, the key-value pair in xxx=yyy '
        'format will be kwargs for dataset.evaluate() function (deprecate), '
        'change to --eval-options instead.')
    parser.add_argument(
        '--eval-options',
        nargs='+',
        action=DictAction,
        default={},
        help='custom options for evaluation, the key-value pair in xxx=yyy '
        'format will be kwargs for dataset.evaluate() function')
    parser.add_argument(
        '--cfg-options',
        nargs='+',
        action=DictAction,
        default={},
        help='override some settings in the used config, the key-value pair '
        'in xxx=yyy format will be merged into config file. For example, '
        "'--cfg-options model.backbone.depth=18 model.backbone.with_cp=True'")
    parser.add_argument(
        '--average-clips',
        choices=['score', 'prob', None],
        default=None,
        help='average type when averaging test clips')
    parser.add_argument(
        '-la', '--launcher',
        choices=['none', 'pytorch', 'slurm', 'mpi'],
        default='pytorch',
        help='job launcher')
    parser.add_argument('--local_rank', type=int, default=0)
    parser.add_argument(
        '--onnx',
        action='store_true',
        help='Whether to test with onnx model or not')
    parser.add_argument(
        '--fuse_conv_bn',
        action='store_true',
        help='Whether to test with TensorRT engine or not')
    parser.add_argument(
        '--tensorrt',
        action='store_true',
        help='Whether to test with TensorRT engine or not')
    args = parser.parse_args()
    if 'LOCAL_RANK' not in os.environ:
        os.environ['LOCAL_RANK'] = str(args.local_rank)

    if args.options and args.eval_options:
        raise ValueError(
            '--options and --eval-options cannot be both '
            'specified, --options is deprecated in favor of --eval-options')
    if args.options:
        warnings.warn('--options is deprecated in favor of --eval-options')
        args.eval_options = args.options
    return args


def turn_off_pretrained(cfg):
    # recursively find all pretrained in the model config,
    # and set them None to avoid redundant pretrain steps for testing
    if 'pretrained' in cfg:
        cfg.pretrained = None

    # recursively turn off pretrained value
    for sub_cfg in cfg.values():
        if isinstance(sub_cfg, dict):
            turn_off_pretrained(sub_cfg)


def inference_pytorch(args, cfg, distributed, data_loader,logger):
    """Get predictions by pytorch models."""
    if args.average_clips is not None:
        # You can set average_clips during testing, it will override the
        # original setting
        if cfg.model.get('test_cfg') is None and cfg.get('test_cfg') is None:
            cfg.model.setdefault('test_cfg',
                                 dict(average_clips=args.average_clips))
        else:
            if cfg.model.get('test_cfg') is not None:
                cfg.model.test_cfg.average_clips = args.average_clips
            else:
                cfg.test_cfg.average_clips = args.average_clips

    # remove redundant pretrain steps for testing
    # turn_off_pretrained(cfg.model)

    # build the model and load checkpoint
    model = build_model(
        cfg.model, train_cfg=None, test_cfg=cfg.get('test_cfg'))

    if len(cfg.module_hooks) > 0:
        register_module_hooks(model, cfg.module_hooks)

    fp16_cfg = cfg.get('fp16', None)
    if fp16_cfg is not None:
        wrap_fp16_model(model)



    load_checkpoint(model, args.checkpoint, map_location='cpu',strict=True )
    logger.info("load checkpoint from {}".format(args.checkpoint))
    

    if args.fuse_conv_bn:
        model = fuse_conv_bn(model)
    logger.info("fuse_conv_bn : {}".format(fuse_conv_bn))
    logger.info("distributed: {}".format(distributed))
    start = time()
    if not distributed:
        model = build_dp(
            model, default_device, default_args=dict(device_ids=cfg.get("gpu_ids", [0,])))
        outputs = single_gpu_test(model, data_loader)
        
    else:
        model = build_ddp(
            model,
            default_device,
            default_args=dict(
                device_ids=[int(os.environ['LOCAL_RANK'])],
                broadcast_buffers=False))
        outputs = multi_gpu_test(model, data_loader, args.tmpdir,
                                 args.gpu_collect)
    end = time()
    logger.info("Inference dataset use {}s".format(int(end - start)))
    return outputs


def setup_env(args):
    if args.tensorrt and args.onnx:
        raise ValueError(
            'Cannot set onnx mode and tensorrt mode at the same time.')

    cfg = Config.fromfile(args.config)

    cfg.merge_from_dict(args.cfg_options)

    # set multi-process settings
    setup_multi_processes(cfg)
    
    # load anno file
    test_config = cfg.data.test
    # if osp.isfile(args.anno_file):
    #     cfg = Config._merge_a_into_b(
    #         dict(ann_file_test=args.anno_file), cfg
    #     )
        
    #     test_config = Config._merge_a_into_b(
    #         dict(ann_file=args.anno_file), test_config 
    #     )
    

    # Load output_config from cfg
    output_config = cfg.get('output_config', {})
    if args.out:
        # Overwrite output_config from args.out
        out_file = osp.join(args.out, "prediction.pkl")
        output_config = Config._merge_a_into_b(
            dict(out=out_file), output_config)
        
        os.makedirs(args.out, exist_ok=True)
        out_log_file = osp.join(args.out, "log.txt")
        logger = get_root_logger(log_file=out_log_file, )
    # Load eval_config from cfg
    eval_config = cfg.get('eval_config', {})
    if args.eval:
        # Overwrite eval_config from args.eval
        eval_config = Config._merge_a_into_b(
            dict(metrics=args.eval), eval_config)
    if args.eval_options:
        # Add options from args.eval_options
        eval_config = Config._merge_a_into_b(args.eval_options, eval_config)

    assert output_config or eval_config, \
        ('Please specify at least one operation (save or eval the '
         'results) with the argument "--out" or "--eval"')
    dataset_type = cfg.data.test.type
    if output_config.get('out', None):
        if 'output_format' in output_config:
            # ugly workround to make recognition and localization the same
            warnings.warn(
                'Skip checking `output_format` in localization task.')
        else:
            out = output_config['out']
            # make sure the dirname of the output path exists
            mmcv.mkdir_or_exist(osp.dirname(out))
            _, suffix = osp.splitext(out)
            if dataset_type == 'AVADataset':
                assert suffix[1:] == 'csv', ('For AVADataset, the format of '
                                             'the output file should be csv')
            else:
                assert suffix[1:] in file_handlers, (
                    'The format of the output '
                    'file should be json, pickle or yaml')

    # set cudnn benchmark
    if cfg.get('cudnn_benchmark', False):
        torch.backends.cudnn.benchmark = True
    cfg.data.test.test_mode = True

    # init distributed env first, since logger depends on the dist info.
    if args.launcher == 'none':
        distributed = False
    else:
        distributed = True
        init_dist(args.launcher, **cfg.dist_params)
    # The flag is used to register module's hooks
    cfg.setdefault('module_hooks', [])
    cfg.distributed = distributed

    logger.info("config: \n{}".format(cfg))
    logger.info("test_config: \n{}".format(test_config))
    logger.info("args: \n{}".format(args))
    return cfg, output_config, eval_config, test_config, logger

def main():
    args = parse_args()
    cfg, output_config, eval_config, test_config, logger  = setup_env(args)
    
    # build the dataloader
    dataset = build_dataset(test_config, dict(test_mode=True))

    logger.info("videos_per_gpu: {}".format(cfg.data.get('videos_per_gpu', 1),))
    logger.info("workers_per_gpu: {}".format(cfg.data.get('workers_per_gpu', 1),))

    dataloader_setting = dict(
        videos_per_gpu=cfg.data.get('videos_per_gpu', 1),
        workers_per_gpu=cfg.data.get('workers_per_gpu', 1),
        dist=cfg.distributed,
        shuffle=False)
    dataloader_setting = dict(dataloader_setting,
                              **cfg.data.get('test_dataloader', {}))
    print(dataloader_setting)
    data_loader = build_dataloader(dataset, **dataloader_setting)

    logger.info("data pipeline: {}".format(dataset.pipeline.transforms))
    logger.info("data loader size: {}".format(len(data_loader)))
        
    outputs = inference_pytorch(args, cfg, cfg.distributed, data_loader, logger)

    rank, _ = get_dist_info()
    if rank == 0:
        eval_res = dataset.evaluate(outputs, **eval_config)
        for name, val in eval_res.items():
            logger.info(f'{name}: {val:.04f}')


if __name__ == '__main__':
    main()


tools/test.py
dun-research/DRCA
# Copyright (c) OpenMMLab. All rights reserved.
import argparse
import copy
import os
import os.path as osp
import time
import warnings

import mmcv
import torch
import torch.distributed as dist
from mmcv import Config, DictAction
from mmcv.runner import get_dist_info, init_dist, set_random_seed
from mmcv.utils import get_git_hash

from mmaction import __version__

from mmaction.apis import init_random_seed, train_model
from mmaction.datasets import build_dataset
from mmaction.models import build_model
from mmaction.utils import (collect_env, get_root_logger,
                            register_module_hooks, setup_multi_processes)
from mmcv.runner.fp16_utils import wrap_fp16_model
from register import register_custom_modules

def parse_args():
    parser = argparse.ArgumentParser(description='Train a recognizer')
    parser.add_argument('config', help='train config file path')
    parser.add_argument('--work-dir', help='the dir to save logs and models')
    parser.add_argument(
        '--resume-from', help='the checkpoint file to resume from')
    parser.add_argument(
        '--checkpoint',  help='the checkpoint file to resume from')
    parser.add_argument(
        '--validate',
        action='store_true',
        help='whether to evaluate the checkpoint during training')
    parser.add_argument(
        '--test-last',
        action='store_true',
        help='whether to test the checkpoint after training')
    parser.add_argument(
        '--test-best',
        action='store_true',
        help=('whether to test the best checkpoint (if applicable) after '
              'training'))
    group_gpus = parser.add_mutually_exclusive_group()
    group_gpus.add_argument(
        '--gpus',
        type=int,
        help='number of gpus to use '
        '(only applicable to non-distributed training)')
    group_gpus.add_argument(
        '--gpu-ids',
        type=int,
        nargs='+',
        help='ids of gpus to use '
        '(only applicable to non-distributed training)')
    parser.add_argument('--seed', type=int, default=42, help='random seed')
    parser.add_argument(
        '--diff-seed',
        action='store_true',
        help='Whether or not set different seeds for different ranks')
    parser.add_argument(
        '--deterministic',
        action='store_true',
        help='whether to set deterministic options for CUDNN backend.')
    parser.add_argument(
        '--cfg-options',
        nargs='+',
        action=DictAction,
        default={},
        help='override some settings in the used config, the key-value pair '
        'in xxx=yyy format will be merged into config file. For example, '
        "'--cfg-options model.backbone.depth=18 model.backbone.with_cp=True'")
    parser.add_argument(
        '-la', '--launcher',
        choices=['none', 'pytorch', 'slurm', 'mpi'],
        default='pytorch',
        help='job launcher')
    parser.add_argument('--local_rank', type=int, default=0)
    args = parser.parse_args()
    if 'LOCAL_RANK' not in os.environ:
        os.environ['LOCAL_RANK'] = str(args.local_rank)

    return args


def main():
    args = parse_args()
    args.validate = True

    cfg = Config.fromfile(args.config)

    cfg.merge_from_dict(args.cfg_options)

    # set multi-process settings
    setup_multi_processes(cfg)

    # set cudnn_benchmark
    if cfg.get('cudnn_benchmark', False):
        torch.backends.cudnn.benchmark = True



    # work_dir is determined in this priority:
    # CLI > config file > default (base filename)
    if args.work_dir is not None:
        # update configs according to CLI args if args.work_dir is not None
        cfg.work_dir = args.work_dir
    elif cfg.get('work_dir', None) is None:
        # use config filename as default work_dir if cfg.work_dir is None
        cfg.work_dir = osp.join('./work_dirs',
                                osp.splitext(osp.basename(args.config))[0])
    if args.resume_from is not None:
        cfg.resume_from = args.resume_from

    if args.gpu_ids is not None or args.gpus is not None:
        warnings.warn(
            'The Args `gpu_ids` and `gpus` are only used in non-distributed '
            'mode and we highly encourage you to use distributed mode, i.e., '
            'launch training with dist_train.sh. The two args will be '
            'deperacted.')
        if args.gpu_ids is not None:
            warnings.warn(
                'Non-distributed training can only use 1 gpu now. We will '
                'use the 1st one in gpu_ids. ')
            cfg.gpu_ids = [args.gpu_ids[0]]
        elif args.gpus is not None:
            warnings.warn('Non-distributed training can only use 1 gpu now. ')
            cfg.gpu_ids = range(1)

    # init distributed env first, since logger depends on the dist info.
    if args.launcher == 'none':
        distributed = False
        num_gpus = torch.cuda.device_count()
    else:
        distributed = True
        init_dist(args.launcher, **cfg.dist_params)
        _, world_size = get_dist_info()
        cfg.gpu_ids = range(world_size)
        num_gpus = world_size

    # set LR
    batch_size_per_gpu = cfg.get("batch_size_per_gpu", cfg.data.videos_per_gpu)
    base_lr = cfg.get("base_lr", -1)
    if base_lr >= 0:
        batch_size = batch_size_per_gpu * num_gpus
        lr = batch_size * base_lr / 256
        cfg.optimizer.lr = lr
        cfg.data.videos_per_gpu = batch_size_per_gpu
    else:
        lr = cfg.optimizer.lr
    

    # The flag is used to determine whether it is omnisource training
    cfg.setdefault('omnisource', False)

    # The flag is used to register module's hooks
    cfg.setdefault('module_hooks', [])

    # create work_dir
    mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir))
    # dump config
    cfg.dump(osp.join(cfg.work_dir, osp.basename(args.config)))
    # init logger before other steps
    timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
    log_file = osp.join(cfg.work_dir, f'{timestamp}.log')
    logger = get_root_logger(log_file=log_file, log_level=cfg.log_level)


    # init the meta dict to record some important information such as
    # environment info and seed, which will be logged
    meta = dict()
    # log env info
    env_info_dict = collect_env()
    env_info = '\n'.join([f'{k}: {v}' for k, v in env_info_dict.items()])
    dash_line = '-' * 60 + '\n'
    logger.info('Environment info:\n' + dash_line + env_info + '\n' +
                dash_line)
    meta['env_info'] = env_info

    # log some basic info
    logger.info(f'Distributed training: {distributed}')
    logger.info(f'Config: {cfg.pretty_text}')

    # set random seeds
    seed = init_random_seed(args.seed, distributed=distributed)
    seed = seed + dist.get_rank() if args.diff_seed else seed
    logger.info(f'Set random seed to {seed}, '
                f'deterministic: {args.deterministic}')
    set_random_seed(seed, deterministic=args.deterministic)

    cfg.seed = seed
    
    meta['seed'] = seed
    meta['config_name'] = osp.basename(args.config)
    meta['work_dir'] = osp.basename(cfg.work_dir.rstrip('/\\'))

    model = build_model(
        cfg.model,
        train_cfg=cfg.get('train_cfg'),
        test_cfg=cfg.get('test_cfg'))

    if len(cfg.module_hooks) > 0:
        register_module_hooks(model, cfg.module_hooks)

    fp16_cfg = cfg.get('fp16', None)
    print(f"fp16_cfg: {fp16_cfg}")
    if fp16_cfg is not None:
        wrap_fp16_model(model)
    if cfg.omnisource:
        # If omnisource flag is set, cfg.data.train should be a list
        assert isinstance(cfg.data.train, list)
        datasets = [build_dataset(dataset) for dataset in cfg.data.train]
    else:
        datasets = [build_dataset(cfg.data.train)]

    if len(cfg.workflow) == 2:
        # For simplicity, omnisource is not compatible with val workflow,
        # we recommend you to use `--validate`
        assert not cfg.omnisource
        if args.validate:
            warnings.warn('val workflow is duplicated with `--validate`, '
                          'it is recommended to use `--validate`. see '
                          'https://github.com/open-mmlab/mmaction2/pull/123')
        val_dataset = copy.deepcopy(cfg.data.val)
        datasets.append(build_dataset(val_dataset))
    if cfg.checkpoint_config is not None:
        # save mmaction version, config file content and class names in
        # checkpoints as meta data
        cfg.checkpoint_config.meta = dict(
            mmaction_version=__version__ + get_git_hash(digits=7),
            config=cfg.pretty_text)
    logger.info(f"lr={lr}, batch_size_per_gpu={batch_size_per_gpu}, num_gpus={num_gpus}")
    test_option = dict(test_last=args.test_last, test_best=args.test_best)
    torch.autograd.set_detect_anomaly(True)

    
    train_model(
        model,
        datasets,
        cfg,
        distributed=distributed,
        validate=args.validate,
        test=test_option,
        timestamp=timestamp,
        meta=meta)


if __name__ == '__main__':
    main()

tools/train.py
dun-research/DRCA