brainscore-vision 2.1__py3-none-any.whl

brainscore_vision/models/temporal_model_mmaction2/__init__.py

@@ -0,0 +1,23 @@
+from brainscore_vision import model_registry
+from brainscore_vision.model_helpers.brain_transformation import ModelCommitment
+from brainscore_vision.model_helpers.activations.temporal.utils import get_specified_layers
+from brainscore_vision.model_interface import BrainModel
+from . import model
+
+
+def commit_model(identifier):
+    activations_model=model.get_model(identifier)
+    layers=get_specified_layers(activations_model)
+    return ModelCommitment(identifier=identifier, activations_model=activations_model, layers=layers)
+
+
+model_registry["I3D"] = lambda: commit_model("I3D")
+model_registry["I3D-nonlocal"] = lambda: commit_model("I3D-nonlocal")
+model_registry["SlowFast"] = lambda: commit_model("SlowFast")
+model_registry["X3D"] = lambda: commit_model("X3D")
+model_registry["TimeSformer"] = lambda: commit_model("TimeSformer")
+model_registry["VideoSwin-B"] = lambda: commit_model("VideoSwin-B")
+model_registry["VideoSwin-L"] = lambda: commit_model("VideoSwin-L")
+model_registry["UniFormer-V1"] = lambda: commit_model("UniFormer-V1")
+model_registry["UniFormer-V2-B"] = lambda: commit_model("UniFormer-V2-B")
+model_registry["UniFormer-V2-L"] = lambda: commit_model("UniFormer-V2-L")

brainscore_vision/models/temporal_model_mmaction2/mmaction2.csv

