Initial commit

app.py

@@ -0,0 +1,81 @@
+import cv2
+import glob
+import gradio as gr
+import mediapy
+import nibabel
+import numpy as np
+import shutil
+import torch
+import torch.nn.functional as F
+
+from omegaconf import OmegaConf
+from skp import builder
+
+
+def window(x, WL=400, WW=2500):
+    lower, upper = WL - WW // 2, WL + WW // 2
+    x = np.clip(x, lower, upper)
+    x = x - lower
+    x = x / (upper - lower)
+    return (x * 255).astype("uint8")
+
+
+def rescale(x):
+    x = x / 255.
+    x = x - 0.5
+    x = x * 2.0
+    return x
+
+
+def generate_segmentation_video(study):
+    img = nibabel.load(study).get_fdata()[:, ::-1, ::-1].transpose(2, 1, 0)
+    img = window(img)
+
+    X = torch.from_numpy(img).float().unsqueeze(0).unsqueeze(0)
+    X = F.interpolate(X, size=(192, 192, 192), mode="nearest")
+    X = rescale(X)
+    with torch.no_grad():
+        seg_output = seg_model(X)
+
+    seg_output = torch.sigmoid(seg_output)
+    p_spine = seg_output[:, :7].sum(1)
+    seg_output = torch.argmax(seg_output, dim=1) + 1
+    seg_output[p_spine < 0.5] = 0
+    seg_output = F.interpolate(seg_output.unsqueeze(0).float(), size=img.shape, mode="nearest")
+    seg_output = seg_output.squeeze(0).squeeze(0).numpy()
+    seg_output = (seg_output * 255 / 7).astype("uint8")
+    seg_output = np.stack([cv2.applyColorMap(_, cv2.COLORMAP_JET) for _ in seg_output])
+
+    frames = []
+    skip = 8
+    for idx in range(0, img.shape[2], skip):
+        i = img[:, :, idx]
+        o = seg_output[:, :, idx]
+        i = cv2.cvtColor(i, cv2.COLOR_GRAY2RGB)
+        frame = np.concatenate((i, o), 1)
+        frames.append(frame)
+    mediapy.write_video("video.mp4", frames, fps=30)
+    return "video.mp4"
+
+
+ffmpeg_path = shutil.which('ffmpeg')
+mediapy.set_ffmpeg(ffmpeg_path)
+
+config = OmegaConf.load("configs/pseudoseg000.yaml")
+config.model.load_pretrained = "seg.ckpt"
+seg_model = builder.build_model(config).eval()
+examples = glob.glob("examples/*.nii.gz")
+
+with gr.Blocks(theme="dark-peach") as demo:
+    select_study = gr.Dropdown(choices=sorted(examples), type="value", label="Select a study")
+    button_predict = gr.Button("Predict")
+    video_output = gr.Video()
+    button_predict.click(fn=generate_segmentation_video,
+                         inputs=select_study,
+                         outputs=video_output)
+
+
+if __name__ == "__main__":
+    demo.launch(debug=True, share=True)
+
+

configs/chunk000.yaml

@@ -0,0 +1,89 @@
+experiment:
+  seed: 88
+  save_dir: ../experiments/
+
+
+data:
+  annotations: ../data/train_vertebra_chunks_kfold.csv
+  data_dir: ../data/train-numpy-vertebra-chunks
+  input: filename
+  target: fracture
+  outer_fold: 0
+  dataset:
+    name: NumpyChunkDataset
+    params:
+      flip: true
+      invert: false
+      channels: grayscale
+      z_lt: resample_resample
+      z_gt: resample_resample
+      num_images: 64
+
+
+transform:
+  resize:
+    name: resize_ignore_3d
+    params:
+      imsize: [64, 288, 288]
+  augment:
+    null
+  crop:
+    null
+  preprocess:
+    name: Preprocessor
+    params:
+      image_range: [0, 255]
+      input_range: [0, 1]
+      mean: [0.5]
+      sdev: [0.5]
+
+
+task:
+  name: ClassificationTask
+  params:
+
+
+model:
+  name: Net3D
+  params:
+    backbone: x3d_l
+    backbone_params:
+      z_strides: [1, 1, 1, 1, 1]
+    pretrained: true
+    num_classes: 1
+    dropout: 0.2
+    pool: avg
+    in_channels: 1
+    multisample_dropout: true
+
+
+loss:
+  name: BCEWithLogitsLoss
+  params:
+
+
+optimizer:
+  name: AdamW
+  params:
+    lr: 3.0e-4
+    weight_decay: 5.0e-4
+
+
+scheduler:
+  name: CosineAnnealingLR
+  params:
+    final_lr: 0.0
+
+
+train:
+  batch_size: 4
+  num_epochs: 10
+
+
+evaluate:
+  metrics: [AUROC]
+  monitor: auc_mean
+  mode: max
+
+
+

configs/chunkseq003.yaml

@@ -0,0 +1,67 @@
+experiment:
+  seed: 88
+  save_dir: ../experiments/
+
+
+data:
+  annotations: ../data/train_chunk_features_kfold.csv
+  data_dir: ../data/train-chunk000-features/foldx
+  input: filename
+  target: [C1, C2, C3, C4, C5, C6, C7, patient_overall]
+  outer_fold: 0
+  dataset:
+    name: FeatureDataset
+    params:
+      seq_len: 7
+      reverse: false
+      normalize: false
+      exam_level_label: true
+
+
+task:
+  name: ClassificationTask
+  params:
+
+
+model:
+  name: DualTransformer
+  params:
+    num_classes: 1
+    embedding_dim: 432
+    hidden_dim: 864
+    n_layers: 3
+    n_heads: 16
+
+
+loss:
+  name: MultilabelWeightedBCE
+  params:
+    weights: [1, 1, 1, 1, 1, 1, 1, 7]
+    pos_weight: 2.0
+
+
+optimizer:
+  name: AdamW
+  params:
+    lr: 1.0e-5
+    weight_decay: 5.0e-4
+
+
+scheduler:
+  name: CosineAnnealingLR
+  params:
+    final_lr: 0
+
+
+train:
+  batch_size: 32
+  num_epochs: 25
+
+
+evaluate:
+  batch_size: 1
+  metrics: [CompetitionMetric, AUROC]
+  monitor: comp_metric
+  mode: min
+
+

configs/pseudoseg000.yaml

@@ -0,0 +1,110 @@
+experiment:
+  seed: 88
+  save_dir: ../experiments/
+
+
+data:
+  annotations: ../data/train_seg_whole_192_kfold_with_pseudo.csv
+  data_dir: ../data/
+  input: filename
+  target: label
+  outer_fold: 0
+  dataset:
+    name: NumpyChunkSegmentDataset
+    params:
+      segmentation_format: numpy
+      channels: grayscale
+      flip: true
+      transpose: true
+      invert: false
+      verbose: true
+      num_images: 192
+      z_lt: resample_resample
+      z_gt: resample_resample
+      one_hot_encode: true
+      num_classes: 8
+      add_foreground_channel: false
+
+
+transform:
+  resize:
+    name: resize_ignore_3d
+    params:
+      imsize: [192, 192, 192]
+  augment:
+    null
+  crop:
+    null
+  preprocess:
+    name: Preprocessor
+    params:
+      image_range: [0, 255]
+      input_range: [0, 1]
+      mean: [0.5]
+      sdev: [0.5]
+
+
+task:
+  name: SegmentationTask3D
+  params:
+    chunk_validation: true
+
+
+model:
+  name: NetSegment3D
+  params:
+    architecture: DeepLabV3Plus_3D
+    encoder_name: x3d_l
+    encoder_params:
+      pretrained: true
+      output_stride: 16
+      z_strides: [2, 2, 2, 2, 2]
+    decoder_params:
+      upsampling: 4
+    deep_supervision: true
+    num_classes: 8
+    in_channels: 1
+    dropout: 0.2
+
+
+loss:
+  name: SupervisorLoss
+  params:
+    segmentation_loss: DiceBCELoss
+    scale_factors: [0.25, 0.25]
+    loss_weights: [1.0, 0.25, 0.25]
+    loss_params:
+      dice_loss_params:
+        mode: multilabel
+        exponent: 2
+        smooth: 1.0
+      bce_loss_params:
+        smooth_factor: 0.01
+        pos_weight: 1.0
+      dice_loss_weight: 1.0
+      bce_loss_weight: 0.2
+
+
+optimizer:
+  name: AdamW
+  params:
+    lr: 3.0e-4
+    weight_decay: 5.0e-4
+
+
+scheduler:
+  name: CosineAnnealingLR
+  params:
+    final_lr: 0.0
+
+
+train:
+  batch_size: 4
+  num_epochs: 10
+
+
+evaluate:
+  batch_size: 1
+  metrics: [DSC]
+  monitor: dsc_ignore_mean
+  mode: max

examples/1.2.826.0.1.3680043.15773.nii.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a316d4cdb9534c662a209dea2b50fd57168398b1a658d14937ce285d3b792917
+size 65868417

packages.txt

@@ -0,0 +1 @@
+ffmpeg

requirements.txt

@@ -0,0 +1,7 @@
+omegaconf
+mediapy
+nibabel
+opencv-python
+timm
+torch
+transformers

seg.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aa6ee0036af98df68621b5cacbed9b4cd290eb1b59c6af7785a7e9c81ed74afa
+size 21569386

seq.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:653637b500e3ae5ffab8b07f34d36662396ab3eacec8024e5ecea952d7c2c07e
+size 18011334

skp/builder.py

@@ -0,0 +1,187 @@
+import numpy as np
+import torch
+
+from . import models 
+
+
+def get_name_and_params(base):
+    name = getattr(base, 'name')
+    params = getattr(base, 'params') or {}
+    return name, params
+
+
+def get_transform(base, transform, mode=None):
+    if not base: return None
+    transform = getattr(base, transform)
+    if not transform: return None
+    name, params = get_name_and_params(transform)
+    if mode:
+        params.update({'mode': mode})
+    return getattr(data.transforms, name)(**params)
+
+
+def build_transforms(cfg, mode):
+    # 1-Resize
+    resizer = get_transform(cfg.transform, 'resize')
+    # 2-(Optional) Data augmentation
+    augmenter = None
+    if mode == "train":
+        augmenter = get_transform(cfg.transform, 'augment')
+    # 3-(Optional) Crop
+    cropper = get_transform(cfg.transform, 'crop', mode=mode)
+    # 4-Preprocess
+    preprocessor = get_transform(cfg.transform, 'preprocess')
+    return {
+        'resize': resizer,
+        'augment': augmenter,
+        'crop': cropper,
+        'preprocess': preprocessor
+    }
+
+
+def build_dataset(cfg, data_info, mode):
+    dataset_class = getattr(data.datasets, cfg.data.dataset.name)
+    dataset_params = cfg.data.dataset.params
+    dataset_params.test_mode = mode != 'train'
+    dataset_params = dict(dataset_params)
+    if "FeatureDataset" not in cfg.data.dataset.name:
+        transforms = build_transforms(cfg, mode)
+        dataset_params.update(transforms)
+    dataset_params.update(data_info)
+    return dataset_class(**dataset_params)
+
+
+def build_dataloader(cfg, dataset, mode):
+
+    def worker_init_fn(worker_id):                                                          
+        np.random.seed(np.random.get_state()[1][0] + worker_id)
+
+    dataloader_params = {}
+    dataloader_params['num_workers'] = cfg.data.num_workers
+    dataloader_params['drop_last'] = mode == 'train'
+    dataloader_params['shuffle'] = mode == 'train'
+    dataloader_params["pin_memory"] = cfg.data.get("pin_memory", True)
+    if mode in ('train', 'valid'):
+        if mode == "train":
+            dataloader_params['batch_size'] = cfg.train.batch_size
+        elif mode == "valid":
+            dataloader_params["batch_size"] = cfg.evaluate.get("batch_size") or cfg.train.batch_size
+        sampler = None
+        if cfg.data.get("sampler") and mode == 'train':
+            name, params = get_name_and_params(cfg.data.sampler)
+            sampler = getattr(data.samplers, name)(dataset, **params)
+        if sampler:
+            dataloader_params['shuffle'] = False
+            if cfg.strategy == 'ddp':
+                sampler = data.samplers.DistributedSamplerWrapper(sampler)
+            dataloader_params['sampler'] = sampler
+            print(f'Using sampler {sampler} for training ...')
+        elif cfg.strategy == 'ddp':
+            dataloader_params["shuffle"] = False
+            dataloader_params['sampler'] = DistributedSampler(dataset, shuffle=mode=="train")
+    else:
+        assert cfg.strategy != "ddp", "DDP currently not supported for inference"
+        dataloader_params['batch_size'] = cfg.evaluate.get("batch_size") or cfg.train.batch_size
+
+    loader = DataLoader(dataset,
+        **dataloader_params,
+        worker_init_fn=worker_init_fn)
+    return loader
+
+
+def build_model(cfg):
+    name, params = get_name_and_params(cfg.model)
+    if cfg.model.params.get("cnn_params", None):
+        cnn_params = cfg.model.params.cnn_params
+        if cnn_params.get("load_pretrained_backbone", None):
+            if "foldx" in cnn_params.load_pretrained_backbone:
+                cfg.model.params.cnn_params.load_pretrained_backbone = cnn_params.load_pretrained_backbone.\
+                    replace("foldx", f"fold{cfg.data.outer_fold}")
+    print(f'Creating model <{name}> ...')
+    model = getattr(models.engine, name)(**params)
+    if 'backbone' in cfg.model.params:
+        print(f'  Using backbone <{cfg.model.params.backbone}> ...')
+    if 'pretrained' in cfg.model.params:
+        print(f'  Pretrained : {cfg.model.params.pretrained}')
+    if "load_pretrained" in cfg.model:
+        import re
+        if "foldx" in cfg.model.load_pretrained:
+            cfg.model.load_pretrained = cfg.model.load_pretrained.replace("foldx", f"fold{cfg.data.outer_fold}")
+        print(f"  Loading pretrained checkpoint from {cfg.model.load_pretrained}")
+        weights = torch.load(cfg.model.load_pretrained, map_location=lambda storage, loc: storage)['state_dict']
+        weights = {re.sub(r'^model.', '', k) : v for k,v in weights.items() if "loss_fn" not in k}
+        model.load_state_dict(weights) 
+    return model 
+
+
+def build_loss(cfg):
+    name, params = get_name_and_params(cfg.loss)
+    print(f'Using loss function <{name}> ...')
+    params = dict(params)
+    if "pos_weight" in params:
+        params["pos_weight"] = torch.tensor(params["pos_weight"])
+    criterion = getattr(losses, name)(**params)
+    return criterion
+
+
+def build_scheduler(cfg, optimizer):
+    # Some schedulers will require manipulation of config params
+    # My specifications were to make it more intuitive for me
+    name, params = get_name_and_params(cfg.scheduler)
+    print(f'Using learning rate schedule <{name}> ...')
+
+    if name == 'CosineAnnealingLR':
+        # eta_min <-> final_lr
+        # Set T_max as 100000 ... this is changed in on_train_start() method
+        # of the LightningModule task 
+
+        params = {
+            'T_max': 100000,
+            'eta_min': max(params.final_lr, 1.0e-8)
+        }
+
+    if name in ('OneCycleLR', 'CustomOneCycleLR'):
+        # Use learning rate from optimizer parameters as initial learning rate
+        lr_0 = cfg.optimizer.params.lr
+        lr_1 = params.max_lr
+        lr_2 = params.final_lr
+        # lr_0 -> lr_1 -> lr_2 
+        pct_start = params.pct_start
+        params = {}
+        params['steps_per_epoch'] = 100000 # see above- will fix in task
+        params['epochs'] = cfg.train.num_epochs
+        params['max_lr'] = lr_1
+        params['pct_start'] = pct_start
+        params['div_factor'] = lr_1 / lr_0 # max/init
+        params['final_div_factor'] = lr_0 / max(lr_2, 1.0e-8) # init/final
+
+    scheduler = getattr(optim, name)(optimizer=optimizer, **params)
+    
+    # Some schedulers might need more manipulation after instantiation
+    if name in ('OneCycleLR', 'CustomOneCycleLR'):
+        scheduler.pct_start = params['pct_start']
+
+    # Set update frequency
+    if name in ('OneCycleLR', 'CustomOneCycleLR', 'CosineAnnealingLR'):
+        scheduler.update_frequency = 'on_batch'
+    elif name in ('ReduceLROnPlateau'):
+        scheduler.update_frequency = 'on_valid'
+    else:
+        scheduler.update_frequency = 'on_epoch'
+
+    return scheduler
+
+
+def build_optimizer(cfg, parameters):
+    name, params = get_name_and_params(cfg.optimizer)
+    print(f'Using optimizer <{name}> ...')
+    optimizer = getattr(optim, name)(parameters, **params)
+    return optimizer
+
+
+def build_task(cfg, model):
+    name, params = get_name_and_params(cfg.task)
+    print(f'Building task <{name}> ...')
+    return getattr(tasks, name)(cfg, model, **params)
+
+