@@ -0,0 +1,24 @@
+name,config,checkpoint,Kinetics400-top1,Kinetics400-top5,FLOPs,params,note
+I3D,https://github.com/open-mmlab/mmaction2/blob/main/configs/recognition/i3d/i3d_imagenet-pretrained-r50_8xb8-32x2x1-100e_kinetics400-rgb.py,https://download.openmmlab.com/mmaction/v1.0/recognition/i3d/i3d_imagenet-pretrained-r50_8xb8-32x2x1-100e_kinetics400-rgb/i3d_imagenet-pretrained-r50_8xb8-32x2x1-100e_kinetics400-rgb_20220812-e213c223.pth,73.47,91.27,43.5G,28.0M,
+I3D-nonlocal,https://github.com/open-mmlab/mmaction2/blob/main/configs/recognition/i3d/i3d_imagenet-pretrained-r50-nl-dot-product_8xb8-32x2x1-100e_kinetics400-rgb.py,https://download.openmmlab.com/mmaction/v1.0/recognition/i3d/i3d_imagenet-pretrained-r50-nl-dot-product_8xb8-32x2x1-100e_kinetics400-rgb/i3d_imagenet-pretrained-r50-nl-dot-product_8xb8-32x2x1-100e_kinetics400-rgb_20220812-8e1f2148.pth,,,,,
+TSM,https://github.com/open-mmlab/mmaction2/blob/main/configs/recognition/tsm/tsm_imagenet-pretrained-r50_8xb16-1x1x8-50e_kinetics400-rgb.py,https://download.openmmlab.com/mmaction/v1.0/recognition/tsm/tsm_imagenet-pretrained-r50_8xb16-1x1x8-50e_kinetics400-rgb/tsm_imagenet-pretrained-r50_8xb16-1x1x8-50e_kinetics400-rgb_20220831-64d69186.pth,75.12,91.55,65.75G,23.87M,
+SlowFast,https://github.com/open-mmlab/mmaction2/blob/main/configs/recognition/slowfast/slowfast_r50_8xb8-8x8x1-256e_kinetics400-rgb.py,https://download.openmmlab.com/mmaction/v1.0/recognition/slowfast/slowfast_r50_8xb8-8x8x1-256e_kinetics400-rgb/slowfast_r50_8xb8-8x8x1-256e_kinetics400-rgb_20220818-1cb6dfc8.pth,,,,,
+X3D,https://github.com/open-mmlab/mmaction2/blob/main/configs/recognition/x3d/x3d_m_16x5x1_facebook-kinetics400-rgb.py,https://download.openmmlab.com/mmaction/v1.0/recognition/x3d/facebook/x3d_m_16x5x1_facebook-kinetics400-rgb_20201027-3f42382a.pth,,,,,
+TimeSformer,https://github.com/open-mmlab/mmaction2/blob/main/configs/recognition/timesformer/timesformer_divST_8xb8-8x32x1-15e_kinetics400-rgb.py,https://download.openmmlab.com/mmaction/v1.0/recognition/timesformer/timesformer_divST_8xb8-8x32x1-15e_kinetics400-rgb/timesformer_divST_8xb8-8x32x1-15e_kinetics400-rgb_20220815-a4d0d01f.pth,,,,,
+VideoSwin-B,https://github.com/open-mmlab/mmaction2/blob/main/configs/recognition/swin/swin-base-p244-w877_in1k-pre_8xb8-amp-32x2x1-30e_kinetics400-rgb.py,https://download.openmmlab.com/mmaction/v1.0/recognition/swin/swin-base-p244-w877_in1k-pre_8xb8-amp-32x2x1-30e_kinetics400-rgb/swin-base-p244-w877_in1k-pre_8xb8-amp-32x2x1-30e_kinetics400-rgb_20220930-182ec6cc.pth,,,,,
+VideoSwin-L,https://github.com/open-mmlab/mmaction2/blob/main/configs/recognition/swin/swin-large-p244-w877_in22k-pre_8xb8-amp-32x2x1-30e_kinetics400-rgb.py,https://download.openmmlab.com/mmaction/v1.0/recognition/swin/swin-large-p244-w877_in22k-pre_8xb8-amp-32x2x1-30e_kinetics400-rgb/swin-large-p244-w877_in22k-pre_8xb8-amp-32x2x1-30e_kinetics400-rgb_20220930-78ad8b11.pth,,,,,
+UniFormer-V1,https://github.com/open-mmlab/mmaction2/blob/main/configs/recognition/uniformer/uniformer-base_imagenet1k-pre_32x4x1_kinetics400-rgb.py,https://download.openmmlab.com/mmaction/v1.0/recognition/uniformerv1/uniformer-base_imagenet1k-pre_32x4x1_kinetics400-rgb_20221219-b776322c.pth,,,,,
+UniFormer-V2-B,https://github.com/open-mmlab/mmaction2/blob/main/configs/recognition/uniformerv2/uniformerv2-base-p16-res224_clip-kinetics710-pre_8xb32-u8_kinetics400-rgb.py,https://download.openmmlab.com/mmaction/v1.0/recognition/uniformerv2/uniformerv2-base-p16-res224_clip-kinetics710-pre_8xb32-u8_kinetics400-rgb/uniformerv2-base-p16-res224_clip-kinetics710-pre_8xb32-u8_kinetics400-rgb_20230313-75be0806.pth,,,,,
+UniFormer-V2-L,https://github.com/open-mmlab/mmaction2/blob/main/configs/recognition/uniformerv2/uniformerv2-large-p14-res224_clip-kinetics710-pre_u32_kinetics400-rgb.py,https://download.openmmlab.com/mmaction/v1.0/recognition/uniformerv2/kinetics400/uniformerv2-large-p14-res224_clip-kinetics710-pre_u32_kinetics400-rgb_20221219-56a46f64.pth,,,,,
+VideoMAE-V2,https://github.com/open-mmlab/mmaction2/blob/main/configs/recognition/videomaev2/vit-base-p16_videomaev2-vit-g-dist-k710-pre_16x4x1_kinetics-400.py,https://download.openmmlab.com/mmaction/v1.0/recognition/videomaev2/vit-base-p16_videomaev2-vit-g-dist-k710-pre_16x4x1_kinetics-400/vit-base-p16_videomaev2-vit-g-dist-k710-pre_16x4x1_kinetics-400_20230510-3e7f93b2.pth,86.6,97.3,180G,87M,
+VideoMAE-V1,https://github.com/open-mmlab/mmaction2/blob/main/configs/recognition/videomae/vit-base-p16_videomae-k400-pre_16x4x1_kinetics-400.py,https://download.openmmlab.com/mmaction/v1.0/recognition/videomae/vit-base-p16_videomae-k400-pre_16x4x1_kinetics-400_20221013-860a3cd3.pth,81.3,95.0,180G,87M,
+R2plus1D,https://github.com/open-mmlab/mmaction2/blob/main/configs/recognition/r2plus1d/r2plus1d_r34_8xb8-32x2x1-180e_kinetics400-rgb.py,https://download.openmmlab.com/mmaction/v1.0/recognition/r2plus1d/r2plus1d_r34_8xb8-32x2x1-180e_kinetics400-rgb/r2plus1d_r34_8xb8-32x2x1-180e_kinetics400-rgb_20220812-4270588c.pth,75.46,92.28,213G,63.8M,
+I3D-nonlocal,https://github.com/open-mmlab/mmaction2/blob/main/configs/recognition/i3d/i3d_imagenet-pretrained-r50-nl-dot-product_8xb8-32x2x1-100e_kinetics400-rgb.py,https://download.openmmlab.com/mmaction/v1.0/recognition/i3d/i3d_imagenet-pretrained-r50-nl-dot-product_8xb8-32x2x1-100e_kinetics400-rgb/i3d_imagenet-pretrained-r50-nl-dot-product_8xb8-32x2x1-100e_kinetics400-rgb_20220812-8e1f2148.pth,74.80,92.07,59.3G,35.4M,
+TSN,https://github.com/open-mmlab/mmaction2/blob/main/configs/recognition/tsn/tsn_imagenet-pretrained-r50_8xb32-1x1x8-100e_kinetics400-rgb.py,https://download.openmmlab.com/mmaction/v1.0/recognition/tsn/tsn_imagenet-pretrained-r50_8xb32-1x1x8-100e_kinetics400-rgb/tsn_imagenet-pretrained-r50_8xb32-1x1x8-100e_kinetics400-rgb_20220906-2692d16c.pth,74.12,91.34,102.7G,24.33M,
+C3D,https://github.com/open-mmlab/mmaction2/blob/main/configs/recognition/c3d/c3d_sports1m-pretrained_8xb30-16x1x1-45e_ucf101-rgb.py,https://download.openmmlab.com/mmaction/v1.0/recognition/c3d/c3d_sports1m-pretrained_8xb30-16x1x1-45e_ucf101-rgb/c3d_sports1m-pretrained_8xb30-16x1x1-45e_ucf101-rgb_20220811-31723200.pth,,,38.5G,78.4M,
+UniFormer-V2,https://github.com/open-mmlab/mmaction2/blob/main/configs/recognition/uniformerv2/uniformerv2-base-p16-res224_clip-kinetics710-pre_8xb32-u8_kinetics400-rgb.py,https://download.openmmlab.com/mmaction/v1.0/recognition/uniformerv2/uniformerv2-base-p16-res224_clip-kinetics710-pre_8xb32-u8_kinetics400-rgb/uniformerv2-base-p16-res224_clip-kinetics710-pre_8xb32-u8_kinetics400-rgb_20230313-75be0806.pth,,,0.1T,115M,
+VideoSwin,https://github.com/open-mmlab/mmaction2/blob/main/configs/recognition/swin/swin-base-p244-w877_in1k-pre_8xb8-amp-32x2x1-30e_kinetics400-rgb.py,https://download.openmmlab.com/mmaction/v1.0/recognition/swin/swin-base-p244-w877_in1k-pre_8xb8-amp-32x2x1-30e_kinetics400-rgb/swin-base-p244-w877_in1k-pre_8xb8-amp-32x2x1-30e_kinetics400-rgb_20220930-182ec6cc.pth,80.57,94.49,282G,88.0M,
+C2D,https://github.com/open-mmlab/mmaction2/blob/main/configs/recognition/c2d/c2d_r50-in1k-pre-nopool_8xb32-8x8x1-100e_kinetics400-rgb.py,https://download.openmmlab.com/mmaction/v1.0/recognition/c2d/c2d_r50-in1k-pre-nopool_8xb32-8x8x1-100e_kinetics400-rgb/c2d_r50-in1k-pre-nopool_8xb32-8x8x1-100e_kinetics400-rgb_20221027-e0227b22.pth,73.44,91.00,33G,24.3M,
+TSM-nonlocal,https://github.com/open-mmlab/mmaction2/blob/main/configs/recognition/tsm/tsm_imagenet-pretrained-r50-nl-dot-product_8xb16-1x1x8-50e_kinetics400-rgb.py,https://download.openmmlab.com/mmaction/v1.0/recognition/tsm/tsm_imagenet-pretrained-r50-nl-dot-product_8xb16-1x1x8-50e_kinetics400-rgb/tsm_imagenet-pretrained-r50-nl-dot-product_8xb16-1x1x8-50e_kinetics400-rgb_20220831-108bfde5.pth,74.49,91.15,61.30G,31.68M,
+CSN,https://github.com/open-mmlab/mmaction2/blob/main/configs/recognition/csn/ircsn_ig65m-pretrained-r50-bnfrozen_8xb12-32x2x1-58e_kinetics400-rgb.py,https://download.openmmlab.com/mmaction/v1.0/recognition/csn/ircsn_ig65m-pretrained-r50-bnfrozen_8xb12-32x2x1-58e_kinetics400-rgb/ircsn_ig65m-pretrained-r50-bnfrozen_8xb12-32x2x1-58e_kinetics400-rgb_20220811-44395bae.pth,79.44,94.26,55.90G,13.13M,
+TPN,https://github.com/open-mmlab/mmaction2/blob/main/configs/recognition/tpn/tpn-slowonly_r50_8xb8-8x8x1-150e_kinetics400-rgb.py,https://download.openmmlab.com/mmaction/v1.0/recognition/tpn/tpn-slowonly_r50_8xb8-8x8x1-150e_kinetics400-rgb/tpn-slowonly_r50_8xb8-8x8x1-150e_kinetics400-rgb_20220913-97d0835d.pth,74.20,91.48,,,

brainscore_vision/models/temporal_model_mmaction2/model.py

@@ -0,0 +1,226 @@
+import os
+import numpy as np
+
+import mmengine
+import mmaction
+from mmaction.apis import init_recognizer
+from mmengine.registry import init_default_scope
+from mmengine.dataset import Compose, pseudo_collate
+
+from brainscore_vision.model_helpers.activations.temporal.model import PytorchWrapper
+
+
+LARGE_MODEL_LAYER_STEP = 2
+HOME = os.path.join(os.path.dirname(mmaction.__file__), "models")
+
+
+class MMActionModelWrapper(PytorchWrapper):
+    meta = None
+
+    def load_meta(self, path=os.path.join(os.path.dirname(__file__), "mmaction2.csv")):
+        if self.meta is None:
+            import pandas as pd
+            self.meta = pd.read_csv(path)
+    
+    def __init__(self, model_name, process_output=None, *args, **kwargs):
+        self.load_meta()
+
+        _num_frames = None
+        num_frames = kwargs.get("num_frames")
+        if isinstance(num_frames, (list, tuple)):
+            if num_frames[0] == num_frames[1]:
+                _num_frames = num_frames
+        elif num_frames is not None:
+            _num_frames = num_frames
+
+        model_data = self.meta[self.meta['name'] == model_name].iloc[0]  # return a Series
+        config = model_data['config']
+        checkpoint = model_data['checkpoint']
+        config = config.replace("https://github.com/open-mmlab/mmaction2/blob/main/", "")
+        config_path = os.path.join(HOME, config)
+        config = mmengine.Config.fromfile(config_path)
+
+        test_pipeline_cfg = config.test_pipeline
+        # SampleFrames: clip_len x frame_interval (sampling interval) x num_clips
+        # change every ThreeCrop and TenCrop to CenterCrop
+        for i, pipeline in enumerate(test_pipeline_cfg):
+            if pipeline['type'] in ['ThreeCrop', 'TenCrop']:
+                test_pipeline_cfg[i] = {'type': 'CenterCrop', 'crop_size': pipeline['crop_size']}
+            if pipeline['type'] in ['SampleFrames']:
+                test_pipeline_cfg[i].update({"num_clips": 1, 'frame_interval': 1})
+
+        model = init_recognizer(config, checkpoint, device="cpu")
+        init_default_scope(model.cfg.get('default_scope', 'mmaction'))
+        test_pipeline = Compose(test_pipeline_cfg[3:])
+
+        def transform_video(video):
+            imgs = video.to_numpy()
+            data = {'imgs': imgs, 'num_clips': 1, 'modality': 'RGB'}
+            if _num_frames is not None:
+                data['clip_len'] = _num_frames
+                assert len(imgs) == _num_frames
+            else:
+                data['clip_len'] = len(imgs)
+
+            data = test_pipeline(data)
+            return data
+        
+        super().__init__(model_name, model, transform_video, process_output, *args, **kwargs)
+
+    def forward(self, inputs):
+        data = pseudo_collate(inputs)
+        data["inputs"] = [d.to(self._device) for d in data["inputs"]]
+        result = self._model.test_step(data)[0]
+        return result
+    
+
+def get_model(identifier):
+    if identifier == "I3D":
+        process_output = None
+        inferencer_kwargs = {
+            "fps": 12.5,
+            "layer_activation_format": {
+                "backbone.conv1": "CTHW",  # too large: (C: 64, T: *, H: 128, W: 128)
+                **{f"backbone.layer{i}": "CTHW" for i in range(1, 5)},
+                "cls_head": "C",
+            },
+            "num_frames": (5, np.inf),
+        }
+
+    if identifier == "I3D-nonlocal":
+        process_output = None
+        inferencer_kwargs = {
+            "fps": 12.5,
+            "layer_activation_format": {
+                "backbone.conv1": "CTHW",  # too large: (C: 64, T: *, H: 128, W: 128)
+                **{f"backbone.layer{i}": "CTHW" for i in range(1, 5)},
+                "cls_head": "C",
+            },
+            "num_frames": (5, np.inf),
+        }
+
+    if identifier == "TSM":
+        process_output = None
+        inferencer_kwargs = {
+            "fps": 25,
+            "layer_activation_format": {},
+        }
+
+    if identifier == "SlowFast":
+        process_output = None
+        inferencer_kwargs = {
+            "fps": 12.5,
+            "layer_activation_format": {
+                "backbone.slow_path.conv1_lateral": "CTHW",
+                **{f"backbone.slow_path.layer{i}_lateral": "CTHW" for i in range(1, 4)},
+                "cls_head": "C",
+            },
+            "num_frames": 32,  # TODO: in fact can be multiple of 4?
+        }
+
+    if identifier == "X3D":
+        process_output = None
+        inferencer_kwargs = {
+            "fps": 30,
+            "layer_activation_format": {
+                "backbone.conv1_t": "CTHW",
+                **{f"backbone.layer{i}": "CTHW" for i in range(1, 5)},
+                "cls_head": "C",
+            },
+        }
+
+    if identifier == "TimeSformer":
+        inferencer_kwargs = {
+            "fps": 8,
+            "layer_activation_format": {
+                "backbone.patch_embed": "CTHW",
+                **{f"backbone.transformer_layers.layers.{i}": "HWTC" for i in range(0, 12, LARGE_MODEL_LAYER_STEP)},
+                "cls_head": "C",
+            },
+            "num_frames": 8
+        }
+        def process_output(layer, layer_name, inputs, output):
+            if layer_name == "backbone.patch_embed":
+                B = inputs[0].shape[0]
+                C = output.shape[-1]
+                output = output.reshape(B, -1, 14, 14, C)
+            if layer_name.startswith("backbone.transformer_layers.layers."):
+                output = output[:, 1:]
+                B = output.shape[0]
+                C = output.shape[-1]
+                output = output.reshape(B, 14, 14, -1, C)
+            return output
+        
+    if identifier in ["VideoSwin-B", "VideoSwin-L"]:
+
+        transformer_layers = {
+            **{f"backbone.layers.0.blocks.{i}": "THWC" for i in range(0, 2, LARGE_MODEL_LAYER_STEP)},
+            **{f"backbone.layers.1.blocks.{i}": "THWC" for i in range(0, 2, LARGE_MODEL_LAYER_STEP)},
+            **{f"backbone.layers.2.blocks.{i}": "THWC" for i in range(0, 18, LARGE_MODEL_LAYER_STEP)},
+            **{f"backbone.layers.3.blocks.{i}": "THWC" for i in range(0, 2, LARGE_MODEL_LAYER_STEP)},
+        }
+
+        inferencer_kwargs = {
+            "fps": 12.5,
+            "layer_activation_format": {
+                "backbone.patch_embed": "CTHW",
+                **transformer_layers,
+                "cls_head": "C",
+            },
+        }
+        process_output = None
+
+    if identifier == "UniFormer-V1":
+
+        transformer_layers = {
+            **{f"backbone.blocks1.{i}": "CTHW" for i in range(0, 5, LARGE_MODEL_LAYER_STEP)},
+            **{f"backbone.blocks2.{i}": "CTHW" for i in range(0, 8, LARGE_MODEL_LAYER_STEP)},
+            **{f"backbone.blocks3.{i}": "CTHW" for i in range(0, 20, LARGE_MODEL_LAYER_STEP)},
+            **{f"backbone.blocks4.{i}": "CTHW" for i in range(0, 7, LARGE_MODEL_LAYER_STEP)},
+        }
+
+        inferencer_kwargs = {
+            "fps": 6.25,
+            "layer_activation_format": {
+                "backbone.pos_drop": "CTHW",
+                **transformer_layers,
+                "cls_head": "C",
+            },
+        }
+        process_output = None
+
+    if identifier.startswith("UniFormer-V2"):
+
+        if identifier == "UniFormer-V2-B":
+            num_frames = 8
+            num_transformer_layers = 12
+            img_size = 14
+        elif identifier == "UniFormer-V2-L":
+            num_frames = 32
+            num_transformer_layers = 24
+            img_size = 16
+
+        transformer_layers = {
+            **{f"backbone.transformer.resblocks.{i}": "HWTC" for i in range(0, num_transformer_layers, LARGE_MODEL_LAYER_STEP)},
+        }
+
+        inferencer_kwargs = {
+            "fps": 25,
+            "layer_activation_format": {
+                "backbone.conv1": "CTHW",
+                **transformer_layers,
+                "backbone": "C",
+                "cls_head": "C",
+            },
+            "num_frames": num_frames
+        }
+        def process_output(layer, layer_name, inputs, output):
+            if layer_name.startswith("backbone.transformer.resblocks."):
+                T = inputs[1]
+                C = output.shape[-1]
+                output = output[1:]  # remove the class token
+                output = output.reshape(img_size, img_size, -1, T, C).permute(2, 0, 1, 3, 4)  # BHWTC 
+            return output
+
+    model = MMActionModelWrapper(identifier, process_output, **inferencer_kwargs)
+    return model