skp/models/__init__.py

@@ -0,0 +1 @@
+from . import engine

skp/models/backbones.py

@@ -0,0 +1,114 @@
+import re
+import timm
+import torch
+
+from functools import partial
+from timm.models.vision_transformer import VisionTransformer
+from timm.models.swin_transformer_v2 import SwinTransformerV2
+
+from .vmz.backbones import *
+
+
+def check_name(name, s):
+    return bool(re.search(s, name))
+
+
+def create_backbone(name, pretrained, features_only=False, **kwargs):
+    try:
+        model = timm.create_model(name, pretrained=pretrained, 
+                                  features_only=features_only, 
+                                  num_classes=0, global_pool="")
+    except Exception as e:
+        assert name in BACKBONES, f"{name} is not a valid backbone"
+        model = BACKBONES[name](pretrained=pretrained, features_only=features_only, **kwargs)
+    with torch.no_grad():
+        if check_name(name, r"x3d|csn|r2plus1d|i3d"):
+            dim_feats = model(torch.randn((2, 3, 64, 64, 64))).size(1)
+        elif isinstance(model, (VisionTransformer, SwinTransformerV2)):
+            dim_feats = model.norm.normalized_shape[0]
+        else:
+            dim_feats = model(torch.randn((2, 3, 128, 128))).size(1)
+    return model, dim_feats
+
+
+def create_csn(name, pretrained, features_only=False, z_strides=[1, 1, 1, 1, 1], **kwargs):
+    if features_only:
+        raise Exception("features_only is currently not supported")
+    if not pretrained:
+        from pytorchvideo.models import hub
+        model = getattr(hub, name)(pretrained=False)
+    else:
+        model = torch.hub.load("facebookresearch/pytorchvideo:main", model=name, pretrained=pretrained)
+    model.blocks[5] = nn.Identity()
+    return model
+
+
+def create_x3d(name, pretrained, features_only=False, z_strides=[1, 1, 1, 1, 1], **kwargs):
+    if not pretrained:
+        from pytorchvideo.models import hub
+        model = getattr(hub, name)(pretrained=False)
+    else:
+        model = torch.hub.load("facebookresearch/pytorchvideo", model=name, pretrained=pretrained)   
+    for idx, z in enumerate(z_strides):
+        assert z in [1, 2], "Only z-strides of 1 or 2 are supported"
+        if z == 2:
+            if idx == 0:
+                stem_layer = model.blocks[0].conv.conv_t
+                w = stem_layer.weight
+                w = w.repeat(1, 1, 3, 1, 1)
+                in_channels, out_channels = stem_layer.in_channels, stem_layer.out_channels
+                model.blocks[0].conv.conv_t = nn.Conv3d(in_channels, out_channels, kernel_size=(3, 3, 3), stride=(2, 2, 2), padding=(1, 1, 1))
+            else:                
+                model.blocks[idx].res_blocks[0].branch1_conv.stride = (2, 2, 2)
+                model.blocks[idx].res_blocks[0].branch2.conv_b.stride = (2, 2, 2)
+
+    if features_only:
+        model.blocks[-1] = nn.Identity()
+        model = X3D_Features(model)
+    else:
+        model.blocks[-1] = nn.Sequential(
+                model.blocks[-1].pool.pre_conv,
+                model.blocks[-1].pool.pre_norm,
+                model.blocks[-1].pool.pre_act,
+            )
+
+    return model
+
+
+def create_i3d(name, pretrained, features_only=False, **kwargs):
+    from pytorchvideo.models import hub
+    model = getattr(hub, name)(pretrained=pretrained)
+    model.blocks[-1] = nn.Identity()
+    return model
+
+
+class X3D_Features(nn.Module):
+
+    def __init__(self, model):
+        super().__init__()
+        self.model = model 
+        self.out_channels = [24, 24, 48, 96, 192]
+
+    def forward(self, x):
+        features = []
+        for idx in range(len(self.model.blocks) - 1):
+            x = self.model.blocks[idx](x)
+            features.append(x)
+        return features
+
+
+BACKBONES = {
+    "x3d_xs": partial(create_x3d, name="x3d_xs"),
+    "x3d_s": partial(create_x3d, name="x3d_s"),
+    "x3d_m": partial(create_x3d, name="x3d_m"),
+    "x3d_l": partial(create_x3d, name="x3d_l"),
+    "i3d_r50": partial(create_i3d, name="i3d_r50"),
+    "csn_r101": partial(create_csn, name="csn_r101"),
+    "ir_csn_50": ir_csn_50,
+    "ir_csn_101": ir_csn_101,
+    "ir_csn_152": ir_csn_152,
+    "ip_csn_50": ip_csn_50,
+    "ip_csn_101": ip_csn_101,
+    "ip_csn_152": ip_csn_152,
+    "r2plus1d_34": r2plus1d_34
+}

skp/models/engine.py

@@ -0,0 +1,257 @@
+import math
+import numpy as np
+import re
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from pytorchvideo.models.x3d import create_x3d_stem
+from timm.models.vision_transformer import VisionTransformer
+from timm.models.swin_transformer_v2 import SwinTransformerV2
+from . import backbones
+from . import segmentation
+from .pooling import create_pool2d_layer, create_pool3d_layer
+from .sequence import Transformer, DualTransformer, DualTransformerV2
+from .tools import change_initial_stride, change_num_input_channels
+
+
+class Net2D(nn.Module):
+
+    def __init__(self,
+                 backbone,
+                 pretrained,
+                 num_classes,
+                 dropout,
+                 pool,
+                 in_channels=3,
+                 change_stride=None,
+                 feature_reduction=None,
+                 multisample_dropout=False,
+                 load_pretrained_backbone=None,
+                 freeze_backbone=False,
+                 backbone_params={},
+                 pool_layer_params={}):
+
+        super().__init__()
+        self.backbone, dim_feats = backbones.create_backbone(name=backbone, pretrained=pretrained, **backbone_params)
+        if isinstance(pool, str):
+            self.pool_layer = create_pool2d_layer(name=pool, **pool_layer_params)
+        else:
+            self.pool_layer = nn.Identity()
+        if pool == "catavgmax": 
+            dim_feats *= 2
+        self.msdo = multisample_dropout
+        if in_channels != 3:
+            self.backbone = change_num_input_channels(self.backbone, in_channels)
+        if change_stride:
+            self.backbone = change_initial_stride(self.backbone, tuple(change_stride), in_channels)
+        self.dropout = nn.Dropout(p=dropout)
+        if isinstance(feature_reduction, int):
+            # Use 1D grouped convolution to reduce # of parameters
+            groups = math.gcd(dim_feats, feature_reduction)
+            self.feature_reduction = nn.Conv1d(dim_feats, feature_reduction, groups=groups, kernel_size=1,
+                                               stride=1, bias=False)
+            dim_feats = feature_reduction
+        self.classifier = nn.Linear(dim_feats, num_classes) 
+
+        if load_pretrained_backbone:
+            # Assumes that model has a `backbone` attribute
+            # Note: if you want to load the entire pretrained model, this is done via the
+            # builder.build_model function
+            print(f"Loading pretrained backbone from {load_pretrained_backbone} ...")
+            weights = torch.load(load_pretrained_backbone, map_location=lambda storage, loc: storage)['state_dict']
+            weights = {re.sub(r'^model.', '', k) : v for k,v in weights.items()}
+            # Get feature_reduction, if present
+            feat_reduce_weight = {re.sub(r"^feature_reduction.", "", k): v
+                                  for k, v in weights.items() if "feature_reduction" in k}
+            # Get backbone only
+            weights = {re.sub(r'^backbone.', '', k) : v for k,v in weights.items() if 'backbone' in k}
+            self.backbone.load_state_dict(weights)
+            if len(feat_reduce_weight) > 0:
+                print("Also loading feature reduction layer ...")
+                self.feature_reduction.load_state_dict(feat_reduce_weight)
+
+        if freeze_backbone:
+            print("Freezing backbone ...")
+            for param in self.backbone.parameters():
+                param.requires_grad = False
+
+    def extract_features(self, x):
+        features = self.backbone(x)
+        features = self.pool_layer(features)
+        if isinstance(self.backbone, VisionTransformer):
+            features = features[:, self.backbone.num_prefix_tokens:].mean(dim=1)
+        if isinstance(self.backbone, SwinTransformerV2):
+            features = features.mean(dim=1)
+        if hasattr(self, "feature_reduction"):
+            features = self.feature_reduction(features.unsqueeze(-1)).squeeze(-1)
+        return features
+
+    def forward(self, x):
+        features = self.extract_features(x)
+        if self.msdo:
+            x = torch.mean(torch.stack([self.classifier(self.dropout(features)) for _ in range(5)]), dim=0)
+        else:
+            x = self.classifier(self.dropout(features))
+        # Important nuance:
+        # For binary classification, the model returns a tensor of shape (N,)
+        # Otherwise, (N,C)
+        return x[:, 0] if self.classifier.out_features == 1 else x
+
+
+class SeqNet2D(Net2D):
+
+    def forward(self, x):
+        # x.shape = (N, C, Z, H, W)
+        features = torch.stack([self.extract_features(x[:, :, _]) for _ in range(x.size(2))], dim=2)
+        features = features.max(2)[0]
+
+        if self.msdo:
+            x = torch.mean(torch.stack([self.classifier(self.dropout(features)) for _ in range(5)]), dim=0)
+        else:
+            x = self.classifier(self.dropout(features))
+        # Important nuance:
+        # For binary classification, the model returns a tensor of shape (N,)
+        # Otherwise, (N,C)
+        return x[:, 0] if self.classifier.out_features == 1 else x
+
+
+class TDCNN(nn.Module):
+
+    def __init__(self, cnn_params, transformer_params, freeze_cnn=False, freeze_transformer=False):
+        super().__init__()
+        self.cnn = Net2D(**cnn_params)
+        del self.cnn.dropout
+        del self.cnn.classifier
+        self.transformer = Transformer(**transformer_params)
+
+        if freeze_cnn:
+            for param in self.cnn.parameters():
+                param.requires_grad = False
+
+        if freeze_transformer:
+            for param in self.transformer.parameters():
+                param.requires_grad = False
+
+    def extract_features(self, x):
+        N, C, Z, H, W = x.size()
+        assert N == 1, "For feature extraction, batch size must be 1"
+        features = self.cnn.extract_features(x.squeeze(0).transpose(0, 1)).unsqueeze(0)
+        # features.shape = (1, Z, dim_feats)
+        return self.transformer.extract_features((features, torch.ones((features.size(0), features.size(1))).to(features.device)))
+
+    def forward(self, x):
+        # BCZHW
+        features = torch.stack([self.cnn.extract_features(x[:, :, i]) for i in range(x.size(2))], dim=1)
+        # B, seq_len, dim_feat
+        return self.transformer((features, torch.ones((features.size(0), features.size(1))).to(features.device)))
+
+
+class Net2DWith3DStem(Net2D):
+
+    def __init__(self, *args, **kwargs):
+        stem_out_channels = kwargs.pop("stem_out_channels", 24)
+        load_pretrained_stem = kwargs.pop("load_pretrained_stem", None)
+        conv_kernel_size = tuple(kwargs.pop("conv_kernel_size", (5, 3, 3)))
+        conv_stride = tuple(kwargs.pop("conv_stride", (1, 2, 2)))
+        in_channels = kwargs.pop("in_channels", 3)
+        kwargs["in_channels"] = stem_out_channels
+        super().__init__(*args, **kwargs)
+        self.stem_layer = create_x3d_stem(in_channels=in_channels,
+                                          out_channels=stem_out_channels,
+                                          conv_kernel_size=conv_kernel_size,
+                                          conv_stride=conv_stride)
+        if kwargs["pretrained"]:
+            from pytorchvideo.models.hub import x3d_l
+            self.stem_layer.load_state_dict(x3d_l(pretrained=True).blocks[0].state_dict())
+
+        if load_pretrained_stem:
+            import re
+            print(f"  Loading pretrained stem from {load_pretrained_stem} ...")
+            weights = torch.load(load_pretrained_stem, map_location=lambda storage, loc: storage)['state_dict']
+            stem_weights = {k.replace("model.backbone.blocks.0.", ""): v for k, v in weights.items() if "backbone.blocks.0" in k}
+            self.stem_layer.load_state_dict(stem_weights)
+
+    def forward(self, x):
+        x = self.stem_layer(x)
+        x = x.mean(3)
+        features = self.extract_features(x)
+        if self.msdo:
+            x = torch.mean(torch.stack([self.classifier(self.dropout(features)) for _ in range(5)]), dim=0)
+        else:
+            x = self.classifier(self.dropout(features))
+        # Important nuance:
+        # For binary classification, the model returns a tensor of shape (N,)
+        # Otherwise, (N,C)
+        return x[:, 0] if self.classifier.out_features == 1 else x
+
+
+class Net3D(Net2D):
+
+    def __init__(self, *args, **kwargs):
+        z_strides = kwargs.pop("z_strides", [1,1,1,1,1])
+        super().__init__(*args, **kwargs)
+        self.pool_layer = create_pool3d_layer(name=kwargs["pool"], **kwargs.pop("pool_layer_params", {}))
+
+
+class NetSegment2D(nn.Module):
+    """ For now, this class essentially servers as a wrapper for the 
+    segmentation model which is mostly defined in the segmentation submodule, 
+    similar to the original segmentation_models.pytorch.
+
+    It may be worth refactoring it in the future, such that you define this as
+    a general class, then select your choice of encoder and decoder. The encoder
+    is pretty much the same across all the segmentation models currently 
+    implemented (DeepLabV3+, FPN, Unet).
+    """
+    def __init__(self,
+                 architecture,
+                 encoder_name,
+                 encoder_params,
+                 decoder_params,
+                 num_classes,
+                 dropout,
+                 in_channels,
+                 load_pretrained_encoder=None,
+                 freeze_encoder=False,
+                 deep_supervision=False,
+                 pool_layer_params={},
+                 aux_head_params={}):
+
+        super().__init__()
+
+        self.segmentation_model = getattr(segmentation, architecture)(
+                encoder_name=encoder_name,
+                encoder_params=encoder_params,
+                dropout=dropout,
+                classes=num_classes,
+                deep_supervision=deep_supervision,
+                in_channels=in_channels,
+                **decoder_params
+            )
+
+
+        if load_pretrained_encoder: 
+            # Assumes that model has a `encoder` attribute
+            # Note: if you want to load the entire pretrained model, this is done via the
+            # builder.build_model function
+            print(f"Loading pretrained encoder from {load_pretrained_encoder} ...")
+            weights = torch.load(load_pretrained_encoder, map_location=lambda storage, loc: storage)['state_dict']
+            weights = {re.sub(r'^model.segmentation_model', '', k) : v for k,v in weights.items()}
+            # Get encoder only
+            weights = {re.sub(r'^encoder.', '', k) : v for k,v in weights.items() if 'backbone' in k}
+            self.segmentation_model.encoder.load_state_dict(weights)
+
+        if freeze_encoder:
+            print("Freezing encoder ...")
+            for param in self.segmentation_model.encoder.parameters():
+                param.requires_grad = False
+
+
+    def forward(self, x):
+        return self.segmentation_model(x)
+
+
+class NetSegment3D(NetSegment2D): 
+
+    pass