brainscore_vision/models/temporal_model_mmaction2/requirements.txt

@@ -0,0 +1,5 @@
+importlib-metadata<5
+mmaction2 @ git+https://github.com/YingtianDt/mmaction2.git@533edc3
+mmengine
+torch
+torchvision

brainscore_vision/models/temporal_model_mmaction2/test.py

@@ -0,0 +1,24 @@
+import pytest
+
+from brainscore_vision import load_model
+
+
+model_list = [
+    "I3D",
+    "I3D-nonlocal",
+    "SlowFast",
+    "X3D",
+    "TimeSformer",
+    "VideoSwin-B",
+    "VideoSwin-L",
+    "UniFormer-V1",
+    "UniFormer-V2-B",
+    "UniFormer-V2-L",
+]
+
+@pytest.mark.private_access
+@pytest.mark.memory_intense
+@pytest.mark.parametrize("model_identifier", model_list)
+def test_load(model_identifier):
+    model = load_model(model_identifier)
+    assert model is not None

brainscore_vision/models/temporal_model_openstl/__init__.py

@@ -0,0 +1,18 @@
+from brainscore_vision import model_registry
+from brainscore_vision.model_helpers.brain_transformation import ModelCommitment
+from brainscore_vision.model_helpers.activations.temporal.utils import get_specified_layers
+from brainscore_vision.model_interface import BrainModel
+from . import model
+
+
+def commit_model(identifier):
+    activations_model=model.get_model(identifier)
+    layers=get_specified_layers(activations_model)
+    return ModelCommitment(identifier=identifier, activations_model=activations_model, layers=layers)
+
+
+model_registry["ConvLSTM"] = lambda: commit_model("ConvLSTM")
+model_registry["PredRNN"] = lambda: commit_model("PredRNN")
+model_registry["SimVP"] = lambda: commit_model("SimVP")
+model_registry["TAU"] = lambda: commit_model("TAU")
+model_registry["MIM"] = lambda: commit_model("MIM")

brainscore_vision/models/temporal_model_openstl/model.py

@@ -0,0 +1,206 @@
+import os
+import imp
+import numpy as np
+from collections import OrderedDict
+
+import torch
+from torchvision import transforms
+import openstl
+from openstl.methods import method_maps
+from openstl.utils import reshape_patch
+
+from brainscore_vision.model_helpers.activations.temporal.model import PytorchWrapper
+from brainscore_vision.model_helpers.activations.temporal.utils import download_weight_file
+
+
+# We only use models trained on KITTI dataset, because it is the most ecological, 
+# diverse, challenging, and widely used dataset for next frame prediction among 
+# the datasets used by OpenSTL repo.
+IMAGE_SIZES = (128, 160)  # for KITTI
+KITTI_CONFIG_DIR = os.path.join(os.path.dirname(openstl.__file__), "configs/kitticaltech")
+KITTI_FPS = 10  # from paper: https://www.cvlibs.net/publications/Geiger2012CVPR.pdf
+LARGE_MODEL_LAYER_STEP = 2
+
+transform_image = transforms.Resize(IMAGE_SIZES)
+
+
+class LSTMWrapper(PytorchWrapper):
+    def _register_hook(self, layer, layer_name, target_dict):
+        def hook_function(_layer, _input, output, name=layer_name, target_dict=target_dict):
+            output = self._process_activation(_layer, name, _input, output)
+            target_dict.setdefault(name, []).append(PytorchWrapper._tensor_to_numpy(output)) 
+
+        hook = layer.register_forward_hook(hook_function)
+        return hook
+    
+    def get_activations(self, inputs, layer_names):
+        self._model.eval()
+        layer_results = OrderedDict()
+        hooks = []
+
+        for layer_name in layer_names:
+            layer = self.get_layer(layer_name)
+            hook = self._register_hook(layer, layer_name, target_dict=layer_results)
+            hooks.append(hook)
+
+        with torch.no_grad():
+            self.forward(inputs)
+
+        for hook in hooks:
+            hook.remove()
+
+        # stack the T dim to be the second dim
+        for layer_name, activations in layer_results.items():
+            layer_results[layer_name] = np.stack(activations, axis=1)
+
+        return layer_results
+
+    def forward(self, inputs):
+        tensor = torch.stack(inputs)
+        tensor = tensor.to(self._device)
+        return self._model(tensor, return_loss=False)
+    
+
+class MIMWrapper(LSTMWrapper):
+    def forward(self, inputs):
+        output = super().forward(inputs)
+        # clear MIMBlock.convlstm_c
+        def _clear_helper(module):
+            if hasattr(module, "convlstm_c"):
+                module.convlstm_c = None
+            for child in module.children():
+                _clear_helper(child)
+        _clear_helper(self._model)
+        return output
+
+
+def _get_config(name, parent_dir):
+    config = imp.load_source(name, os.path.join(parent_dir, f"{name}.py")).__dict__
+    config = {k: v for k, v in config.items() if not k.startswith("__")}
+    return config
+
+
+def get_model(identifier):
+    config = _get_config(identifier, KITTI_CONFIG_DIR)
+    config["method"] = config["method"].lower()
+    config['dataname'] = "kitticaltech"
+    config['dataname'] = "kitticaltech"
+    config['metrics'] = ['mse', 'mae']  # not in use, just to initialize the model
+    config['in_shape'] = [None, 3, *IMAGE_SIZES]
+
+    if identifier == "PredRNN":
+        layer_activation_format = {
+            **{f"cell_list.{i}": "TCHW" for i in range(4)},
+            "conv_last": "TCHW"
+        }
+
+        def process_output(layer, layer_name, inputs, output):
+            if layer_name.startswith("cell_list"):
+                h, c, m = output
+                return m
+            else:
+                return output
+        
+        wrapper_cls = LSTMWrapper
+        kwargs = {}
+        weight_name = "kitticaltech_predrnn_one_ep100.pth"
+
+    elif identifier == "ConvLSTM":
+        layer_activation_format = {
+            **{f"cell_list.{i}": "TCHW" for i in range(4)},
+            "conv_last": "TCHW"
+        }
+
+        def process_output(layer, layer_name, inputs, output):
+            if layer_name.startswith("cell_list"):
+                h, c = output
+                return c
+            else:
+                return output
+        
+        wrapper_cls = LSTMWrapper
+        kwargs = {}
+        weight_name = "kitticaltech_convlstm_one_ep100.pth"
+
+    elif identifier in ["SimVP", "TAU"]:
+        num_frames = 10
+        layer_activation_format = {
+            **{f"enc.enc.{i}": "TCHW" for i in range(0, 2, LARGE_MODEL_LAYER_STEP)},
+            **{f"hid.enc.{i}": "TCHW" for i in range(0, 6, LARGE_MODEL_LAYER_STEP)},
+            **{f"dec.dec.{i}": "TCHW" for i in range(0, 2, LARGE_MODEL_LAYER_STEP)},
+        }
+
+        config['in_shape'] = [num_frames, 3, *IMAGE_SIZES]
+        wrapper_cls = PytorchWrapper
+        kwargs = {"num_frames": num_frames}
+
+        def process_output(layer, layer_name, inputs, output):
+            if layer_name.startswith("enc") or layer_name.startswith("dec"):
+                output = output.view(-1, num_frames, *output.shape[1:])
+            elif layer_name.startswith("hid"):
+                output = output[:, None]  # time-compressed layers
+            return output
+        if identifier == "SimVP":
+            weight_name = "kitticaltech_simvp_gsta_one_ep100.pth"
+        elif identifier == "TAU":
+            weight_name = "kitticaltech_tau_one_ep100.pth"
+
+    elif identifier == "MIM":
+        layer_activation_format = {
+            **{f"stlstm_layer.{i}": "TCHW" for i in range(0, 4, LARGE_MODEL_LAYER_STEP)},
+            **{f"stlstm_layer_diff.{i}": "TCHW" for i in range(0, 3, LARGE_MODEL_LAYER_STEP)},
+            "conv_last": "TCHW"
+        }
+
+        def process_output(layer, layer_name, inputs, output):
+            if layer_name.startswith("stlstm_layer."):
+                h, c, m = output
+                ret = m
+            elif layer_name.startswith("stlstm_layer_diff."):
+                h, c = output
+                ret = c
+            else:
+                ret = output
+            return ret
+        
+        wrapper_cls = MIMWrapper
+        kwargs = {}
+        weight_name = "kitticaltech_mim_one_ep100.pth"
+
+
+    model = method_maps[config["method"]](**config).model
+    url = f"https://github.com/chengtan9907/OpenSTL/releases/download/kitti-weights/{weight_name}"
+    weight_path = download_weight_file(url, folder="temporal_model_openstl")
+    model.load_state_dict(torch.load(weight_path, map_location="cpu"))
+
+    def transform_video_lstm(video):
+        frames = torch.Tensor(video.to_numpy() / 255.0).permute(0, 3, 1, 2)
+        frames = transform_image(frames)
+        frames = frames.permute(0, 2, 3, 1)[None, :]  # BTHWC
+        patch_size = config["patch_size"]
+        assert 5 == frames.ndim
+        batch_size, seq_length, img_height, img_width, num_channels = frames.shape
+        a = frames.reshape(batch_size, seq_length,
+                                    img_height//patch_size, patch_size,
+                                    img_width//patch_size, patch_size,
+                                    num_channels)
+        b = a.transpose(3, 4)
+        patches = b.reshape(batch_size, seq_length,
+                                    img_height//patch_size,
+                                    img_width//patch_size,
+                                    patch_size*patch_size*num_channels)[0]
+        return patches
+    
+    def transform_video_simvp(video):
+        frames = torch.Tensor(video.to_numpy() / 255.0).permute(0, 3, 1, 2)
+        frames = transform_image(frames)
+        return frames
+    
+    if identifier in ("PredRNN", "ConvLSTM", "MIM"):
+        transform_video = transform_video_lstm  
+    else: 
+        transform_video = transform_video_simvp
+
+    return wrapper_cls(identifier, model, transform_video, fps=KITTI_FPS, 
+                          layer_activation_format=layer_activation_format,
+                          process_output=process_output, **kwargs)