skp/models/pooling/__init__.py

@@ -0,0 +1,3 @@
+from .pool3d import create_pool3d_layer 
+from .pool2d import create_pool2d_layer 
+from .pool1d import create_pool1d_layer

skp/models/pooling/gem.py

@@ -0,0 +1,35 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# From: https://github.com/filipradenovic/cnnimageretrieval-pytorch/blob/master/cirtorch/layers/pooling.py
+def gem_1d(x, p=3, eps=1e-6):
+    return F.avg_pool1d(x.clamp(min=eps).pow(p), (x.size(-1),)).pow(1./p)
+
+
+def gem_2d(x, p=3, eps=1e-6):
+    return F.avg_pool2d(x.clamp(min=eps).pow(p), (x.size(-2), x.size(-1))).pow(1./p)
+
+
+def gem_3d(x, p=3, eps=1e-6):
+    return F.avg_pool3d(x.clamp(min=eps).pow(p), (x.size(-3), x.size(-2), x.size(-1))).pow(1./p)
+
+
+_GEM_FN = {
+    1: gem_1d, 2: gem_2d, 3: gem_3d
+}
+
+
+class GeM(nn.Module):
+
+    def __init__(self, p=3, eps=1e-6, dim=2):
+        super().__init__()
+        self.p = nn.Parameter(torch.ones(1)*p)
+        self.eps = eps
+        self.dim = dim
+        self.flatten = nn.Flatten(1)
+
+    def forward(self, x):
+        pooled = _GEM_FN[self.dim](x, p=self.p, eps=self.eps)
+        return self.flatten(pooled)

skp/models/pooling/pool1d.py

@@ -0,0 +1,107 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .gem import GeM
+
+
+def adaptive_avgmax_pool1d(x, output_size=1):
+    x_avg = F.adaptive_avg_pool1d(x, output_size)
+    x_max = F.adaptive_max_pool1d(x, output_size)
+    return 0.5 * (x_avg + x_max)
+
+
+def adaptive_catavgmax_pool1d(x, output_size=1):
+    x_avg = F.adaptive_avg_pool1d(x, output_size)
+    x_max = F.adaptive_max_pool1d(x, output_size)
+    return torch.cat((x_avg, x_max), 1)
+
+
+def select_adaptive_pool1d(x, pool_type='avg', output_size=1):
+    """Selectable global pooling function with dynamic input kernel size
+    """
+    if pool_type == 'avg':
+        x = F.adaptive_avg_pool1d(x, output_size)
+    elif pool_type == 'avgmax':
+        x = adaptive_avgmax_pool1d(x, output_size)
+    elif pool_type == 'catavgmax':
+        x = adaptive_catavgmax_pool1d(x, output_size)
+    elif pool_type == 'max':
+        x = F.adaptive_max_pool1d(x, output_size)
+    else:
+        assert False, 'Invalid pool type: %s' % pool_type
+    return x
+
+
+class FastAdaptiveAvgPool1d(nn.Module):
+    def __init__(self, flatten=False):
+        super(FastAdaptiveAvgPool1d, self).__init__()
+        self.flatten = flatten
+
+    def forward(self, x):
+        return x.mean(2, keepdim=not self.flatten)
+
+
+class AdaptiveAvgMaxPool1d(nn.Module):
+    def __init__(self, output_size=1):
+        super(AdaptiveAvgMaxPool1d, self).__init__()
+        self.output_size = output_size
+
+    def forward(self, x):
+        return adaptive_avgmax_pool1d(x, self.output_size)
+
+
+class AdaptiveCatAvgMaxPool1d(nn.Module):
+    def __init__(self, output_size=1):
+        super(AdaptiveCatAvgMaxPool1d, self).__init__()
+        self.output_size = output_size
+
+    def forward(self, x):
+        return adaptive_catavgmax_pool1d(x, self.output_size)
+
+
+class SelectAdaptivePool1d(nn.Module):
+    """Selectable global pooling layer with dynamic input kernel size
+    """
+    def __init__(self, output_size=1, pool_type='fast', flatten=False):
+        super(SelectAdaptivePool1d, self).__init__()
+        self.pool_type = pool_type or ''  # convert other falsy values to empty string for consistent TS typing
+        self.flatten = nn.Flatten(1) if flatten else nn.Identity()
+        if pool_type == '':
+            self.pool = nn.Identity()  # pass through
+        elif pool_type == 'fast':
+            assert output_size == 1
+            self.pool = FastAdaptiveAvgPool1d(flatten)
+            self.flatten = nn.Identity()
+        elif pool_type == 'avg':
+            self.pool = nn.AdaptiveAvgPool1d(output_size)
+        elif pool_type == 'avgmax':
+            self.pool = AdaptiveAvgMaxPool1d(output_size)
+        elif pool_type == 'catavgmax':
+            self.pool = AdaptiveCatAvgMaxPool1d(output_size)
+        elif pool_type == 'max':
+            self.pool = nn.AdaptiveMaxPool1d(output_size)
+        else:
+            assert False, 'Invalid pool type: %s' % pool_type
+
+    def is_identity(self):
+        return not self.pool_type
+
+    def forward(self, x):
+        x = self.pool(x)
+        x = self.flatten(x)
+        return x
+
+    def __repr__(self):
+        return self.__class__.__name__ + ' (' \
+               + 'pool_type=' + self.pool_type \
+               + ', flatten=' + str(self.flatten) + ')'
+
+
+def create_pool1d_layer(name, **kwargs):
+    assert name in ["avg", "max", "fast", "avgmax", "catavgmax", "gem"]
+    if name != "gem":
+        pool1d_layer = SelectAdaptivePool1d(pool_type=name, flatten=True)
+    elif name == "gem": 
+        pool1d_layer = GeM(dim=1, **kwargs)
+    return pool1d_layer

skp/models/pooling/pool2d.py

@@ -0,0 +1,16 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from timm.models.layers import SelectAdaptivePool2d
+
+from .gem import GeM
+
+
+def create_pool2d_layer(name, **kwargs):
+    assert name in ["avg", "max", "fast", "avgmax", "catavgmax", "gem"]
+    if name != "gem":
+        pool2d_layer = SelectAdaptivePool2d(pool_type=name, flatten=True)
+    elif name == "gem": 
+        pool2d_layer = GeM(dim=2, **kwargs)
+    return pool2d_layer

skp/models/pooling/pool3d.py

@@ -0,0 +1,107 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .gem import GeM
+
+
+def adaptive_avgmax_pool3d(x, output_size=1):
+    x_avg = F.adaptive_avg_pool3d(x, output_size)
+    x_max = F.adaptive_max_pool3d(x, output_size)
+    return 0.5 * (x_avg + x_max)
+
+
+def adaptive_catavgmax_pool3d(x, output_size=1):
+    x_avg = F.adaptive_avg_pool3d(x, output_size)
+    x_max = F.adaptive_max_pool3d(x, output_size)
+    return torch.cat((x_avg, x_max), 1)
+
+
+def select_adaptive_pool3d(x, pool_type='avg', output_size=1):
+    """Selectable global pooling function with dynamic input kernel size
+    """
+    if pool_type == 'avg':
+        x = F.adaptive_avg_pool3d(x, output_size)
+    elif pool_type == 'avgmax':
+        x = adaptive_avgmax_pool3d(x, output_size)
+    elif pool_type == 'catavgmax':
+        x = adaptive_catavgmax_pool3d(x, output_size)
+    elif pool_type == 'max':
+        x = F.adaptive_max_pool3d(x, output_size)
+    else:
+        assert False, 'Invalid pool type: %s' % pool_type
+    return x
+
+
+class FastAdaptiveAvgPool3d(nn.Module):
+    def __init__(self, flatten=False):
+        super(FastAdaptiveAvgPool3d, self).__init__()
+        self.flatten = flatten
+
+    def forward(self, x):
+        return x.mean((2,3,4), keepdim=not self.flatten)
+
+
+class AdaptiveAvgMaxPool3d(nn.Module):
+    def __init__(self, output_size=1):
+        super(AdaptiveAvgMaxPool3d, self).__init__()
+        self.output_size = output_size
+
+    def forward(self, x):
+        return adaptive_avgmax_pool3d(x, self.output_size)
+
+
+class AdaptiveCatAvgMaxPool3d(nn.Module):
+    def __init__(self, output_size=1):
+        super(AdaptiveCatAvgMaxPool3d, self).__init__()
+        self.output_size = output_size
+
+    def forward(self, x):
+        return adaptive_catavgmax_pool3d(x, self.output_size)
+
+
+class SelectAdaptivePool3d(nn.Module):
+    """Selectable global pooling layer with dynamic input kernel size
+    """
+    def __init__(self, output_size=1, pool_type='fast', flatten=False):
+        super(SelectAdaptivePool3d, self).__init__()
+        self.pool_type = pool_type or ''  # convert other falsy values to empty string for consistent TS typing
+        self.flatten = nn.Flatten(1) if flatten else nn.Identity()
+        if pool_type == '':
+            self.pool = nn.Identity()  # pass through
+        elif pool_type == 'fast':
+            assert output_size == 1
+            self.pool = FastAdaptiveAvgPool3d(flatten)
+            self.flatten = nn.Identity()
+        elif pool_type == 'avg':
+            self.pool = nn.AdaptiveAvgPool3d(output_size)
+        elif pool_type == 'avgmax':
+            self.pool = AdaptiveAvgMaxPool3d(output_size)
+        elif pool_type == 'catavgmax':
+            self.pool = AdaptiveCatAvgMaxPool3d(output_size)
+        elif pool_type == 'max':
+            self.pool = nn.AdaptiveMaxPool3d(output_size)
+        else:
+            assert False, 'Invalid pool type: %s' % pool_type
+
+    def is_identity(self):
+        return not self.pool_type
+
+    def forward(self, x):
+        x = self.pool(x)
+        x = self.flatten(x)
+        return x
+
+    def __repr__(self):
+        return self.__class__.__name__ + ' (' \
+               + 'pool_type=' + self.pool_type \
+               + ', flatten=' + str(self.flatten) + ')'
+
+
+def create_pool3d_layer(name, **kwargs):
+    assert name in ["avg", "max", "fast", "avgmax", "catavgmax", "gem"]
+    if name != "gem":
+        pool1d_layer = SelectAdaptivePool3d(pool_type=name, flatten=True)
+    elif name == "gem": 
+        pool1d_layer = GeM(dim=3, **kwargs)
+    return pool1d_layer

skp/models/rev_mvit/REV_MVIT_B_16_CONV.yaml