brainscore_vision/models/temporal_model_openstl/requirements.txt

@@ -0,0 +1,3 @@
+torch
+torchvision
+openstl @ git+https://github.com/YingtianDt/OpenSTL.git

brainscore_vision/models/temporal_model_openstl/test.py

@@ -0,0 +1,19 @@
+import pytest
+
+from brainscore_vision import load_model
+
+
+model_list = [
+    "ConvLSTM",
+    "PredRNN",
+    "SimVP",
+    "TAU",
+    "MIM"
+]
+
+@pytest.mark.private_access
+@pytest.mark.memory_intense
+@pytest.mark.parametrize("model_identifier", model_list)
+def test_load(model_identifier):
+    model = load_model(model_identifier)
+    assert model is not None

brainscore_vision/models/temporal_model_torchvision/__init__.py

@@ -0,0 +1,19 @@
+from brainscore_vision import model_registry
+from brainscore_vision.model_helpers.brain_transformation import ModelCommitment
+from brainscore_vision.model_helpers.activations.temporal.utils import get_specified_layers
+from brainscore_vision.model_interface import BrainModel
+from . import model
+
+
+def commit_model(identifier):
+    activations_model=model.get_model(identifier)
+    layers=get_specified_layers(activations_model)
+    return ModelCommitment(identifier=identifier, activations_model=activations_model, layers=layers)
+
+
+model_registry['r3d_18'] = lambda: commit_model('r3d_18')
+model_registry['r2plus1d_18'] = lambda: commit_model('r2plus1d_18')
+model_registry['mc3_18'] = lambda: commit_model('mc3_18')
+model_registry['s3d'] = lambda: commit_model('s3d')
+model_registry['mvit_v1_b'] = lambda: commit_model('mvit_v1_b')
+model_registry['mvit_v2_s'] = lambda: commit_model('mvit_v2_s')