@@ -0,0 +1,109 @@
+TRAIN:
+  ENABLE: True
+  DATASET: imagenet
+  BATCH_SIZE: 256
+  EVAL_PERIOD: 10
+  CHECKPOINT_PERIOD: 1
+  AUTO_RESUME: True
+
+DATA:
+  # PATH_TO_DATA_DIR: path-to-imagenet-dir
+  MEAN: [0.485, 0.456, 0.406]
+  STD: [0.229, 0.224, 0.225]
+  NUM_FRAMES: 64
+  TRAIN_CROP_SIZE: 224
+  TEST_CROP_SIZE: 224
+  INPUT_CHANNEL_NUM: [3]
+MVIT:
+  PATCH_2D: False
+  ZERO_DECAY_POS_CLS: False
+  MODE: "conv"
+  CLS_EMBED_ON: False
+  PATCH_KERNEL: [3, 7, 7]
+  PATCH_STRIDE: [2, 4, 4]
+  PATCH_PADDING: [1, 3, 3]
+  EMBED_DIM: 96
+  NUM_HEADS: 1
+  MLP_RATIO: 4.0
+  QKV_BIAS: True
+  DROPPATH_RATE: 0.1
+  DROPOUT_RATE: 0.0
+  DEPTH: 16
+  LAYER_SCALE_INIT_VALUE: 0.0
+  HEAD_INIT_SCALE: 1.0
+  USE_MEAN_POOLING: False
+  USE_ABS_POS: True
+  USE_FIXED_SINCOS_POS: False
+  SEP_POS_EMBED: False
+  REL_POS_SPATIAL: False
+  REL_POS_TEMPORAL: False
+  REL_POS_ZERO_INIT: False
+  RESIDUAL_POOLING: False
+  NORM: "layernorm"
+  NORM_STEM: False
+  DIM_MUL: [[1, 2.0], [3, 2.0], [14, 2.0]]
+  HEAD_MUL: [[1, 2.0], [3, 2.0], [14, 2.0]]
+  POOL_FIRST: null
+  POOL_KVQ_KERNEL: [1, 3, 3]
+  POOL_KV_STRIDE_ADAPTIVE: [1, 4, 4]
+  POOL_Q_STRIDE: [[1, 1, 2, 2], [3, 1, 2, 2], [14, 1, 2, 2]]
+  SEPARATE_QKV : True
+  REV:
+    ENABLE: True
+    RESPATH_FUSE: "concat"
+    BUFFER_LAYERS : [1,3, 14]
+    RES_PATH : "conv"
+    PRE_Q_FUSION: "concat_linear_2"
+DETECTION:
+  ENABLE: False
+AUG:
+  ENABLE: True
+  COLOR_JITTER: 0.4
+  AA_TYPE: rand-m9-n6-mstd0.5-inc1
+  INTERPOLATION: bicubic
+  RE_PROB: 0.25
+  RE_MODE: pixel
+  RE_COUNT: 1
+  RE_SPLIT: False
+MIXUP:
+  ENABLE: True
+  ALPHA: 0.8
+  CUTMIX_ALPHA: 1.0
+  PROB: 1.0
+  SWITCH_PROB: 0.5
+  LABEL_SMOOTH_VALUE: 0.1
+SOLVER:
+  BASE_LR_SCALE_NUM_SHARDS: True
+  BASE_LR: 0.00025
+  LR_POLICY: cosine
+  MAX_EPOCH: 300
+  MOMENTUM: 0.9
+  WEIGHT_DECAY: 0.05
+  WARMUP_EPOCHS: 70.0
+  WARMUP_START_LR: 1e-8
+  OPTIMIZING_METHOD: adamw
+  COSINE_AFTER_WARMUP: True
+  COSINE_END_LR: 1e-6
+  ZERO_WD_1D_PARAM: True
+  CLIP_GRAD_L2NORM: 1.0
+MODEL:
+  NUM_CLASSES: 1000
+  ARCH: mvit
+  MODEL_NAME: MViT
+  LOSS_FUNC: soft_cross_entropy
+  DROPOUT_RATE: 0.0
+  HEAD_ACT: "softmax"
+  DETACH_FINAL_FC: False
+CONTRASTIVE:
+  NUM_MLP_LAYERS: 1
+TEST:
+  ENABLE: False
+  DATASET: imagenet
+  BATCH_SIZE: 256
+DATA_LOADER:
+  NUM_WORKERS: 8
+  PIN_MEMORY: True
+NUM_GPUS: 2
+NUM_SHARDS: 1
+RNG_SEED: 0
+OUTPUT_DIR: .

skp/models/rev_mvit/attention.py

@@ -0,0 +1,568 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+
+
+import numpy
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.init import trunc_normal_
+
+from .common import DropPath, Mlp
+
+
+def attention_pool(tensor, pool, thw_shape, has_cls_embed=True, norm=None):
+    if pool is None:
+        return tensor, thw_shape
+    tensor_dim = tensor.ndim
+    if tensor_dim == 4:
+        pass
+    elif tensor_dim == 3:
+        tensor = tensor.unsqueeze(1)
+    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
+    T, H, W = thw_shape
+    tensor = (
+        tensor.reshape(B * N, T, H, W, C).permute(0, 4, 1, 2, 3).contiguous()
+    )
+
+    tensor = pool(tensor)
+
+    thw_shape = [tensor.shape[2], tensor.shape[3], tensor.shape[4]]
+    L_pooled = tensor.shape[2] * tensor.shape[3] * tensor.shape[4]
+    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)
+    # Assert tensor_dim in [3, 4]
+    if tensor_dim == 4:
+        pass
+    else:  #  tensor_dim == 3:
+        tensor = tensor.squeeze(1)
+    return tensor, thw_shape
+
+
+def get_rel_pos(rel_pos, d):
+    if isinstance(d, int):
+        ori_d = rel_pos.shape[0]
+        if ori_d == d:
+            return rel_pos
+        else:
+            # Interpolate rel pos.
+            new_pos_embed = F.interpolate(
+                rel_pos.reshape(1, ori_d, -1).permute(0, 2, 1),
+                size=d,
+                mode="linear",
+            )
+
+            return new_pos_embed.reshape(-1, d).permute(1, 0)
+
+
+def cal_rel_pos_spatial(
+    attn, q, k, has_cls_embed, q_shape, k_shape, rel_pos_h, rel_pos_w
+):
+    """
+    Decomposed Spatial Relative Positional Embeddings.
+    """
+    sp_idx = 1 if has_cls_embed else 0
+    q_t, q_h, q_w = q_shape
+    k_t, k_h, k_w = k_shape
+    dh = int(2 * max(q_h, k_h) - 1)
+    dw = int(2 * max(q_w, k_w) - 1)
+
+    # Scale up rel pos if shapes for q and k are different.
+    q_h_ratio = max(k_h / q_h, 1.0)
+    k_h_ratio = max(q_h / k_h, 1.0)
+    dist_h = (
+        torch.arange(q_h)[:, None] * q_h_ratio
+        - torch.arange(k_h)[None, :] * k_h_ratio
+    )
+    dist_h += (k_h - 1) * k_h_ratio
+    q_w_ratio = max(k_w / q_w, 1.0)
+    k_w_ratio = max(q_w / k_w, 1.0)
+    dist_w = (
+        torch.arange(q_w)[:, None] * q_w_ratio
+        - torch.arange(k_w)[None, :] * k_w_ratio
+    )
+    dist_w += (k_w - 1) * k_w_ratio
+
+    # Intepolate rel pos if needed.
+    rel_pos_h = get_rel_pos(rel_pos_h, dh)
+    rel_pos_w = get_rel_pos(rel_pos_w, dw)
+    Rh = rel_pos_h[dist_h.long()]
+    Rw = rel_pos_w[dist_w.long()]
+
+    B, n_head, q_N, dim = q.shape
+
+    r_q = q[:, :, sp_idx:].reshape(B, n_head, q_t, q_h, q_w, dim)
+    rel_h_q = torch.einsum(
+        "bythwc,hkc->bythwk", r_q, Rh
+    )  # [B, H, q_t, qh, qw, k_h]
+    rel_w_q = torch.einsum(
+        "bythwc,wkc->bythwk", r_q, Rw
+    )  # [B, H, q_t, qh, qw, k_w]
+
+    attn[:, :, sp_idx:, sp_idx:] = (
+        attn[:, :, sp_idx:, sp_idx:].view(B, -1, q_t, q_h, q_w, k_t, k_h, k_w)
+        + rel_h_q[:, :, :, :, :, None, :, None]
+        + rel_w_q[:, :, :, :, :, None, None, :]
+    ).view(B, -1, q_t * q_h * q_w, k_t * k_h * k_w)
+
+    return attn
+
+
+def cal_rel_pos_temporal(attn, q, has_cls_embed, q_shape, k_shape, rel_pos_t):
+    """
+    Temporal Relative Positional Embeddings.
+    """
+    sp_idx = 1 if has_cls_embed else 0
+    q_t, q_h, q_w = q_shape
+    k_t, k_h, k_w = k_shape
+    dt = int(2 * max(q_t, k_t) - 1)
+    # Intepolate rel pos if needed.
+    rel_pos_t = get_rel_pos(rel_pos_t, dt)
+
+    # Scale up rel pos if shapes for q and k are different.
+    q_t_ratio = max(k_t / q_t, 1.0)
+    k_t_ratio = max(q_t / k_t, 1.0)
+    dist_t = (
+        torch.arange(q_t)[:, None] * q_t_ratio
+        - torch.arange(k_t)[None, :] * k_t_ratio
+    )
+    dist_t += (k_t - 1) * k_t_ratio
+    Rt = rel_pos_t[dist_t.long()]
+
+    B, n_head, q_N, dim = q.shape
+
+    r_q = q[:, :, sp_idx:].reshape(B, n_head, q_t, q_h, q_w, dim)
+    # [B, H, q_t, q_h, q_w, dim] -> [q_t, B, H, q_h, q_w, dim] -> [q_t, B*H*q_h*q_w, dim]
+    r_q = r_q.permute(2, 0, 1, 3, 4, 5).reshape(
+        q_t, B * n_head * q_h * q_w, dim
+    )
+
+    # [q_t, B*H*q_h*q_w, dim] * [q_t, dim, k_t] = [q_t, B*H*q_h*q_w, k_t] -> [B*H*q_h*q_w, q_t, k_t]
+    rel = torch.matmul(r_q, Rt.transpose(1, 2)).transpose(0, 1)
+    # [B*H*q_h*q_w, q_t, k_t] -> [B, H, q_t, q_h, q_w, k_t]
+    rel = rel.view(B, n_head, q_h, q_w, q_t, k_t).permute(0, 1, 4, 2, 3, 5)
+
+    attn[:, :, sp_idx:, sp_idx:] = (
+        attn[:, :, sp_idx:, sp_idx:].view(B, -1, q_t, q_h, q_w, k_t, k_h, k_w)
+        + rel[:, :, :, :, :, :, None, None]
+    ).view(B, -1, q_t * q_h * q_w, k_t * k_h * k_w)
+
+    return attn
+
+
+class MultiScaleAttention(nn.Module):
+    def __init__(
+        self,
+        dim,
+        dim_out,
+        input_size,
+        num_heads=8,
+        qkv_bias=False,
+        drop_rate=0.0,
+        kernel_q=(1, 1, 1),
+        kernel_kv=(1, 1, 1),
+        stride_q=(1, 1, 1),
+        stride_kv=(1, 1, 1),
+        norm_layer=nn.LayerNorm,
+        has_cls_embed=True,
+        # Options include `conv`, `avg`, and `max`.
+        mode="conv",
+        # If True, perform pool before projection.
+        pool_first=False,
+        rel_pos_spatial=False,
+        rel_pos_temporal=False,
+        rel_pos_zero_init=False,
+        residual_pooling=False,
+        separate_qkv=False,
+    ):
+        super().__init__()
+        self.pool_first = pool_first
+        self.separate_qkv = separate_qkv
+        self.drop_rate = drop_rate
+        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
+        self.mode = mode
+        padding_q = [int(q // 2) for q in kernel_q]
+        padding_kv = [int(kv // 2) for kv in kernel_kv]
+
+        if pool_first or separate_qkv:
+            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)
+        else:
+            self.qkv = nn.Linear(dim, dim_out * 3, bias=qkv_bias)
+
+        self.proj = nn.Linear(dim_out, dim_out)
+        if drop_rate > 0.0:
+            self.proj_drop = nn.Dropout(drop_rate)
+
+        # 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 = ()
+
+        if mode in ("avg", "max"):
+            pool_op = nn.MaxPool3d if mode == "max" else nn.AvgPool3d
+            self.pool_q = (
+                pool_op(kernel_q, stride_q, padding_q, ceil_mode=False)
+                if len(kernel_q) > 0
+                else None
+            )
+            self.pool_k = (
+                pool_op(kernel_kv, stride_kv, padding_kv, ceil_mode=False)
+                if len(kernel_kv) > 0
+                else None
+            )
+            self.pool_v = (
+                pool_op(kernel_kv, stride_kv, padding_kv, ceil_mode=False)
+                if len(kernel_kv) > 0
+                else None
+            )
+        elif mode == "conv" or mode == "conv_unshared":
+            if pool_first:
+                dim_conv = dim // num_heads if mode == "conv" else dim
+            else:
+                dim_conv = dim_out // num_heads if mode == "conv" else dim_out
+            self.pool_q = (
+                nn.Conv3d(
+                    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.Conv3d(
+                    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.Conv3d(
+                    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
+        else:
+            raise NotImplementedError(f"Unsupported model {mode}")
+
+        self.rel_pos_spatial = rel_pos_spatial
+        self.rel_pos_temporal = rel_pos_temporal
+        if self.rel_pos_spatial:
+            assert input_size[1] == input_size[2]
+            size = input_size[1]
+            q_size = size // stride_q[1] if len(stride_q) > 0 else size
+            kv_size = size // stride_kv[1] if len(stride_kv) > 0 else size
+            rel_sp_dim = 2 * max(q_size, kv_size) - 1
+
+            self.rel_pos_h = nn.Parameter(torch.zeros(rel_sp_dim, head_dim))
+            self.rel_pos_w = nn.Parameter(torch.zeros(rel_sp_dim, head_dim))
+            if not rel_pos_zero_init:
+                trunc_normal_(self.rel_pos_h, std=0.02)
+                trunc_normal_(self.rel_pos_w, std=0.02)
+        if self.rel_pos_temporal:
+            self.rel_pos_t = nn.Parameter(
+                torch.zeros(2 * input_size[0] - 1, head_dim)
+            )
+            if not rel_pos_zero_init:
+                trunc_normal_(self.rel_pos_t, std=0.02)
+
+        self.residual_pooling = residual_pooling
+
+    def forward(self, x, thw_shape):
+        B, N, _ = x.shape
+
+        if self.pool_first:
+            if self.mode == "conv_unshared":
+                fold_dim = 1
+            else:
+                fold_dim = self.num_heads
+            x = x.reshape(B, N, fold_dim, -1).permute(0, 2, 1, 3)
+            q = k = v = x
+        else:
+            assert self.mode != "conv_unshared"
+            if not self.separate_qkv:
+                qkv = (
+                    self.qkv(x)
+                    .reshape(B, N, 3, self.num_heads, -1)
+                    .permute(2, 0, 3, 1, 4)
+                )
+                q, k, v = qkv[0], qkv[1], qkv[2]
+            else:
+                q = k = v = x
+                q = (
+                    self.q(q)
+                    .reshape(B, N, self.num_heads, -1)
+                    .permute(0, 2, 1, 3)
+                )
+                k = (
+                    self.k(k)
+                    .reshape(B, N, self.num_heads, -1)
+                    .permute(0, 2, 1, 3)
+                )
+                v = (
+                    self.v(v)
+                    .reshape(B, N, self.num_heads, -1)
+                    .permute(0, 2, 1, 3)
+                )
+
+        q, q_shape = attention_pool(
+            q,
+            self.pool_q,
+            thw_shape,
+            has_cls_embed=self.has_cls_embed,
+            norm=self.norm_q if hasattr(self, "norm_q") else None,
+        )
+        k, k_shape = attention_pool(
+            k,
+            self.pool_k,
+            thw_shape,
+            has_cls_embed=self.has_cls_embed,
+            norm=self.norm_k if hasattr(self, "norm_k") else None,
+        )
+        v, v_shape = attention_pool(
+            v,
+            self.pool_v,
+            thw_shape,
+            has_cls_embed=self.has_cls_embed,
+            norm=self.norm_v if hasattr(self, "norm_v") else None,
+        )
+
+        if self.pool_first:
+            q_N = (
+                numpy.prod(q_shape) + 1
+                if self.has_cls_embed
+                else numpy.prod(q_shape)
+            )
+            k_N = (
+                numpy.prod(k_shape) + 1
+                if self.has_cls_embed
+                else numpy.prod(k_shape)
+            )
+            v_N = (
+                numpy.prod(v_shape) + 1
+                if self.has_cls_embed
+                else numpy.prod(v_shape)
+            )
+
+            q = q.permute(0, 2, 1, 3).reshape(B, q_N, -1)
+            q = (
+                self.q(q)
+                .reshape(B, q_N, self.num_heads, -1)
+                .permute(0, 2, 1, 3)
+            )
+
+            v = v.permute(0, 2, 1, 3).reshape(B, v_N, -1)
+            v = (
+                self.v(v)
+                .reshape(B, v_N, self.num_heads, -1)
+                .permute(0, 2, 1, 3)
+            )
+
+            k = k.permute(0, 2, 1, 3).reshape(B, k_N, -1)
+            k = (
+                self.k(k)
+                .reshape(B, k_N, self.num_heads, -1)
+                .permute(0, 2, 1, 3)
+            )
+
+        N = q.shape[2]
+        attn = (q * self.scale) @ k.transpose(-2, -1)
+        if self.rel_pos_spatial:
+            attn = cal_rel_pos_spatial(
+                attn,
+                q,
+                k,
+                self.has_cls_embed,
+                q_shape,
+                k_shape,
+                self.rel_pos_h,
+                self.rel_pos_w,
+            )
+
+        if self.rel_pos_temporal:
+            attn = cal_rel_pos_temporal(
+                attn,
+                q,
+                self.has_cls_embed,
+                q_shape,
+                k_shape,
+                self.rel_pos_t,
+            )
+        attn = attn.softmax(dim=-1)
+
+        x = attn @ v
+
+        if self.residual_pooling:
+            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)
+
+        if self.drop_rate > 0.0:
+            x = self.proj_drop(x)
+        return x, q_shape
+
+
+class MultiScaleBlock(nn.Module):
+    def __init__(
+        self,
+        dim,
+        dim_out,
+        num_heads,
+        input_size,
+        mlp_ratio=4.0,
+        qkv_bias=False,
+        qk_scale=None,
+        drop_rate=0.0,
+        drop_path=0.0,
+        layer_scale_init_value=0.0,
+        act_layer=nn.GELU,
+        norm_layer=nn.LayerNorm,
+        up_rate=None,
+        kernel_q=(1, 1, 1),
+        kernel_kv=(1, 1, 1),
+        stride_q=(1, 1, 1),
+        stride_kv=(1, 1, 1),
+        mode="conv",
+        has_cls_embed=True,
+        pool_first=False,
+        rel_pos_spatial=False,
+        rel_pos_temporal=False,
+        rel_pos_zero_init=False,
+        residual_pooling=False,
+        dim_mul_in_att=False,
+        separate_qkv=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
+        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]
+        att_dim = dim_out if dim_mul_in_att else dim
+        self.attn = MultiScaleAttention(
+            dim,
+            att_dim,
+            num_heads=num_heads,
+            input_size=input_size,
+            qkv_bias=qkv_bias,
+            drop_rate=drop_rate,
+            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,
+            mode=mode,
+            pool_first=pool_first,
+            rel_pos_spatial=rel_pos_spatial,
+            rel_pos_temporal=rel_pos_temporal,
+            rel_pos_zero_init=rel_pos_zero_init,
+            residual_pooling=residual_pooling,
+            separate_qkv=separate_qkv,
+        )
+        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
+        # TODO: check the use case for up_rate, and merge the following lines
+        if up_rate is not None and up_rate > 1:
+            mlp_dim_out = dim * up_rate
+        else:
+            mlp_dim_out = dim_out
+        self.mlp = Mlp(
+            in_features=att_dim,
+            hidden_features=mlp_hidden_dim,
+            out_features=mlp_dim_out,
+            act_layer=act_layer,
+            drop_rate=drop_rate,
+        )
+        if layer_scale_init_value > 0:
+            self.gamma_1 = nn.Parameter(
+                layer_scale_init_value * torch.ones((dim)), requires_grad=True
+            )
+            self.gamma_2 = nn.Parameter(
+                layer_scale_init_value * torch.ones((dim_out)),
+                requires_grad=True,
+            )
+        else:
+            self.gamma_1, self.gamma_2 = None, None
+
+        if dim != dim_out:
+            self.proj = nn.Linear(dim, dim_out)
+
+        self.pool_skip = (
+            nn.MaxPool3d(
+                kernel_skip, stride_skip, padding_skip, ceil_mode=False
+            )
+            if len(stride_skip) > 0 and numpy.prod(stride_skip) > 1
+            else None
+        )
+
+    def forward(self, x, thw_shape=None):
+        x_norm = self.norm1(x)
+        x_block, thw_shape_new = self.attn(x_norm, thw_shape)
+        if self.dim_mul_in_att and self.dim != self.dim_out:
+            x = self.proj(x_norm)
+        x_res, _ = attention_pool(
+            x, self.pool_skip, thw_shape, has_cls_embed=self.has_cls_embed
+        )
+        if self.gamma_1 is not None:
+            x = x_res + self.drop_path(self.gamma_1 * x_block)
+        else:
+            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)
+        if self.gamma_2 is not None:
+            x = x + self.drop_path(self.gamma_2 * x_mlp)
+        else:
+            x = x + self.drop_path(x_mlp)
+        if thw_shape:
+            return x, thw_shape_new
+        else:
+            return x

skp/models/rev_mvit/batchnorm_helper.py

@@ -0,0 +1,112 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+
+"""BatchNorm (BN) utility functions and custom batch-size BN implementations"""
+
+from functools import partial
+import torch
+import torch.nn as nn
+
+from pytorchvideo.layers.batch_norm import (
+    NaiveSyncBatchNorm1d,
+    NaiveSyncBatchNorm3d,
+)  # noqa
+
+
+def get_norm(cfg):
+    """
+    Args:
+        cfg (CfgNode): model building configs, details are in the comments of
+            the config file.
+    Returns:
+        nn.Module: the normalization layer.
+    """
+    if cfg.BN.NORM_TYPE in {"batchnorm", "sync_batchnorm_apex"}:
+        return nn.BatchNorm3d
+    elif cfg.BN.NORM_TYPE == "sub_batchnorm":
+        return partial(SubBatchNorm3d, num_splits=cfg.BN.NUM_SPLITS)
+    elif cfg.BN.NORM_TYPE == "sync_batchnorm":
+        return partial(
+            NaiveSyncBatchNorm3d,
+            num_sync_devices=cfg.BN.NUM_SYNC_DEVICES,
+            global_sync=cfg.BN.GLOBAL_SYNC,
+        )
+    else:
+        raise NotImplementedError(
+            "Norm type {} is not supported".format(cfg.BN.NORM_TYPE)
+        )
+
+
+class SubBatchNorm3d(nn.Module):
+    """
+    The standard BN layer computes stats across all examples in a GPU. In some
+    cases it is desirable to compute stats across only a subset of examples
+    (e.g., in multigrid training https://arxiv.org/abs/1912.00998).
+    SubBatchNorm3d splits the batch dimension into N splits, and run BN on
+    each of them separately (so that the stats are computed on each subset of
+    examples (1/N of batch) independently. During evaluation, it aggregates
+    the stats from all splits into one BN.
+    """
+
+    def __init__(self, num_splits, **args):
+        """
+        Args:
+            num_splits (int): number of splits.
+            args (list): other arguments.
+        """
+        super(SubBatchNorm3d, self).__init__()
+        self.num_splits = num_splits
+        num_features = args["num_features"]
+        # Keep only one set of weight and bias.
+        if args.get("affine", True):
+            self.affine = True
+            args["affine"] = False
+            self.weight = torch.nn.Parameter(torch.ones(num_features))
+            self.bias = torch.nn.Parameter(torch.zeros(num_features))
+        else:
+            self.affine = False
+        self.bn = nn.BatchNorm3d(**args)
+        args["num_features"] = num_features * num_splits
+        self.split_bn = nn.BatchNorm3d(**args)
+
+    def _get_aggregated_mean_std(self, means, stds, n):
+        """
+        Calculate the aggregated mean and stds.
+        Args:
+            means (tensor): mean values.
+            stds (tensor): standard deviations.
+            n (int): number of sets of means and stds.
+        """
+        mean = means.view(n, -1).sum(0) / n
+        std = (
+            stds.view(n, -1).sum(0) / n
+            + ((means.view(n, -1) - mean) ** 2).view(n, -1).sum(0) / n
+        )
+        return mean.detach(), std.detach()
+
+    def aggregate_stats(self):
+        """
+        Synchronize running_mean, and running_var. Call this before eval.
+        """
+        if self.split_bn.track_running_stats:
+            (
+                self.bn.running_mean.data,
+                self.bn.running_var.data,
+            ) = self._get_aggregated_mean_std(
+                self.split_bn.running_mean,
+                self.split_bn.running_var,
+                self.num_splits,
+            )
+
+    def forward(self, x):
+        if self.training:
+            n, c, t, h, w = x.shape
+            x = x.view(n // self.num_splits, c * self.num_splits, t, h, w)
+            x = self.split_bn(x)
+            x = x.view(n, c, t, h, w)
+        else:
+            x = self.bn(x)
+        if self.affine:
+            x = x * self.weight.view((-1, 1, 1, 1))
+            x = x + self.bias.view((-1, 1, 1, 1))
+        return x

skp/models/rev_mvit/common.py

@@ -0,0 +1,154 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+import torch
+import torch.nn as nn
+
+
+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)
+
+
+class TwoStreamFusion(nn.Module):
+    def __init__(self, mode, dim=None, kernel=3, padding=1):
+        """
+        A general constructor for neural modules fusing two equal sized tensors
+        in forward. Following options are supported:
+
+        "add" / "max" / "min" / "avg"             : respective operations on the two halves.
+        "concat"                                  : NOOP.
+        "concat_linear_{dim_mult}_{drop_rate}"    : MLP to fuse with hidden dim "dim_mult"
+                                                    (optional, def 1.) higher than input dim
+                                                    with optional dropout "drop_rate" (def: 0.)
+        "ln+concat_linear_{dim_mult}_{drop_rate}" : perform MLP after layernorm on the input.
+
+        """
+        super().__init__()
+        self.mode = mode
+        if mode == "add":
+            self.fuse_fn = lambda x: torch.stack(torch.chunk(x, 2, dim=2)).sum(
+                dim=0
+            )
+        elif mode == "max":
+            self.fuse_fn = (
+                lambda x: torch.stack(torch.chunk(x, 2, dim=2))
+                .max(dim=0)
+                .values
+            )
+        elif mode == "min":
+            self.fuse_fn = (
+                lambda x: torch.stack(torch.chunk(x, 2, dim=2))
+                .min(dim=0)
+                .values
+            )
+        elif mode == "avg":
+            self.fuse_fn = lambda x: torch.stack(torch.chunk(x, 2, dim=2)).mean(
+                dim=0
+            )
+        elif mode == "concat":
+            # x itself is the channel concat version
+            self.fuse_fn = lambda x: x
+        elif "concat_linear" in mode:
+            if len(mode.split("_")) == 2:
+                dim_mult = 1.0
+                drop_rate = 0.0
+            elif len(mode.split("_")) == 3:
+                dim_mult = float(mode.split("_")[-1])
+                drop_rate = 0.0
+
+            elif len(mode.split("_")) == 4:
+                dim_mult = float(mode.split("_")[-2])
+                drop_rate = float(mode.split("_")[-1])
+            else:
+                raise NotImplementedError
+
+            if mode.split("+")[0] == "ln":
+                self.fuse_fn = nn.Sequential(
+                    nn.LayerNorm(dim),
+                    Mlp(
+                        in_features=dim,
+                        hidden_features=int(dim * dim_mult),
+                        act_layer=nn.GELU,
+                        out_features=dim,
+                        drop_rate=drop_rate,
+                    ),
+                )
+            else:
+                self.fuse_fn = Mlp(
+                    in_features=dim,
+                    hidden_features=int(dim * dim_mult),
+                    act_layer=nn.GELU,
+                    out_features=dim,
+                    drop_rate=drop_rate,
+                )
+
+        else:
+            raise NotImplementedError
+
+    def forward(self, x):
+        if "concat_linear" in self.mode:
+            return self.fuse_fn(x) + x
+
+        else:
+            return self.fuse_fn(x)

skp/models/rev_mvit/head_helper.py

@@ -0,0 +1,140 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+
+"""ResNe(X)t Head helper."""
+
+import torch
+import torch.nn as nn
+
+from .batchnorm_helper import (
+    NaiveSyncBatchNorm1d as NaiveSyncBatchNorm1d,
+)
+
+
+class MLPHead(nn.Module):
+    def __init__(
+        self,
+        dim_in,
+        dim_out,
+        mlp_dim,
+        num_layers,
+        bn_on=False,
+        bias=True,
+        flatten=False,
+        xavier_init=True,
+        bn_sync_num=1,
+        global_sync=False,
+    ):
+        super(MLPHead, self).__init__()
+        self.flatten = flatten
+        b = False if bn_on else bias
+        # assert bn_on or bn_sync_num=1
+        mlp_layers = [nn.Linear(dim_in, mlp_dim, bias=b)]
+        mlp_layers[-1].xavier_init = xavier_init
+        for i in range(1, num_layers):
+            if bn_on:
+                if global_sync or bn_sync_num > 1:
+                    mlp_layers.append(
+                        NaiveSyncBatchNorm1d(
+                            num_sync_devices=bn_sync_num,
+                            global_sync=global_sync,
+                            num_features=mlp_dim,
+                        )
+                    )
+                else:
+                    mlp_layers.append(nn.BatchNorm1d(num_features=mlp_dim))
+            mlp_layers.append(nn.ReLU(inplace=True))
+            if i == num_layers - 1:
+                d = dim_out
+                b = bias
+            else:
+                d = mlp_dim
+            mlp_layers.append(nn.Linear(mlp_dim, d, bias=b))
+            mlp_layers[-1].xavier_init = xavier_init
+        self.projection = nn.Sequential(*mlp_layers)
+
+    def forward(self, x):
+        if x.ndim == 5:
+            x = x.permute((0, 2, 3, 4, 1))
+        if self.flatten:
+            x = x.reshape(-1, x.shape[-1])
+
+        return self.projection(x)
+
+
+class TransformerBasicHead(nn.Module):
+    """
+    BasicHead. No pool.
+    """
+
+    def __init__(
+        self,
+        dim_in,
+        num_classes,
+        dropout_rate=0.0,
+        act_func="softmax",
+        cfg=None,
+    ):
+        """
+        Perform linear projection and activation as head for tranformers.
+        Args:
+            dim_in (int): the channel dimension of the input to the head.
+            num_classes (int): the channel dimensions of the output to the head.
+            dropout_rate (float): dropout rate. If equal to 0.0, perform no
+                dropout.
+            act_func (string): activation function to use. 'softmax': applies
+                softmax on the output. 'sigmoid': applies sigmoid on the output.
+        """
+        super(TransformerBasicHead, self).__init__()
+        if dropout_rate > 0.0:
+            self.dropout = nn.Dropout(dropout_rate)
+        self.projection = nn.Linear(dim_in, num_classes, bias=True)
+
+        if cfg.CONTRASTIVE.NUM_MLP_LAYERS == 1:
+            self.projection = nn.Linear(dim_in, num_classes, bias=True)
+        else:
+            self.projection = MLPHead(
+                dim_in,
+                num_classes,
+                cfg.CONTRASTIVE.MLP_DIM,
+                cfg.CONTRASTIVE.NUM_MLP_LAYERS,
+                bn_on=cfg.CONTRASTIVE.BN_MLP,
+                bn_sync_num=cfg.BN.NUM_SYNC_DEVICES
+                if cfg.CONTRASTIVE.BN_SYNC_MLP
+                else 1,
+                global_sync=(
+                    cfg.CONTRASTIVE.BN_SYNC_MLP and cfg.BN.GLOBAL_SYNC
+                ),
+            )
+        self.detach_final_fc = cfg.MODEL.DETACH_FINAL_FC
+
+        # Softmax for evaluation and testing.
+        if act_func == "softmax":
+            self.act = nn.Softmax(dim=1)
+        elif act_func == "sigmoid":
+            self.act = nn.Sigmoid()
+        elif act_func == "none":
+            self.act = None
+        else:
+            raise NotImplementedError(
+                "{} is not supported as an activation"
+                "function.".format(act_func)
+            )
+
+    def forward(self, x):
+        if hasattr(self, "dropout"):
+            x = self.dropout(x)
+        if self.detach_final_fc:
+            x = x.detach()
+        x = self.projection(x)
+
+        if not self.training:
+            if self.act is not None:
+                x = self.act(x)
+            # Performs fully convolutional inference.
+            if x.ndim == 5 and x.shape[1:4] > torch.Size([1, 1, 1]):
+                x = x.mean([1, 2, 3])
+
+        x = x.view(x.shape[0], -1)
+
+        return x

skp/models/rev_mvit/reversible_mvit.py

@@ -0,0 +1,696 @@
+import sys
+from functools import partial
+import torch
+from torch import nn
+from torch.autograd import Function as Function
+
+from .attention import MultiScaleAttention, attention_pool
+from .common import Mlp, TwoStreamFusion, drop_path
+from .utils import round_width
+
+
+class ReversibleMViT(nn.Module):
+    """
+    Reversible model builder. This builds the reversible transformer encoder
+    and allows reversible training.
+
+    Karttikeya Mangalam, Haoqi Fan, Yanghao Li, Chao-Yuan Wu, Bo Xiong,
+    Christoph Feichtenhofer, Jitendra Malik
+    "Reversible Vision Transformers"
+
+    https://openaccess.thecvf.com/content/CVPR2022/papers/Mangalam_Reversible_Vision_Transformers_CVPR_2022_paper.pdf
+    """
+
+    def __init__(self, config, model):
+        """
+        The `__init__` method of any subclass should also contain these
+            arguments.
+        Args:
+            cfg (CfgNode): model building configs, details are in the
+                comments of the config file.
+            model (nn.Module): parent MViT module this module forms
+                a reversible encoder in.
+        """
+
+        super().__init__()
+        self.cfg = config
+
+        embed_dim = self.cfg.MVIT.EMBED_DIM
+        depth = self.cfg.MVIT.DEPTH
+        num_heads = self.cfg.MVIT.NUM_HEADS
+        mlp_ratio = self.cfg.MVIT.MLP_RATIO
+        qkv_bias = self.cfg.MVIT.QKV_BIAS
+
+        drop_path_rate = self.cfg.MVIT.DROPPATH_RATE
+        self.dropout = config.MVIT.DROPOUT_RATE
+        self.pre_q_fusion = self.cfg.MVIT.REV.PRE_Q_FUSION
+        dpr = [
+            x.item() for x in torch.linspace(0, drop_path_rate, depth)
+        ]  # stochastic depth decay rule
+
+        input_size = model.patch_dims
+
+        self.layers = nn.ModuleList([])
+        self.no_custom_backward = False
+
+        if self.cfg.MVIT.NORM == "layernorm":
+            norm_layer = partial(nn.LayerNorm, eps=1e-6)
+        else:
+            raise NotImplementedError("Only supports layernorm.")
+
+        dim_mul, head_mul = torch.ones(depth + 1), torch.ones(depth + 1)
+        for i in range(len(self.cfg.MVIT.DIM_MUL)):
+            dim_mul[self.cfg.MVIT.DIM_MUL[i][0]] = self.cfg.MVIT.DIM_MUL[i][1]
+        for i in range(len(self.cfg.MVIT.HEAD_MUL)):
+            head_mul[self.cfg.MVIT.HEAD_MUL[i][0]] = self.cfg.MVIT.HEAD_MUL[i][
+                1
+            ]
+
+        pool_q = model.pool_q
+        pool_kv = model.pool_kv
+        stride_q = model.stride_q
+        stride_kv = model.stride_kv
+
+        for i in range(depth):
+
+            num_heads = round_width(num_heads, head_mul[i])
+
+            # Upsampling inside the MHPA, input to the Q-pooling block is lower C dimension
+            # This localizes the feature changes in a single block, making more computation reversible.
+            embed_dim = round_width(
+                embed_dim, dim_mul[i - 1] if i > 0 else 1.0, divisor=num_heads
+            )
+            dim_out = round_width(
+                embed_dim,
+                dim_mul[i],
+                divisor=round_width(num_heads, head_mul[i + 1]),
+            )
+
+            if i in self.cfg.MVIT.REV.BUFFER_LAYERS:
+                layer_type = StageTransitionBlock
+                input_mult = 2 if "concat" in self.pre_q_fusion else 1
+            else:
+                layer_type = ReversibleBlock
+                input_mult = 1
+
+            dimout_correction = (
+                2 if (input_mult == 2 and "concat" in self.pre_q_fusion) else 1
+            )
+
+            self.layers.append(
+                layer_type(
+                    dim=embed_dim
+                    * input_mult,  # added only for concat fusion before Qpooling layers
+                    input_size=input_size,
+                    dim_out=dim_out * input_mult // dimout_correction,
+                    num_heads=num_heads,
+                    cfg=self.cfg,
+                    mlp_ratio=mlp_ratio,
+                    qkv_bias=qkv_bias,
+                    drop_path=dpr[i],
+                    norm_layer=norm_layer,
+                    kernel_q=pool_q[i] if len(pool_q) > i else [],
+                    kernel_kv=pool_kv[i] if len(pool_kv) > i else [],
+                    stride_q=stride_q[i] if len(stride_q) > i else [],
+                    stride_kv=stride_kv[i] if len(stride_kv) > i else [],
+                    layer_id=i,
+                    pre_q_fusion=self.pre_q_fusion,
+                )
+            )
+            # F is the attention block
+            self.layers[-1].F.thw = input_size
+
+            if len(stride_q[i]) > 0:
+                input_size = [
+                    size // stride
+                    for size, stride in zip(input_size, stride_q[i])
+                ]
+
+        embed_dim = dim_out
+
+    @staticmethod
+    def vanilla_backward(h, layers, buffer):
+        """
+        Using rev layers without rev backpropagation. Debugging purposes only.
+        Activated with self.no_custom_backward.
+        """
+
+        # split into hidden states (h) and attention_output (a)
+        h, a = torch.chunk(h, 2, dim=-1)
+        for _, layer in enumerate(layers):
+            a, h = layer(a, h)
+
+        return torch.cat([a, h], dim=-1)
+
+    def forward(self, x):
+
+        # process the layers in a reversible stack and an irreversible stack.
+        stack = []
+        for l_i in range(len(self.layers)):
+            if isinstance(self.layers[l_i], StageTransitionBlock):
+                stack.append(("StageTransition", l_i))
+            else:
+                if len(stack) == 0 or stack[-1][0] == "StageTransition":
+                    stack.append(("Reversible", []))
+                stack[-1][1].append(l_i)
+
+        for layer_seq in stack:
+
+            if layer_seq[0] == "StageTransition":
+                x = self.layers[layer_seq[1]](x)
+
+            else:
+                x = torch.cat([x, x], dim=-1)
+
+                # no need for custom backprop in eval/model stat log
+                if not self.training or self.no_custom_backward:
+                    executing_fn = ReversibleMViT.vanilla_backward
+                else:
+                    executing_fn = RevBackProp.apply
+
+                x = executing_fn(
+                    x,
+                    self.layers[layer_seq[1][0] : layer_seq[1][-1] + 1],
+                    [],  # buffer activations
+                )
+
+        # Apply dropout
+        x = nn.functional.dropout(x, p=self.dropout, training=self.training)
+
+        return x
+
+
+class RevBackProp(Function):
+    """
+    Custom Backpropagation function to allow (A) flusing memory in foward
+    and (B) activation recomputation reversibly in backward for gradient calculation.
+
+    Inspired by https://github.com/RobinBruegger/RevTorch/blob/master/revtorch/revtorch.py
+    """
+
+    @staticmethod
+    def forward(
+        ctx,
+        x,
+        layers,
+        buffer_layers,  # List of layer ids for int activation to buffer
+    ):
+        """
+        Reversible Forward pass. Any intermediate activations from `buffer_layers` are
+        cached in ctx for forward pass. This is not necessary for standard usecases.
+        Each reversible layer implements its own forward pass logic.
+        """
+        buffer_layers.sort()
+
+        X_1, X_2 = torch.chunk(x, 2, dim=-1)
+
+        intermediate = []
+
+        for layer in layers:
+
+            X_1, X_2 = layer(X_1, X_2)
+
+            if layer.layer_id in buffer_layers:
+                intermediate.extend([X_1.detach(), X_2.detach()])
+
+        if len(buffer_layers) == 0:
+            all_tensors = [X_1.detach(), X_2.detach()]
+        else:
+            intermediate = [torch.LongTensor(buffer_layers), *intermediate]
+            all_tensors = [X_1.detach(), X_2.detach(), *intermediate]
+
+        ctx.save_for_backward(*all_tensors)
+        ctx.layers = layers
+
+        return torch.cat([X_1, X_2], dim=-1)
+
+    @staticmethod
+    def backward(ctx, dx):
+        """
+        Reversible Backward pass. Any intermediate activations from `buffer_layers` are
+        recovered from ctx. Each layer implements its own loic for backward pass (both
+        activation recomputation and grad calculation).
+        """
+        dX_1, dX_2 = torch.chunk(dx, 2, dim=-1)
+
+        # retrieve params from ctx for backward
+        X_1, X_2, *int_tensors = ctx.saved_tensors
+
+        # no buffering
+        if len(int_tensors) != 0:
+            buffer_layers = int_tensors[0].tolist()
+
+        else:
+            buffer_layers = []
+
+        layers = ctx.layers
+
+        for _, layer in enumerate(layers[::-1]):
+
+            if layer.layer_id in buffer_layers:
+
+                X_1, X_2, dX_1, dX_2 = layer.backward_pass(
+                    Y_1=int_tensors[
+                        buffer_layers.index(layer.layer_id) * 2 + 1
+                    ],
+                    Y_2=int_tensors[
+                        buffer_layers.index(layer.layer_id) * 2 + 2
+                    ],
+                    dY_1=dX_1,
+                    dY_2=dX_2,
+                )
+
+            else:
+
+                X_1, X_2, dX_1, dX_2 = layer.backward_pass(
+                    Y_1=X_1,
+                    Y_2=X_2,
+                    dY_1=dX_1,
+                    dY_2=dX_2,
+                )
+
+        dx = torch.cat([dX_1, dX_2], dim=-1)
+
+        del int_tensors
+        del dX_1, dX_2, X_1, X_2
+
+        return dx, None, None
+
+
+class StageTransitionBlock(nn.Module):
+    """
+    Blocks for changing the feature dimensions in MViT (using Q-pooling).
+    See Section 3.3.1 in paper for details.
+    """
+
+    def __init__(
+        self,
+        dim,
+        input_size,
+        dim_out,
+        num_heads,
+        mlp_ratio,
+        qkv_bias,
+        drop_path,
+        kernel_q,
+        kernel_kv,
+        stride_q,
+        stride_kv,
+        cfg,
+        norm_layer=nn.LayerNorm,
+        pre_q_fusion=None,
+        layer_id=0,
+    ):
+        """
+        Uses the same structure of F and G functions as Reversible Block except
+        without using reversible forward (and backward) pass.
+        """
+        super().__init__()
+
+        self.drop_path_rate = drop_path
+
+        embed_dim = dim
+
+        self.F = AttentionSubBlock(
+            dim=embed_dim,
+            input_size=input_size,
+            num_heads=num_heads,
+            cfg=cfg,
+            dim_out=dim_out,
+            kernel_q=kernel_q,
+            kernel_kv=kernel_kv,
+            stride_q=stride_q,
+            stride_kv=stride_kv,
+            norm_layer=norm_layer,
+        )
+
+        self.G = MLPSubblock(
+            dim=dim_out,
+            mlp_ratio=mlp_ratio,
+            norm_layer=norm_layer,
+        )
+
+        self.layer_id = layer_id
+
+        self.is_proj = False
+        self.has_cls_embed = cfg.MVIT.CLS_EMBED_ON
+
+        self.is_conv = False
+        self.pool_first = cfg.MVIT.POOL_FIRST
+        self.mode = cfg.MVIT.MODE
+        self.pre_q_fuse = TwoStreamFusion(pre_q_fusion, dim=dim)
+
+        if cfg.MVIT.REV.RES_PATH == "max":
+            self.res_conv = False
+            self.pool_skip = nn.MaxPool3d(
+                # self.attention.attn.pool_q.kernel_size,
+                [s + 1 if s > 1 else s for s in self.F.attn.pool_q.stride],
+                self.F.attn.pool_q.stride,
+                [int(k // 2) for k in self.F.attn.pool_q.stride],
+                # self.attention.attn.pool_q.padding,
+                ceil_mode=False,
+            )
+
+        elif cfg.MVIT.REV.RES_PATH == "conv":
+            self.res_conv = True
+        else:
+            raise NotImplementedError
+
+        # Add a linear projection in residual branch
+        if embed_dim != dim_out:
+            self.is_proj = True
+            self.res_proj = nn.Linear(embed_dim, dim_out, bias=True)
+
+    def forward(
+        self,
+        x,
+    ):
+        """
+        Forward logic is similar to MultiScaleBlock with Q-pooling.
+        """
+        x = self.pre_q_fuse(x)
+
+        # fork tensor for residual connections
+        x_res = x
+
+        # This uses conv to pool the residual hidden features
+        # but done before pooling only if not pool_first
+        if self.is_proj and not self.pool_first:
+            x_res = self.res_proj(x_res)
+
+        if self.res_conv:
+
+            # Pooling the hidden features with the same conv as Q
+            N, L, C = x_res.shape
+
+            # This handling is the same as that of q in MultiScaleAttention
+            if self.mode == "conv_unshared":
+                fold_dim = 1
+            else:
+                fold_dim = self.F.attn.num_heads
+
+            # Output is (B, N, L, C)
+            x_res = x_res.reshape(N, L, fold_dim, C // fold_dim).permute(
+                0, 2, 1, 3
+            )
+
+            x_res, _ = attention_pool(
+                x_res,
+                self.F.attn.pool_q,
+                # thw_shape = self.attention.attn.thw,
+                thw_shape=self.F.thw,
+                has_cls_embed=self.has_cls_embed,
+                norm=self.F.attn.norm_q
+                if hasattr(self.F.attn, "norm_q")
+                else None,
+            )
+            x_res = x_res.permute(0, 2, 1, 3).reshape(N, x_res.shape[2], C)
+
+        else:
+            # Pooling the hidden features with max op
+            x_res, _ = attention_pool(
+                x_res,
+                self.pool_skip,
+                thw_shape=self.F.attn.thw,
+                has_cls_embed=self.has_cls_embed,
+            )
+
+        # If pool_first then project to higher dim now
+        if self.is_proj and self.pool_first:
+            x_res = self.res_proj(x_res)
+
+        x = self.F(x)
+        x = x_res + x
+        x = x + self.G(x)
+
+        x = drop_path(x, drop_prob=self.drop_path_rate, training=self.training)
+
+        return x
+
+
+class ReversibleBlock(nn.Module):
+    """
+    Reversible Blocks for Reversible Vision Transformer and also
+    for state-preserving blocks in Reversible MViT. See Section
+    3.3.2 in paper for details.
+    """
+
+    def __init__(
+        self,
+        dim,
+        input_size,
+        dim_out,
+        num_heads,
+        mlp_ratio,
+        qkv_bias,
+        drop_path,
+        kernel_q,
+        kernel_kv,
+        stride_q,
+        stride_kv,
+        cfg,
+        norm_layer=nn.LayerNorm,
+        layer_id=0,
+        **kwargs
+    ):
+        """
+        Block is composed entirely of function F (Attention
+        sub-block) and G (MLP sub-block) including layernorm.
+        """
+        super().__init__()
+
+        self.drop_path_rate = drop_path
+
+        self.F = AttentionSubBlock(
+            dim=dim,
+            input_size=input_size,
+            num_heads=num_heads,
+            cfg=cfg,
+            dim_out=dim_out,
+            kernel_q=kernel_q,
+            kernel_kv=kernel_kv,
+            stride_q=stride_q,
+            stride_kv=stride_kv,
+            norm_layer=norm_layer,
+        )
+
+        self.G = MLPSubblock(
+            dim=dim,
+            mlp_ratio=mlp_ratio,
+            norm_layer=norm_layer,
+        )
+
+        self.layer_id = layer_id
+
+        self.seeds = {}
+
+    def seed_cuda(self, key):
+        """
+        Fix seeds to allow for stochastic elements such as
+        dropout to be reproduced exactly in activation
+        recomputation in the backward pass.
+        """
+
+        # randomize seeds
+        # use cuda generator if available
+        if (
+            hasattr(torch.cuda, "default_generators")
+            and len(torch.cuda.default_generators) > 0
+        ):
+            # GPU
+            device_idx = torch.cuda.current_device()
+            seed = torch.cuda.default_generators[device_idx].seed()
+        else:
+            # CPU
+            seed = int(torch.seed() % sys.maxsize)
+
+        self.seeds[key] = seed
+        torch.manual_seed(self.seeds[key])
+
+    def forward(self, X_1, X_2):
+        """
+        forward pass equations:
+        Y_1 = X_1 + Attention(X_2), F = Attention
+        Y_2 = X_2 + MLP(Y_1), G = MLP
+        """
+
+        self.seed_cuda("attn")
+        # Y_1 : attn_output
+        f_X_2 = self.F(X_2)
+
+        self.seed_cuda("droppath")
+        f_X_2_dropped = drop_path(
+            f_X_2, drop_prob=self.drop_path_rate, training=self.training
+        )
+
+        # Y_1 = X_1 + f(X_2)
+        Y_1 = X_1 + f_X_2_dropped
+
+        # free memory
+        del X_1
+
+        self.seed_cuda("FFN")
+        g_Y_1 = self.G(Y_1)
+
+        torch.manual_seed(self.seeds["droppath"])
+        g_Y_1_dropped = drop_path(
+            g_Y_1, drop_prob=self.drop_path_rate, training=self.training
+        )
+
+        # Y_2 = X_2 + g(Y_1)
+        Y_2 = X_2 + g_Y_1_dropped
+
+        del X_2
+
+        return Y_1, Y_2
+
+    def backward_pass(
+        self,
+        Y_1,
+        Y_2,
+        dY_1,
+        dY_2,
+    ):
+        """
+        equation for activation recomputation:
+        X_2 = Y_2 - G(Y_1), G = MLP
+        X_1 = Y_1 - F(X_2), F = Attention
+        """
+
+        # temporarily record intermediate activation for G
+        # and use them for gradient calculcation of G
+        with torch.enable_grad():
+
+            Y_1.requires_grad = True
+
+            torch.manual_seed(self.seeds["FFN"])
+            g_Y_1 = self.G(Y_1)
+
+            torch.manual_seed(self.seeds["droppath"])
+            g_Y_1 = drop_path(
+                g_Y_1, drop_prob=self.drop_path_rate, training=self.training
+            )
+
+            g_Y_1.backward(dY_2, retain_graph=True)
+
+        # activation recomputation is by design and not part of
+        # the computation graph in forward pass.
+        with torch.no_grad():
+
+            X_2 = Y_2 - g_Y_1
+            del g_Y_1
+
+            dY_1 = dY_1 + Y_1.grad
+            Y_1.grad = None
+
+        # record F activations and calc gradients on F
+        with torch.enable_grad():
+            X_2.requires_grad = True
+
+            torch.manual_seed(self.seeds["attn"])
+            f_X_2 = self.F(X_2)
+
+            torch.manual_seed(self.seeds["droppath"])
+            f_X_2 = drop_path(
+                f_X_2, drop_prob=self.drop_path_rate, training=self.training
+            )
+
+            f_X_2.backward(dY_1, retain_graph=True)
+
+        # propagate reverse computed acitvations at the start of
+        # the previou block for backprop.s
+        with torch.no_grad():
+
+            X_1 = Y_1 - f_X_2
+
+            del f_X_2, Y_1
+            dY_2 = dY_2 + X_2.grad
+
+            X_2.grad = None
+            X_2 = X_2.detach()
+
+        return X_1, X_2, dY_1, dY_2
+
+
+class MLPSubblock(nn.Module):
+    """
+    This creates the function G such that the entire block can be
+    expressed as F(G(X)). Includes pre-LayerNorm.
+    """
+
+    def __init__(
+        self,
+        dim,
+        mlp_ratio,
+        norm_layer=nn.LayerNorm,
+    ):
+
+        super().__init__()
+        self.norm = norm_layer(dim, eps=1e-6, elementwise_affine=True)
+
+        mlp_hidden_dim = int(dim * mlp_ratio)
+
+        self.mlp = Mlp(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            act_layer=nn.GELU,
+        )
+
+    def forward(self, x):
+        return self.mlp(self.norm(x))
+
+
+class AttentionSubBlock(nn.Module):
+    """
+    This creates the function F such that the entire block can be
+    expressed as F(G(X)). Includes pre-LayerNorm.
+    """
+
+    def __init__(
+        self,
+        dim,
+        input_size,
+        num_heads,
+        cfg,
+        dim_out=None,
+        kernel_q=(1, 1, 1),
+        kernel_kv=(1, 1, 1),
+        stride_q=(1, 1, 1),
+        stride_kv=(1, 1, 1),
+        norm_layer=nn.LayerNorm,
+    ):
+
+        super().__init__()
+        self.norm = norm_layer(dim, eps=1e-6, elementwise_affine=True)
+
+        # This will be set externally during init
+        self.thw = None
+
+        # the actual attention details are the same as Multiscale
+        # attention for MViTv2 (with channel up=projection inside block)
+        # can also implement no upprojection attention for vanilla ViT
+        self.attn = MultiScaleAttention(
+            dim,
+            dim_out,
+            input_size=input_size,
+            num_heads=num_heads,
+            kernel_q=kernel_q,
+            kernel_kv=kernel_kv,
+            stride_q=stride_q,
+            stride_kv=stride_kv,
+            norm_layer=norm_layer,
+            drop_rate=cfg.MVIT.DROPOUT_RATE,
+            qkv_bias=cfg.MVIT.QKV_BIAS,
+            has_cls_embed=cfg.MVIT.CLS_EMBED_ON,
+            mode=cfg.MVIT.MODE,
+            pool_first=cfg.MVIT.POOL_FIRST,
+            rel_pos_spatial=cfg.MVIT.REL_POS_SPATIAL,
+            rel_pos_temporal=cfg.MVIT.REL_POS_TEMPORAL,
+            rel_pos_zero_init=cfg.MVIT.REL_POS_ZERO_INIT,
+            residual_pooling=cfg.MVIT.RESIDUAL_POOLING,
+            separate_qkv=cfg.MVIT.SEPARATE_QKV,
+        )
+
+    def forward(self, x):
+        out, _ = self.attn(self.norm(x), self.thw)
+        return out

skp/models/rev_mvit/stem_helper.py

@@ -0,0 +1,325 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+
+"""ResNe(X)t 3D stem helper."""
+
+import torch
+import torch.nn as nn
+
+
+def get_stem_func(name):
+    """
+    Retrieves the stem module by name.
+    """
+    trans_funcs = {"x3d_stem": X3DStem, "basic_stem": ResNetBasicStem}
+    assert (
+        name in trans_funcs.keys()
+    ), "Transformation function '{}' not supported".format(name)
+    return trans_funcs[name]
+
+
+class VideoModelStem(nn.Module):
+    """
+    Video 3D stem module. Provides stem operations of Conv, BN, ReLU, MaxPool
+    on input data tensor for one or multiple pathways.
+    """
+
+    def __init__(
+        self,
+        dim_in,
+        dim_out,
+        kernel,
+        stride,
+        padding,
+        inplace_relu=True,
+        eps=1e-5,
+        bn_mmt=0.1,
+        norm_module=nn.BatchNorm3d,
+        stem_func_name="basic_stem",
+    ):
+        """
+        The `__init__` method of any subclass should also contain these
+        arguments. List size of 1 for single pathway models (C2D, I3D, Slow
+        and etc), list size of 2 for two pathway models (SlowFast).
+
+        Args:
+            dim_in (list): the list of channel dimensions of the inputs.
+            dim_out (list): the output dimension of the convolution in the stem
+                layer.
+            kernel (list): the kernels' size of the convolutions in the stem
+                layers. Temporal kernel size, height kernel size, width kernel
+                size in order.
+            stride (list): the stride sizes of the convolutions in the stem
+                layer. Temporal kernel stride, height kernel size, width kernel
+                size in order.
+            padding (list): the paddings' sizes of the convolutions in the stem
+                layer. Temporal padding size, height padding size, width padding
+                size in order.
+            inplace_relu (bool): calculate the relu on the original input
+                without allocating new memory.
+            eps (float): epsilon for batch norm.
+            bn_mmt (float): momentum for batch norm. Noted that BN momentum in
+                PyTorch = 1 - BN momentum in Caffe2.
+            norm_module (nn.Module): nn.Module for the normalization layer. The
+                default is nn.BatchNorm3d.
+            stem_func_name (string): name of the the stem function applied on
+                input to the network.
+        """
+        super(VideoModelStem, self).__init__()
+
+        assert (
+            len(
+                {
+                    len(dim_in),
+                    len(dim_out),
+                    len(kernel),
+                    len(stride),
+                    len(padding),
+                }
+            )
+            == 1
+        ), "Input pathway dimensions are not consistent. {} {} {} {} {}".format(
+            len(dim_in),
+            len(dim_out),
+            len(kernel),
+            len(stride),
+            len(padding),
+        )
+
+        self.num_pathways = len(dim_in)
+        self.kernel = kernel
+        self.stride = stride
+        self.padding = padding
+        self.inplace_relu = inplace_relu
+        self.eps = eps
+        self.bn_mmt = bn_mmt
+        # Construct the stem layer.
+        self._construct_stem(dim_in, dim_out, norm_module, stem_func_name)
+
+    def _construct_stem(self, dim_in, dim_out, norm_module, stem_func_name):
+        trans_func = get_stem_func(stem_func_name)
+
+        for pathway in range(len(dim_in)):
+            stem = trans_func(
+                dim_in[pathway],
+                dim_out[pathway],
+                self.kernel[pathway],
+                self.stride[pathway],
+                self.padding[pathway],
+                self.inplace_relu,
+                self.eps,
+                self.bn_mmt,
+                norm_module,
+            )
+            self.add_module("pathway{}_stem".format(pathway), stem)
+
+    def forward(self, x):
+        assert (
+            len(x) == self.num_pathways
+        ), "Input tensor does not contain {} pathway".format(self.num_pathways)
+        # use a new list, don't modify in-place the x list, which is bad for activation checkpointing.
+        y = []
+        for pathway in range(len(x)):
+            m = getattr(self, "pathway{}_stem".format(pathway))
+            y.append(m(x[pathway]))
+        return y
+
+
+class ResNetBasicStem(nn.Module):
+    """
+    ResNe(X)t 3D stem module.
+    Performs spatiotemporal Convolution, BN, and Relu following by a
+        spatiotemporal pooling.
+    """
+
+    def __init__(
+        self,
+        dim_in,
+        dim_out,
+        kernel,
+        stride,
+        padding,
+        inplace_relu=True,
+        eps=1e-5,
+        bn_mmt=0.1,
+        norm_module=nn.BatchNorm3d,
+    ):
+        """
+        The `__init__` method of any subclass should also contain these arguments.
+
+        Args:
+            dim_in (int): the channel dimension of the input. Normally 3 is used
+                for rgb input, and 2 or 3 is used for optical flow input.
+            dim_out (int): the output dimension of the convolution in the stem
+                layer.
+            kernel (list): the kernel size of the convolution in the stem layer.
+                temporal kernel size, height kernel size, width kernel size in
+                order.
+            stride (list): the stride size of the convolution in the stem layer.
+                temporal kernel stride, height kernel size, width kernel size in
+                order.
+            padding (int): the padding size of the convolution in the stem
+                layer, temporal padding size, height padding size, width
+                padding size in order.
+            inplace_relu (bool): calculate the relu on the original input
+                without allocating new memory.
+            eps (float): epsilon for batch norm.
+            bn_mmt (float): momentum for batch norm. Noted that BN momentum in
+                PyTorch = 1 - BN momentum in Caffe2.
+            norm_module (nn.Module): nn.Module for the normalization layer. The
+                default is nn.BatchNorm3d.
+        """
+        super(ResNetBasicStem, self).__init__()
+        self.kernel = kernel
+        self.stride = stride
+        self.padding = padding
+        self.inplace_relu = inplace_relu
+        self.eps = eps
+        self.bn_mmt = bn_mmt
+        # Construct the stem layer.
+        self._construct_stem(dim_in, dim_out, norm_module)
+
+    def _construct_stem(self, dim_in, dim_out, norm_module):
+        self.conv = nn.Conv3d(
+            dim_in,
+            dim_out,
+            self.kernel,
+            stride=self.stride,
+            padding=self.padding,
+            bias=False,
+        )
+        self.bn = norm_module(
+            num_features=dim_out, eps=self.eps, momentum=self.bn_mmt
+        )
+        self.relu = nn.ReLU(self.inplace_relu)
+        self.pool_layer = nn.MaxPool3d(
+            kernel_size=[1, 3, 3], stride=[1, 2, 2], padding=[0, 1, 1]
+        )
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        x = self.relu(x)
+        x = self.pool_layer(x)
+        return x
+
+
+class X3DStem(nn.Module):
+    """
+    X3D's 3D stem module.
+    Performs a spatial followed by a depthwise temporal Convolution, BN, and Relu following by a
+        spatiotemporal pooling.
+    """
+
+    def __init__(
+        self,
+        dim_in,
+        dim_out,
+        kernel,
+        stride,
+        padding,
+        inplace_relu=True,
+        eps=1e-5,
+        bn_mmt=0.1,
+        norm_module=nn.BatchNorm3d,
+    ):
+        """
+        The `__init__` method of any subclass should also contain these arguments.
+
+        Args:
+            dim_in (int): the channel dimension of the input. Normally 3 is used
+                for rgb input, and 2 or 3 is used for optical flow input.
+            dim_out (int): the output dimension of the convolution in the stem
+                layer.
+            kernel (list): the kernel size of the convolution in the stem layer.
+                temporal kernel size, height kernel size, width kernel size in
+                order.
+            stride (list): the stride size of the convolution in the stem layer.
+                temporal kernel stride, height kernel size, width kernel size in
+                order.
+            padding (int): the padding size of the convolution in the stem
+                layer, temporal padding size, height padding size, width
+                padding size in order.
+            inplace_relu (bool): calculate the relu on the original input
+                without allocating new memory.
+            eps (float): epsilon for batch norm.
+            bn_mmt (float): momentum for batch norm. Noted that BN momentum in
+                PyTorch = 1 - BN momentum in Caffe2.
+            norm_module (nn.Module): nn.Module for the normalization layer. The
+                default is nn.BatchNorm3d.
+        """
+        super(X3DStem, self).__init__()
+        self.kernel = kernel
+        self.stride = stride
+        self.padding = padding
+        self.inplace_relu = inplace_relu
+        self.eps = eps
+        self.bn_mmt = bn_mmt
+        # Construct the stem layer.
+        self._construct_stem(dim_in, dim_out, norm_module)
+
+    def _construct_stem(self, dim_in, dim_out, norm_module):
+        self.conv_xy = nn.Conv3d(
+            dim_in,
+            dim_out,
+            kernel_size=(1, self.kernel[1], self.kernel[2]),
+            stride=(1, self.stride[1], self.stride[2]),
+            padding=(0, self.padding[1], self.padding[2]),
+            bias=False,
+        )
+        self.conv = nn.Conv3d(
+            dim_out,
+            dim_out,
+            kernel_size=(self.kernel[0], 1, 1),
+            stride=(self.stride[0], 1, 1),
+            padding=(self.padding[0], 0, 0),
+            bias=False,
+            groups=dim_out,
+        )
+
+        self.bn = norm_module(
+            num_features=dim_out, eps=self.eps, momentum=self.bn_mmt
+        )
+        self.relu = nn.ReLU(self.inplace_relu)
+
+    def forward(self, x):
+        x = self.conv_xy(x)
+        x = self.conv(x)
+        x = self.bn(x)
+        x = self.relu(x)
+        return x
+
+
+class PatchEmbed(nn.Module):
+    """
+    PatchEmbed.
+    """
+
+    def __init__(
+        self,
+        dim_in=3,
+        dim_out=768,
+        kernel=(1, 16, 16),
+        stride=(1, 4, 4),
+        padding=(1, 7, 7),
+        conv_2d=False,
+    ):
+        super().__init__()
+        if conv_2d:
+            conv = nn.Conv2d
+        else:
+            conv = nn.Conv3d
+        self.proj = conv(
+            dim_in,
+            dim_out,
+            kernel_size=kernel,
+            stride=stride,
+            padding=padding,
+        )
+
+    def forward(self, x, keep_spatial=False):
+        x = self.proj(x)
+        if keep_spatial:
+            return x, x.shape
+        # B C (T) H W -> B (T)HW C
+        return x.flatten(2).transpose(1, 2), x.shape

skp/models/rev_mvit/utils.py

@@ -0,0 +1,221 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+
+import numpy as np
+import torch
+
+
+def round_width(width, multiplier, min_width=1, divisor=1, verbose=False):
+    if not multiplier:
+        return width
+    width *= multiplier
+    min_width = min_width or divisor
+    if verbose:
+        print(f"min width {min_width}")
+        print(f"width {width} divisor {divisor}")
+        print(f"other {int(width + divisor / 2) // divisor * divisor}")
+
+    width_out = max(min_width, int(width + divisor / 2) // divisor * divisor)
+    if width_out < 0.9 * width:
+        width_out += divisor
+    return int(width_out)
+
+
+def validate_checkpoint_wrapper_import(checkpoint_wrapper):
+    """
+    Check if checkpoint_wrapper is imported.
+    """
+    if checkpoint_wrapper is None:
+        raise ImportError("Please install fairscale.")
+
+
+def get_gkern(kernlen, std):
+    """Returns a 2D Gaussian kernel array."""
+
+    def _gaussian_fn(kernlen, std):
+        n = torch.arange(0, kernlen).float()
+        n -= n.mean()
+        n /= std
+        w = torch.exp(-0.5 * n**2)
+        return w
+
+    gkern1d = _gaussian_fn(kernlen, std)
+    gkern2d = torch.outer(gkern1d, gkern1d)
+    return gkern2d / gkern2d.sum()
+
+
+# --------------------------------------------------------
+# 2D sine-cosine position embedding
+# References:
+# Transformer: https://github.com/tensorflow/models/blob/master/official/nlp/transformer/model_utils.py
+# MoCo v3: https://github.com/facebookresearch/moco-v3
+# --------------------------------------------------------
+def get_3d_sincos_pos_embed(embed_dim, grid_size, t_size, cls_token=False):
+    """
+    grid_size: int of the grid height and width
+    t_size: int of the temporal size
+    return:
+    pos_embed: [t_size*grid_size*grid_size, embed_dim] or [1+t_size*grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+    """
+    assert embed_dim % 4 == 0
+    embed_dim_spatial = embed_dim // 4 * 3
+    embed_dim_temporal = embed_dim // 4
+
+    # spatial
+    grid_h = np.arange(grid_size, dtype=np.float32)
+    grid_w = np.arange(grid_size, dtype=np.float32)
+    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = np.stack(grid, axis=0)
+
+    grid = grid.reshape([2, 1, grid_size, grid_size])
+    pos_embed_spatial = get_2d_sincos_pos_embed_from_grid(
+        embed_dim_spatial, grid
+    )
+
+    # temporal
+    grid_t = np.arange(t_size, dtype=np.float32)
+    pos_embed_temporal = get_1d_sincos_pos_embed_from_grid(
+        embed_dim_temporal, grid_t
+    )
+
+    # concate: [T, H, W] order
+    pos_embed_temporal = pos_embed_temporal[:, np.newaxis, :]
+    pos_embed_temporal = np.repeat(
+        pos_embed_temporal, grid_size**2, axis=1
+    )  # [T, H*W, D // 4]
+    pos_embed_spatial = pos_embed_spatial[np.newaxis, :, :]
+    pos_embed_spatial = np.repeat(
+        pos_embed_spatial, t_size, axis=0
+    )  # [T, H*W, D // 4 * 3]
+
+    pos_embed = np.concatenate([pos_embed_temporal, pos_embed_spatial], axis=-1)
+    pos_embed = pos_embed.reshape([-1, embed_dim])  # [T*H*W, D]
+
+    if cls_token:
+        pos_embed = np.concatenate(
+            [np.zeros([1, embed_dim]), pos_embed], axis=0
+        )
+    return pos_embed
+
+
+def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False):
+    """
+    grid_size: int of the grid height and width
+    return:
+    pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+    """
+    grid_h = np.arange(grid_size, dtype=np.float32)
+    grid_w = np.arange(grid_size, dtype=np.float32)
+    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = np.stack(grid, axis=0)
+
+    grid = grid.reshape([2, 1, grid_size, grid_size])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    if cls_token:
+        pos_embed = np.concatenate(
+            [np.zeros([1, embed_dim]), pos_embed], axis=0
+        )
+    return pos_embed
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    assert embed_dim % 2 == 0
+
+    # use half of dimensions to encode grid_h
+    emb_h = get_1d_sincos_pos_embed_from_grid(
+        embed_dim // 2, grid[0]
+    )  # (H*W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid(
+        embed_dim // 2, grid[1]
+    )  # (H*W, D/2)
+
+    emb = np.concatenate([emb_h, emb_w], axis=1)  # (H*W, D)
+    return emb
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """
+    embed_dim: output dimension for each position
+    pos: a list of positions to be encoded: size (M,)
+    out: (M, D)
+    """
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=np.float)
+    omega /= embed_dim / 2.0
+    omega = 1.0 / 10000**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = np.einsum("m,d->md", pos, omega)  # (M, D/2), outer product
+
+    emb_sin = np.sin(out)  # (M, D/2)
+    emb_cos = np.cos(out)  # (M, D/2)
+
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+    return emb
+
+
+# --------------------------------------------------------
+# Interpolate position embeddings for high-resolution
+# References:
+# DeiT: https://github.com/facebookresearch/deit
+# --------------------------------------------------------
+def interpolate_pos_embed(model, checkpoint_model):
+    if "pos_embed" in checkpoint_model:
+        pos_embed_checkpoint = checkpoint_model["pos_embed"]
+        embedding_size = pos_embed_checkpoint.shape[-1]
+        num_patches = model.patch_embed.num_patches
+        num_extra_tokens = model.pos_embed.shape[-2] - num_patches
+        # height (== width) for the checkpoint position embedding
+        orig_size = int(
+            (pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5
+        )
+        # height (== width) for the new position embedding
+        new_size = int(num_patches**0.5)
+        # class_token and dist_token are kept unchanged
+        if orig_size != new_size:
+            print(
+                "Position interpolate from %dx%d to %dx%d"
+                % (orig_size, orig_size, new_size, new_size)
+            )
+            extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
+            # only the position tokens are interpolated
+            pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
+            pos_tokens = pos_tokens.reshape(
+                -1, orig_size, orig_size, embedding_size
+            ).permute(0, 3, 1, 2)
+            pos_tokens = torch.nn.functional.interpolate(
+                pos_tokens,
+                size=(new_size, new_size),
+                mode="bicubic",
+                align_corners=False,
+            )
+            pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
+            new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
+            checkpoint_model["pos_embed"] = new_pos_embed
+
+
+def calc_mvit_feature_geometry(cfg):
+    feat_size = [
+        [
+            cfg.DATA.NUM_FRAMES // cfg.MVIT.PATCH_STRIDE[0]
+            if len(cfg.MVIT.PATCH_STRIDE) > 2
+            else 1,
+            cfg.DATA.TRAIN_CROP_SIZE // cfg.MVIT.PATCH_STRIDE[-2],
+            cfg.DATA.TRAIN_CROP_SIZE // cfg.MVIT.PATCH_STRIDE[-1],
+        ]
+        for i in range(cfg.MVIT.DEPTH)
+    ]
+    feat_stride = [
+        [
+            cfg.MVIT.PATCH_STRIDE[0] if len(cfg.MVIT.PATCH_STRIDE) > 2 else 1,
+            cfg.MVIT.PATCH_STRIDE[-2],
+            cfg.MVIT.PATCH_STRIDE[-1],
+        ]
+        for i in range(cfg.MVIT.DEPTH)
+    ]
+    for _, x in enumerate(cfg.MVIT.POOL_Q_STRIDE):
+        for i in range(cfg.MVIT.DEPTH):
+            if i >= x[0]:
+                for j in range(len(feat_size[i])):
+                    feat_size[i][j] = feat_size[i][j] // x[j + 1]
+                    feat_stride[i][j] = feat_stride[i][j] * x[j + 1]
+    return feat_size, feat_stride