Add files using upload-large-folder tool

fMRI-foundation-model-old/fMRI-MAE-mamba/accel.slurm

@@ -0,0 +1,65 @@
+#!/bin/bash
+#SBATCH --account=fmri
+#SBATCH --partition=a40
+#SBATCH --job-name=fmrimamba
+#SBATCH --ntasks-per-node=1
+#SBATCH --nodes=1
+#SBATCH --gres=gpu:1
+#SBATCH --gpus-per-task=1       # Set to equal gres=gpu:#!
+#SBATCH --time=20:00:00         # total run time limit (HH:MM:SS)
+#SBATCH -e slurms/%j.err        # first create a "slurms" folder in current directory to store logs
+#SBATCH -o slurms/%j.out
+#SBATCH --comment=medarc
+#SBATCH --no-requeue
+#SBATCH --exclusive
+
+export NUM_GPUS=1 # Set to equal gres=gpu:#!
+echo NUM_GPUS=$NUM_GPUS
+
+source ~/.bashrc
+
+cd /weka/proj-fmri/paulscotti/fMRI-foundation-model/fMRI-MAE/
+jupyter nbconvert mamba.ipynb --to python
+if [ $? -ne 0 ]; then
+  echo "Error: Conversion of ipynb to Python failed. Exiting."
+  exit 1
+fi
+
+# Make sure another job doesnt use same port, here using random number
+export MASTER_PORT=$((RANDOM % (19000 - 11000 + 1) + 11000))
+export HOSTNAMES=$(scontrol show hostnames "$SLURM_JOB_NODELIST")
+export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
+echo MASTER_ADDR=${MASTER_ADDR}
+echo MASTER_PORT=${MASTER_PORT}
+echo WORLD_SIZE=$((${SLURM_NNODES} * ${NUM_GPUS}))
+
+# Other settings relevant for multi-node
+export NCCL_DEBUG=WARN #INFO
+export NCCL_PROTO=simple
+export FI_EFA_USE_DEVICE_RDMA=1
+export FI_EFA_FORK_SAFE=1
+export FI_LOG_LEVEL=1
+export PYTHONFAULTHANDLER=1
+export CUDA_LAUNCH_BLOCKING=0
+export OMPI_MCA_mtl_base_verbose=1
+export FI_EFA_ENABLE_SHM_TRANSFER=0
+export FI_PROVIDER=efa
+export FI_EFA_TX_MIN_CREDITS=64
+export NCCL_TREE_THRESHOLD=0
+export TORCH_DISTRIBUTED_DEBUG=DETAIL
+export NCCL_P2P_DISABLE=1
+module load openmpi cuda/12.1
+
+srun torchrun \
+--nproc_per_node=$NUM_GPUS \
+--nnodes=$SLURM_NNODES \
+--rdzv_id=$SLURM_JOBID \
+--rdzv_backend=c10d \
+--rdzv_endpoint=${MASTER_ADDR}:$MASTER_PORT \
+--rdzv_conf=timeout=90 \
+mamba.py
+
+if [ $? -ne 0 ]; then
+  echo "Error: srun command failed. Please check the logs for more details."
+  exit 1
+fi

fMRI-foundation-model-old/fMRI-MAE-mamba/config.yaml

@@ -0,0 +1,53 @@
+# Model Config
+model_name: "framepatchsize4_bs8_300ep"
+use_cls_token: False
+use_contrastive_loss: False
+constrastive_loss_weight: 0.01
+
+# Training Configs
+batch_size: 8
+num_workers: 10
+num_epochs: 300
+seed: 42
+max_lr: 3.0e-6 # Keep the x.0 else will be converted to string
+num_samples_per_epoch: 1024
+
+# Saving progress
+ckpt_saving: True
+ckpt_interval: 50
+resume_from_ckpt: True
+wandb_log: True
+
+# MAE Config
+tube_start_masking_ratio: 0.95
+tube_end_masking_ratio: 0.95
+decoder_mask_ratio: 0.85
+
+# Model Config
+patch_size: 8
+frame_patch_size: 4
+use_rope_emb: False
+masking_strategy: "MNI"
+
+# ViT Config
+encoder_model: "vit_base"
+decoder_model: "vit_small"
+
+# Mamba config
+#embed_dim: 512
+#depth: 32
+
+# Data Config
+img_size: [88, 104, 72] # Image Size
+num_frames: 4
+is_s3: False
+# 000495.tar is missing??
+train_urls: ["/weka/proj-fmri/shared/NSD_MNI_wds/{000000..000738}.tar"]
+# train_urls: ["s3://proj-fmri/fmri_foundation_datasets/NSD_MNI_wds/{000000..000494}.tar","s3://proj-fmri/fmri_foundation_datasets/NSD_MNI_wds/{000496..000740}.tar"]
+
+#"/weka/proj-fmri/paulscotti/old_fMRI-foundation-model/dataset_creation/wds_creation/wds/000000.tar"
+#"s3://proj-fmri/fmri_foundation_datasets/NSD_MNI_wds/{000000..000740}.tar"
+# test_urls: "s3://proj-fmri/fmri_foundation_datasets/openneuro_MNI_wds/000000.tar"
+# train_urls: "/weka/proj-fmri/paulscotti/old_fMRI-foundation-model/dataset_creation/wds_creation/wds/{000001..000240}.tar"
+# train_urls: "/weka/proj-fmri/paulscotti/fMRI-foundation-model/dataset_creation/wds_creation/wds/{000000..000382}.tar"
+# test_urls: "/weka/proj-fmri/paulscotti/old_fMRI-foundation-model/dataset_creation/wds_creation/wds/000000.tar"

fMRI-foundation-model-old/fMRI-MAE-mamba/dataloader.py

@@ -0,0 +1,772 @@
+import numpy as np
+from torch import Tensor, stack
+from torch.utils.data import IterDataPipe, DataLoader
+from torch.utils.data.datapipes._decorator import functional_datapipe
+from torch.utils.data.datapipes.iter.callable import MapperIterDataPipe
+from torch.utils.data.datapipes.utils.common import StreamWrapper, match_masks
+from torchdata.datapipes.iter import IterableWrapper, IterDataPipe, TarArchiveLoader
+from torchdata.datapipes.iter.load.s3io import S3FileLoaderIterDataPipe
+from torchdata.datapipes.utils.common import validate_pathname_binary_tuple
+import warnings
+import subprocess
+import tarfile
+import tempfile
+import time
+import random
+from omegaconf import DictConfig, ListConfig
+import webdataset as wds
+import re
+import os
+import copy
+from io import BufferedIOBase, BytesIO, RawIOBase
+from typing import (
+    IO,
+    Any,
+    Callable,
+    Dict,
+    Iterable,
+    Iterator,
+    List,
+    Optional,
+    Sequence,
+    Tuple,
+    TypeVar,
+    Union,
+    cast,
+)
+T_co = TypeVar("T_co", covariant=True)
+
+class ShuffledListDataPipe(IterDataPipe[T_co]):
+    def __init__(
+        self,
+        source_list: List[T_co],
+        *,
+        shuffle: bool = True,
+        cycle: Union[bool, int] = True,
+    ):
+        super().__init__()
+        self.source = source_list
+        self._enabled = shuffle
+        self._seed = None
+        self._rng = random.Random()
+        if isinstance(cycle, bool):
+            if cycle:
+                self._cycle = -1
+            else:
+                self._cycle = 1
+        else:
+            assert isinstance(cycle, int)
+            self._cycle = cycle
+
+    def __iter__(self) -> Iterator[T_co]:
+        source = copy.copy(self.source)
+        cycle = self._cycle
+        epochs = 0
+        while cycle == -1 or epochs < cycle:
+            if self._enabled:
+                self._rng.shuffle(source)
+            yield from source
+            epochs += 1
+
+    def __len__(self):
+        if self.count == -1:
+            raise TypeError(
+                f"This {type(self).__name__} instance cycles forever, and "
+                f"therefore doesn't have valid length."
+            )
+        else:
+            return self.count * len(self.source)
+
+    def __getstate__(self):
+        state = (
+            self.source,
+            self._enabled,
+            self._seed,
+            self._rng.getstate(),
+            self._valid_iterator_id,
+            self._number_of_samples_yielded,
+        )
+        if IterDataPipe.getstate_hook is not None:
+            return IterDataPipe.getstate_hook(state)
+        return state
+
+    def __setstate__(self, state):
+        (
+            self.source,
+            self._enabled,
+            self._seed,
+            rng_state,
+            self._valid_iterator_id,
+            self._number_of_samples_yielded,
+        ) = state
+        self._rng = random.Random()
+        self._rng.setstate(rng_state)
+
+def get_job_temp_dir(dl_root: str) -> str:
+    try:
+        job_or_array_id = (
+            os.environ.get("SLURM_ARRAY_JOB_ID", "") or os.environ["SLURM_JOB_ID"]
+        )
+        task_id = os.environ.get("SLURM_ARRAY_TASK_ID", "") or "0"
+        return os.path.join(dl_root, f"JOB_{job_or_array_id}/TASK_{task_id}")
+    except KeyError as e:
+        print("SLURM_JOB_ID env var not set: You sure this job is healthy?")
+        raise e
+
+def _shard_expand(s: str) -> List[str]:
+    expansion = r"\{[0-9]+\.\.[0-9]+\}"
+    m = re.search(expansion, s)
+    if not m:
+        return [s]
+    prefix = s[: m.start()]
+    rest = _shard_expand(s[m.end() :])
+    rng = s[m.start() + 1 : m.end() - 1]
+    lohi = rng.split("..")
+    if len(lohi[0]) == len(lohi[1]) and lohi[0].startswith("0"):
+        fmt = "{prefix}{i:0>{l}d}{r}"
+    elif len(lohi[0]) <= len(lohi[1]):
+        if lohi[0].startswith("0") and lohi[0] != "0":
+            raise ValueError(
+                "shard_expand: low bound must not start with 0 if low bound is shorter"
+            )
+        fmt = "{prefix}{i}{r}"
+    else:
+        raise ValueError("shard_expand: low bound must be shorter than high bound")
+    lo, hi = (int(x) for x in lohi)
+    if lo >= hi:
+        raise ValueError(f"shard_expand: bad range in in shard spec {s}.")
+    result = []
+    for i in range(lo, hi + 1):
+        for r in rest:
+            expanded: str = fmt.format(prefix=prefix, i=i, r=r, l=len(lohi[1]))
+            result.append(expanded)
+    return result
+
+@functional_datapipe("custom_shard_expand")
+class CustomShardExpanderIterDataPipe(IterDataPipe[str]):
+    def __init__(self, source_datapipe: IterDataPipe[str]) -> None:
+        super().__init__()
+        self.source_datapipe: IterDataPipe[str] = source_datapipe
+
+    def __iter__(self) -> Iterator[str]:
+        for path in self.source_datapipe:
+            yield from _shard_expand(path)
+
+def warn_and_continue(exn):
+    """Call in an exception handler to ignore any exception, issue a warning, and continue."""
+    print(exn)
+    warnings.warn(repr(exn))
+    time.sleep(0.05)
+    return True
+error_handler = warn_and_continue
+
+def is_stream_handle(data):
+    obj_to_check = data.file_obj if isinstance(data, StreamWrapper) else data
+    return isinstance(obj_to_check, (BufferedIOBase, RawIOBase))
+
+def read_stream_handles(data):
+    if not is_stream_handle(data):
+        return data
+    else:
+        ds = data
+        data = b"".join(data)
+        ds.close()
+        del ds
+        return data
+        
+def stream_reader(sample):
+    return {k: read_stream_handles(v) for k, v in sample.items()}
+
+def to_our_format(sample):
+    wds_key = sample.pop("__key__")
+    sample = {k.lstrip("."): v for k, v in sample.items()}
+    sample["__key__"] = wds_key.split("/")[-1]
+    sample["__url__"] = "/".join(wds_key.split("/")[:-1])
+    return sample
+
+def add_processors(
+    datapipeline,
+    processors: Optional[ListConfig],
+    description: str,
+    error_handler: Callable = warn_and_continue,
+):
+    if not processors:
+        return datapipeline
+    else:
+        for i, processor_config in enumerate(processors):
+            processor = instantiate(processor_config)
+            if isinstance(processor, AbstractFilter):
+                print(
+                    f"Adding filter {processor.__class__.__name__} as {description} #{i} "
+                    f"to the datapipeline"
+                )
+                datapipeline = datapipeline.filter(processor.filter)
+            elif isinstance(processor, AbstractMapper):
+                print(
+                    f"Adding mapper {processor.__class__.__name__} as {description} #{i} "
+                    f"to the datapipeline"
+                )
+                datapipeline = datapipeline.map_with_handler(
+                    processor.map,
+                    handler=error_handler,
+                    called_cls_name=processor.__class__.__name__,
+                )
+            else:
+                raise TypeError(
+                    f"chosen {description} {processor.__class__.__name__} should be either subclass"
+                    "AbstractMapper or AbstractFilter"
+                )
+        return datapipeline
+
+@functional_datapipe("map_with_handler")
+class MapperWithErrorHandlingIterDataPipe(MapperIterDataPipe):
+    def __init__(
+        self,
+        datapipe: IterDataPipe,
+        fn: Callable,
+        handler: Callable = wds.reraise_exception,
+        input_col: Optional[Union[str, int]] = None,
+        output_col: Optional[Union[str, int]] = None,
+        called_cls_name: Optional[str] = None,
+    ):
+        # for now, disbable input and output col since this is never used anyways
+        if input_col is not None:
+            raise NotImplementedError("`input_col` argument currently not supported")
+
+        if output_col is not None:
+            raise NotImplementedError("`output_col` argument currently not supported")
+
+        super().__init__(datapipe, fn)
+        self.handler = handler
+        self._apply_fn_ = self._apply_fn
+
+    def __iter__(self) -> Iterator[T_co]:
+        for data in self.datapipe:
+            try:
+                res = self._apply_fn_(data)
+                if res is None:
+                    continue
+                yield res
+            except Exception as e:
+                if self.handler(e):
+                    continue
+                else:
+                    raise e
+
+__S3_TOOLS__ = {
+    "s3": ["/usr/local/bin/aws", "s3"],
+}
+
+def is_tar(x: str) -> bool:
+    return x.endswith(".tar")
+
+def ls_aws(
+    path: str,
+    tool: str = "s3",
+    recursive: bool = True,
+    raise_errors: bool = True,
+    skip_files: bool = True,
+):
+    assert path.startswith("s3://"), path
+
+    # in case we have a file, set recursive to false
+    isfile = bool(os.path.splitext(path)[1])
+    if isfile and skip_files:
+        return [path]
+
+    if not path.endswith("/") and not isfile:
+        path = path + "/"
+    cmd = [*__S3_TOOLS__[tool], "ls", f"{path}"]
+
+    if recursive and not isfile:
+        cmd += ["--recursive"]
+    try:
+        result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+        result.check_returncode()
+
+        stdout = result.stdout.decode().split("\n")
+        # return empty list in case path is not a real file/directory
+        if not stdout[0]:
+            return []
+
+        # return path if it is an exisiting file
+        if isfile:
+            return [path]
+
+        # strip timestamp and object size
+        out = [line.split(" ")[-1] for line in stdout if line]
+        if recursive:
+            # aws s3 ls returns prefix+filename
+            bucket = path[: path.find("/", len("s3://"))]  # s3://<bucket>/...
+            out = [os.path.join(bucket, o) for o in out]
+        else:
+            # aws s3 ls returns filename only
+            out = [os.path.join(path, o) for i in out]
+
+        return out
+    except subprocess.CalledProcessError as e:
+        print(f"Got exception while trying to load data! {e.__class__.__name__}: {e}")
+        if raise_errors:
+            raise e
+        else:
+            return []
+
+@functional_datapipe("wrapped_load_files_by_s3")
+class WrappedS3FileLoaderIterDataPipe(S3FileLoaderIterDataPipe):
+    def __init__(
+        self,
+        *args,
+        ignore_missing_files: bool = False,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        self.ignore_missing_files = ignore_missing_files
+
+    def __iter__(self) -> Iterator[Tuple[str, StreamWrapper]]:
+        # same as parent but maybe ignoring errors
+        for url in self.source_datapipe:
+            try:
+                yield url, StreamWrapper(BytesIO(self.handler.s3_read(url)))
+            except ValueError as e:
+                if not self.ignore_missing_files:
+                    raise FileNotFoundError(url)
+                else:
+                    print(f"Warning: Could not download {url}")
+
+@functional_datapipe("list_files_with_aws_cli")
+class S3CLIFileListerIterDataPipe(IterDataPipe[T_co]):
+    def __init__(
+        self,
+        source_datapipe: Union[str, Sequence[str], IterDataPipe],
+        s3_tool: str = "s3",
+        n_retries: int = 10,
+        sleep_interval: float = 0.01,
+    ):
+        if isinstance(source_datapipe, str):
+            source_datapipe = [
+                source_datapipe,
+            ]
+        if not isinstance(source_datapipe, IterDataPipe):
+            self.datapipe: IterDataPipe = IterableWrapper(source_datapipe)  # type: ignore[assignment]
+        else:
+            self.datapipe = source_datapipe
+
+        assert s3_tool in __S3_TOOLS__, f"`s3_tool` has to be in {list(__S3_TOOLS__)}"
+
+        self.s3_tool = s3_tool
+
+        self.n_retries = n_retries
+        self.sleep = sleep_interval
+
+    def __iter__(self) -> Iterator[str]:
+        for root in self.datapipe:
+            files = None
+            for _ in range(self.n_retries):
+                try:
+                    files = ls_aws(
+                        path=root, tool=self.s3_tool, recursive=True, raise_errors=True
+                    )
+                    break
+                except subprocess.CalledProcessError:
+                    time.sleep(self.sleep)
+
+            if files is None:
+                print(
+                    f"Could not ls data expected under {root} in {self.n_retries} tries. Not yielding ..."
+                )
+                continue
+
+            for file in files:
+                yield file
+
+def _download_tar(url, scratch, verbose=False, n_retries=100):
+    start = time.perf_counter()
+    if verbose:
+        _log(f"downloading {url}")
+    uid = "".join(filter(str.isalnum, os.path.splitext(url)[0]))
+    idx = 0
+    path = os.path.join(scratch, f"{os.getpid()}.{uid}.{idx}.tar")
+    while os.path.exists(path):
+        idx += 1
+        path = os.path.join(scratch, f"{os.getpid()}.{uid}.{idx}.tar")
+
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    for i_try in range(n_retries):
+        p = subprocess.call(
+            [
+                "/usr/local/bin/aws",
+                "s3",
+                "cp",
+                url,
+                path,
+                "--quiet",
+            ]
+        )
+
+        if p != 0:
+            # mostly to work around
+            # Error when retrieving credentials from Ec2InstanceMetadata: No
+            # credentials found in credential_source referenced in profile
+            # default
+            if i_try == 0:
+                # check if the requested file actually exists
+                if not ls_aws(url, skip_files=False, raise_errors=False):
+                    raise FileNotFoundError(url)
+
+            if i_try + 1 < n_retries:
+                time.sleep(random.uniform(0.0, 0.1))
+                continue
+            else:
+                raise Exception(
+                    f"File {url} should exist but failed to download after trying {i_try+1} times."
+                )
+        break
+
+    if verbose:
+        _log(
+            f"wrote {path} in {time.perf_counter() - start} secs after {i_try+1} tries."
+        )
+    return path
+
+
+@functional_datapipe("download_with_s3_cli")
+class S3CLITarDownloader(IterDataPipe[T_co]):
+    def __init__(
+        self,
+        source_datapipe,
+        n_retries: int = 10,
+        verbose: bool = False,
+        dl_root: str = "/scratch",
+        mode: str = "r:*",
+        aws_kwargs: Optional[dict] = None,
+        ignore_missing_files: bool = False,
+    ):
+        self.source_datapipe = source_datapipe
+        assert os.path.isdir(dl_root), f"`dl_root` {dl_root} is not a valid directory"
+        # get slurm job id and create a subdir in `dl_root` to download all the tars to
+        dl_root = get_job_temp_dir(dl_root)
+        os.makedirs(dl_root, exist_ok=True)
+        self.dl_root = dl_root
+        self.mode = mode
+        self.n_retries = n_retries
+        self.ignore_missing_files = ignore_missing_files
+
+        self.verbose = verbose
+
+    def _yield_next(self, url: str, local_path: str):
+        if self.verbose:
+            _log("popping queue")
+
+        if self.verbose:
+            _log(f"loading {local_path}")
+        tarstream = tarfile.open(local_path, self.mode)
+
+        if self.verbose:
+            _log(f"yielding {url}")
+
+        yield (url, StreamWrapper(tarstream))
+        if self.verbose:
+            _log("new tar request")
+
+        if self.verbose:
+            _log(f"removing previous tar at {local_path}")
+        try:
+            os.remove(local_path)
+        except FileNotFoundError:
+            _log(
+                f"WARNING: Could not find previous tar for deletion. Unless a clean-up was triggered this is unexpected. The location was {local_path}"
+            )
+
+    def __iter__(self):
+        with tempfile.TemporaryDirectory(dir=self.dl_root) as scratch:
+            for url in self.source_datapipe:
+                try:
+                    local_path = _download_tar(
+                        url, scratch, verbose=self.verbose, n_retries=self.n_retries
+                    )
+                except FileNotFoundError as e:
+                    if not self.ignore_missing_files:
+                        raise e
+                else:
+                    yield from self._yield_next(url, local_path)
+
+
+@functional_datapipe("load_from_tar_and_handle_error")
+class TarArchiveLoaderWithErrorHandlingIterDataPipe(TarArchiveLoader):
+    def __init__(
+        self,
+        datapipe: Iterable[Tuple[str, BufferedIOBase]],
+        mode: str = "r:*",
+        length: int = -1,
+        handler: Callable = wds.reraise_exception,
+    ):
+        super().__init__(datapipe=datapipe, mode=mode, length=length)
+        self.handler = handler
+
+    def __iter__(self) -> Iterator[Tuple[str, BufferedIOBase]]:
+        for data in self.datapipe:
+            validate_pathname_binary_tuple(data)
+            pathname, data_stream = data
+            try:
+                if isinstance(data_stream, StreamWrapper) and isinstance(
+                    data_stream.file_obj, tarfile.TarFile
+                ):
+                    tar = data_stream.file_obj
+                else:
+                    reading_mode = (
+                        self.mode
+                        if hasattr(data_stream, "seekable") and data_stream.seekable()
+                        else self.mode.replace(":", "|")
+                    )
+                    # typing.cast is used here to silence mypy's type checker
+                    tar = tarfile.open(
+                        fileobj=cast(Optional[IO[bytes]], data_stream),
+                        mode=reading_mode,
+                    )
+                for tarinfo in tar:
+                    if not tarinfo.isfile():
+                        continue
+                    extracted_fobj = tar.extractfile(tarinfo)
+                    if extracted_fobj is None:
+                        warnings.warn(
+                            f"failed to extract file {tarinfo.name} from source tarfile {pathname}"
+                        )
+                        raise tarfile.ExtractError
+                    inner_pathname = os.path.normpath(
+                        os.path.join(pathname, tarinfo.name)
+                    )
+
+                    yield inner_pathname, StreamWrapper(extracted_fobj, data_stream, name=inner_pathname)  # type: ignore[misc]
+            except Exception as e:
+                warnings.warn(
+                    f"Unable to extract files from corrupted tarfile stream {pathname} due to: {e}, abort!"
+                )
+                if self.handler(e):
+                    if hasattr(e, "args") and len(e.args) > 0:
+                        e.args = (e.args[0] + " @ " + str(pathname),) + e.args[1:]
+            finally:
+                if isinstance(data_stream, StreamWrapper):
+                    data_stream.autoclose()
+
+class DictCollator:
+    def __init__(
+        self,
+        combine_tensors: bool = True,
+        combine_scalars: bool = True,
+        timeout=None,
+    ):
+        self.combine_tensors = combine_tensors
+        self.combine_scalars = combine_scalars
+
+        collate = self._collate
+        if timeout is not None:
+            collate = timeout_wrapper(collate, timeout=timeout)
+        self.collate = collate
+
+    def __call__(self, samples):
+        return self.collate(samples)
+
+    def _collate(self, samples):
+        keys = set.intersection(*[set(sample.keys()) for sample in samples])
+        batched = {key: [] for key in keys}
+
+        for s in samples:
+            [batched[key].append(s[key]) for key in batched]
+
+        result = {}
+        for key in batched:
+            if isinstance(batched[key][0], (int, float)):
+                if self.combine_scalars:
+                    result[key] = np.array(list(batched[key]))
+            elif isinstance(batched[key][0], Tensor):
+                if self.combine_tensors:
+                    result[key] = stack(list(batched[key]))
+            elif isinstance(batched[key][0], np.ndarray):
+                if self.combine_tensors:
+                    result[key] = np.array(list(batched[key]))
+            else:
+                result[key] = list(batched[key])
+
+        del samples
+        del batched
+        return result
+
+def timeout_wrapper(func: Callable, timeout: Optional[float] = None) -> Callable:
+    if timeout is None or timeout <= 0.0:
+        return func
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        result = [None]
+        exception = [None]
+        event = threading.Event()
+
+        def wrapped_func():
+            try:
+                result[0] = func(*args, **kwargs)
+            except Exception as e:
+                exception[0] = e
+            finally:
+                event.set()
+
+        thread = threading.Thread(target=wrapped_func)
+        thread.start()
+        event.wait(timeout)
+
+        if not event.is_set():
+            raise TimeoutError(f"Function call timed out (longer than {timeout} secs).")
+
+        thread.join()
+
+        if exception[0] is not None:
+            err = exception[0]
+            del exception
+            raise err
+
+        del thread
+        del exception
+        del wrapped_func
+        del event
+        del args
+        del kwargs
+
+        ret = result[0]
+        del result
+        return ret
+
+    return wrapper
+
+class DecoderWithTimeout(wds.Decoder):
+    """Decode samples using a list of handlers.
+
+    For each key/data item, this iterates through the list of
+    handlers until some handler returns something other than None.
+    """
+
+    def __init__(self, *args, timeout: Optional[float] = None, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        if timeout is not None:
+            self.decode = timeout_wrapper(self.decode, timeout)
+
+def setup_decoder(
+    decoders,
+    additional_decoder_kwargs: Optional[Union[DictConfig, Dict]] = None,
+):
+    if not additional_decoder_kwargs:
+        additional_decoder_kwargs = {}
+
+    if not isinstance(decoders, (List, ListConfig)):
+        decoders = [decoders]
+    handlers = []
+    for decoder_spec in decoders:
+        if isinstance(decoder_spec, (Dict, DictConfig)):
+            decoder = instantiate(decoder_spec)
+        elif isinstance(decoder_spec, str):
+            decoder = wds.autodecode.ImageHandler(decoder_spec)
+        else:
+            raise TypeError(f"{decoder_spec} not a thing for decoders.")
+        handlers.append(decoder)
+        print(f"Adding decoder {decoder.__class__.__name__} to decoders.")
+
+    decoder = DecoderWithTimeout(
+        handlers,
+        partial=additional_decoder_kwargs.pop("partial", True),
+        **additional_decoder_kwargs,  # todo
+    )
+    return decoder
+
+def create_dataset(urls, is_s3=False, prefetch=None, 
+                     s3_buffer_size=None, s3_ignore_missing_files=True,
+                     sample_shuffle=1, shard_shuffle=1000, cycle=True,
+                     split_workers_along_tars=True, decoders="torch"):
+    if isinstance(urls, str):
+        urls = [urls]
+    dp = IterableWrapper(urls)
+    dp = dp.custom_shard_expand()
+    if is_s3: 
+        dp = dp.list_files_with_aws_cli().filter(is_tar)
+    else:
+        dp = dp.list_files(masks="*.tar", recursive=True)
+    
+    dp = ShuffledListDataPipe(list(dp), shuffle=shard_shuffle > 1, cycle=cycle)
+
+    if split_workers_along_tars:
+        dp = dp.sharding_filter()
+
+    if is_s3:
+        # if prefetch is None:
+        #     s3_buffer_size = int(s3_buffer_size) if s3_buffer_size is not None else None
+        #     dp = dp.wrapped_load_files_by_s3(
+        #         buffer_size=s3_buffer_size,
+        #         ignore_missing_files=s3_ignore_missing_files,
+        #     )
+        dp = dp.download_with_s3_cli(ignore_missing_files=s3_ignore_missing_files)
+    else:
+        dp = dp.open_files(mode="b")
+
+    if prefetch is not None:
+        dp = dp.custom_prefetch(buffer_size=prefetch)
+
+    dp = dp.load_from_tar_and_handle_error(handler=error_handler)
+    dp = dp.webdataset()
+
+    if not split_workers_along_tars:
+        dp = dp.sharding_filter()
+
+    dp = dp.map_with_handler(
+            stream_reader,
+            handler=error_handler,
+            called_cls_name="StreamReader",
+        )
+
+    dp = dp.shuffle(buffer_size=sample_shuffle) if sample_shuffle > 1 else dp
+
+    dp = dp.map_with_handler(
+            to_our_format,
+            handler=error_handler,
+            called_cls_name="ToOurFormat",
+        )
+
+    decoder: DecoderWithTimeout = setup_decoder(decoders) # ,additional_decoder_kwargs)
+
+    dp = dp.map_with_handler(
+        decoder.decode,
+        handler=error_handler,
+        called_cls_name="Decoder",
+    )
+    
+    return dp
+
+def create_loader(
+    datapipeline: IterDataPipe,
+    batch_size: int,
+    num_workers: int,
+    partial: bool = False,
+    collation_fn: Optional[Union[Callable, Dict, DictConfig]] = DictCollator(),
+    batched_transforms: Optional[ListConfig] = None,
+    loader_kwargs: Optional[Union[Dict, DictConfig]] = None,
+) -> DataLoader:
+    if not loader_kwargs:
+        loader_kwargs = {}
+
+    loader_kwargs.pop("shuffle", None)
+    if not batched_transforms:
+        batched_transforms = []
+
+    print("#" * 100)
+    print("Building dataloader with the following parameters")
+    print(f"batch_size: {batch_size}, num_workers: {num_workers}")
+    for key in loader_kwargs:
+        print(key, ": ", loader_kwargs[key])
+    print("#" * 100)
+
+    datapipeline = datapipeline.batch(batch_size, drop_last=not partial)
+
+    if isinstance(collation_fn, (Dict, DictConfig)):
+        collation_fn = instantiate(collation_fn)
+    datapipeline = datapipeline.collate(collate_fn=collation_fn)
+    loader = DataLoader(
+        datapipeline, batch_size=None, num_workers=num_workers, **loader_kwargs
+    )
+    return loader

fMRI-foundation-model-old/fMRI-MAE-mamba/mindeye_config.yaml

@@ -0,0 +1,30 @@
+# Model Config
+model_name: "pretrain_framepatchsize4"
+mae_model_name: "framepatchsize4_bs8_300ep"
+
+# Training Configs
+global_batch_size: 32
+mixed_precision: "fp16"
+num_epochs: 12
+seed: 42
+max_lr: 3.0e-4 # Keep the x.0 else will be converted to string
+multi_subject: False
+multisubject_ckpt: None
+
+# Saving progress
+ckpt_saving: False
+ckpt_interval: 99
+resume_from_ckpt: False
+wandb_log: True
+
+# Model Config
+in_dim: 370176
+hidden_dim: 512
+drop: .15
+mixup_pct: 0.33
+
+# Data Config
+nsd_wds_path: "/weka/proj-fmri/shared/mindeyev2_dataset/wds"
+nsd_raw_path: "/weka/proj-fmri/shared/mindeyev2_dataset"
+nsd_image_path: "/weka/proj-fmri/shared/mindeyev2_dataset"
+num_sessions: 40

fMRI-foundation-model-old/fMRI-MAE-mamba/mindeye_models.py

@@ -0,0 +1,276 @@
+from typing import Optional
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class MindEyeModule(nn.Module):
+    def __init__(self):
+        super().__init__()
+    def forward(self, x):
+        return x
+class RidgeRegression(nn.Module):
+    # make sure to add weight_decay when initializing optimizer
+    def __init__(self, input_sizes, out_features, seq_len=1): 
+        super().__init__()
+        self.out_features = out_features
+        self.linears = torch.nn.ModuleList([
+                torch.nn.Linear(input_size, out_features) for input_size in input_sizes
+            ])
+        self.seq_len = seq_len
+    def forward(self, x, subj_idx):
+        out = torch.cat([self.linears[subj_idx](x[:,seq]).unsqueeze(1) for seq in range(self.seq_len)], dim=1)
+        return out
+class BrainNetwork(nn.Module):
+    def __init__(self, h=1024, in_dim=15724, out_dim=768, seq_len=1, n_blocks=4, drop=.15, clip_size=768, clip_scale=1):
+        super().__init__()
+        self.seq_len = seq_len
+        self.h = h
+        self.clip_size = clip_size
+        self.clip_scale = clip_scale
+        self.mixer_blocks1 = nn.ModuleList([
+            self.mixer_block1(h, drop) for _ in range(n_blocks)
+        ])
+        self.mixer_blocks2 = nn.ModuleList([
+            self.mixer_block2(seq_len, drop) for _ in range(n_blocks)
+        ])
+        
+        # Output linear layer
+        self.backbone_linear = nn.Linear(h * seq_len, out_dim, bias=True) 
+        self.clip_proj = self.projector(clip_size, clip_size, h=clip_size)
+            
+    def projector(self, in_dim, out_dim, h=2048):
+        return nn.Sequential(
+            nn.LayerNorm(in_dim),
+            nn.GELU(),
+            nn.Linear(in_dim, h),
+            nn.LayerNorm(h),
+            nn.GELU(),
+            nn.Linear(h, h),
+            nn.LayerNorm(h),
+            nn.GELU(),
+            nn.Linear(h, out_dim)
+        )
+    
+    def mlp(self, in_dim, out_dim, drop):
+        return nn.Sequential(
+            nn.Linear(in_dim, out_dim),
+            nn.GELU(),
+            nn.Dropout(drop),
+            nn.Linear(out_dim, out_dim),
+        )
+    
+    def mixer_block1(self, h, drop):
+        return nn.Sequential(
+            nn.LayerNorm(h),
+            self.mlp(h, h, drop),  # Token mixing
+        )
+
+    def mixer_block2(self, seq_len, drop):
+        return nn.Sequential(
+            nn.LayerNorm(seq_len),
+            self.mlp(seq_len, seq_len, drop)  # Channel mixing
+        )
+        
+    def forward(self, x):
+        # make empty tensors
+        c = torch.Tensor([0.])
+        
+        # Mixer blocks
+        residual1 = x
+        residual2 = x.permute(0,2,1)
+        for block1, block2 in zip(self.mixer_blocks1,self.mixer_blocks2):
+            x = block1(x) + residual1
+            residual1 = x
+            x = x.permute(0,2,1)
+            
+            x = block2(x) + residual2
+            residual2 = x
+            x = x.permute(0,2,1)
+            
+        x = x.reshape(x.size(0), -1)
+        backbone = self.backbone_linear(x).reshape(len(x), -1, self.clip_size)
+        if self.clip_scale>0:
+            c = self.clip_proj(backbone)
+        
+        return backbone, c
+    
+    
+import open_clip
+import kornia
+class FrozenOpenCLIPImageEmbedder(nn.Module):
+    """
+    Uses the OpenCLIP vision transformer encoder for images
+    """
+
+    def __init__(
+        self,
+        arch="ViT-H-14",
+        version="laion2b_s32b_b79k",
+        device="cuda",
+        init_device="cpu",
+        max_length=77,
+        freeze=True,
+        antialias=True,
+        ucg_rate=0.0,
+        unsqueeze_dim=False,
+        repeat_to_max_len=False,
+        num_image_crops=0,
+        output_tokens=False,
+        l2_norm_tokens=False,
+        only_tokens=False,
+        cache_dir: Optional[str] = None,
+    ):
+        super().__init__()
+        model, _, _ = open_clip.create_model_and_transforms(
+            arch,
+            device=torch.device(init_device),
+            pretrained=version,
+            cache_dir=cache_dir,
+        )
+        del model.transformer
+        self.model = model
+        self.max_crops = num_image_crops
+        self.pad_to_max_len = self.max_crops > 0
+        self.repeat_to_max_len = repeat_to_max_len and (not self.pad_to_max_len)
+        self.device = device
+        self.max_length = max_length
+        if freeze:
+            self.freeze()
+
+        self.antialias = antialias
+
+        self.register_buffer(
+            "mean", torch.Tensor([0.48145466, 0.4578275, 0.40821073]), persistent=False
+        )
+        self.register_buffer(
+            "std", torch.Tensor([0.26862954, 0.26130258, 0.27577711]), persistent=False
+        )
+        self.ucg_rate = ucg_rate
+        self.unsqueeze_dim = unsqueeze_dim
+        self.stored_batch = None
+        self.model.visual.output_tokens = output_tokens
+        self.output_tokens = output_tokens
+        if only_tokens:
+            assert output_tokens
+        self.only_tokens = only_tokens
+        self.l2_norm_tokens = l2_norm_tokens
+        if l2_norm_tokens:
+            assert output_tokens
+
+    def preprocess(self, x):
+        # normalize to [0,1]
+        x = kornia.geometry.resize(
+            x,
+            (224, 224),
+            interpolation="bicubic",
+            align_corners=True,
+            antialias=self.antialias,
+        )
+        x = (x + 1.0) / 2.0
+        # renormalize according to clip
+        x = kornia.enhance.normalize(x, self.mean, self.std)
+        return x
+
+    def freeze(self):
+        self.model = self.model.eval()
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def forward(self, image, no_dropout=False):
+        z = self.encode_with_vision_transformer(image)
+        tokens = None
+        if self.output_tokens:
+            z, tokens = z[0], z[1]
+        z = z.to(image.dtype)
+        if self.ucg_rate > 0.0 and not no_dropout and not (self.max_crops > 0):
+            z = (
+                torch.bernoulli(
+                    (1.0 - self.ucg_rate) * torch.ones(z.shape[0], device=z.device)
+                )[:, None]
+                * z
+            )
+            if tokens is not None:
+                tokens = (
+                    expand_dims_like(
+                        torch.bernoulli(
+                            (1.0 - self.ucg_rate)
+                            * torch.ones(tokens.shape[0], device=tokens.device)
+                        ),
+                        tokens,
+                    )
+                    * tokens
+                )
+        if self.unsqueeze_dim:
+            z = z[:, None, :]
+        if self.output_tokens:
+            assert not self.repeat_to_max_len
+            assert not self.pad_to_max_len
+            if self.only_tokens:
+                return tokens
+            return tokens, z
+        if self.repeat_to_max_len:
+            if z.dim() == 2:
+                z_ = z[:, None, :]
+            else:
+                z_ = z
+            return repeat(z_, "b 1 d -> b n d", n=self.max_length), z
+        elif self.pad_to_max_len:
+            assert z.dim() == 3
+            z_pad = torch.cat(
+                (
+                    z,
+                    torch.zeros(
+                        z.shape[0],
+                        self.max_length - z.shape[1],
+                        z.shape[2],
+                        device=z.device,
+                    ),
+                ),
+                1,
+            )
+            return z_pad, z_pad[:, 0, ...]
+        return z
+
+    def encode_with_vision_transformer(self, img):
+        # if self.max_crops > 0:
+        #    img = self.preprocess_by_cropping(img)
+        if img.dim() == 5:
+            assert self.max_crops == img.shape[1]
+            img = rearrange(img, "b n c h w -> (b n) c h w")
+        img = self.preprocess(img)
+        if not self.output_tokens:
+            assert not self.model.visual.output_tokens
+            x = self.model.visual(img)
+            tokens = None
+        else:
+            assert self.model.visual.output_tokens
+            x, tokens = self.model.visual(img)
+            if self.l2_norm_tokens:
+                token_shape = tokens.shape
+                tokens = tokens.flatten(1)
+                tokens = torch.nn.functional.normalize(tokens, dim=-1)
+                tokens = (tokens - .0002) / .0015
+                tokens = tokens.view(token_shape)
+                tokens = (tokens * 1.0957) + .1598
+        if self.max_crops > 0:
+            x = rearrange(x, "(b n) d -> b n d", n=self.max_crops)
+            # drop out between 0 and all along the sequence axis
+            x = (
+                torch.bernoulli(
+                    (1.0 - self.ucg_rate)
+                    * torch.ones(x.shape[0], x.shape[1], 1, device=x.device)
+                )
+                * x
+            )
+            if tokens is not None:
+                tokens = rearrange(tokens, "(b n) t d -> b t (n d)", n=self.max_crops)
+                logpy.warning(
+                    f"You are running very experimental token-concat in {self.__class__.__name__}. "
+                    f"Check what you are doing, and then remove this message."
+                )
+        if self.output_tokens:
+            return x, tokens
+        return x
+
+    def encode(self, text):
+        return self(text)

fMRI-foundation-model-old/fMRI-MAE-mamba/models.py

@@ -0,0 +1,504 @@
+# Copyright (c) 2015-present, Facebook, Inc.
+# All rights reserved.
+import os
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from functools import partial
+from torch import Tensor
+from typing import Optional
+import torch.utils.checkpoint as checkpoint
+
+from einops import rearrange
+from einops.layers.torch import Rearrange
+import numpy as np
+from timm.models.vision_transformer import _cfg
+from timm.models.layers import trunc_normal_
+
+from timm.models.layers import DropPath, to_2tuple
+from timm.models.registry import register_model
+from timm.models.vision_transformer import _load_weights
+
+import math
+
+from mamba_ssm.modules.mamba_simple import Mamba
+
+try:
+    from mamba_ssm.ops.triton.layernorm import RMSNorm, layer_norm_fn, rms_norm_fn
+except ImportError:
+    RMSNorm, layer_norm_fn, rms_norm_fn = None, None, None
+
+def posemb_sincos_4d(patches, temperature=10000, dtype=torch.float32):
+    _, f, d, h, w, dim, device, dtype = (*patches.shape, patches.device, patches.dtype)
+
+    z, y, x, t = torch.meshgrid(
+        torch.arange(f, device=device),
+        torch.arange(d, device=device),
+        torch.arange(h, device=device),
+        torch.arange(w, device=device),
+        indexing="ij",
+    )
+
+    fourier_dim = dim // 8
+
+    omega = torch.arange(fourier_dim, device=device) / (fourier_dim - 1)
+    omega = 1.0 / (temperature**omega)
+
+    z, y, x, t = [v.flatten()[:, None] * omega[None, :] for v in [z, y, x, t]]
+
+    pe = torch.cat(
+        (z.sin(), z.cos(), y.sin(), y.cos(), x.sin(), x.cos(), t.sin(), t.cos()), dim=1
+    )
+    pe = F.pad(pe, (0, dim - (fourier_dim * 8)))
+    return pe.type(dtype)
+
+
+
+class Block(nn.Module):
+    def __init__(
+        self, dim, mixer_cls, norm_cls=nn.LayerNorm, fused_add_norm=False, residual_in_fp32=False,drop_path=0.,
+    ):
+        """
+        Simple block wrapping a mixer class with LayerNorm/RMSNorm and residual connection"
+
+        This Block has a slightly different structure compared to a regular
+        prenorm Transformer block.
+        The standard block is: LN -> MHA/MLP -> Add.
+        [Ref: https://arxiv.org/abs/2002.04745]
+        Here we have: Add -> LN -> Mixer, returning both
+        the hidden_states (output of the mixer) and the residual.
+        This is purely for performance reasons, as we can fuse add and LayerNorm.
+        The residual needs to be provided (except for the very first block).
+        """
+        super().__init__()
+        self.residual_in_fp32 = residual_in_fp32
+        self.fused_add_norm = fused_add_norm
+        self.mixer = mixer_cls(dim)
+        self.norm = norm_cls(dim)
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        if self.fused_add_norm:
+            assert RMSNorm is not None, "RMSNorm import fails"
+            assert isinstance(
+                self.norm, (nn.LayerNorm, RMSNorm)
+            ), "Only LayerNorm and RMSNorm are supported for fused_add_norm"
+
+    def forward(
+        self, hidden_states: Tensor, residual: Optional[Tensor] = None, inference_params=None,
+        use_checkpoint=False
+    ):
+        r"""Pass the input through the encoder layer.
+
+        Args:
+            hidden_states: the sequence to the encoder layer (required).
+            residual: hidden_states = Mixer(LN(residual))
+        """
+        if not self.fused_add_norm:
+            residual = (residual + self.drop_path(hidden_states)) if residual is not None else hidden_states
+            hidden_states = self.norm(residual.to(dtype=self.norm.weight.dtype))
+            if self.residual_in_fp32:
+                residual = residual.to(torch.float32)
+        else:
+            fused_add_norm_fn = rms_norm_fn if isinstance(self.norm, RMSNorm) else layer_norm_fn
+            hidden_states, residual = fused_add_norm_fn(
+                hidden_states if residual is None else self.drop_path(hidden_states),
+                self.norm.weight,
+                self.norm.bias,
+                residual=residual,
+                prenorm=True,
+                residual_in_fp32=self.residual_in_fp32,
+                eps=self.norm.eps,
+            )
+        if use_checkpoint:
+            hidden_states = checkpoint.checkpoint(self.mixer, hidden_states, inference_params)
+        else:
+            hidden_states = self.mixer(hidden_states, inference_params=inference_params)
+        return hidden_states, residual
+
+    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
+        return self.mixer.allocate_inference_cache(batch_size, max_seqlen, dtype=dtype, **kwargs)
+
+
+def create_block(
+    d_model,
+    ssm_cfg=None,
+    norm_epsilon=1e-5,
+    drop_path=0.,
+    rms_norm=True,
+    residual_in_fp32=True,
+    fused_add_norm=True,
+    layer_idx=None,
+    bimamba=True,
+    device=None,
+    dtype=None,
+):
+    import inspect 
+    inspect.getsourcefile(Mamba)
+    factory_kwargs = {"device": device, "dtype": dtype}
+    if ssm_cfg is None:
+        ssm_cfg = {}
+    mixer_cls = partial(Mamba, layer_idx=layer_idx, bimamba=bimamba, **ssm_cfg, **factory_kwargs)
+    norm_cls = partial(nn.LayerNorm if not rms_norm else RMSNorm, eps=norm_epsilon)
+    block = Block(
+        d_model,
+        mixer_cls,
+        norm_cls=norm_cls,
+        drop_path=drop_path,
+        fused_add_norm=fused_add_norm,
+        residual_in_fp32=residual_in_fp32,
+    )
+    block.layer_idx = layer_idx
+    return block
+
+
+# https://github.com/huggingface/transformers/blob/c28d04e9e252a1a099944e325685f14d242ecdcd/src/transformers/models/gpt2/modeling_gpt2.py#L454
+def _init_weights(
+    module,
+    n_layer,
+    initializer_range=0.02,  # Now only used for embedding layer.
+    rescale_prenorm_residual=True,
+    n_residuals_per_layer=1,  # Change to 2 if we have MLP
+):
+    if isinstance(module, nn.Linear):
+        if module.bias is not None:
+            if not getattr(module.bias, "_no_reinit", False):
+                nn.init.zeros_(module.bias)
+    elif isinstance(module, nn.Embedding):
+        nn.init.normal_(module.weight, std=initializer_range)
+
+    if rescale_prenorm_residual:
+        # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme:
+        #   > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale
+        #   > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers.
+        #   >   -- GPT-2 :: https://openai.com/blog/better-language-models/
+        #
+        # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py
+        for name, p in module.named_parameters():
+            if name in ["out_proj.weight", "fc2.weight"]:
+                # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block
+                # Following Pytorch init, except scale by 1/sqrt(2 * n_layer)
+                # We need to reinit p since this code could be called multiple times
+                # Having just p *= scale would repeatedly scale it down
+                nn.init.kaiming_uniform_(p, a=math.sqrt(5))
+                with torch.no_grad():
+                    p /= math.sqrt(n_residuals_per_layer * n_layer)
+
+
+def segm_init_weights(m):
+    if isinstance(m, nn.Linear):
+        trunc_normal_(m.weight, std=0.02)
+        if isinstance(m, nn.Linear) and m.bias is not None:
+            nn.init.constant_(m.bias, 0)
+    elif isinstance(m, nn.LayerNorm):
+        nn.init.constant_(m.bias, 0)
+        nn.init.constant_(m.weight, 1.0)
+
+
+class PatchEmbed(nn.Module):
+    """ Image to Patch Embedding
+    """
+    def __init__(self, img_size=224, patch_size=16, kernel_size=1, in_chans=3, embed_dim=768):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.num_patches = num_patches
+        self.tubelet_size = kernel_size
+
+        self.proj = nn.Conv3d(
+            in_chans, embed_dim, 
+            kernel_size=(kernel_size, patch_size[0], patch_size[1]),
+            stride=(kernel_size, patch_size[0], patch_size[1])
+        )
+
+    def forward(self, x):
+        x = self.proj(x)
+        return x
+
+ 
+# sin-cos position encoding
+# https://github.com/jadore801120/attention-is-all-you-need-pytorch/blob/master/transformer/Models.py#L31
+def get_sinusoid_encoding_table(n_position, d_hid): 
+    ''' Sinusoid position encoding table ''' 
+    # TODO: make it with torch instead of numpy 
+    def get_position_angle_vec(position): 
+        return [position / np.power(10000, 2 * (hid_j // 2) / d_hid) for hid_j in range(d_hid)] 
+
+    sinusoid_table = np.array([get_position_angle_vec(pos_i) for pos_i in range(n_position)]) 
+    sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2]) # dim 2i 
+    sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2]) # dim 2i+1 
+
+    return  torch.tensor(sinusoid_table, dtype=torch.float, requires_grad=False).unsqueeze(0) 
+
+
+class VisionMamba(nn.Module):
+    def __init__(
+            self, 
+            img_size=224, 
+            patch_size=16, 
+            frame_patch_size=4,
+            depth=24, 
+            embed_dim=192, 
+            channels=3, 
+            drop_rate=0.,
+            drop_path_rate=0.1,
+            ssm_cfg=None, 
+            norm_epsilon=1e-5, 
+            initializer_cfg=None,
+            fused_add_norm=True,
+            rms_norm=True, 
+            residual_in_fp32=True,
+            bimamba=True,
+            # video
+            kernel_size=1, 
+            num_frames=8, 
+            use_rope_emb=False,
+            device=None,
+            dtype=None,
+            # checkpoint
+            use_checkpoint=False,
+            checkpoint_num=0,
+        ):
+        factory_kwargs = {"device": device, "dtype": dtype} # follow MambaLMHeadModel
+        super().__init__()
+        self.residual_in_fp32 = residual_in_fp32
+        self.fused_add_norm = fused_add_norm
+        self.use_checkpoint = use_checkpoint
+        self.checkpoint_num = checkpoint_num
+        print(f'Use checkpoint: {use_checkpoint}')
+        print(f'Checkpoint number: {checkpoint_num}')
+       
+
+        # pretrain parameters
+        self.d_model = self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
+        patch_depth, patch_height, patch_width = patch_size
+        self.patchify = Rearrange(
+            "b c (f pf) (d pd) (h ph) (w pw) -> b f d h w (pd ph pw pf c)",
+            pd=patch_depth,
+            ph=patch_height,
+            pw=patch_width,
+            pf=frame_patch_size,
+        )
+        # self.patch_embed = PatchEmbed(
+        #     img_size=img_size, patch_size=patch_size, 
+        #     kernel_size=kernel_size,
+        #     in_chans=channels, embed_dim=embed_dim
+        # )
+        self.patch_dim = channels * patch_depth * patch_height * patch_width * frame_patch_size
+        self.patch_to_emb = nn.Sequential(
+            nn.LayerNorm(self.patch_dim),
+            nn.Linear(self.patch_dim, self.embed_dim),
+            nn.LayerNorm(self.embed_dim),
+        )
+        
+        self.encoder_proj = nn.Sequential(
+                                nn.LayerNorm(self.embed_dim), 
+                                nn.GELU(), 
+                                nn.Linear(self.embed_dim, self.embed_dim),
+                            )
+        
+        self.decoder_proj = nn.Sequential(
+                                nn.LayerNorm(self.embed_dim), 
+                                nn.GELU(), 
+                                nn.Linear(self.embed_dim, self.patch_dim),
+                            )
+
+        self.cls_token = nn.Parameter(torch.randn(1, 1, self.embed_dim))
+        image_depth, image_height, image_width = img_size
+        self.use_rope_emb = use_rope_emb
+        self.posemb_sincos_4d = posemb_sincos_4d(
+                torch.zeros(
+                    1,
+                    num_frames,
+                    image_depth // patch_depth,
+                    image_height // patch_height,
+                    image_width // patch_width,
+                    self.embed_dim,
+                )
+            ).to(device)
+        print ("posemb_sincos_4d", self.posemb_sincos_4d.shape)
+        self.encoder_to_decoder = nn.Linear(self.embed_dim, self.embed_dim, bias=False)
+        self.mask_token = nn.Parameter(torch.zeros(1, 1, self.embed_dim))
+
+        # self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, self.embed_dim))
+        # self.temporal_pos_embedding = nn.Parameter(torch.zeros(1, num_frames // kernel_size, embed_dim))
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+        inter_dpr = [0.0] + dpr
+        self.drop_path = DropPath(drop_path_rate) if drop_path_rate > 0. else nn.Identity()
+        # mamba blocks
+        print ("depth", depth, "d_model", embed_dim, "rms_norm", rms_norm, "residual_in_fp32", residual_in_fp32, "fused_add_norm", fused_add_norm, "bimamba", bimamba, "ssm_cfg", ssm_cfg)
+        self.encoder_layers = nn.ModuleList(
+            [
+                create_block(
+                    embed_dim,
+                    ssm_cfg=ssm_cfg,
+                    norm_epsilon=norm_epsilon,
+                    rms_norm=rms_norm,
+                    residual_in_fp32=residual_in_fp32,
+                    fused_add_norm=fused_add_norm,
+                    layer_idx=i,
+                    bimamba=bimamba,
+                    drop_path=inter_dpr[i],
+                    **factory_kwargs,
+                )
+                for i in range(depth)
+            ]
+        )
+        self.decoder_layers = nn.ModuleList(
+            [
+                create_block(
+                    embed_dim,
+                    ssm_cfg=ssm_cfg,
+                    norm_epsilon=norm_epsilon,
+                    rms_norm=rms_norm,
+                    residual_in_fp32=residual_in_fp32,
+                    fused_add_norm=fused_add_norm,
+                    layer_idx=i,
+                    bimamba=bimamba,
+                    drop_path=inter_dpr[i],
+                    **factory_kwargs,
+                )
+                for i in range(depth)
+            ]
+        )
+        
+        # output head
+        self.norm_f = (nn.LayerNorm if not rms_norm else RMSNorm)(embed_dim, eps=norm_epsilon, **factory_kwargs)
+  
+
+        # original init
+        self.apply(segm_init_weights)
+        # trunc_normal_(self.pos_embed, std=.02)
+
+        # mamba init
+        self.apply(
+            partial(
+                _init_weights,
+                n_layer=depth,
+                **(initializer_cfg if initializer_cfg is not None else {}),
+            )
+        )
+
+    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
+        return {
+            i: layer.allocate_inference_cache(batch_size, max_seqlen, dtype=dtype, **kwargs)
+            for i, layer in enumerate(self.layers)
+        }
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {"pos_embed", "cls_token", "temporal_pos_embedding"}
+    
+    def get_num_layers(self):
+        return (len(self.encoder_layers), len(self.decoder_layers))
+
+    @torch.jit.ignore()
+    def load_pretrained(self, checkpoint_path, prefix=""):
+        _load_weights(self, checkpoint_path, prefix)
+
+    def forward(self, x, inference_params=None, encoder_mask=None, decoder_mask=None, verbose=False):
+        # x = self.forward_features(x, inference_params, encoder_mask, decoder_mask, verbose)
+        # # x = self.head(self.head_drop(x))
+        
+        # ENCODER
+        if decoder_mask is None:
+            if verbose: print(x.shape)
+            x = self.patchify(x)
+            if verbose: print("patched", x.shape)
+            x = self.patch_to_emb(x)
+            if verbose: print("patched_emb", x.shape)
+            x = rearrange(x, "b ... d -> b (...) d")
+            if verbose: print("reshaped", x.shape)
+            if not self.use_rope_emb:
+                if verbose: print("pe", self.posemb_sincos_4d.shape)
+                x = x + self.posemb_sincos_4d
+            if verbose: print("x", x.shape)
+            x = x[:, encoder_mask]
+            # if self.use_cls_token:
+            cls_tokens = self.cls_token.expand(len(x), -1, -1)
+            x = torch.cat((cls_tokens, x), dim=1)
+            if verbose: print("masked", x.shape)
+            
+            # Mamba below
+            residual = None
+            hidden_states = x
+            for idx, layer in enumerate(self.encoder_layers):
+                if self.use_checkpoint and idx < self.checkpoint_num:
+                    hidden_states, residual = layer(
+                        hidden_states, residual, inference_params=inference_params,
+                        use_checkpoint=True
+                    )
+                else:
+                    hidden_states, residual = layer(
+                        hidden_states, residual, inference_params=inference_params
+                    )
+            # assuming fused_add_norm is True
+            fused_add_norm_fn = rms_norm_fn if isinstance(self.norm_f, RMSNorm) else layer_norm_fn
+            hidden_states = fused_add_norm_fn(
+                self.drop_path(hidden_states),
+                self.norm_f.weight,
+                self.norm_f.bias,
+                eps=self.norm_f.eps,
+                residual=residual,
+                prenorm=False, # Set prenorm=False here since we don't need the residual
+                residual_in_fp32=self.residual_in_fp32,
+            )
+        else: # DECODER
+            if verbose: print(x.shape)
+            x = self.encoder_to_decoder(x)
+            B, _, _ = x.shape
+            N = decoder_mask.sum()
+            mask = None
+            if not self.use_rope_emb:
+                pos_embed = self.posemb_sincos_4d
+                if verbose: print("pe", pos_embed.shape)
+                pos_emd_encoder = pos_embed[encoder_mask]
+                pos_emd_decoder = pos_embed[decoder_mask]
+                if verbose: print("pos_emd_encoder", pos_emd_encoder.shape)
+                if verbose: print("pos_emd_decoder", pos_emd_decoder.shape)
+                # if self.use_cls_token:
+                cls_tokens = x[:,:1,:]
+                x = x[:,1:,:]
+                x = torch.cat([x + pos_emd_encoder, 
+                               self.mask_token.repeat(B, N, 1) + pos_emd_decoder], 
+                              dim=1)
+                # if self.use_cls_token:
+                x = torch.cat([cls_tokens, x], dim=1)
+            else:
+                mask = torch.cat((torch.where(encoder_mask)[0], torch.where(decoder_mask)[0]))
+                # No abs positional embeddings for RoPE
+                x = torch.cat([x,self.mask_token.repeat(B, N-1, 1)],dim=1) # N-1 due to cls token
+            
+            # Mamba below
+            residual = None
+            hidden_states = x
+            for idx, layer in enumerate(self.decoder_layers):
+                if self.use_checkpoint and idx < self.checkpoint_num:
+                    hidden_states, residual = layer(
+                        hidden_states, residual, inference_params=inference_params,
+                        use_checkpoint=True
+                    )
+                else:
+                    hidden_states, residual = layer(
+                        hidden_states, residual, inference_params=inference_params
+                    )
+            # assuming fused_add_norm is True
+            fused_add_norm_fn = rms_norm_fn if isinstance(self.norm_f, RMSNorm) else layer_norm_fn
+            hidden_states = fused_add_norm_fn(
+                self.drop_path(hidden_states),
+                self.norm_f.weight,
+                self.norm_f.bias,
+                eps=self.norm_f.eps,
+                residual=residual,
+                prenorm=False, # Set prenorm=False here since we don't need the residual
+                residual_in_fp32=self.residual_in_fp32,
+            )
+        
+        if decoder_mask is None:
+            return self.encoder_proj(hidden_states)
+        else:
+            return self.decoder_proj(hidden_states)

fMRI-foundation-model-old/fMRI-MAE-mamba/utils.py

@@ -0,0 +1,267 @@
+from io import BytesIO
+import os
+import random
+import numpy as np
+import torch
+from einops import rearrange
+from nilearn import plotting
+from PIL import Image
+from skimage import filters
+from torchvision import transforms
+import nibabel as nib
+from nilearn import plotting
+import matplotlib.pyplot as plt
+import re
+
+def my_split_by_node(urls): return urls
+
+def is_interactive():
+    import __main__ as main
+
+    return not hasattr(main, "__file__")
+
+def my_split_by_node(urls): return urls
+
+def seed_everything(seed=0, cudnn_deterministic=True):
+    random.seed(seed)
+    os.environ['PYTHONHASHSEED'] = str(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    
+
+def grayscale_decoder(image_data):
+    return np.array(Image.open(BytesIO(image_data))).astype(np.float32) / 65535
+
+
+def numpy_decoder(npy_data):
+    return np.load(BytesIO(npy_data))
+
+
+def reshape_to_2d(tensor):
+    if tensor.ndim == 5:
+        tensor = tensor[0]
+    assert tensor.ndim == 4
+    return rearrange(tensor, "b h w c -> (b h) (c w)")
+
+
+def reshape_to_original(tensor_2d, h=64, w=64, c=48):
+    # print(tensor_2d.shape) # torch.Size([1, 256, 3072])
+    return rearrange(tensor_2d, "(tr h) (c w) -> tr h w c", h=h, w=w, c=c)
+
+
+def plot_numpy_nii(image):
+    while image.ndim > 3:
+        image = image[0]
+    nii = nib.Nifti1Image(image.astype(np.float32), np.eye(4))  # noqa
+    plotting.plot_epi(nii, cmap="gray")
+
+
+def threshold_based_masking(org_images):
+    thresholds = filters.threshold_multiotsu(org_images.numpy(), classes=3)
+    brain_segmentation = org_images > thresholds.min()
+    return brain_segmentation
+
+
+def get_brain_pos_patches(
+    func,
+    patch_depth=8,
+    patch_height=8,
+    patch_width=8,
+    frame_patch_size=1,
+    masking_strategy="conservative",
+):
+    _, _, depth = func.shape
+    if masking_strategy == "conservative":
+        func = func.sum(axis=(-1), keepdim=True).repeat(1, 1, depth)
+    else:
+        raise Exception("Not implemented other masking strategies than conservative.")
+
+    return func
+
+
+class DataPrepper:
+    def __init__(
+        self,
+        num_frames=4,
+        masking_strategy="MNI",
+        patch_depth=8,
+        patch_height=8,
+        patch_width=8,
+        frame_patch_size=1,
+    ):
+        self.num_frames = num_frames
+        self.masking_strategy = masking_strategy
+        self.patch_depth = 8
+        self.patch_height = 8
+        self.patch_width = 8
+        self.frame_patch_size = 1
+
+    def __call__(self, func):
+        start_timepoint = np.random.choice(np.arange(func.shape[1] - self.num_frames))
+        timepoints = np.arange(start_timepoint, start_timepoint + self.num_frames)
+
+        func = func[:,timepoints]
+
+        if self.masking_strategy=="MNI":
+            return func, None
+        
+        brain_segmentation = threshold_based_masking(func.mean(1))
+        pos_patches = None
+        for brain in brain_segmentation:
+            output = get_brain_pos_patches(
+                brain,
+                patch_depth=self.patch_depth,
+                patch_height=self.patch_height,
+                patch_width=self.patch_width,
+                frame_patch_size=self.frame_patch_size,
+                masking_strategy=self.masking_strategy,
+            )
+            if pos_patches is None:
+                pos_patches = output[None]
+            else:
+                pos_patches = torch.vstack((pos_patches, output[None]))
+        return func, pos_patches
+
+
+def plot_slices(unpatches):
+    if unpatches.ndim == 5:
+        unpatches = unpatches[0]
+    return transforms.ToPILImage()(reshape_to_2d(unpatches))
+
+
+def check_loss(loss):
+    if loss.isnan().any():
+        raise ValueError('NaN loss')
+        
+
+def count_params(model):
+    total = sum(p.numel() for p in model.parameters())
+    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print("param counts:\n{:,} total\n{:,} trainable".format(total, trainable))
+    return trainable
+
+
+def contrastive_loss(
+    cls_token1: torch.Tensor, cls_token2: torch.Tensor, temperature: torch.Tensor
+):
+    feat1 = cls_token1 / cls_token1.norm(dim=1, keepdim=True)
+    feat2 = cls_token2 / cls_token2.norm(dim=1, keepdim=True)
+
+    cosine_sim = feat1 @ feat2.T
+    logit_scale = temperature.exp()  # log scale, learned during training
+    feat1 = cosine_sim * logit_scale
+    feat2 = feat1.T
+
+    labels = torch.arange(feat1.shape[0]).to(feat1.device)
+    loss = (
+        torch.nn.functional.cross_entropy(feat1, labels)
+        + torch.nn.functional.cross_entropy(feat2, labels)
+    ) / 2
+    return loss
+
+### MindEye functions ###
+
+def soft_clip_loss(preds, targs, temp=0.006):
+    clip_clip = (targs @ targs.T)/temp
+    brain_clip = (preds @ targs.T)/temp
+    loss1 = -(brain_clip.log_softmax(-1) * clip_clip.softmax(-1)).sum(-1).mean()
+    loss2 = -(brain_clip.T.log_softmax(-1) * clip_clip.softmax(-1)).sum(-1).mean()
+    
+    loss = (loss1 + loss2)/2
+    return loss
+
+def cosine_anneal(start, end, steps):
+    return end + (start - end)/2 * (1 + torch.cos(torch.pi*torch.arange(steps)/(steps-1)))
+
+def mixco(voxels, beta=0.15, s_thresh=0.5, perm=None, betas=None, select=None):
+    if perm is None:
+        perm = torch.randperm(voxels.shape[0])
+    voxels_shuffle = voxels[perm].to(voxels.device,dtype=voxels.dtype)
+    if betas is None:
+        betas = torch.distributions.Beta(beta, beta).sample([voxels.shape[0]]).to(voxels.device,dtype=voxels.dtype)
+    if select is None:
+        select = (torch.rand(voxels.shape[0]) <= s_thresh).to(voxels.device)
+    betas_shape = [-1] + [1]*(len(voxels.shape)-1)
+    voxels[select] = voxels[select] * betas[select].reshape(*betas_shape) + \
+        voxels_shuffle[select] * (1 - betas[select]).reshape(*betas_shape)
+    betas[~select] = 1
+    return voxels, perm, betas, select
+
+def mixco_nce(preds, targs, temp=0.1, perm=None, betas=None, select=None, distributed=False, 
+              accelerator=None, local_rank=None, bidirectional=True):
+    brain_clip = (preds @ targs.T)/temp
+    
+    if perm is not None and betas is not None and select is not None:
+        probs = torch.diag(betas)
+        probs[torch.arange(preds.shape[0]).to(preds.device), perm] = 1 - betas
+
+        loss = -(brain_clip.log_softmax(-1) * probs).sum(-1).mean()
+        if bidirectional:
+            loss2 = -(brain_clip.T.log_softmax(-1) * probs.T).sum(-1).mean()
+            loss = (loss + loss2)/2
+        return loss
+    else:
+        loss =  F.cross_entropy(brain_clip, torch.arange(brain_clip.shape[0]).to(brain_clip.device))
+        if bidirectional:
+            loss2 = F.cross_entropy(brain_clip.T, torch.arange(brain_clip.shape[0]).to(brain_clip.device))
+            loss = (loss + loss2)/2
+        return loss
+    
+def topk(similarities,labels,k=5):
+    if k > similarities.shape[0]:
+        k = similarities.shape[0]
+    topsum=0
+    for i in range(k):
+        topsum += torch.sum(torch.argsort(similarities,axis=1)[:,-(i+1)] == labels)/len(labels)
+    return topsum
+
+def batchwise_cosine_similarity(Z,B):
+    Z = Z.flatten(1)
+    B = B.flatten(1).T
+    Z_norm = torch.linalg.norm(Z, dim=1, keepdim=True)  # Size (n, 1).
+    B_norm = torch.linalg.norm(B, dim=0, keepdim=True)  # Size (1, b).
+    cosine_similarity = ((Z @ B) / (Z_norm @ B_norm)).T
+    return cosine_similarity
+
+def prenormed_batchwise_cosine_similarity(Z,B):
+    return (Z @ B.T).T
+
+def torch_to_Image(x):
+    if x.ndim==4:
+        x=x[0]
+    return transforms.ToPILImage()(x)
+
+def get_masking_ratio(current_epoch, total_epochs, start_masking_ratio, end_masking_ratio):
+    """Returns the masking ratio for the current epochs. Linearly increase the masking ratio over the span of the training"""
+    return start_masking_ratio + (end_masking_ratio-start_masking_ratio) * ((current_epoch+1)/total_epochs)
+
+def view_brain(data):
+    if torch.is_tensor(data):
+        data = data.numpy()
+    if data.ndim==5:
+        new_nii = nib.Nifti1Image((data[0,0].astype(np.float32)-.5)*2, np.eye(4))
+    elif data.ndim==4:
+        new_nii = nib.Nifti1Image((data[0].astype(np.float32)-.5)*2, np.eye(4))
+    elif data.ndim==3:
+        new_nii = nib.Nifti1Image((data.astype(np.float32)-.5)*2, np.eye(4))
+    else:
+        raise Exception("Check dimensionality of your brain data")
+    return plotting.view_img(new_nii, bg_img=None, vmax=1, cmap=plt.cm.gray, threshold=None)
+
+def get_first_tar(train_urls):
+    if isinstance(train_urls, list):
+        # If train_urls is a list, get the first element
+        url = train_urls[0]
+    else:
+        # If train_urls is a string, treat it as the only element
+        url = train_urls
+
+    # Extract the first tar file using regular expression
+    match = re.search(r'\{(\d+)\.\.', url)
+    if match:
+        first_tar = match.group(1)
+        return f"/scratch/fmri_foundation_datasets/NSD_MNI_wds/{first_tar}.tar"
+    else:
+        return None

fMRI-foundation-model-old/fMRI-MAE/.ipynb_checkpoints/main-checkpoint.ipynb

@@ -0,0 +1,1350 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "b8e236f1-385a-4d93-bb39-bea3ee384d76",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Number of available CUDA devices: 1\n",
+      "LOCAL RANK=0\n",
+      "NUM GPUS=1\n",
+      "NODE=0\n",
+      "GLOBAL RANK=0\n",
+      "WORLD_SIZE=1\n",
+      "PID of this process = 1339581\n",
+      "device = cuda distributed = False num_devices = 1 local rank = 0 world size = 1 data_type = torch.float16\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Import packages and setup gpu configuration.\n",
+    "# This code block shouldnt need to be adjusted!\n",
+    "import os\n",
+    "import sys\n",
+    "import json\n",
+    "import yaml\n",
+    "import numpy as np\n",
+    "import pandas as pd\n",
+    "import copy\n",
+    "import math\n",
+    "from einops import rearrange\n",
+    "from einops.layers.torch import Rearrange\n",
+    "import time\n",
+    "import random\n",
+    "import h5py\n",
+    "import webdataset as wds\n",
+    "import gc\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.nn.functional as F\n",
+    "from torchvision import transforms\n",
+    "import utils\n",
+    "from models import *\n",
+    "import nibabel as nib\n",
+    "from nilearn import plotting\n",
+    "\n",
+    "import schedulefree\n",
+    "\n",
+    "# tf32 data type is faster than standard float32\n",
+    "torch.backends.cuda.matmul.allow_tf32 = True\n",
+    "\n",
+    "### Multi-GPU config ###\n",
+    "device_count = torch.cuda.device_count()\n",
+    "print(f\"Number of available CUDA devices: {device_count}\")\n",
+    "\n",
+    "local_rank = os.getenv('LOCAL_RANK')\n",
+    "if local_rank is None: \n",
+    "    local_rank = 0\n",
+    "else:\n",
+    "    local_rank = int(local_rank)\n",
+    "print(f\"LOCAL RANK={local_rank}\")\n",
+    "\n",
+    "num_devices = os.getenv('NUM_GPUS')\n",
+    "if num_devices is None: \n",
+    "    num_devices = 1\n",
+    "else:\n",
+    "    num_devices = int(num_devices)\n",
+    "print(f\"NUM GPUS={num_devices}\")\n",
+    "distributed = True if num_devices>1 else False\n",
+    "if distributed: assert device_count==num_devices\n",
+    "\n",
+    "node = os.getenv('SLURM_NODEID')\n",
+    "if node is None:\n",
+    "    node = 0\n",
+    "else:\n",
+    "    node = int(node)\n",
+    "print(f\"NODE={node}\")\n",
+    "\n",
+    "global_rank = os.getenv('RANK')\n",
+    "if global_rank is None:\n",
+    "    global_rank = 0\n",
+    "else:\n",
+    "    global_rank = int(global_rank)\n",
+    "print(f\"GLOBAL RANK={global_rank}\")\n",
+    "\n",
+    "world_size = os.getenv('WORLD_SIZE')\n",
+    "if world_size is None: \n",
+    "    world_size = 1\n",
+    "else:\n",
+    "    world_size = int(world_size)\n",
+    "print(f\"WORLD_SIZE={world_size}\")\n",
+    "\n",
+    "if utils.is_interactive():\n",
+    "    # Following allows you to change functions in models.py or utils.py and \n",
+    "    # have this notebook automatically update with your revisions\n",
+    "    %load_ext autoreload\n",
+    "    %autoreload 2\n",
+    "    from tqdm.notebook import tqdm\n",
+    "else:\n",
+    "    from tqdm import tqdm\n",
+    "\n",
+    "# Load parameters from yaml config\n",
+    "config = yaml.load(open('config.yaml', 'r'), Loader=yaml.FullLoader)\n",
+    "\n",
+    "# create global variables from the config\n",
+    "for attribute_name in config.keys():\n",
+    "    globals()[attribute_name] = config[f'{attribute_name}']\n",
+    "\n",
+    "data_type = torch.float16 # change depending on your mixed_precision\n",
+    "# batch_size = global_batch_size // num_devices\n",
+    "global_batch_size = batch_size * world_size\n",
+    "\n",
+    "# FSDP Setup\n",
+    "if distributed:\n",
+    "    import torch.distributed as dist\n",
+    "    import torch.multiprocessing as mp\n",
+    "    from torch.distributed.fsdp import FullyShardedDataParallel as FSDP\n",
+    "    from torch.distributed.fsdp.api import BackwardPrefetch, CPUOffload, ShardingStrategy\n",
+    "    import functools\n",
+    "    from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy, transformer_auto_wrap_policy\n",
+    "    print(\"starting init_process_group...\")\n",
+    "    dist.init_process_group(\"nccl\", rank=global_rank, world_size=world_size)\n",
+    "    print(f\"setting device to cuda:{local_rank}\")\n",
+    "    try:\n",
+    "        torch.cuda.set_device(local_rank)\n",
+    "        device = torch.device('cuda',local_rank)\n",
+    "        print(f\"\\nSuccessfully set cuda:{local_rank} | global_rank{global_rank} | node{node}\")\n",
+    "    except Exception as error:        \n",
+    "        print(f\"\\nFAILED TO SET DEVICE cuda:{local_rank} | global_rank{global_rank} | node{node}\")\n",
+    "        print(\"An exception occurred:\", error)\n",
+    "        \n",
+    "else:\n",
+    "    device = torch.device('cuda')\n",
+    "\n",
+    "print(\"PID of this process =\",os.getpid())\n",
+    "print(\"device =\", device, \"distributed =\",distributed, \"num_devices =\", num_devices, \"local rank =\", local_rank, \"world size =\", world_size, \"data_type =\", data_type)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d08179db-9c6a-4bc6-a245-79fae6884ca2",
+   "metadata": {},
+   "source": [
+    "# Configuration"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "e6b1c3fe-28ab-40b7-8906-c6a9c8070d00",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'model_name': 'mini_nomask_logitsCLS_downstream_40ep_l', 'use_cls_token': True, 'use_contrastive_loss': True, 'contrastive_loss_weight': 0.1, 'batch_size': 256, 'num_workers': 10, 'num_epochs': 20, 'seed': 42, 'max_lr': 3e-06, 'num_samples_per_epoch': 1024, 'test_num_samples_per_epoch': 384, 'ckpt_saving': True, 'ckpt_interval': 50, 'resume_from_ckpt': True, 'wandb_log': True, 'tube_start_masking_ratio': 0.75, 'tube_end_masking_ratio': 0.75, 'decoder_mask_ratio': 0.75, 'patch_size': [8, 8, 8], 'frame_patch_size': 4, 'use_rope_emb': False, 'masking_strategy': 'None', 'encoder_model': 'vit_mini', 'decoder_model': 'vit_mini', 'img_size': [88, 104, 72], 'num_frames': 4, 'is_s3': False, 'train_urls': ['/weka/proj-fmri/shared/NSD_MNI_wds/{000000..000699}.tar'], 'test_urls': ['/weka/proj-fmri/shared/NSD_MNI_wds/{000700..000738}.tar']}\n",
+      "outdir /weka/proj-fmri/paulscotti/fMRI-foundation-model/ckpts/mini_nomask_logitsCLS_downstream_40ep_l\n",
+      "global_batch_size 256\n",
+      "use_cls_token True\n",
+      "num_patches 5148\n",
+      "num_patches_per_timepoint 1287\n",
+      "num_encoder_patches 321\n",
+      "num_decoder_patches 321\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(config)\n",
+    "\n",
+    "# seed all random functions\n",
+    "utils.seed_everything(seed)\n",
+    "\n",
+    "outdir = os.path.abspath(f'../ckpts/{model_name}')\n",
+    "os.makedirs(outdir,exist_ok=True)\n",
+    "print(\"outdir\", outdir)\n",
+    "print(\"global_batch_size\", global_batch_size)\n",
+    "print(\"use_cls_token\", use_cls_token)\n",
+    "\n",
+    "if type(patch_size) == int:\n",
+    "    patch_size = [patch_size,patch_size,patch_size]\n",
+    "patch_depth = patch_size[0]\n",
+    "patch_height = patch_size[1]\n",
+    "patch_width = patch_size[2]\n",
+    "\n",
+    "num_patches = int(\n",
+    "    (img_size[0] / patch_depth)\n",
+    "    * (img_size[1] / patch_height)\n",
+    "    * (img_size[2] / patch_width)\n",
+    "    * num_frames\n",
+    ")\n",
+    "num_patches_per_timepoint = num_patches // frame_patch_size\n",
+    "num_encoder_patches = int(np.floor((num_patches_per_timepoint * num_frames // frame_patch_size) * (1 - tube_start_masking_ratio)))\n",
+    "num_decoder_patches = int(np.floor((num_patches_per_timepoint * num_frames  // frame_patch_size) * (1 - decoder_mask_ratio)))\n",
+    "print(\"num_patches\", num_patches)\n",
+    "print(\"num_patches_per_timepoint\", num_patches_per_timepoint)\n",
+    "print(\"num_encoder_patches\", num_encoder_patches)\n",
+    "print(\"num_decoder_patches\", num_decoder_patches)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1ae8419d-988f-42c6-acb6-b258e0694eee",
+   "metadata": {},
+   "source": [
+    "# Prep models"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "5e37b6b9-5b91-4c4a-af85-ac2af69704e7",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "param counts:\n",
+      "1,693,008 total\n",
+      "1,693,008 trainable\n",
+      "\n",
+      "encoder\n",
+      "torch.Size([256, 1, 4, 88, 104, 72])\n",
+      "patched torch.Size([256, 1, 11, 13, 9, 2048])\n",
+      "reshaped torch.Size([256, 1287, 2048])\n",
+      "masked torch.Size([256, 321, 2048])\n",
+      "patched_emb torch.Size([256, 321, 48])\n",
+      "pe torch.Size([1287, 48])\n",
+      "masked torch.Size([256, 322, 48])\n",
+      "torch.Size([256, 322, 48])\n",
+      "\n",
+      "decoder\n",
+      "torch.Size([256, 322, 48])\n",
+      "pe torch.Size([1287, 48])\n",
+      "pos_emd_encoder torch.Size([321, 48])\n",
+      "pos_emd_decoder torch.Size([321, 48])\n",
+      "x_concat torch.Size([256, 643, 48])\n",
+      "torch.Size([256, 643, 48])\n",
+      "proj torch.Size([256, 643, 2048])\n",
+      "\n",
+      "enc_cls_token torch.Size([256, 1, 48])\n",
+      "encoder_patches torch.Size([256, 321, 48])\n",
+      "dec_cls_token torch.Size([256, 1, 2048])\n",
+      "decoder_patches torch.Size([256, 642, 2048])\n"
+     ]
+    }
+   ],
+   "source": [
+    "vit_size = {\n",
+    "    \"encoder\": encoder_model,\n",
+    "    \"decoder\": decoder_model\n",
+    "}\n",
+    "    \n",
+    "model = get_vit(\n",
+    "    size=vit_size,\n",
+    "    image_size=img_size,  # depth, height, width\n",
+    "    image_patch_size=(patch_depth,patch_height,patch_width),  # depth, height, width patch size\n",
+    "    frames=num_frames,\n",
+    "    frame_patch_size=frame_patch_size,\n",
+    "    channels=1,\n",
+    "    use_rope_emb=use_rope_emb,\n",
+    "    use_cls_token=use_cls_token,\n",
+    ")\n",
+    "utils.count_params(model)\n",
+    "\n",
+    "# function to select random num_frames from sample and obtain brain-positive patches\n",
+    "aug_transform = utils.DataPrepper(\n",
+    "    num_frames=num_frames*2,\n",
+    "    masking_strategy=masking_strategy,\n",
+    "    patch_depth=patch_depth,\n",
+    "    patch_height=patch_height,\n",
+    "    patch_width=patch_width,\n",
+    "    frame_patch_size=frame_patch_size,\n",
+    ")\n",
+    "\n",
+    "# test that the model works without error\n",
+    "model = model.to(device)\n",
+    "encoder_mask = torch.zeros(num_patches_per_timepoint).to(torch.bool)\n",
+    "encoder_mask[:num_encoder_patches] = True\n",
+    "decoder_mask = torch.zeros(num_patches_per_timepoint).to(torch.bool)\n",
+    "decoder_mask[-num_decoder_patches:] = True\n",
+    "decoder_mask[encoder_mask] = False\n",
+    "with torch.no_grad():\n",
+    "    print(\"\\nencoder\")\n",
+    "    encoder_out = model(\n",
+    "                torch.randn(batch_size, 1, num_frames, img_size[0], img_size[1], img_size[2]).to(device),\n",
+    "                encoder_mask=encoder_mask,\n",
+    "                verbose=True)\n",
+    "    print(\"\\ndecoder\")\n",
+    "    decoder_out = model(\n",
+    "                encoder_out, \n",
+    "                encoder_mask=encoder_mask, \n",
+    "                decoder_mask=decoder_mask, \n",
+    "                verbose=True)\n",
+    "    if use_cls_token:\n",
+    "        enc_cls_token = encoder_out[:, :1, :]\n",
+    "        encoder_patches = encoder_out[:, 1:, :]\n",
+    "        dec_cls_token = decoder_out[:, :1, :]\n",
+    "        decoder_patches = decoder_out[:, 1:, :]\n",
+    "        print(\"\\nenc_cls_token\", enc_cls_token.shape)\n",
+    "        print(\"encoder_patches\", encoder_patches.shape)\n",
+    "        print(\"dec_cls_token\", dec_cls_token.shape)\n",
+    "        print(\"decoder_patches\", decoder_patches.shape)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "414cb007-5552-4c5a-8fbd-9b158565770c",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "### Add \"linear\" probe"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "bb3df855-81c4-4d2c-aff2-42a19dd0e67e",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "class LinearProbe(nn.Module):\n",
+    "    def __init__(self, input_dim, h=256, num_classes=8):\n",
+    "        super(LinearProbe, self).__init__()\n",
+    "        # self.classifier = nn.Linear(input_dim, num_classes)\n",
+    "        self.classifier = nn.Sequential(\n",
+    "            nn.LayerNorm(input_dim),\n",
+    "            nn.GELU(),\n",
+    "            nn.Linear(input_dim, h),\n",
+    "            nn.LayerNorm(h),\n",
+    "            nn.GELU(),\n",
+    "            nn.Linear(h, h),\n",
+    "            nn.LayerNorm(h),\n",
+    "            nn.GELU(),\n",
+    "            nn.Linear(h, num_classes)\n",
+    "        )\n",
+    "    def forward(self, x):\n",
+    "        x = self.classifier(x)\n",
+    "        return x"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "2e6a865b-8348-486e-ab88-040a87fa2ee5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# if use_cls_token:\n",
+    "#     model.cont = LinearProbe((num_encoder_patches+1)*model.encoder_embed_dim,h=768,num_classes=768)\n",
+    "# else:\n",
+    "# model.cont = LinearProbe(model.encoder_embed_dim,h=256,num_classes=256)\n",
+    "# model = model.to(device)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3dd51ddf-fb71-48f4-bdbd-88753b44d2aa",
+   "metadata": {},
+   "source": [
+    "## Create dataset and data loaders"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "75dd0398-0eb8-464e-8bb8-13e7ce1a8480",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# from dataloader import create_dataset, create_loader\n",
+    "# train_urls = train_urls[0]\n",
+    "# print(train_urls)\n",
+    "\n",
+    "# train_dp = create_dataset(train_urls, \n",
+    "#                           is_s3=train_urls[:2]==\"s3\", \n",
+    "#                           sample_shuffle=100, shard_shuffle=100)\n",
+    "# train_dl = create_loader(train_dp, batch_size=batch_size, num_workers=num_workers)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "cadb1eb6-dc41-416b-ace5-70bd7c7d887c",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "['/weka/proj-fmri/shared/NSD_MNI_wds/{000000..000699}.tar']\n",
+      "['/weka/proj-fmri/shared/NSD_MNI_wds/{000700..000738}.tar']\n"
+     ]
+    }
+   ],
+   "source": [
+    "def log_and_continue(exn):\n",
+    "    \"\"\"Call in an exception handler to ignore any exception, issue a warning, and continue.\"\"\"\n",
+    "    print(f'Handling webdataset error ({repr(exn)}). Ignoring.')\n",
+    "    return True\n",
+    "\n",
+    "def filter_corrupted_images(sample):\n",
+    "    \"\"\"If all the required files are not present don't use them.\"\"\"\n",
+    "    correct_data = (\"func.npy\" in sample)\n",
+    "    return correct_data\n",
+    "\n",
+    "### ================      Train Dataset and DataLoader    ====================\n",
+    "from braceexpand import braceexpand\n",
+    "print(train_urls)\n",
+    "if is_s3:\n",
+    "    expanded_urls = [f\"pipe:aws s3 cp {url} -\" for pattern in train_urls for url in braceexpand(pattern)]\n",
+    "else:\n",
+    "    expanded_urls = [str(url) for pattern in train_urls for url in braceexpand(pattern)]\n",
+    "\n",
+    "train_data = (\n",
+    "    wds.WebDataset(expanded_urls, resampled=True, nodesplitter=wds.split_by_node, handler=log_and_continue)\n",
+    "    .shuffle(100, initial=100, rng=random.Random(seed))\n",
+    "    .select(filter_corrupted_images)\n",
+    "    .decode(\"torch\")\n",
+    ")\n",
+    "train_dl = torch.utils.data.DataLoader(train_data, batch_size=batch_size, shuffle=False, drop_last=True, pin_memory=True)\n",
+    "\n",
+    "### ================      Test Dataset and DataLoader    ====================\n",
+    "print(test_urls)\n",
+    "if is_s3:\n",
+    "    expanded_urls = [f\"pipe:aws s3 cp {url} -\" for pattern in test_urls for url in braceexpand(pattern)]\n",
+    "else:\n",
+    "    expanded_urls = [str(url) for pattern in train_urls for url in braceexpand(pattern)]\n",
+    "\n",
+    "test_data = (\n",
+    "    wds.WebDataset(expanded_urls, resampled=True, nodesplitter=wds.split_by_node, handler=log_and_continue)\n",
+    "    .shuffle(100, initial=100, rng=random.Random(seed))\n",
+    "    .select(filter_corrupted_images)\n",
+    "    .decode(\"torch\")\n",
+    ")\n",
+    "test_dl = torch.utils.data.DataLoader(test_data, batch_size=batch_size, shuffle=False, drop_last=True, pin_memory=True)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "15d05a32-12eb-494a-82df-dce4c7e9924c",
+   "metadata": {},
+   "source": [
+    "### Check data loaders work"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "480e7c0b-f58e-4c35-80fe-1284ee0e3966",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# if utils.is_interactive():\n",
+    "#     start_time = time.time() \n",
+    "#     num_it = 2\n",
+    "#     print(f\"Yielding {num_it} batches\")\n",
+    "    \n",
+    "#     for i, batch in enumerate(test_dl):\n",
+    "#         print(\"iter\",i)\n",
+    "#         input_func = batch['func.npy']\n",
+    "#         subject_id = batch['subject_id.txt']\n",
+    "#         subject_id = [int(subject[-2:]) for subject in subject_id]\n",
+    "#         # session_id = batch['session_id.txt']\n",
+    "#         # session_id = [int(session[-2:]) for session in session_id]\n",
+    "#         func, brain_pos_pats = aug_transform(input_func)\n",
+    "#         if i >= (num_it-1):\n",
+    "#             break\n",
+    "    \n",
+    "#     print(\"Done!\")\n",
+    "#     print(\"input_func\", input_func.shape)\n",
+    "#     print(\"func\", func.shape)\n",
+    "#     print(\"subject_id\", subject_id)\n",
+    "\n",
+    "#     end_time = time.time()  \n",
+    "#     execution_time = end_time - start_time  \n",
+    "#     print(f\"Execution time: {execution_time} seconds\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "349c6fba-5494-4964-b03d-e01c3afe47db",
+   "metadata": {},
+   "source": [
+    "### Playing with the data, visualization of patching + masking"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "fea69a0b-f0a0-4309-b689-553b911a8da5",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# if utils.is_interactive():\n",
+    "#     func, brain_pos_pats = aug_transform(input_func)\n",
+    "#     print(func.shape)\n",
+    "#     display(utils.view_brain(func,cut_coords=(44,44,44)))\n",
+    "# # plt.hist(func[0,0].flatten().clamp(.25,3),bins=100)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2b8e6baa-4b1c-4f38-b078-70b2b092d14d",
+   "metadata": {},
+   "source": [
+    "# Set up optimizer and saving functions"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "9553ad47-e76d-4ee1-a238-aad13a4be043",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "if distributed:    \n",
+    "    # my_auto_wrap_policy = functools.partial(\n",
+    "    #     size_based_auto_wrap_policy, min_num_params=200000\n",
+    "    # )\n",
+    "    my_auto_wrap_policy = functools.partial(\n",
+    "        transformer_auto_wrap_policy, \n",
+    "        transformer_layer_cls={\n",
+    "            Attention, # <--- Your Transformer layer class\n",
+    "        },\n",
+    "    )\n",
+    "    print(f\"\\nPrepping FSDP on {global_rank} {node}...\\n\")\n",
+    "    model = model.to(device)\n",
+    "    model = FSDP(\n",
+    "        model,\n",
+    "        sharding_strategy=ShardingStrategy.HYBRID_SHARD,\n",
+    "        auto_wrap_policy=my_auto_wrap_policy,\n",
+    "        use_orig_params=False,\n",
+    "        cpu_offload=None, #CPUOffload(offload_params=True)\n",
+    "        sync_module_states=True,\n",
+    "        limit_all_gathers=True, # See https://github.com/pytorch/pytorch/issues/91165\n",
+    "        device_id=device,\n",
+    "    )\n",
+    "    print(f\"\\nSuccessfully loaded FSDP model to device on global_rank {global_rank}\\n\")\n",
+    "    dist.barrier()\n",
+    "    print(f\"\\nSuccessfully loaded FSDP model to device on global_rank {global_rank}\\n\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "4da73c08-ca61-48ef-9e63-b70db6f07a59",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "num_iterations_per_epoch 4\n",
+      "probe_num_iterations_per_epoch 1\n",
+      "total_steps 80\n",
+      "\n",
+      "Done with model preparations!\n",
+      "param counts:\n",
+      "1,838,256 total\n",
+      "1,838,256 trainable\n"
+     ]
+    }
+   ],
+   "source": [
+    "no_decay = [\"bias\", \"LayerNorm.bias\", \"LayerNorm.weight\"]\n",
+    "opt_grouped_parameters = [\n",
+    "    {'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': 1e-2},\n",
+    "    {'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0},\n",
+    "]\n",
+    "\n",
+    "if distributed:\n",
+    "    max_lr = max_lr * global_batch_size\n",
+    "    print(f\"multiply lr {max_lr} by global batch size: max_lr={max_lr}\")\n",
+    "\n",
+    "# optimizer = torch.optim.AdamW(opt_grouped_parameters, lr=max_lr)\n",
+    "optimizer = schedulefree.AdamWScheduleFree(opt_grouped_parameters, lr=max_lr)\n",
+    "\n",
+    "num_iterations_per_epoch = num_samples_per_epoch // global_batch_size\n",
+    "print(\"num_iterations_per_epoch\", num_iterations_per_epoch)\n",
+    "\n",
+    "probe_num_iterations_per_epoch = test_num_samples_per_epoch // global_batch_size\n",
+    "print(\"probe_num_iterations_per_epoch\", probe_num_iterations_per_epoch)\n",
+    "\n",
+    "total_steps = num_epochs * num_iterations_per_epoch * num_devices\n",
+    "print(\"total_steps\", total_steps)\n",
+    "\n",
+    "# lr_scheduler = torch.optim.lr_scheduler.OneCycleLR(\n",
+    "#     optimizer,\n",
+    "#     max_lr=max_lr,\n",
+    "#     total_steps=total_steps,\n",
+    "# )\n",
+    "\n",
+    "print(\"\\nDone with model preparations!\")\n",
+    "num_params = utils.count_params(model)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "f51f674d-7cd4-49d1-bd8d-c697d9614f65",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "def save_ckpt(model,tag=\"last\"):\n",
+    "    if distributed: dist.barrier()\n",
+    "    model_states = model.state_dict()\n",
+    "    if global_rank == 0:\n",
+    "        os.makedirs(outdir,exist_ok=True)\n",
+    "        ckpt_path = outdir+f'/{tag}.pth'\n",
+    "        torch.save({\n",
+    "            'epoch': epoch,\n",
+    "            'model_state_dict': model_states,\n",
+    "            'optimizer_state_dict': optimizer.state_dict(),\n",
+    "        }, ckpt_path)\n",
+    "        print(f\"\\n---saved {ckpt_path}!---\\n\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ef9b4262-435b-4872-ab4a-424cb9dfd37a",
+   "metadata": {},
+   "source": [
+    "# Start wandb (if enabled)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "56431733-4fb5-4072-840e-22536608f9f1",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "wandb found run mini_nomask_logitsCLS_downstream_40ep_l\n",
+      "wandb_config:\n",
+      " {'model_name': 'mini_nomask_logitsCLS_downstream_40ep_l', 'global_batch_size': 256, 'batch_size': 256, 'num_epochs': 20, 'num_samples_per_epoch': 1024, 'test_num_samples_per_epoch': 384, 'num_iterations_per_epoch': 4, 'encoder_model': 'vit_mini', 'decoder_model': 'vit_mini', 'tube_start_masking_ratio': 0.75, 'tube_end_masking_ratio': 0.75, 'decoder_mask_ratio': 0.75, 'num_frames': 4, 'patch_size': [8, 8, 8], 'frame_patch_size': 4, 'use_contrastive_loss': True, 'use_cls_token': True, 'contrastive_loss_weight': 0.1, 'num_params': 1838256, 'max_lr': 3e-06, 'ckpt_interval': 50, 'ckpt_saving': False, 'seed': 42, 'distributed': False, 'num_devices': 1, 'world_size': 1, 'train_urls': ['/weka/proj-fmri/shared/NSD_MNI_wds/{000000..000699}.tar']}\n",
+      "wandb_id: mini_nomask_logitsCLS_downstream_40ep_l\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mpaul-scotti\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
+     ]
+    },
+    {
+     "data": {
+      "text/html": [
+       "wandb version 0.16.6 is available!  To upgrade, please run:\n",
+       " $ pip install wandb --upgrade"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "Tracking run with wandb version 0.16.4"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "Run data is saved locally in <code>/weka/proj-fmri/paulscotti/MindEyeV2/wandb/run-20240421_203216-mini_nomask_logitsCLS_downstream_40ep_l</code>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "Syncing run <strong><a href='https://wandb.ai/paul-scotti/found/runs/mini_nomask_logitsCLS_downstream_40ep_l' target=\"_blank\">mini_nomask_logitsCLS_downstream_40ep_l</a></strong> to <a href='https://wandb.ai/paul-scotti/found' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       " View project at <a href='https://wandb.ai/paul-scotti/found' target=\"_blank\">https://wandb.ai/paul-scotti/found</a>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       " View run at <a href='https://wandb.ai/paul-scotti/found/runs/mini_nomask_logitsCLS_downstream_40ep_l' target=\"_blank\">https://wandb.ai/paul-scotti/found/runs/mini_nomask_logitsCLS_downstream_40ep_l</a>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "if utils.is_interactive():\n",
+    "#     wandb_log = False\n",
+    "    ckpt_saving = False\n",
+    "if local_rank==0 and wandb_log: # only use main process for wandb logging\n",
+    "    import wandb\n",
+    "    wandb_project = 'found'\n",
+    "    print(f\"wandb {wandb_project} run {model_name}\")\n",
+    "    # need to configure wandb beforehand in terminal with \"wandb init\"!\n",
+    "    wandb_config = {\n",
+    "      \"model_name\": model_name,\n",
+    "      \"global_batch_size\": global_batch_size,\n",
+    "      \"batch_size\": batch_size,\n",
+    "      \"num_epochs\": num_epochs,\n",
+    "      \"num_samples_per_epoch\": num_samples_per_epoch,\n",
+    "      \"test_num_samples_per_epoch\": test_num_samples_per_epoch,\n",
+    "      \"num_iterations_per_epoch\": num_iterations_per_epoch,\n",
+    "      \"encoder_model\": encoder_model,\n",
+    "      \"decoder_model\": decoder_model,\n",
+    "      \"tube_start_masking_ratio\": tube_start_masking_ratio,\n",
+    "      \"tube_end_masking_ratio\": tube_end_masking_ratio,\n",
+    "      \"decoder_mask_ratio\": decoder_mask_ratio,\n",
+    "      \"num_frames\": num_frames,\n",
+    "      \"patch_size\": patch_size,\n",
+    "      \"frame_patch_size\": frame_patch_size,\n",
+    "      \"use_contrastive_loss\": use_contrastive_loss,\n",
+    "      \"use_cls_token\": use_cls_token,\n",
+    "      \"contrastive_loss_weight\": contrastive_loss_weight,\n",
+    "      \"num_params\": num_params,\n",
+    "      \"max_lr\": max_lr,\n",
+    "      \"ckpt_interval\": ckpt_interval,\n",
+    "      \"ckpt_saving\": ckpt_saving,\n",
+    "      \"seed\": seed,\n",
+    "      \"distributed\": distributed,\n",
+    "      \"num_devices\": num_devices,\n",
+    "      \"world_size\": world_size,\n",
+    "      \"train_urls\": train_urls,\n",
+    "    }\n",
+    "    print(\"wandb_config:\\n\",wandb_config)\n",
+    "    print(\"wandb_id:\",model_name)\n",
+    "    wandb.init(\n",
+    "        id=model_name,\n",
+    "        project=wandb_project,\n",
+    "        name=model_name,\n",
+    "        config=wandb_config,\n",
+    "        resume=\"allow\",\n",
+    "    )\n",
+    "else:\n",
+    "    wandb_log = False"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c43a5055-8afd-468a-93bf-32f94bd1d042",
+   "metadata": {},
+   "source": [
+    "# Start training"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "f5e394dd-2745-41fa-a5aa-54b7b1373a8f",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "epoch = 0\n",
+    "lrs, train_losses, recon_losses, contrastive_losses = [], [], [], []\n",
+    "cos_sim_encoder_output, cos_sim_decoder_output, cos_sim_encoder_output_patchwise = [], [], []\n",
+    "probe_losses, probe_accs, test_losses, test_accs = [], [], [], []"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "id": "0b9a5605-fd81-41d4-964e-1752ed6e1289",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# # resume from ckpt (e.g., if you are resuming from a run that got pre-empted)\n",
+    "# load_progress = False\n",
+    "# if wandb_log:\n",
+    "#     if wandb.run.resumed:\n",
+    "#         load_checkpoint_in_model(model, outdir+\"/last\")\n",
+    "#         load_progress = True\n",
+    "# elif resume_from_ckpt: # if resuming without using wandb\n",
+    "#     load_checkpoint_in_model(model, outdir+\"/last\")\n",
+    "#     load_progress = True"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "id": "154e8f51-6918-45df-b6b8-ca46c89ba34e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "if masking_strategy==\"MNI\":\n",
+    "    from einops.layers.torch import Rearrange\n",
+    "    MNI_brain = nib.load(\"/weka/proj-fmri/paulscotti/fMRI-foundation-model/dataset_creation/afni_conversion/tpl-MNI152NLin2009cAsym_res-02_T1w_brain.nii.gz\").get_fdata()\n",
+    "    brain_pos_voxels = MNI_brain[6:94,8:112,10:82]\n",
+    "    brain_pos_pats = Rearrange(\n",
+    "            \"b c (f pf) (d pd) (h ph) (w pw) -> b f d h w (pd ph pw pf c)\",\n",
+    "            pd=patch_depth,\n",
+    "            ph=patch_height,\n",
+    "            pw=patch_width,\n",
+    "            pf=1,\n",
+    "        )(torch.Tensor(brain_pos_voxels)[None,None,None])\n",
+    "    brain_pos_pats_vit = rearrange(brain_pos_pats, \"b ... d -> b (...) d\").mean(-1)[0]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "id": "c3461199-e805-4e9c-8c91-894e83cf8bc3",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "2778883b6d5b442dba4d66aa5a0be678",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Overall:   0%|          | 0/20 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "0 0.078125 2.080991744995117\n",
+      "test 0 0.146484375 2.087881088256836\n",
+      "test 1 0.033203125 2.1139793395996094\n",
+      "{'train/loss': 0.6546683609485626, 'train/recon_losses': nan, 'train/contrastive_losses': 6.546683311462402, 'train/num_steps': 0, 'train/cos_sim_encoder_output': 0.6011962890625, 'train/cos_sim_decoder_output': nan, 'train/cos_sim_encoder_output_patchwise': 0.2939453125, 'train/probe_losses': 2.080991744995117, 'train/probe_accs': 0.078125, 'test/probe_losses': 2.1009302139282227, 'test/probe_accs': 0.08984375, 'lr': 3e-06, 'epoch': 0, 'tube_mask_ratio': 0.75, 'decoder_mask_ratio': 0.75}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/admin/home-paulscotti/found/lib/python3.10/site-packages/numpy/core/fromnumeric.py:3504: RuntimeWarning: Mean of empty slice.\n",
+      "  return _methods._mean(a, axis=axis, dtype=dtype,\n",
+      "/admin/home-paulscotti/found/lib/python3.10/site-packages/numpy/core/_methods.py:129: RuntimeWarning: invalid value encountered in scalar divide\n",
+      "  ret = ret.dtype.type(ret / rcount)\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "0 0.2265625 2.132457733154297\n",
+      "test 0 0.10546875 2.0733165740966797\n",
+      "test 1 0.1875 1.9566822052001953\n",
+      "{'train/loss': 0.646026149392128, 'train/recon_losses': nan, 'train/contrastive_losses': 6.460261344909668, 'train/num_steps': 0, 'train/cos_sim_encoder_output': 0.6077880859375, 'train/cos_sim_decoder_output': nan, 'train/cos_sim_encoder_output_patchwise': 0.294189453125, 'train/probe_losses': 2.132457733154297, 'train/probe_accs': 0.2265625, 'test/probe_losses': 2.0149993896484375, 'test/probe_accs': 0.146484375, 'lr': 3e-06, 'epoch': 1, 'tube_mask_ratio': 0.75, 'decoder_mask_ratio': 0.75}\n",
+      "0 0.123046875 2.136058807373047\n",
+      "test 0 0.119140625 2.041532516479492\n",
+      "test 1 0.275390625 1.9539203643798828\n",
+      "{'train/loss': 0.6089650392532349, 'train/recon_losses': nan, 'train/contrastive_losses': 6.0896501541137695, 'train/num_steps': 0, 'train/cos_sim_encoder_output': 0.609375, 'train/cos_sim_decoder_output': nan, 'train/cos_sim_encoder_output_patchwise': 0.3001708984375, 'train/probe_losses': 2.136058807373047, 'train/probe_accs': 0.123046875, 'test/probe_losses': 1.9977264404296875, 'test/probe_accs': 0.197265625, 'lr': 3e-06, 'epoch': 2, 'tube_mask_ratio': 0.75, 'decoder_mask_ratio': 0.75}\n",
+      "0 0.294921875 1.9231395721435547\n",
+      "test 0 0.123046875 1.9924964904785156\n",
+      "test 1 0.232421875 2.002117156982422\n",
+      "{'train/loss': 0.6081273704767227, 'train/recon_losses': nan, 'train/contrastive_losses': 6.081273555755615, 'train/num_steps': 0, 'train/cos_sim_encoder_output': 0.6173095703125, 'train/cos_sim_decoder_output': nan, 'train/cos_sim_encoder_output_patchwise': 0.29730224609375, 'train/probe_losses': 1.9231395721435547, 'train/probe_accs': 0.294921875, 'test/probe_losses': 1.9973068237304688, 'test/probe_accs': 0.177734375, 'lr': 3e-06, 'epoch': 3, 'tube_mask_ratio': 0.75, 'decoder_mask_ratio': 0.75}\n",
+      "0 0.15625 1.9835872650146484\n",
+      "test 0 0.162109375 2.0044002532958984\n",
+      "test 1 0.181640625 1.993692398071289\n",
+      "{'train/loss': 0.6120030879974365, 'train/recon_losses': nan, 'train/contrastive_losses': 6.120030879974365, 'train/num_steps': 0, 'train/cos_sim_encoder_output': 0.61279296875, 'train/cos_sim_decoder_output': nan, 'train/cos_sim_encoder_output_patchwise': 0.30224609375, 'train/probe_losses': 1.9835872650146484, 'train/probe_accs': 0.15625, 'test/probe_losses': 1.9990463256835938, 'test/probe_accs': 0.171875, 'lr': 3e-06, 'epoch': 4, 'tube_mask_ratio': 0.75, 'decoder_mask_ratio': 0.75}\n",
+      "0 0.08203125 2.095348358154297\n",
+      "test 0 0.052734375 2.236387252807617\n",
+      "test 1 0.0703125 2.2432994842529297\n",
+      "{'train/loss': 0.6117484122514725, 'train/recon_losses': nan, 'train/contrastive_losses': 6.117484092712402, 'train/num_steps': 0, 'train/cos_sim_encoder_output': 0.6046142578125, 'train/cos_sim_decoder_output': nan, 'train/cos_sim_encoder_output_patchwise': 0.3111572265625, 'train/probe_losses': 2.095348358154297, 'train/probe_accs': 0.08203125, 'test/probe_losses': 2.2398433685302734, 'test/probe_accs': 0.0615234375, 'lr': 3e-06, 'epoch': 5, 'tube_mask_ratio': 0.75, 'decoder_mask_ratio': 0.75}\n",
+      "0 0.15625 2.082639694213867\n",
+      "test 0 0.01171875 2.237241744995117\n",
+      "test 1 0.05859375 2.2262916564941406\n",
+      "{'train/loss': 0.610412523150444, 'train/recon_losses': nan, 'train/contrastive_losses': 6.104125022888184, 'train/num_steps': 0, 'train/cos_sim_encoder_output': 0.6123046875, 'train/cos_sim_decoder_output': nan, 'train/cos_sim_encoder_output_patchwise': 0.30657958984375, 'train/probe_losses': 2.082639694213867, 'train/probe_accs': 0.15625, 'test/probe_losses': 2.231766700744629, 'test/probe_accs': 0.03515625, 'lr': 3e-06, 'epoch': 6, 'tube_mask_ratio': 0.75, 'decoder_mask_ratio': 0.75}\n",
+      "0 0.00390625 2.2390995025634766\n",
+      "test 0 0.1015625 2.056934356689453\n",
+      "test 1 0.064453125 1.9691143035888672\n",
+      "{'train/loss': 0.6085265427827835, 'train/recon_losses': nan, 'train/contrastive_losses': 6.085265159606934, 'train/num_steps': 0, 'train/cos_sim_encoder_output': 0.63623046875, 'train/cos_sim_decoder_output': nan, 'train/cos_sim_encoder_output_patchwise': 0.31298828125, 'train/probe_losses': 2.2390995025634766, 'train/probe_accs': 0.00390625, 'test/probe_losses': 2.01302433013916, 'test/probe_accs': 0.0830078125, 'lr': 3e-06, 'epoch': 7, 'tube_mask_ratio': 0.75, 'decoder_mask_ratio': 0.75}\n",
+      "0 0.494140625 1.8124122619628906\n",
+      "test 0 0.392578125 1.8616065979003906\n",
+      "test 1 0.3828125 1.8176498413085938\n",
+      "{'train/loss': 0.6146479845046997, 'train/recon_losses': nan, 'train/contrastive_losses': 6.146479606628418, 'train/num_steps': 0, 'train/cos_sim_encoder_output': 0.613525390625, 'train/cos_sim_decoder_output': nan, 'train/cos_sim_encoder_output_patchwise': 0.31793212890625, 'train/probe_losses': 1.8124122619628906, 'train/probe_accs': 0.494140625, 'test/probe_losses': 1.8396282196044922, 'test/probe_accs': 0.3876953125, 'lr': 3e-06, 'epoch': 8, 'tube_mask_ratio': 0.75, 'decoder_mask_ratio': 0.75}\n",
+      "0 0.021484375 2.2320098876953125\n",
+      "test 0 0.025390625 2.181182861328125\n",
+      "test 1 0.033203125 2.175107955932617\n",
+      "{'train/loss': 0.6042157411575317, 'train/recon_losses': nan, 'train/contrastive_losses': 6.042157173156738, 'train/num_steps': 0, 'train/cos_sim_encoder_output': 0.6165771484375, 'train/cos_sim_decoder_output': nan, 'train/cos_sim_encoder_output_patchwise': 0.31536865234375, 'train/probe_losses': 2.2320098876953125, 'train/probe_accs': 0.021484375, 'test/probe_losses': 2.178145408630371, 'test/probe_accs': 0.029296875, 'lr': 3e-06, 'epoch': 9, 'tube_mask_ratio': 0.75, 'decoder_mask_ratio': 0.75}\n",
+      "0 0.04296875 2.068845748901367\n",
+      "test 0 0.095703125 2.147918701171875\n",
+      "test 1 0.275390625 2.0189361572265625\n",
+      "{'train/loss': 0.6057839542627335, 'train/recon_losses': nan, 'train/contrastive_losses': 6.057839393615723, 'train/num_steps': 0, 'train/cos_sim_encoder_output': 0.6259765625, 'train/cos_sim_decoder_output': nan, 'train/cos_sim_encoder_output_patchwise': 0.30718994140625, 'train/probe_losses': 2.068845748901367, 'train/probe_accs': 0.04296875, 'test/probe_losses': 2.0834274291992188, 'test/probe_accs': 0.185546875, 'lr': 3e-06, 'epoch': 10, 'tube_mask_ratio': 0.75, 'decoder_mask_ratio': 0.75}\n",
+      "0 0.248046875 2.012929916381836\n",
+      "test 0 0.1875 2.1199779510498047\n",
+      "test 1 0.177734375 2.0317916870117188\n",
+      "{'train/loss': 0.6008661240339279, 'train/recon_losses': nan, 'train/contrastive_losses': 6.008661270141602, 'train/num_steps': 0, 'train/cos_sim_encoder_output': 0.624267578125, 'train/cos_sim_decoder_output': nan, 'train/cos_sim_encoder_output_patchwise': 0.3172607421875, 'train/probe_losses': 2.012929916381836, 'train/probe_accs': 0.248046875, 'test/probe_losses': 2.0758848190307617, 'test/probe_accs': 0.1826171875, 'lr': 3e-06, 'epoch': 11, 'tube_mask_ratio': 0.75, 'decoder_mask_ratio': 0.75}\n",
+      "0 0.15625 2.024129867553711\n",
+      "test 0 0.361328125 1.8665599822998047\n",
+      "test 1 0.22265625 1.9156570434570312\n",
+      "{'train/loss': 0.6274944394826889, 'train/recon_losses': nan, 'train/contrastive_losses': 6.274944305419922, 'train/num_steps': 0, 'train/cos_sim_encoder_output': 0.62109375, 'train/cos_sim_decoder_output': nan, 'train/cos_sim_encoder_output_patchwise': 0.3148193359375, 'train/probe_losses': 2.024129867553711, 'train/probe_accs': 0.15625, 'test/probe_losses': 1.891108512878418, 'test/probe_accs': 0.2919921875, 'lr': 3e-06, 'epoch': 12, 'tube_mask_ratio': 0.75, 'decoder_mask_ratio': 0.75}\n",
+      "0 0.046875 2.0176239013671875\n",
+      "test 0 0.140625 1.9483413696289062\n",
+      "test 1 0.287109375 1.853952407836914\n",
+      "{'train/loss': 0.5929619073867798, 'train/recon_losses': nan, 'train/contrastive_losses': 5.929618835449219, 'train/num_steps': 0, 'train/cos_sim_encoder_output': 0.6199951171875, 'train/cos_sim_decoder_output': nan, 'train/cos_sim_encoder_output_patchwise': 0.31878662109375, 'train/probe_losses': 2.0176239013671875, 'train/probe_accs': 0.046875, 'test/probe_losses': 1.9011468887329102, 'test/probe_accs': 0.2138671875, 'lr': 3e-06, 'epoch': 13, 'tube_mask_ratio': 0.75, 'decoder_mask_ratio': 0.75}\n",
+      "0 0.025390625 2.3530330657958984\n",
+      "test 0 0.22265625 1.9583473205566406\n",
+      "test 1 0.26171875 2.0058231353759766\n",
+      "{'train/loss': 0.6034502238035202, 'train/recon_losses': nan, 'train/contrastive_losses': 6.034502029418945, 'train/num_steps': 0, 'train/cos_sim_encoder_output': 0.636962890625, 'train/cos_sim_decoder_output': nan, 'train/cos_sim_encoder_output_patchwise': 0.31683349609375, 'train/probe_losses': 2.3530330657958984, 'train/probe_accs': 0.025390625, 'test/probe_losses': 1.9820852279663086, 'test/probe_accs': 0.2421875, 'lr': 3e-06, 'epoch': 14, 'tube_mask_ratio': 0.75, 'decoder_mask_ratio': 0.75}\n",
+      "0 0.037109375 2.5336132049560547\n",
+      "test 0 0.06640625 2.2813949584960938\n",
+      "test 1 0.345703125 1.9364681243896484\n",
+      "{'train/loss': 0.5937714576721191, 'train/recon_losses': nan, 'train/contrastive_losses': 5.937714576721191, 'train/num_steps': 0, 'train/cos_sim_encoder_output': 0.6195068359375, 'train/cos_sim_decoder_output': nan, 'train/cos_sim_encoder_output_patchwise': 0.31719970703125, 'train/probe_losses': 2.5336132049560547, 'train/probe_accs': 0.037109375, 'test/probe_losses': 2.108931541442871, 'test/probe_accs': 0.2060546875, 'lr': 3e-06, 'epoch': 15, 'tube_mask_ratio': 0.75, 'decoder_mask_ratio': 0.75}\n",
+      "0 0.080078125 2.155559539794922\n",
+      "test 0 0.056640625 2.157787322998047\n",
+      "test 1 0.146484375 2.017892837524414\n",
+      "{'train/loss': 0.6005350202322006, 'train/recon_losses': nan, 'train/contrastive_losses': 6.005350112915039, 'train/num_steps': 0, 'train/cos_sim_encoder_output': 0.6202392578125, 'train/cos_sim_decoder_output': nan, 'train/cos_sim_encoder_output_patchwise': 0.3182373046875, 'train/probe_losses': 2.155559539794922, 'train/probe_accs': 0.080078125, 'test/probe_losses': 2.0878400802612305, 'test/probe_accs': 0.1015625, 'lr': 3e-06, 'epoch': 16, 'tube_mask_ratio': 0.75, 'decoder_mask_ratio': 0.75}\n",
+      "0 0.111328125 2.099821090698242\n",
+      "test 0 0.4296875 1.8980712890625\n",
+      "test 1 0.294921875 1.9661064147949219\n",
+      "{'train/loss': 0.6011319160461426, 'train/recon_losses': nan, 'train/contrastive_losses': 6.011319160461426, 'train/num_steps': 0, 'train/cos_sim_encoder_output': 0.6202392578125, 'train/cos_sim_decoder_output': nan, 'train/cos_sim_encoder_output_patchwise': 0.31549072265625, 'train/probe_losses': 2.099821090698242, 'train/probe_accs': 0.111328125, 'test/probe_losses': 1.932088851928711, 'test/probe_accs': 0.3623046875, 'lr': 3e-06, 'epoch': 17, 'tube_mask_ratio': 0.75, 'decoder_mask_ratio': 0.75}\n",
+      "0 0.119140625 2.0135326385498047\n",
+      "test 0 0.265625 1.861032485961914\n",
+      "test 1 0.158203125 1.972625732421875\n",
+      "{'train/loss': 0.5891992747783661, 'train/recon_losses': nan, 'train/contrastive_losses': 5.891992568969727, 'train/num_steps': 0, 'train/cos_sim_encoder_output': 0.6190185546875, 'train/cos_sim_decoder_output': nan, 'train/cos_sim_encoder_output_patchwise': 0.316650390625, 'train/probe_losses': 2.0135326385498047, 'train/probe_accs': 0.119140625, 'test/probe_losses': 1.9168291091918945, 'test/probe_accs': 0.2119140625, 'lr': 3e-06, 'epoch': 18, 'tube_mask_ratio': 0.75, 'decoder_mask_ratio': 0.75}\n",
+      "0 0.181640625 2.0815296173095703\n",
+      "test 0 0.154296875 2.060464859008789\n",
+      "test 1 0.46875 1.8413829803466797\n",
+      "{'train/loss': 0.5968149304389954, 'train/recon_losses': nan, 'train/contrastive_losses': 5.968149185180664, 'train/num_steps': 0, 'train/cos_sim_encoder_output': 0.61767578125, 'train/cos_sim_decoder_output': nan, 'train/cos_sim_encoder_output_patchwise': 0.31536865234375, 'train/probe_losses': 2.0815296173095703, 'train/probe_accs': 0.181640625, 'test/probe_losses': 1.9509239196777344, 'test/probe_accs': 0.3115234375, 'lr': 3e-06, 'epoch': 19, 'tube_mask_ratio': 0.75, 'decoder_mask_ratio': 0.75}\n"
+     ]
+    }
+   ],
+   "source": [
+    "mse = nn.MSELoss()\n",
+    "l1 = nn.L1Loss()\n",
+    "crossentropy = nn.CrossEntropyLoss()\n",
+    "if use_contrastive_loss:\n",
+    "    contrastive_temps = utils.cosine_anneal(0.004, 0.0075, num_epochs)\n",
+    "progress_bar = tqdm(range(epoch, num_epochs), disable=local_rank!=0, desc=\"Overall\")\n",
+    "for epoch in progress_bar:\n",
+    "    # get the masking ratio for the current epoch\n",
+    "    tube_mask_ratio = utils.get_masking_ratio(\n",
+    "        current_epoch=epoch, \n",
+    "        total_epochs=num_epochs, \n",
+    "        start_masking_ratio=tube_start_masking_ratio, \n",
+    "        end_masking_ratio=tube_end_masking_ratio\n",
+    "    )\n",
+    "    with torch.cuda.amp.autocast(dtype=data_type):\n",
+    "        model.train()\n",
+    "        optimizer.train()\n",
+    "        for train_i, batch in enumerate(train_dl):\n",
+    "            optimizer.zero_grad()\n",
+    "\n",
+    "            input_func = batch['func.npy']\n",
+    "\n",
+    "            subject_id = batch['subject_id.txt']\n",
+    "            subject_id = torch.Tensor([int(subject[-2:]) for subject in subject_id]).long()\n",
+    "            subject_id = torch.repeat_interleave(subject_id.long(), 2).to(device)\n",
+    "            # session_id = batch['session_id.txt']\n",
+    "            # session_id = torch.Tensor([int(session[-2:]) for session in session_id]).long().repeat(2).to(device)\n",
+    "            # session_id = torch.repeat_interleave(session_id.long(), 2)\n",
+    "\n",
+    "            if masking_strategy==\"None\":\n",
+    "                func, _ = aug_transform(input_func)\n",
+    "                brain_pos_pats_vit = torch.ones(num_patches_per_timepoint)\n",
+    "            elif masking_strategy==\"MNI\":\n",
+    "                func, _ = aug_transform(input_func)\n",
+    "            else:\n",
+    "                func, brain_pos_voxels = aug_transform(input_func)\n",
+    "                brain_pos_pats = model.patchify(torch.Tensor(brain_pos_voxels)[None,None,None])\n",
+    "                brain_pos_pats_vit = rearrange(brain_pos_pats, \"b ... d -> b (...) d\").mean(-1)[0]\n",
+    "\n",
+    "            func = func.reshape(-1, num_frames, \n",
+    "                                func.shape[-3], func.shape[-2], func.shape[-1])\n",
+    "            func = func.unsqueeze(1).clamp(0,1)\n",
+    "            \n",
+    "            # create encoder and decoder masks\n",
+    "            rand_patches = torch.randperm(num_patches_per_timepoint)\n",
+    "            \n",
+    "            encoder_mask = torch.zeros(num_patches_per_timepoint).to(torch.bool)\n",
+    "            encoder_mask[rand_patches[:num_encoder_patches]] = True\n",
+    "            encoder_mask = encoder_mask.tile(num_frames//frame_patch_size)\n",
+    "            \n",
+    "            decoder_mask = torch.zeros(num_patches_per_timepoint).to(torch.bool)\n",
+    "            decoder_mask[rand_patches[num_encoder_patches:num_encoder_patches+num_decoder_patches]] = True\n",
+    "            decoder_mask = decoder_mask.tile(num_frames//frame_patch_size)\n",
+    "\n",
+    "            # encode the tube patches\n",
+    "            encoder_out = model(func, encoder_mask=encoder_mask, device=device)\n",
+    "            if use_cls_token:\n",
+    "                enc_cls_token = encoder_out[:,:1,:]\n",
+    "\n",
+    "            # decode both the encoder_out patches and masked decoder patches\n",
+    "            decoder_out = model(encoder_out, encoder_mask=encoder_mask, decoder_mask=decoder_mask, device=device)\n",
+    "            # subset only the reconstructed decoder patches\n",
+    "            output = decoder_out[:, -decoder_mask.sum():]\n",
+    "\n",
+    "            # compare to ground truth and calculate loss\n",
+    "            target_patches = model.patchify(func)\n",
+    "            target_patches_vit = rearrange(target_patches, \"b ... d -> b (...) d\")\n",
+    "            target = target_patches_vit.to(device)[:, decoder_mask]\n",
+    "\n",
+    "            target_mean = target.mean(0)\n",
+    "            target_std = target.std(0)\n",
+    "            target_normed = (target - target_mean) / (target_std + 1e-6)\n",
+    "\n",
+    "            recon_loss = mse(output, target_normed)\n",
+    "            recon_losses.append(recon_loss.item())\n",
+    "            loss = recon_loss\n",
+    "\n",
+    "            # contrastive loss\n",
+    "            if use_contrastive_loss:\n",
+    "                # encode the decoder patches\n",
+    "                encoder_out2 = model(func, encoder_mask=decoder_mask, device=device)\n",
+    "                enc_cls_token2 = encoder_out2[:,:1,:]\n",
+    "                \n",
+    "                temp = contrastive_temps[epoch]\n",
+    "                \n",
+    "                logits = (nn.functional.normalize(enc_cls_token.flatten(1),dim=-1) @\n",
+    "                            nn.functional.normalize(enc_cls_token2.flatten(1),dim=-1).T) / temp\n",
+    "                \n",
+    "                # logits = (nn.functional.normalize(model.cont(encoder_out.flatten(1)),dim=-1) @\n",
+    "                #             nn.functional.normalize(model.cont(encoder_out2.flatten(1)),dim=-1).T) / temp\n",
+    "                \n",
+    "                labels = torch.arange(len(logits)).long().to(device)\n",
+    "                loss1 = crossentropy(logits, labels)\n",
+    "                # loss1 = -(logits.log_softmax(-1) * labels.softmax(-1)).sum(-1).mean()\n",
+    "                loss2 = crossentropy(logits.T, labels)\n",
+    "                contr_loss = (loss1 + loss2)/2\n",
+    "                \n",
+    "                contrastive_losses.append(contr_loss.item())\n",
+    "                loss += (contr_loss * contrastive_loss_weight)\n",
+    "\n",
+    "            cos_sim_encoder_output_patchwise.append(utils.patchwise_cosine_similarity(encoder_out).mean().item())\n",
+    "            cos_sim_encoder_output.append(utils.batchwise_cosine_similarity(encoder_out.flatten(1)/1e3,encoder_out.flatten(1)/1e3)[~torch.eye(len(encoder_out),dtype=torch.bool)].mean().item())\n",
+    "            cos_sim_decoder_output.append(utils.batchwise_cosine_similarity(output,output)[~torch.eye(len(output),dtype=torch.bool)].mean().item())\n",
+    "\n",
+    "            loss.backward()\n",
+    "            optimizer.step()\n",
+    "            lrs.append(optimizer.param_groups[0][\"lr\"])\n",
+    "            train_losses.append(loss.item())\n",
+    "\n",
+    "            if train_i >= (num_iterations_per_epoch-1):\n",
+    "                break\n",
+    "\n",
+    "        # reset linear_probe\n",
+    "        # if use_cls_token:\n",
+    "        #     linear_probe = LinearProbe((num_patches_per_timepoint+1)*model.encoder_embed_dim)\n",
+    "        # else:\n",
+    "        #     linear_probe = LinearProbe(num_patches_per_timepoint*model.encoder_embed_dim)\n",
+    "        linear_probe = LinearProbe(model.encoder_embed_dim)\n",
+    "        linear_probe = linear_probe.to(device)\n",
+    "        probe_opt_grouped_parameters = [\n",
+    "            {'params': [p for n, p in linear_probe.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': 1e-2},\n",
+    "            {'params': [p for n, p in linear_probe.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0},\n",
+    "        ]\n",
+    "        probe_optimizer = torch.optim.AdamW(probe_opt_grouped_parameters, lr=3e-3)\n",
+    "\n",
+    "        if True:#(epoch % 5 == 0) or (epoch == num_epochs-1):\n",
+    "            model.eval()\n",
+    "            optimizer.eval()\n",
+    "            linear_probe.train()\n",
+    "            for probe_i, batch in enumerate(train_dl):\n",
+    "                probe_optimizer.zero_grad()\n",
+    "\n",
+    "                input_func = batch['func.npy']\n",
+    "\n",
+    "                subject_id = batch['subject_id.txt']\n",
+    "                subject_id = torch.Tensor([int(subject[-2:]) for subject in subject_id]).long()\n",
+    "                subject_id = torch.repeat_interleave(subject_id.long(), 2).to(device)\n",
+    "\n",
+    "                func, _ = aug_transform(input_func)\n",
+    "                func = func.reshape(-1, num_frames, \n",
+    "                                    func.shape[-3], func.shape[-2], func.shape[-1])\n",
+    "                func = func.unsqueeze(1).clamp(0,1)\n",
+    "\n",
+    "                encoder_mask = torch.ones(num_patches_per_timepoint).to(torch.bool)\n",
+    "                encoder_mask = encoder_mask.tile(num_frames//frame_patch_size)\n",
+    "\n",
+    "                # encode the tube patches\n",
+    "                with torch.no_grad():\n",
+    "                    encoder_out = model(func, encoder_mask=encoder_mask, device=device)\n",
+    "                    encoder_out = encoder_out[:,:1,:]\n",
+    "                    encoder_out = nn.functional.normalize(encoder_out,dim=-1)\n",
+    "\n",
+    "                # linear probe\n",
+    "                subject_pred = linear_probe(encoder_out.flatten(1).to(device))\n",
+    "                probe_loss = crossentropy(subject_pred, subject_id-1) # minus 1 because subject_id is 1-indexed\n",
+    "\n",
+    "                probe_accuracy = (torch.max(subject_pred,1).indices == (subject_id-1)).sum() / len(subject_id)\n",
+    "                probe_accs.append(probe_accuracy.item())\n",
+    "                probe_losses.append(probe_loss.item())\n",
+    "\n",
+    "                print(probe_i, probe_accuracy.item(), probe_loss.item())\n",
+    "\n",
+    "                probe_loss.backward()\n",
+    "                probe_optimizer.step()\n",
+    "\n",
+    "                if probe_i >= (probe_num_iterations_per_epoch-1):\n",
+    "                    break\n",
+    "\n",
+    "            for test_i, batch in enumerate(test_dl):\n",
+    "                input_func = batch['func.npy']\n",
+    "\n",
+    "                subject_id = batch['subject_id.txt']\n",
+    "                subject_id = torch.Tensor([int(subject[-2:]) for subject in subject_id]).long()\n",
+    "                subject_id = torch.repeat_interleave(subject_id.long(), 2).to(device)\n",
+    "\n",
+    "                func, _ = aug_transform(input_func)\n",
+    "                func = func.reshape(-1, num_frames, \n",
+    "                                    func.shape[-3], func.shape[-2], func.shape[-1])\n",
+    "                func = func.unsqueeze(1).clamp(0,1)\n",
+    "\n",
+    "                encoder_mask = torch.ones(num_patches_per_timepoint).to(torch.bool)\n",
+    "                encoder_mask = encoder_mask.tile(num_frames//frame_patch_size)\n",
+    "\n",
+    "                # encode the tube patches\n",
+    "                with torch.no_grad():\n",
+    "                    encoder_out = model(func, encoder_mask=encoder_mask, device=device)\n",
+    "                    encoder_out = encoder_out[:,:1,:]\n",
+    "                    encoder_out = nn.functional.normalize(encoder_out,dim=-1)\n",
+    "\n",
+    "                # linear probe\n",
+    "                subject_pred = linear_probe(encoder_out.flatten(1).to(device))\n",
+    "                test_loss = crossentropy(subject_pred, subject_id-1) # minus 1 because subject_id is 1-indexed\n",
+    "\n",
+    "                test_accuracy = (torch.max(subject_pred,1).indices == (subject_id-1)).sum() / len(subject_id)\n",
+    "                test_accs.append(test_accuracy.item())\n",
+    "                test_losses.append(test_loss.item())\n",
+    "\n",
+    "                print(\"test\", test_i, test_accuracy.item(), test_loss.item())\n",
+    "\n",
+    "                if test_i >= 1:\n",
+    "                    break\n",
+    "\n",
+    "        logs = {\n",
+    "            \"train/loss\": np.mean(train_losses[-(train_i + 1) :]),\n",
+    "            \"train/recon_losses\": np.mean(recon_losses[-(train_i + 1) :]),\n",
+    "            \"train/contrastive_losses\": np.mean(contrastive_losses[-(train_i + 1) :]),\n",
+    "            \"train/num_steps\": len(recon_losses),\n",
+    "            \"train/cos_sim_encoder_output\": np.mean(cos_sim_encoder_output[-(train_i + 1) :]),\n",
+    "            \"train/cos_sim_decoder_output\": np.mean(cos_sim_decoder_output[-(train_i + 1) :]),\n",
+    "            \"train/cos_sim_encoder_output_patchwise\": np.mean(cos_sim_encoder_output_patchwise[-(train_i + 1) :]),\n",
+    "            \"train/probe_losses\": np.mean(probe_losses[-(probe_i + 1) :]),\n",
+    "            \"train/probe_accs\": np.mean(probe_accs[-(probe_i + 1) :]),\n",
+    "            \"test/probe_losses\": np.mean(test_losses[-(test_i + 1) :]),\n",
+    "            \"test/probe_accs\": np.mean(test_accs[-(test_i + 1) :]),\n",
+    "            \"lr\": np.mean(lrs[-(train_i + 1) :]),\n",
+    "            \"epoch\": epoch,\n",
+    "            \"tube_mask_ratio\": tube_mask_ratio,\n",
+    "            \"decoder_mask_ratio\": decoder_mask_ratio,\n",
+    "        }\n",
+    "        progress_bar.set_postfix(**logs)\n",
+    "        if utils.is_interactive(): print(logs)\n",
+    "\n",
+    "        # Plot progress (first sample in batch)\n",
+    "        with torch.no_grad():\n",
+    "            if utils.is_interactive() or wandb_log:\n",
+    "                if epoch % 50 == 0:\n",
+    "                    output = (output * target_std) + target_mean\n",
+    "                    idx = 0\n",
+    "                    \n",
+    "                    decode_vis = torch.zeros_like(target_patches_vit)\n",
+    "                    decode_vis[:, decoder_mask] = output.to(decode_vis.device).to(decode_vis.dtype)\n",
+    "                    decoder_unpatches = rearrange(\n",
+    "                        decode_vis,\n",
+    "                        \"b (f d h w) c -> b f d h w c\",\n",
+    "                        d=img_size[0]//patch_depth,\n",
+    "                        h=img_size[1]//patch_height,\n",
+    "                        w=img_size[2]//patch_width,\n",
+    "                    )\n",
+    "                    decoder_func = rearrange(\n",
+    "                        decoder_unpatches,\n",
+    "                        \"b f d h w (pd ph pw pf c) -> b c (f pf) (d pd) (h ph) (w pw)\",\n",
+    "                        b=batch_size*2,\n",
+    "                        f=num_frames//frame_patch_size,\n",
+    "                        d=img_size[0]//patch_depth,\n",
+    "                        h=img_size[1]//patch_height,\n",
+    "                        w=img_size[2]//patch_width,\n",
+    "                        pd=patch_depth,\n",
+    "                        ph=patch_height,\n",
+    "                        pw=patch_width,\n",
+    "                        pf=frame_patch_size,\n",
+    "                    )\n",
+    "                    orig_image = utils.reshape_to_2d(func[idx])\n",
+    "                    recon_image = utils.reshape_to_2d(decoder_func[idx])\n",
+    "\n",
+    "                    combined_image = orig_image.clone()\n",
+    "                    combined_image[recon_image!=0] = recon_image[recon_image!=0]\n",
+    "\n",
+    "                    random_start = np.arange(3100,3450)\n",
+    "                    orig_image = transforms.ToPILImage()(orig_image[:,random_start])\n",
+    "                    recon_image = transforms.ToPILImage()(recon_image[:,random_start])\n",
+    "                    combined_image = transforms.ToPILImage()(combined_image[:,random_start])\n",
+    "\n",
+    "                    if wandb_log:\n",
+    "                        logs[f\"train/orig\"] = wandb.Image(orig_image, caption=f\"epoch{epoch:03d}\")\n",
+    "                        logs[f\"train/recon\"] = wandb.Image(recon_image, caption=f\"epoch{epoch:03d}\")\n",
+    "                        logs[f\"train/combined\"] = wandb.Image(combined_image, caption=f\"epoch{epoch:03d}\")\n",
+    "                    else:\n",
+    "                        if epoch==0:\n",
+    "                            print(\"orig_image\")\n",
+    "                            display(orig_image)\n",
+    "                            print(\"recon_image\")\n",
+    "                            display(recon_image)\n",
+    "                            print(\"combined_image\")\n",
+    "                        display(combined_image)\n",
+    "\n",
+    "    if wandb_log: wandb.log(logs)\n",
+    "\n",
+    "    # Save model checkpoint\n",
+    "    if (ckpt_saving) and ((epoch % ckpt_interval == 0) or (epoch==num_epochs-1)):\n",
+    "        save_ckpt(model,\"last\")\n",
+    "\n",
+    "    # wait for other GPUs to catch up if needed\n",
+    "    if distributed: dist.barrier()\n",
+    "    torch.cuda.empty_cache()\n",
+    "        \n",
+    "if distributed:\n",
+    "    dist.destroy_process_group()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a6779720-0d08-4aec-87e2-439e0c1e9dca",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "encoder_mask.device"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1840c426-d7f5-4d0c-acd1-5ea47d453b91",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "decoder_mask.sum()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fff6e847-2511-4f69-aa70-f1471a5b7d07",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "plt.figure(figsize=(8, 3))\n",
+    "plt.plot(probe_losses)\n",
+    "# plt.title(\"Training re-construction losses\")\n",
+    "plt.show()\n",
+    "if use_contrastive_loss:\n",
+    "    plt.figure(figsize=(8, 3))\n",
+    "    plt.plot(contrastive_losses)\n",
+    "    plt.title(\"Training contrastive losses\")\n",
+    "    plt.show()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "fmri_new",
+   "language": "python",
+   "name": "fmri_new"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.10"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

fMRI-foundation-model-old/fMRI-MAE/accel.slurm

@@ -0,0 +1,65 @@
+#!/bin/bash
+#SBATCH --account=fmri
+#SBATCH --partition=h80i
+#SBATCH --job-name=found
+#SBATCH --ntasks-per-node=1
+#SBATCH --nodes=1
+#SBATCH --gres=gpu:1
+#SBATCH --gpus-per-task=1       # Set to equal gres=gpu:#!
+#SBATCH --time=35:00:00         # total run time limit (HH:MM:SS)
+#SBATCH -e slurms/%j.err        # first create a "slurms" folder in current directory to store logs
+#SBATCH -o slurms/%j.out
+#SBATCH --comment=medarc
+#SBATCH --no-requeue
+#SBATCH --exclude=ip-10-0-134-244,ip-10-0-139-120,ip-10-0-139-20,ip-10-0-168-162
+
+export NUM_GPUS=1 # Set to equal gres=gpu:#!
+echo NUM_GPUS=$NUM_GPUS
+
+source ~/.bashrc
+
+cd /weka/proj-fmri/paulscotti/fMRI-foundation-model/fMRI-MAE/
+jupyter nbconvert main.ipynb --to python
+if [ $? -ne 0 ]; then
+  echo "Error: Conversion of ipynb to Python failed. Exiting."
+  exit 1
+fi
+
+# Make sure another job doesnt use same port, here using random number
+export MASTER_PORT=$((RANDOM % (19000 - 11000 + 1) + 11000))
+export HOSTNAMES=$(scontrol show hostnames "$SLURM_JOB_NODELIST")
+export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
+echo MASTER_ADDR=${MASTER_ADDR}
+echo MASTER_PORT=${MASTER_PORT}
+echo WORLD_SIZE=$((${SLURM_NNODES} * ${NUM_GPUS}))
+
+# Other settings relevant for multi-node
+export NCCL_DEBUG=WARN #INFO
+export NCCL_PROTO=simple
+export FI_EFA_USE_DEVICE_RDMA=1
+export FI_EFA_FORK_SAFE=1
+export FI_LOG_LEVEL=1
+export PYTHONFAULTHANDLER=1
+export CUDA_LAUNCH_BLOCKING=0
+export OMPI_MCA_mtl_base_verbose=1
+export FI_EFA_ENABLE_SHM_TRANSFER=0
+export FI_PROVIDER=efa
+export FI_EFA_TX_MIN_CREDITS=64
+export NCCL_TREE_THRESHOLD=0
+export TORCH_DISTRIBUTED_DEBUG=DETAIL
+export NCCL_P2P_DISABLE=1
+module load openmpi cuda/12.1
+
+srun torchrun \
+--nproc_per_node=$NUM_GPUS \
+--nnodes=$SLURM_NNODES \
+--rdzv_id=$SLURM_JOBID \
+--rdzv_backend=c10d \
+--rdzv_endpoint=${MASTER_ADDR}:$MASTER_PORT \
+--rdzv_conf=timeout=90 \
+main.py
+
+if [ $? -ne 0 ]; then
+  echo "Error: srun command failed. Please check the logs for more details."
+  exit 1
+fi

fMRI-foundation-model-old/fMRI-MAE/dataloader.py

@@ -0,0 +1,772 @@
+import numpy as np
+from torch import Tensor, stack
+from torch.utils.data import IterDataPipe, DataLoader
+from torch.utils.data.datapipes._decorator import functional_datapipe
+from torch.utils.data.datapipes.iter.callable import MapperIterDataPipe
+from torch.utils.data.datapipes.utils.common import StreamWrapper, match_masks
+from torchdata.datapipes.iter import IterableWrapper, IterDataPipe, TarArchiveLoader
+from torchdata.datapipes.iter.load.s3io import S3FileLoaderIterDataPipe
+from torchdata.datapipes.utils.common import validate_pathname_binary_tuple
+import warnings
+import subprocess
+import tarfile
+import tempfile
+import time
+import random
+from omegaconf import DictConfig, ListConfig
+import webdataset as wds
+import re
+import os
+import copy
+from io import BufferedIOBase, BytesIO, RawIOBase
+from typing import (
+    IO,
+    Any,
+    Callable,
+    Dict,
+    Iterable,
+    Iterator,
+    List,
+    Optional,
+    Sequence,
+    Tuple,
+    TypeVar,
+    Union,
+    cast,
+)
+T_co = TypeVar("T_co", covariant=True)
+
+class ShuffledListDataPipe(IterDataPipe[T_co]):
+    def __init__(
+        self,
+        source_list: List[T_co],
+        *,
+        shuffle: bool = True,
+        cycle: Union[bool, int] = True,
+    ):
+        super().__init__()
+        self.source = source_list
+        self._enabled = shuffle
+        self._seed = None
+        self._rng = random.Random()
+        if isinstance(cycle, bool):
+            if cycle:
+                self._cycle = -1
+            else:
+                self._cycle = 1
+        else:
+            assert isinstance(cycle, int)
+            self._cycle = cycle
+
+    def __iter__(self) -> Iterator[T_co]:
+        source = copy.copy(self.source)
+        cycle = self._cycle
+        epochs = 0
+        while cycle == -1 or epochs < cycle:
+            if self._enabled:
+                self._rng.shuffle(source)
+            yield from source
+            epochs += 1
+
+    def __len__(self):
+        if self.count == -1:
+            raise TypeError(
+                f"This {type(self).__name__} instance cycles forever, and "
+                f"therefore doesn't have valid length."
+            )
+        else:
+            return self.count * len(self.source)
+
+    def __getstate__(self):
+        state = (
+            self.source,
+            self._enabled,
+            self._seed,
+            self._rng.getstate(),
+            self._valid_iterator_id,
+            self._number_of_samples_yielded,
+        )
+        if IterDataPipe.getstate_hook is not None:
+            return IterDataPipe.getstate_hook(state)
+        return state
+
+    def __setstate__(self, state):
+        (
+            self.source,
+            self._enabled,
+            self._seed,
+            rng_state,
+            self._valid_iterator_id,
+            self._number_of_samples_yielded,
+        ) = state
+        self._rng = random.Random()
+        self._rng.setstate(rng_state)
+
+def get_job_temp_dir(dl_root: str) -> str:
+    try:
+        job_or_array_id = (
+            os.environ.get("SLURM_ARRAY_JOB_ID", "") or os.environ["SLURM_JOB_ID"]
+        )
+        task_id = os.environ.get("SLURM_ARRAY_TASK_ID", "") or "0"
+        return os.path.join(dl_root, f"JOB_{job_or_array_id}/TASK_{task_id}")
+    except KeyError as e:
+        print("SLURM_JOB_ID env var not set: You sure this job is healthy?")
+        raise e
+
+def _shard_expand(s: str) -> List[str]:
+    expansion = r"\{[0-9]+\.\.[0-9]+\}"
+    m = re.search(expansion, s)
+    if not m:
+        return [s]
+    prefix = s[: m.start()]
+    rest = _shard_expand(s[m.end() :])
+    rng = s[m.start() + 1 : m.end() - 1]
+    lohi = rng.split("..")
+    if len(lohi[0]) == len(lohi[1]) and lohi[0].startswith("0"):
+        fmt = "{prefix}{i:0>{l}d}{r}"
+    elif len(lohi[0]) <= len(lohi[1]):
+        if lohi[0].startswith("0") and lohi[0] != "0":
+            raise ValueError(
+                "shard_expand: low bound must not start with 0 if low bound is shorter"
+            )
+        fmt = "{prefix}{i}{r}"
+    else:
+        raise ValueError("shard_expand: low bound must be shorter than high bound")
+    lo, hi = (int(x) for x in lohi)
+    if lo >= hi:
+        raise ValueError(f"shard_expand: bad range in in shard spec {s}.")
+    result = []
+    for i in range(lo, hi + 1):
+        for r in rest:
+            expanded: str = fmt.format(prefix=prefix, i=i, r=r, l=len(lohi[1]))
+            result.append(expanded)
+    return result
+
+@functional_datapipe("custom_shard_expand")
+class CustomShardExpanderIterDataPipe(IterDataPipe[str]):
+    def __init__(self, source_datapipe: IterDataPipe[str]) -> None:
+        super().__init__()
+        self.source_datapipe: IterDataPipe[str] = source_datapipe
+
+    def __iter__(self) -> Iterator[str]:
+        for path in self.source_datapipe:
+            yield from _shard_expand(path)
+
+def warn_and_continue(exn):
+    """Call in an exception handler to ignore any exception, issue a warning, and continue."""
+    print(exn)
+    warnings.warn(repr(exn))
+    time.sleep(0.05)
+    return True
+error_handler = warn_and_continue
+
+def is_stream_handle(data):
+    obj_to_check = data.file_obj if isinstance(data, StreamWrapper) else data
+    return isinstance(obj_to_check, (BufferedIOBase, RawIOBase))
+
+def read_stream_handles(data):
+    if not is_stream_handle(data):
+        return data
+    else:
+        ds = data
+        data = b"".join(data)
+        ds.close()
+        del ds
+        return data
+        
+def stream_reader(sample):
+    return {k: read_stream_handles(v) for k, v in sample.items()}
+
+def to_our_format(sample):
+    wds_key = sample.pop("__key__")
+    sample = {k.lstrip("."): v for k, v in sample.items()}
+    sample["__key__"] = wds_key.split("/")[-1]
+    sample["__url__"] = "/".join(wds_key.split("/")[:-1])
+    return sample
+
+def add_processors(
+    datapipeline,
+    processors: Optional[ListConfig],
+    description: str,
+    error_handler: Callable = warn_and_continue,
+):
+    if not processors:
+        return datapipeline
+    else:
+        for i, processor_config in enumerate(processors):
+            processor = instantiate(processor_config)
+            if isinstance(processor, AbstractFilter):
+                print(
+                    f"Adding filter {processor.__class__.__name__} as {description} #{i} "
+                    f"to the datapipeline"
+                )
+                datapipeline = datapipeline.filter(processor.filter)
+            elif isinstance(processor, AbstractMapper):
+                print(
+                    f"Adding mapper {processor.__class__.__name__} as {description} #{i} "
+                    f"to the datapipeline"
+                )
+                datapipeline = datapipeline.map_with_handler(
+                    processor.map,
+                    handler=error_handler,
+                    called_cls_name=processor.__class__.__name__,
+                )
+            else:
+                raise TypeError(
+                    f"chosen {description} {processor.__class__.__name__} should be either subclass"
+                    "AbstractMapper or AbstractFilter"
+                )
+        return datapipeline
+
+@functional_datapipe("map_with_handler")
+class MapperWithErrorHandlingIterDataPipe(MapperIterDataPipe):
+    def __init__(
+        self,
+        datapipe: IterDataPipe,
+        fn: Callable,
+        handler: Callable = wds.reraise_exception,
+        input_col: Optional[Union[str, int]] = None,
+        output_col: Optional[Union[str, int]] = None,
+        called_cls_name: Optional[str] = None,
+    ):
+        # for now, disbable input and output col since this is never used anyways
+        if input_col is not None:
+            raise NotImplementedError("`input_col` argument currently not supported")
+
+        if output_col is not None:
+            raise NotImplementedError("`output_col` argument currently not supported")
+
+        super().__init__(datapipe, fn)
+        self.handler = handler
+        self._apply_fn_ = self._apply_fn
+
+    def __iter__(self) -> Iterator[T_co]:
+        for data in self.datapipe:
+            try:
+                res = self._apply_fn_(data)
+                if res is None:
+                    continue
+                yield res
+            except Exception as e:
+                if self.handler(e):
+                    continue
+                else:
+                    raise e
+
+__S3_TOOLS__ = {
+    "s3": ["/usr/local/bin/aws", "s3"],
+}
+
+def is_tar(x: str) -> bool:
+    return x.endswith(".tar")
+
+def ls_aws(
+    path: str,
+    tool: str = "s3",
+    recursive: bool = True,
+    raise_errors: bool = True,
+    skip_files: bool = True,
+):
+    assert path.startswith("s3://"), path
+
+    # in case we have a file, set recursive to false
+    isfile = bool(os.path.splitext(path)[1])
+    if isfile and skip_files:
+        return [path]
+
+    if not path.endswith("/") and not isfile:
+        path = path + "/"
+    cmd = [*__S3_TOOLS__[tool], "ls", f"{path}"]
+
+    if recursive and not isfile:
+        cmd += ["--recursive"]
+    try:
+        result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+        result.check_returncode()
+
+        stdout = result.stdout.decode().split("\n")
+        # return empty list in case path is not a real file/directory
+        if not stdout[0]:
+            return []
+
+        # return path if it is an exisiting file
+        if isfile:
+            return [path]
+
+        # strip timestamp and object size
+        out = [line.split(" ")[-1] for line in stdout if line]
+        if recursive:
+            # aws s3 ls returns prefix+filename
+            bucket = path[: path.find("/", len("s3://"))]  # s3://<bucket>/...
+            out = [os.path.join(bucket, o) for o in out]
+        else:
+            # aws s3 ls returns filename only
+            out = [os.path.join(path, o) for i in out]
+
+        return out
+    except subprocess.CalledProcessError as e:
+        print(f"Got exception while trying to load data! {e.__class__.__name__}: {e}")
+        if raise_errors:
+            raise e
+        else:
+            return []
+
+@functional_datapipe("wrapped_load_files_by_s3")
+class WrappedS3FileLoaderIterDataPipe(S3FileLoaderIterDataPipe):
+    def __init__(
+        self,
+        *args,
+        ignore_missing_files: bool = False,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        self.ignore_missing_files = ignore_missing_files
+
+    def __iter__(self) -> Iterator[Tuple[str, StreamWrapper]]:
+        # same as parent but maybe ignoring errors
+        for url in self.source_datapipe:
+            try:
+                yield url, StreamWrapper(BytesIO(self.handler.s3_read(url)))
+            except ValueError as e:
+                if not self.ignore_missing_files:
+                    raise FileNotFoundError(url)
+                else:
+                    print(f"Warning: Could not download {url}")
+
+@functional_datapipe("list_files_with_aws_cli")
+class S3CLIFileListerIterDataPipe(IterDataPipe[T_co]):
+    def __init__(
+        self,
+        source_datapipe: Union[str, Sequence[str], IterDataPipe],
+        s3_tool: str = "s3",
+        n_retries: int = 10,
+        sleep_interval: float = 0.01,
+    ):
+        if isinstance(source_datapipe, str):
+            source_datapipe = [
+                source_datapipe,
+            ]
+        if not isinstance(source_datapipe, IterDataPipe):
+            self.datapipe: IterDataPipe = IterableWrapper(source_datapipe)  # type: ignore[assignment]
+        else:
+            self.datapipe = source_datapipe
+
+        assert s3_tool in __S3_TOOLS__, f"`s3_tool` has to be in {list(__S3_TOOLS__)}"
+
+        self.s3_tool = s3_tool
+
+        self.n_retries = n_retries
+        self.sleep = sleep_interval
+
+    def __iter__(self) -> Iterator[str]:
+        for root in self.datapipe:
+            files = None
+            for _ in range(self.n_retries):
+                try:
+                    files = ls_aws(
+                        path=root, tool=self.s3_tool, recursive=True, raise_errors=True
+                    )
+                    break
+                except subprocess.CalledProcessError:
+                    time.sleep(self.sleep)
+
+            if files is None:
+                print(
+                    f"Could not ls data expected under {root} in {self.n_retries} tries. Not yielding ..."
+                )
+                continue
+
+            for file in files:
+                yield file
+
+def _download_tar(url, scratch, verbose=False, n_retries=100):
+    start = time.perf_counter()
+    if verbose:
+        _log(f"downloading {url}")
+    uid = "".join(filter(str.isalnum, os.path.splitext(url)[0]))
+    idx = 0
+    path = os.path.join(scratch, f"{os.getpid()}.{uid}.{idx}.tar")
+    while os.path.exists(path):
+        idx += 1
+        path = os.path.join(scratch, f"{os.getpid()}.{uid}.{idx}.tar")
+
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    for i_try in range(n_retries):
+        p = subprocess.call(
+            [
+                "/usr/local/bin/aws",
+                "s3",
+                "cp",
+                url,
+                path,
+                "--quiet",
+            ]
+        )
+
+        if p != 0:
+            # mostly to work around
+            # Error when retrieving credentials from Ec2InstanceMetadata: No
+            # credentials found in credential_source referenced in profile
+            # default
+            if i_try == 0:
+                # check if the requested file actually exists
+                if not ls_aws(url, skip_files=False, raise_errors=False):
+                    raise FileNotFoundError(url)
+
+            if i_try + 1 < n_retries:
+                time.sleep(random.uniform(0.0, 0.1))
+                continue
+            else:
+                raise Exception(
+                    f"File {url} should exist but failed to download after trying {i_try+1} times."
+                )
+        break
+
+    if verbose:
+        _log(
+            f"wrote {path} in {time.perf_counter() - start} secs after {i_try+1} tries."
+        )
+    return path
+
+
+@functional_datapipe("download_with_s3_cli")
+class S3CLITarDownloader(IterDataPipe[T_co]):
+    def __init__(
+        self,
+        source_datapipe,
+        n_retries: int = 10,
+        verbose: bool = False,
+        dl_root: str = "/scratch",
+        mode: str = "r:*",
+        aws_kwargs: Optional[dict] = None,
+        ignore_missing_files: bool = False,
+    ):
+        self.source_datapipe = source_datapipe
+        assert os.path.isdir(dl_root), f"`dl_root` {dl_root} is not a valid directory"
+        # get slurm job id and create a subdir in `dl_root` to download all the tars to
+        dl_root = get_job_temp_dir(dl_root)
+        os.makedirs(dl_root, exist_ok=True)
+        self.dl_root = dl_root
+        self.mode = mode
+        self.n_retries = n_retries
+        self.ignore_missing_files = ignore_missing_files
+
+        self.verbose = verbose
+
+    def _yield_next(self, url: str, local_path: str):
+        if self.verbose:
+            _log("popping queue")
+
+        if self.verbose:
+            _log(f"loading {local_path}")
+        tarstream = tarfile.open(local_path, self.mode)
+
+        if self.verbose:
+            _log(f"yielding {url}")
+
+        yield (url, StreamWrapper(tarstream))
+        if self.verbose:
+            _log("new tar request")
+
+        if self.verbose:
+            _log(f"removing previous tar at {local_path}")
+        try:
+            os.remove(local_path)
+        except FileNotFoundError:
+            _log(
+                f"WARNING: Could not find previous tar for deletion. Unless a clean-up was triggered this is unexpected. The location was {local_path}"
+            )
+
+    def __iter__(self):
+        with tempfile.TemporaryDirectory(dir=self.dl_root) as scratch:
+            for url in self.source_datapipe:
+                try:
+                    local_path = _download_tar(
+                        url, scratch, verbose=self.verbose, n_retries=self.n_retries
+                    )
+                except FileNotFoundError as e:
+                    if not self.ignore_missing_files:
+                        raise e
+                else:
+                    yield from self._yield_next(url, local_path)
+
+
+@functional_datapipe("load_from_tar_and_handle_error")
+class TarArchiveLoaderWithErrorHandlingIterDataPipe(TarArchiveLoader):
+    def __init__(
+        self,
+        datapipe: Iterable[Tuple[str, BufferedIOBase]],
+        mode: str = "r:*",
+        length: int = -1,
+        handler: Callable = wds.reraise_exception,
+    ):
+        super().__init__(datapipe=datapipe, mode=mode, length=length)
+        self.handler = handler
+
+    def __iter__(self) -> Iterator[Tuple[str, BufferedIOBase]]:
+        for data in self.datapipe:
+            validate_pathname_binary_tuple(data)
+            pathname, data_stream = data
+            try:
+                if isinstance(data_stream, StreamWrapper) and isinstance(
+                    data_stream.file_obj, tarfile.TarFile
+                ):
+                    tar = data_stream.file_obj
+                else:
+                    reading_mode = (
+                        self.mode
+                        if hasattr(data_stream, "seekable") and data_stream.seekable()
+                        else self.mode.replace(":", "|")
+                    )
+                    # typing.cast is used here to silence mypy's type checker
+                    tar = tarfile.open(
+                        fileobj=cast(Optional[IO[bytes]], data_stream),
+                        mode=reading_mode,
+                    )
+                for tarinfo in tar:
+                    if not tarinfo.isfile():
+                        continue
+                    extracted_fobj = tar.extractfile(tarinfo)
+                    if extracted_fobj is None:
+                        warnings.warn(
+                            f"failed to extract file {tarinfo.name} from source tarfile {pathname}"
+                        )
+                        raise tarfile.ExtractError
+                    inner_pathname = os.path.normpath(
+                        os.path.join(pathname, tarinfo.name)
+                    )
+
+                    yield inner_pathname, StreamWrapper(extracted_fobj, data_stream, name=inner_pathname)  # type: ignore[misc]
+            except Exception as e:
+                warnings.warn(
+                    f"Unable to extract files from corrupted tarfile stream {pathname} due to: {e}, abort!"
+                )
+                if self.handler(e):
+                    if hasattr(e, "args") and len(e.args) > 0:
+                        e.args = (e.args[0] + " @ " + str(pathname),) + e.args[1:]
+            finally:
+                if isinstance(data_stream, StreamWrapper):
+                    data_stream.autoclose()
+
+class DictCollator:
+    def __init__(
+        self,
+        combine_tensors: bool = True,
+        combine_scalars: bool = True,
+        timeout=None,
+    ):
+        self.combine_tensors = combine_tensors
+        self.combine_scalars = combine_scalars
+
+        collate = self._collate
+        if timeout is not None:
+            collate = timeout_wrapper(collate, timeout=timeout)
+        self.collate = collate
+
+    def __call__(self, samples):
+        return self.collate(samples)
+
+    def _collate(self, samples):
+        keys = set.intersection(*[set(sample.keys()) for sample in samples])
+        batched = {key: [] for key in keys}
+
+        for s in samples:
+            [batched[key].append(s[key]) for key in batched]
+
+        result = {}
+        for key in batched:
+            if isinstance(batched[key][0], (int, float)):
+                if self.combine_scalars:
+                    result[key] = np.array(list(batched[key]))
+            elif isinstance(batched[key][0], Tensor):
+                if self.combine_tensors:
+                    result[key] = stack(list(batched[key]))
+            elif isinstance(batched[key][0], np.ndarray):
+                if self.combine_tensors:
+                    result[key] = np.array(list(batched[key]))
+            else:
+                result[key] = list(batched[key])
+
+        del samples
+        del batched
+        return result
+
+def timeout_wrapper(func: Callable, timeout: Optional[float] = None) -> Callable:
+    if timeout is None or timeout <= 0.0:
+        return func
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        result = [None]
+        exception = [None]
+        event = threading.Event()
+
+        def wrapped_func():
+            try:
+                result[0] = func(*args, **kwargs)
+            except Exception as e:
+                exception[0] = e
+            finally:
+                event.set()
+
+        thread = threading.Thread(target=wrapped_func)
+        thread.start()
+        event.wait(timeout)
+
+        if not event.is_set():
+            raise TimeoutError(f"Function call timed out (longer than {timeout} secs).")
+
+        thread.join()
+
+        if exception[0] is not None:
+            err = exception[0]
+            del exception
+            raise err
+
+        del thread
+        del exception
+        del wrapped_func
+        del event
+        del args
+        del kwargs
+
+        ret = result[0]
+        del result
+        return ret
+
+    return wrapper
+
+class DecoderWithTimeout(wds.Decoder):
+    """Decode samples using a list of handlers.
+
+    For each key/data item, this iterates through the list of
+    handlers until some handler returns something other than None.
+    """
+
+    def __init__(self, *args, timeout: Optional[float] = None, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        if timeout is not None:
+            self.decode = timeout_wrapper(self.decode, timeout)
+
+def setup_decoder(
+    decoders,
+    additional_decoder_kwargs: Optional[Union[DictConfig, Dict]] = None,
+):
+    if not additional_decoder_kwargs:
+        additional_decoder_kwargs = {}
+
+    if not isinstance(decoders, (List, ListConfig)):
+        decoders = [decoders]
+    handlers = []
+    for decoder_spec in decoders:
+        if isinstance(decoder_spec, (Dict, DictConfig)):
+            decoder = instantiate(decoder_spec)
+        elif isinstance(decoder_spec, str):
+            decoder = wds.autodecode.ImageHandler(decoder_spec)
+        else:
+            raise TypeError(f"{decoder_spec} not a thing for decoders.")
+        handlers.append(decoder)
+        print(f"Adding decoder {decoder.__class__.__name__} to decoders.")
+
+    decoder = DecoderWithTimeout(
+        handlers,
+        partial=additional_decoder_kwargs.pop("partial", True),
+        **additional_decoder_kwargs,  # todo
+    )
+    return decoder
+
+def create_dataset(urls, is_s3=False, prefetch=None, 
+                     s3_buffer_size=None, s3_ignore_missing_files=True,
+                     sample_shuffle=1, shard_shuffle=1000, cycle=True,
+                     split_workers_along_tars=True, decoders="torch"):
+    if isinstance(urls, str):
+        urls = [urls]
+    dp = IterableWrapper(urls)
+    dp = dp.custom_shard_expand()
+    if is_s3: 
+        dp = dp.list_files_with_aws_cli().filter(is_tar)
+    else:
+        dp = dp.list_files(masks="*.tar", recursive=True)
+    
+    dp = ShuffledListDataPipe(list(dp), shuffle=shard_shuffle > 1, cycle=cycle)
+
+    if split_workers_along_tars:
+        dp = dp.sharding_filter()
+
+    if is_s3:
+        # if prefetch is None:
+        #     s3_buffer_size = int(s3_buffer_size) if s3_buffer_size is not None else None
+        #     dp = dp.wrapped_load_files_by_s3(
+        #         buffer_size=s3_buffer_size,
+        #         ignore_missing_files=s3_ignore_missing_files,
+        #     )
+        dp = dp.download_with_s3_cli(ignore_missing_files=s3_ignore_missing_files)
+    else:
+        dp = dp.open_files(mode="b")
+
+    if prefetch is not None:
+        dp = dp.custom_prefetch(buffer_size=prefetch)
+
+    dp = dp.load_from_tar_and_handle_error(handler=error_handler)
+    dp = dp.webdataset()
+
+    if not split_workers_along_tars:
+        dp = dp.sharding_filter()
+
+    dp = dp.map_with_handler(
+            stream_reader,
+            handler=error_handler,
+            called_cls_name="StreamReader",
+        )
+
+    dp = dp.shuffle(buffer_size=sample_shuffle) if sample_shuffle > 1 else dp
+
+    dp = dp.map_with_handler(
+            to_our_format,
+            handler=error_handler,
+            called_cls_name="ToOurFormat",
+        )
+
+    decoder: DecoderWithTimeout = setup_decoder(decoders) # ,additional_decoder_kwargs)
+
+    dp = dp.map_with_handler(
+        decoder.decode,
+        handler=error_handler,
+        called_cls_name="Decoder",
+    )
+    
+    return dp
+
+def create_loader(
+    datapipeline: IterDataPipe,
+    batch_size: int,
+    num_workers: int,
+    partial: bool = False,
+    collation_fn: Optional[Union[Callable, Dict, DictConfig]] = DictCollator(),
+    batched_transforms: Optional[ListConfig] = None,
+    loader_kwargs: Optional[Union[Dict, DictConfig]] = None,
+) -> DataLoader:
+    if not loader_kwargs:
+        loader_kwargs = {}
+
+    loader_kwargs.pop("shuffle", None)
+    if not batched_transforms:
+        batched_transforms = []
+
+    print("#" * 100)
+    print("Building dataloader with the following parameters")
+    print(f"batch_size: {batch_size}, num_workers: {num_workers}")
+    for key in loader_kwargs:
+        print(key, ": ", loader_kwargs[key])
+    print("#" * 100)
+
+    datapipeline = datapipeline.batch(batch_size, drop_last=not partial)
+
+    if isinstance(collation_fn, (Dict, DictConfig)):
+        collation_fn = instantiate(collation_fn)
+    datapipeline = datapipeline.collate(collate_fn=collation_fn)
+    loader = DataLoader(
+        datapipeline, batch_size=None, num_workers=num_workers, **loader_kwargs
+    )
+    return loader

fMRI-foundation-model-old/fMRI-MAE/downstream.ipynb

@@ -0,0 +1,1162 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "0859a181-0681-489e-9939-8d47614b9fa8",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "LOCAL RANK=0\n",
+      "NUM GPUS=1\n",
+      "GLOBAL RANK=0\n",
+      "batch_size 32\n",
+      "skipping deepspeed reconfiguration...\n",
+      "PID of this process = 617152\n",
+      "device: cuda\n",
+      "Distributed environment: NO\n",
+      "Num processes: 1\n",
+      "Process index: 0\n",
+      "Local process index: 0\n",
+      "Device: cuda\n",
+      "\n",
+      "Mixed precision type: fp16\n",
+      "\n",
+      "distributed = False num_devices = 1 local rank = 0 world size = 1 data_type = torch.float16\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/admin/home-paulscotti/found/lib/python3.10/site-packages/accelerate/accelerator.py:432: FutureWarning: Passing the following arguments to `Accelerator` is deprecated and will be removed in version 1.0 of Accelerate: dict_keys(['split_batches']). Please pass an `accelerate.DataLoaderConfiguration` instead: \n",
+      "dataloader_config = DataLoaderConfiguration(split_batches=False)\n",
+      "  warnings.warn(\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Import packages and setup gpu configuration.\n",
+    "# This code block shouldnt need to be adjusted!\n",
+    "import os\n",
+    "import sys\n",
+    "import json\n",
+    "import yaml\n",
+    "import numpy as np\n",
+    "import pandas as pd\n",
+    "import copy\n",
+    "import math\n",
+    "from einops import rearrange\n",
+    "from einops.layers.torch import Rearrange\n",
+    "import time\n",
+    "import random\n",
+    "import h5py\n",
+    "import webdataset as wds\n",
+    "import gc\n",
+    "import matplotlib.pyplot as plt\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.nn.functional as F\n",
+    "from torchvision import transforms\n",
+    "import utils\n",
+    "from models import *\n",
+    "from mindeye_models import *\n",
+    "import nibabel as nib\n",
+    "from nilearn import plotting\n",
+    "\n",
+    "# tf32 data type is faster than standard float32\n",
+    "torch.backends.cuda.matmul.allow_tf32 = True\n",
+    "\n",
+    "### Multi-GPU config ###\n",
+    "local_rank = os.getenv('LOCAL_RANK')\n",
+    "if local_rank is None: \n",
+    "    local_rank = 0\n",
+    "else:\n",
+    "    local_rank = int(local_rank)\n",
+    "print(f\"LOCAL RANK={local_rank}\")\n",
+    "\n",
+    "num_devices = os.getenv('NUM_GPUS')\n",
+    "if num_devices is None: \n",
+    "    num_devices = 1\n",
+    "else:\n",
+    "    num_devices = int(num_devices)\n",
+    "print(f\"NUM GPUS={num_devices}\")\n",
+    "distributed = True if num_devices>1 else False\n",
+    "\n",
+    "global_rank = os.getenv('RANK')\n",
+    "if global_rank is None:\n",
+    "    global_rank = 0\n",
+    "else:\n",
+    "    global_rank = int(global_rank)\n",
+    "print(f\"GLOBAL RANK={global_rank}\")\n",
+    "\n",
+    "if utils.is_interactive():\n",
+    "    # Following allows you to change functions in models.py or utils.py and \n",
+    "    # have this notebook automatically update with your revisions\n",
+    "    %load_ext autoreload\n",
+    "    %autoreload 2\n",
+    "    from tqdm.notebook import tqdm\n",
+    "else:\n",
+    "    from tqdm import tqdm\n",
+    "\n",
+    "# Load parameters from yaml config\n",
+    "config = yaml.load(open('config.yaml', 'r'), Loader=yaml.FullLoader)\n",
+    "\n",
+    "# create global variables from the config\n",
+    "for attribute_name in config.keys():\n",
+    "    globals()[attribute_name] = config[f'{attribute_name}']\n",
+    "    \n",
+    "# Load MindEye parameters from yaml config (will override any params with same name)\n",
+    "mindeye_config = yaml.load(open('mindeye_config.yaml', 'r'), Loader=yaml.FullLoader)\n",
+    "\n",
+    "# create global variables from the config\n",
+    "for attribute_name in mindeye_config.keys():\n",
+    "    globals()[attribute_name] = mindeye_config[f'{attribute_name}']\n",
+    "\n",
+    "data_type = torch.float16 # change depending on your mixed_precision\n",
+    "\n",
+    "batch_size = global_batch_size // num_devices\n",
+    "print(\"batch_size\", batch_size)\n",
+    "    \n",
+    "# First use \"accelerate config\" in terminal and setup using deepspeed stage 2 with CPU offloading\n",
+    "from accelerate import Accelerator\n",
+    "accelerator = Accelerator(split_batches=False, mixed_precision=\"fp16\")\n",
+    "\n",
+    "from accelerate.state import AcceleratorState\n",
+    "try:\n",
+    "    AcceleratorState().deepspeed_plugin.deepspeed_config['train_micro_batch_size_per_gpu'] = batch_size\n",
+    "    print(\"deepspeed reconfigured, train_micro_batch_size_per_gpu = \", batch_size)\n",
+    "except:\n",
+    "    print(\"skipping deepspeed reconfiguration...\")\n",
+    "\n",
+    "print(\"PID of this process =\",os.getpid())\n",
+    "device = accelerator.device\n",
+    "print(\"device:\",device)\n",
+    "world_size = accelerator.state.num_processes\n",
+    "num_workers = num_devices\n",
+    "print(accelerator.state)\n",
+    "\n",
+    "print(\"distributed =\",distributed, \"num_devices =\", num_devices, \"local rank =\", local_rank, \"world size =\", world_size, \"data_type =\", data_type)\n",
+    "print = accelerator.print # only print if local_rank=0"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c7b35bc5-8674-4061-a606-f44249d167df",
+   "metadata": {},
+   "source": [
+    "# Configuration"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "d9fa256a-cc44-4346-8d6c-1bed3779d5b0",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "mae config\n",
+      "\n",
+      " {'model_name': 'framepatchsize1_4gpu_bs8_3e7_norm_rope_atomcos_001contrastiveweight', 'use_cls_token': False, 'use_contrastive_loss': True, 'contrastive_loss_weight': 0.001, 'batch_size': 8, 'num_workers': 10, 'num_epochs': 100, 'seed': 42, 'max_lr': 3e-07, 'num_samples_per_epoch': 1024, 'ckpt_saving': True, 'ckpt_interval': 50, 'resume_from_ckpt': True, 'wandb_log': True, 'tube_start_masking_ratio': 0.9, 'tube_end_masking_ratio': 0.9, 'decoder_mask_ratio': 0.85, 'patch_size': [8, 8, 8], 'frame_patch_size': 1, 'use_rope_emb': True, 'masking_strategy': 'MNI', 'encoder_model': 'vit_base', 'decoder_model': 'vit_base', 'img_size': [88, 104, 72], 'num_frames': 4, 'is_s3': False, 'train_urls': ['/weka/proj-fmri/shared/NSD_MNI_wds/{000000..000738}.tar']}\n",
+      "mindeye_config\n",
+      " {'model_name': 'downstream_framepatchsize4_bs8_300ep_contr_nolinear', 'mae_model_name': 'framepatchsize1_4gpu_bs8_300ep_NEWcont_3e8_norm_rope_LONG', 'global_batch_size': 32, 'mixed_precision': 'fp16', 'num_epochs': 12, 'seed': 42, 'max_lr': 0.0003, 'multi_subject': False, 'multisubject_ckpt': 'None', 'ckpt_saving': False, 'ckpt_interval': 99, 'resume_from_ckpt': False, 'wandb_log': False, 'in_dim': 393216, 'hidden_dim': 2056, 'drop': 0.15, 'mixup_pct': 0.0, 'nsd_wds_path': '/weka/proj-fmri/shared/mindeyev2_dataset/wds', 'nsd_raw_path': '/weka/proj-fmri/shared/mindeyev2_dataset', 'nsd_image_path': '/weka/proj-fmri/shared/mindeyev2_dataset', 'num_sessions': 40}\n",
+      "mae_ckpt_pth /weka/proj-fmri/paulscotti/fMRI-foundation-model/ckpts/framepatchsize1_4gpu_bs8_300ep_NEWcont_3e8_norm_rope_LONG/last.pth\n",
+      "outdir /weka/proj-fmri/paulscotti/fMRI-foundation-model/ckpts/downstream_framepatchsize4_bs8_300ep_contr_nolinear\n",
+      "use_cls_token False\n",
+      "num_patches 5148\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(\"mae config\\n\\n\",config)\n",
+    "print(\"mindeye_config\\n\",mindeye_config)\n",
+    "\n",
+    "# if utils.is_interactive():\n",
+    "#     ckpt_saving = False\n",
+    "#     wandb_log = False\n",
+    "\n",
+    "# seed all random functions\n",
+    "utils.seed_everything(seed)\n",
+    "\n",
+    "mae_ckpt_pth = os.path.abspath(f'../ckpts/{mae_model_name}/last.pth')\n",
+    "print(\"mae_ckpt_pth\", mae_ckpt_pth)\n",
+    "\n",
+    "outdir = os.path.abspath(f'../ckpts/{model_name}')\n",
+    "print(\"outdir\", outdir)\n",
+    "\n",
+    "use_cls_token = False #True if use_contrastive_loss else use_cls_token\n",
+    "print(\"use_cls_token\", use_cls_token)\n",
+    "\n",
+    "if type(patch_size) == int:\n",
+    "    patch_size = [patch_size,patch_size,patch_size]\n",
+    "patch_depth = patch_size[0]\n",
+    "patch_height = patch_size[1]\n",
+    "patch_width = patch_size[2]\n",
+    "\n",
+    "num_patches = int(\n",
+    "    (img_size[0] / patch_depth)\n",
+    "    * (img_size[1] / patch_height)\n",
+    "    * (img_size[2] / patch_width)\n",
+    "    * num_frames\n",
+    ")\n",
+    "print(\"num_patches\", num_patches)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8115109a-30a1-448e-a1ea-1b115ed8cb6d",
+   "metadata": {},
+   "source": [
+    "# Load pretrained foundation model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "1b40287e-bab2-4791-9407-e63ad5b0a3b4",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "param counts:\n",
+      "170,830,848 total\n",
+      "170,830,848 trainable\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "170830848"
+      ]
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "vit_size = {\n",
+    "    \"encoder\": encoder_model,\n",
+    "    \"decoder\": decoder_model\n",
+    "}\n",
+    "    \n",
+    "model = get_vit(\n",
+    "    size=vit_size,\n",
+    "    image_size=img_size,  # depth, height, width\n",
+    "    image_patch_size=(patch_depth,patch_height,patch_width),  # depth, height, width patch size\n",
+    "    frames=num_frames,\n",
+    "    frame_patch_size=frame_patch_size,\n",
+    "    channels=1,\n",
+    "    use_rope_emb=use_rope_emb,\n",
+    "    use_cls_token=use_cls_token,\n",
+    ")\n",
+    "utils.count_params(model)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "97b45ad3-e8d5-40cd-b041-1defdbf15155",
+   "metadata": {},
+   "source": [
+    "## Load pretrained ckpt for MAE foundation model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "62bb3f5f-cf7f-42c5-ab8c-d4bf143ba727",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "def save_ckpt(tag):\n",
+    "    ckpt_path = outdir+f'/{tag}.pth'\n",
+    "    if accelerator.is_main_process:\n",
+    "        unwrapped_model = accelerator.unwrap_model(model)\n",
+    "        torch.save({\n",
+    "            'epoch': epoch,\n",
+    "            'model_state_dict': unwrapped_model.state_dict(),\n",
+    "            'optimizer_state_dict': optimizer.state_dict(),\n",
+    "            'lr_scheduler': lr_scheduler.state_dict(),\n",
+    "            'train_losses': losses,\n",
+    "            'test_losses': test_losses,\n",
+    "            'lrs': lrs,\n",
+    "            }, ckpt_path)\n",
+    "    print(f\"\\n---saved {outdir}/{tag} ckpt!---\\n\")\n",
+    "\n",
+    "def load_ckpt(tag,load_lr=True,load_optimizer=True,load_epoch=True,strict=True,outdir=outdir,multisubj_loading=False): \n",
+    "    print(f\"\\n---loading {outdir}/{tag}.pth ckpt---\\n\")\n",
+    "    checkpoint = torch.load(outdir+'/last.pth', map_location='cpu')\n",
+    "    state_dict = checkpoint['model_state_dict']\n",
+    "    if multisubj_loading: # remove incompatible ridge layer that will otherwise error\n",
+    "        state_dict.pop('ridge.linears.0.weight',None)\n",
+    "    model.load_state_dict(state_dict, strict=strict)\n",
+    "    if load_epoch:\n",
+    "        globals()[\"epoch\"] = checkpoint['epoch']\n",
+    "        print(\"Epoch\",epoch)\n",
+    "    if load_optimizer:\n",
+    "        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n",
+    "    if load_lr:\n",
+    "        lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])\n",
+    "    del checkpoint"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "6c426853-d268-4daf-ac19-faa40c2cd362",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "checkpoint = torch.load(mae_ckpt_pth, map_location=device)\n",
+    "model.load_state_dict(checkpoint['model_state_dict'])#, strict=False)\n",
+    "\n",
+    "# set foundation model to evaluation\n",
+    "model.eval()\n",
+    "model.requires_grad_(False)\n",
+    "model.to(device)\n",
+    "pass"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "12fd89a4-955f-439d-b3c9-8f261c06eed2",
+   "metadata": {},
+   "source": [
+    "# Setup MindEye model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "1dae53d1-abff-4ddb-98cd-2cab04b9f0cb",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "nsddata_raw_stimuli = pd.read_csv(f\"{nsd_raw_path}/nsddata_rawdata.csv\")\n",
+    "TR_delay = 3 # to account for bold hrf\n",
+    "train_TRs = np.round(nsddata_raw_stimuli[nsddata_raw_stimuli['shared1000'] == False]['global_TR_onsets'].values + TR_delay).astype(np.int32)\n",
+    "test_TRs = np.round(nsddata_raw_stimuli[nsddata_raw_stimuli['shared1000'] == True]['global_TR_onsets'].values + TR_delay).astype(np.int32)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "4d5015e9-120e-4240-886c-6959796e8faf",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Loaded all 73k possible NSD images! torch.Size([73000, 3, 224, 224])\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Load 73k NSD images\n",
+    "f = h5py.File(f'{nsd_image_path}/coco_images_224_float16.hdf5', 'r')\n",
+    "images = f['images'][:] \n",
+    "images = torch.Tensor(images).to(\"cpu\").to(data_type)\n",
+    "print(\"Loaded all 73k possible NSD images!\", images.shape)\n",
+    "\n",
+    "# Load MindEye hdf5\n",
+    "f = h5py.File(f'{nsd_raw_path}/subj01_mnidata.h5', 'r') #subj01_rawdata_old.h5\n",
+    "mindeye_global_trs = f['global_trs'][:]\n",
+    "mindeye_funcs = f['funcs']"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "d79bcf35-45a6-4dc5-90da-c96c3a596401",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "FrozenOpenCLIPImageEmbedder(\n",
+       "  (model): CLIP(\n",
+       "    (visual): VisionTransformer(\n",
+       "      (conv1): Conv2d(3, 1664, kernel_size=(14, 14), stride=(14, 14), bias=False)\n",
+       "      (patch_dropout): Identity()\n",
+       "      (ln_pre): LayerNorm((1664,), eps=1e-05, elementwise_affine=True)\n",
+       "      (transformer): Transformer(\n",
+       "        (resblocks): ModuleList(\n",
+       "          (0-47): 48 x ResidualAttentionBlock(\n",
+       "            (ln_1): LayerNorm((1664,), eps=1e-05, elementwise_affine=True)\n",
+       "            (attn): MultiheadAttention(\n",
+       "              (out_proj): NonDynamicallyQuantizableLinear(in_features=1664, out_features=1664, bias=True)\n",
+       "            )\n",
+       "            (ls_1): Identity()\n",
+       "            (ln_2): LayerNorm((1664,), eps=1e-05, elementwise_affine=True)\n",
+       "            (mlp): Sequential(\n",
+       "              (c_fc): Linear(in_features=1664, out_features=8192, bias=True)\n",
+       "              (gelu): GELU(approximate='none')\n",
+       "              (c_proj): Linear(in_features=8192, out_features=1664, bias=True)\n",
+       "            )\n",
+       "            (ls_2): Identity()\n",
+       "          )\n",
+       "        )\n",
+       "      )\n",
+       "      (ln_post): LayerNorm((1664,), eps=1e-05, elementwise_affine=True)\n",
+       "    )\n",
+       "    (token_embedding): Embedding(49408, 1280)\n",
+       "    (ln_final): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
+       "  )\n",
+       ")"
+      ]
+     },
+     "execution_count": 8,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# clip_img_embedder = FrozenOpenCLIPImageEmbedder(\n",
+    "#     arch=\"ViT-bigG-14\",\n",
+    "#     version=\"laion2b_s39b_b160k\",\n",
+    "#     output_tokens=True,\n",
+    "#     only_tokens=True,\n",
+    "# )\n",
+    "# clip_seq_dim = 256\n",
+    "# clip_emb_dim = 1664\n",
+    "# clip_img_embedder.to(device)\n",
+    "\n",
+    "clip_img_embedder = FrozenOpenCLIPImageEmbedder(\n",
+    "    arch=\"ViT-bigG-14\",\n",
+    "    version=\"laion2b_s39b_b160k\",\n",
+    "    output_tokens=False,\n",
+    "    only_tokens=False,\n",
+    "    init_device=device,\n",
+    "    device=device,\n",
+    ")\n",
+    "clip_seq_dim = 1\n",
+    "clip_emb_dim = 1280\n",
+    "clip_img_embedder.to(device)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "d89d95f0-b52e-439f-ba5d-c13e2897fe31",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "dividing batch size by subj_list, which will then be concatenated across subj during training...\n",
+      "batch_size = 32 num_iterations_per_epoch = 11 num_samples_per_epoch = 375\n",
+      "Training with 40 sessions\n",
+      "/weka/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/{0..39}.tar\n",
+      "Loaded all subj train dls and betas!\n",
+      "\n",
+      "/weka/proj-fmri/shared/mindeyev2_dataset/wds/subj01/new_test/0.tar\n",
+      "Loaded test dl for subj1!\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "subj = s = 1\n",
+    "subj_list = [subj]\n",
+    "\n",
+    "# if multi_subject:\n",
+    "#     nsessions_allsubj=np.array([40, 40, 32, 30, 40, 32, 40, 30])\n",
+    "#     num_samples_per_epoch = (750*40) // num_devices \n",
+    "# else:\n",
+    "#     num_samples_per_epoch = (750*num_sessions) // num_devices \n",
+    "\n",
+    "num_samples_per_epoch = 375 #(750*num_sessions) // num_devices\n",
+    "\n",
+    "print(\"dividing batch size by subj_list, which will then be concatenated across subj during training...\") \n",
+    "batch_size = batch_size // len(subj_list)\n",
+    "num_iterations_per_epoch = num_samples_per_epoch // (batch_size*len(subj_list))\n",
+    "print(\"batch_size =\", batch_size, \"num_iterations_per_epoch =\",num_iterations_per_epoch, \"num_samples_per_epoch =\",num_samples_per_epoch)\n",
+    "\n",
+    "train_data = {}\n",
+    "train_dl = {}\n",
+    "\n",
+    "print(f\"Training with {num_sessions} sessions\")\n",
+    "train_url = f\"{nsd_wds_path}/subj0{s}/train/\" + \"{0..\" + f\"{num_sessions-1}\" + \"}.tar\"\n",
+    "print(train_url)\n",
+    "    \n",
+    "train_data[f'subj0{s}'] = wds.WebDataset(train_url,resampled=True,nodesplitter=utils.my_split_by_node)\\\n",
+    "                    .shuffle(750, initial=1500, rng=random.Random(42))\\\n",
+    "                    .decode(\"torch\")\\\n",
+    "                    .rename(behav=\"behav.npy\", past_behav=\"past_behav.npy\", future_behav=\"future_behav.npy\", olds_behav=\"olds_behav.npy\")\\\n",
+    "                    .to_tuple(*[\"behav\", \"past_behav\", \"future_behav\", \"olds_behav\"])\n",
+    "# train_dl[f'subj0{s}'] = torch.utils.data.DataLoader(train_data[f'subj0{s}'], batch_size=batch_size, shuffle=False, drop_last=True, pin_memory=True)\n",
+    "train_dl[f'subj0{s}'] = wds.WebLoader(\n",
+    "    train_data[f'subj0{s}'].batched(batch_size), \n",
+    "    pin_memory=True,\n",
+    "    shuffle=False,\n",
+    "    batch_size=None,\n",
+    "    num_workers=num_workers, \n",
+    "    persistent_workers=num_workers>0,\n",
+    ").with_epoch(num_iterations_per_epoch)\n",
+    "\n",
+    "print(\"Loaded all subj train dls and betas!\\n\")\n",
+    "if subj==3:\n",
+    "    num_test=2371\n",
+    "elif subj==4:\n",
+    "    num_test=2188\n",
+    "elif subj==6:\n",
+    "    num_test=2371\n",
+    "elif subj==8:\n",
+    "    num_test=2188\n",
+    "else:\n",
+    "    num_test=3000\n",
+    "test_url = f\"{nsd_wds_path}/subj0{subj}/new_test/\" + \"0.tar\"\n",
+    "print(test_url)\n",
+    "test_data = wds.WebDataset(test_url,resampled=True,nodesplitter=utils.my_split_by_node)\\\n",
+    "                    .decode(\"torch\")\\\n",
+    "                    .rename(behav=\"behav.npy\", past_behav=\"past_behav.npy\", future_behav=\"future_behav.npy\", olds_behav=\"olds_behav.npy\")\\\n",
+    "                    .to_tuple(*[\"behav\", \"past_behav\", \"future_behav\", \"olds_behav\"])\n",
+    "# test_dl = torch.utils.data.DataLoader(test_data, batch_size=num_test, shuffle=False, drop_last=True, pin_memory=True)\n",
+    "test_dl = wds.WebLoader(\n",
+    "    test_data.batched(num_test),\n",
+    "    pin_memory=True,\n",
+    "    shuffle=False,\n",
+    "    batch_size=None,\n",
+    "    num_workers=num_workers,\n",
+    "    persistent_workers=num_workers>0,\n",
+    ").with_epoch(10)\n",
+    "print(f\"Loaded test dl for subj{subj}!\\n\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "e100e116-4b98-4c15-bd8e-252871e72517",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "param counts:\n",
+      "507,388,672 total\n",
+      "507,388,672 trainable\n",
+      "param counts:\n",
+      "507,388,672 total\n",
+      "507,388,672 trainable\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "507388672"
+      ]
+     },
+     "execution_count": 10,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "class MindEyeModule(nn.Module):\n",
+    "    def __init__(self):\n",
+    "        super(MindEyeModule, self).__init__()\n",
+    "    def forward(self, x):\n",
+    "        return x\n",
+    "\n",
+    "# class RidgeRegression(torch.nn.Module):\n",
+    "#     # make sure to add weight_decay when initializing optimizer\n",
+    "#     def __init__(self, input_sizes, out_features, seq_len=1): \n",
+    "#         super(RidgeRegression, self).__init__()\n",
+    "#         self.seq_len = seq_len\n",
+    "#         self.out_features = out_features\n",
+    "#         self.linears = torch.nn.ModuleList([\n",
+    "#                 torch.nn.Linear(input_size, out_features) for input_size in input_sizes\n",
+    "#             ])\n",
+    "#     def forward(self, x, subj_idx):\n",
+    "#         out = torch.cat([self.linears[subj_idx](x[:,seq]).unsqueeze(1) for seq in range(self.seq_len)], dim=1)\n",
+    "#         return out\n",
+    "    \n",
+    "class MLP(torch.nn.Module):\n",
+    "    def __init__(self, input_sizes, out_features, seq_len=0): \n",
+    "        super(MLP, self).__init__()\n",
+    "        self.input_sizes = input_sizes[0]\n",
+    "        self.out_features = out_features\n",
+    "        self.mlp = nn.Sequential(\n",
+    "            nn.LayerNorm(self.input_sizes),\n",
+    "            nn.GELU(),\n",
+    "            nn.Linear(self.input_sizes, out_features),\n",
+    "            nn.LayerNorm(out_features),\n",
+    "            nn.GELU(),\n",
+    "            nn.Linear(out_features, out_features),\n",
+    "            nn.LayerNorm(out_features),\n",
+    "            nn.GELU(),\n",
+    "            nn.Linear(out_features, out_features)\n",
+    "        )\n",
+    "    def forward(self, x, z=None):\n",
+    "        out = self.mlp(x[:,0])\n",
+    "        return out\n",
+    "\n",
+    "\n",
+    "mindeye = MindEyeModule()\n",
+    "mindeye.ridge = MLP(np.array([in_dim]), out_features=clip_emb_dim*clip_seq_dim)\n",
+    "# mindeye.backbone = BrainNetwork(h=hidden_dim, in_dim=hidden_dim, n_blocks=4, drop=drop,\n",
+    "#                           clip_size=clip_emb_dim, out_dim=clip_emb_dim*clip_seq_dim, clip_scale=1)\n",
+    "utils.count_params(mindeye.ridge)\n",
+    "# utils.count_params(mindeye.backbone)\n",
+    "utils.count_params(mindeye)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "47ad9ec0-e0b5-4d2c-9d9b-027a85d176b1",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "total_steps 132\n",
+      "\n",
+      "Done with model preparations!\n",
+      "param counts:\n",
+      "507,388,672 total\n",
+      "507,388,672 trainable\n"
+     ]
+    }
+   ],
+   "source": [
+    "no_decay = [\"bias\", \"LayerNorm.bias\", \"LayerNorm.weight\"]\n",
+    "opt_grouped_parameters = [\n",
+    "    {'params': [p for n, p in mindeye.ridge.named_parameters()], 'weight_decay': 1e-2},\n",
+    "    # {'params': [p for n, p in mindeye.backbone.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': 1e-2},\n",
+    "    # {'params': [p for n, p in mindeye.backbone.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0},\n",
+    "]\n",
+    "\n",
+    "optimizer = torch.optim.AdamW(opt_grouped_parameters, lr=max_lr)\n",
+    "\n",
+    "total_steps = num_epochs * num_iterations_per_epoch\n",
+    "print(\"total_steps\", total_steps)\n",
+    "pct_start = 2/num_epochs if num_epochs>1 else 1.\n",
+    "lr_scheduler = torch.optim.lr_scheduler.OneCycleLR(\n",
+    "    optimizer,\n",
+    "    max_lr=max_lr,\n",
+    "    total_steps=total_steps,\n",
+    ")\n",
+    "\n",
+    "print(\"\\nDone with model preparations!\")\n",
+    "num_params = utils.count_params(mindeye)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2661e670-89b3-4598-a0ba-774e376d0047",
+   "metadata": {},
+   "source": [
+    "# Start wandb (if enabled)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "eae9f959-4457-428a-afb5-45e852490ffe",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "if accelerator.is_main_process and wandb_log: # only use main process for wandb logging\n",
+    "    import wandb\n",
+    "    wandb_project = 'found_downstream'\n",
+    "    print(f\"wandb {wandb_project} run {model_name}\")\n",
+    "    # need to configure wandb beforehand in terminal with \"wandb init\"!\n",
+    "    wandb_config = {\n",
+    "      \"model_name\": model_name,\n",
+    "      \"mae_model_name\": mae_model_name,\n",
+    "      \"global_batch_size\": global_batch_size,\n",
+    "      \"batch_size\": batch_size,\n",
+    "      \"num_epochs\": num_epochs,\n",
+    "      \"num_sessions\": num_sessions,\n",
+    "      \"num_samples_per_epoch\": num_samples_per_epoch,\n",
+    "      \"in_dim\": in_dim,\n",
+    "      \"hidden_dim\": hidden_dim,\n",
+    "      \"mixup_pct\": mixup_pct,\n",
+    "      \"num_params\": num_params,\n",
+    "      \"max_lr\": max_lr,\n",
+    "      \"ckpt_interval\": ckpt_interval,\n",
+    "      \"ckpt_saving\": ckpt_saving,\n",
+    "      \"seed\": seed,\n",
+    "      \"distributed\": distributed,\n",
+    "      \"num_devices\": num_devices,\n",
+    "      \"world_size\": world_size,\n",
+    "      \"train_url\": train_url,\n",
+    "      \"test_url\": test_url,\n",
+    "    }\n",
+    "    print(\"wandb_config:\\n\",wandb_config)\n",
+    "    print(\"wandb_id:\",model_name)\n",
+    "    wandb.init(\n",
+    "        id=model_name,\n",
+    "        project=wandb_project,\n",
+    "        name=model_name,\n",
+    "        config=wandb_config,\n",
+    "        resume=\"allow\",\n",
+    "    )\n",
+    "else:\n",
+    "    wandb_log = False"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7230f232-1700-4967-951d-fb566ae637b3",
+   "metadata": {},
+   "source": [
+    "# Train MindEye model using foundation model inputs"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "a8b849b8-04b2-4908-9b46-a4fa9a565d5a",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "epoch = 0\n",
+    "losses, test_losses, lrs = [], [], []\n",
+    "best_test_loss = 1e9\n",
+    "torch.cuda.empty_cache()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "b91b0e3c-c939-4717-8abc-2f3b8c39d012",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# if masking_strategy==\"MNI\":\n",
+    "from einops.layers.torch import Rearrange\n",
+    "\n",
+    "MNI_brain = nib.load(\"/weka/proj-fmri/paulscotti/fMRI-foundation-model/dataset_creation/afni_conversion/tpl-MNI152NLin2009cAsym_res-02_T1w_brain.nii.gz\").get_fdata()\n",
+    "brain_pos_voxels = MNI_brain[6:94,8:112,10:82]\n",
+    "\n",
+    "# brain_pos_voxels = brain_pos_voxels[:,30:31,:]\n",
+    "\n",
+    "brain_pos_pats = Rearrange(\n",
+    "        \"b c (f pf) (d pd) (h ph) (w pw) -> b f d h w (pd ph pw pf c)\",\n",
+    "        pd=patch_depth,\n",
+    "        ph=patch_height,\n",
+    "        pw=patch_width,\n",
+    "        pf=1,\n",
+    "    )(torch.Tensor(brain_pos_voxels)[None,None,None])\n",
+    "\n",
+    "brain_pos_pats_vit = rearrange(brain_pos_pats, \"b ... d -> b (...) d\").mean(-1)[0]\n",
+    "    \n",
+    "tube_mask = torch.zeros(num_patches // num_frames).to(torch.bool)\n",
+    "batch_positive_approx = (brain_pos_pats_vit > 0)\n",
+    "mask_idx_candidates = torch.where(batch_positive_approx)[0]\n",
+    "mask_idx_candidates = mask_idx_candidates[torch.randperm(len(mask_idx_candidates))]\n",
+    "tube_idx = mask_idx_candidates[:int(num_patches / num_frames * (1 - tube_end_masking_ratio))]\n",
+    "tube_mask[tube_idx] = True\n",
+    "tube_mask = tube_mask.tile(num_frames//frame_patch_size)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "id": "cac228e7-bdf8-4cc7-ab9d-891b5bd5abbd",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# load multisubject stage1 ckpt if set\n",
+    "if multisubject_ckpt!=\"None\" and not resume_from_ckpt:\n",
+    "    load_ckpt(\"last\",outdir=multisubject_ckpt,load_lr=False,load_optimizer=False,load_epoch=False,strict=False,multisubj_loading=True)\n",
+    "    \n",
+    "# load saved ckpt model weights into current model\n",
+    "if resume_from_ckpt:\n",
+    "    load_ckpt(\"last\",load_lr=True,load_optimizer=True,load_epoch=True)\n",
+    "elif wandb_log:\n",
+    "    if wandb.run.resumed:\n",
+    "        load_ckpt(\"last\",load_lr=True,load_optimizer=True,load_epoch=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "id": "b0a23e2f-eb06-4940-9b41-ca0e36dd210f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_dls = [train_dl[f'subj0{s}'] for s in subj_list]\n",
+    "\n",
+    "mindeye, optimizer, *train_dls, lr_scheduler = accelerator.prepare(mindeye, optimizer, *train_dls, lr_scheduler)\n",
+    "# leaving out test_dl since we will only have local_rank 0 device do evals"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "id": "dd08879b-0177-4cac-9438-a82ce0805bc1",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "downstream_framepatchsize4_bs8_300ep_contr_nolinear starting with epoch 0 / 12\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "60a889b1ee2d4d7a8acff2ffaf3a1ca1",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/12 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "6cc02dd1cf08486892e038df7de47cf8",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/11 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "ename": "NameError",
+     "evalue": "name 'err' is not defined",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[17], line 51\u001b[0m\n\u001b[1;32m     49\u001b[0m \u001b[38;5;66;03m## Process it through pretrained MAE ##\u001b[39;00m\n\u001b[1;32m     50\u001b[0m encoder_out \u001b[38;5;241m=\u001b[39m model(voxels_raw, encoder_mask\u001b[38;5;241m=\u001b[39mtube_mask)\n\u001b[0;32m---> 51\u001b[0m \u001b[43merr\u001b[49m\n\u001b[1;32m     52\u001b[0m \u001b[38;5;66;03m# encoder_out = bn(encoder_out)\u001b[39;00m\n\u001b[1;32m     54\u001b[0m voxel0 \u001b[38;5;241m=\u001b[39m encoder_out\u001b[38;5;241m.\u001b[39mflatten(\u001b[38;5;241m1\u001b[39m)\u001b[38;5;241m.\u001b[39munsqueeze(\u001b[38;5;241m1\u001b[39m)\n",
+      "\u001b[0;31mNameError\u001b[0m: name 'err' is not defined"
+     ]
+    }
+   ],
+   "source": [
+    "print(f\"{model_name} starting with epoch {epoch} / {num_epochs}\")\n",
+    "progress_bar = tqdm(range(epoch, num_epochs), disable=not accelerator.is_main_process)\n",
+    "mse = nn.MSELoss()\n",
+    "soft_loss_temps = utils.cosine_anneal(0.004, 0.0075, num_epochs - int(mixup_pct * num_epochs))\n",
+    "\n",
+    "bn = nn.BatchNorm1d(512,affine=False).to(device)\n",
+    "\n",
+    "test_image=None\n",
+    "num_test_eval=batch_size # should instead be average same-image 300 to mimic MindEye2 retrieval evaluation\n",
+    "\n",
+    "for epoch in progress_bar:\n",
+    "    mindeye.train()\n",
+    "\n",
+    "    fwd_percent_correct = 0.\n",
+    "    bwd_percent_correct = 0.\n",
+    "    test_fwd_percent_correct = 0.\n",
+    "    test_bwd_percent_correct = 0.\n",
+    "    loss_clip_total = 0.\n",
+    "    test_loss_clip_total = 0.\n",
+    "\n",
+    "    # pre-load all batches for this epoch (it's MUCH faster to pre-load in bulk than to separate loading per batch)\n",
+    "    voxel_iters = {} # empty dict because diff subjects have differing # of voxels\n",
+    "    image_iters = torch.zeros(num_iterations_per_epoch, batch_size*len(subj_list), 3, 224, 224).float()\n",
+    "    annot_iters = {}\n",
+    "    perm_iters, betas_iters, select_iters = {}, {}, {}\n",
+    "    for s, train_dl in enumerate(train_dls):\n",
+    "        with torch.cuda.amp.autocast(dtype=data_type):\n",
+    "            for iter, (behav0, past_behav0, future_behav0, old_behav0) in enumerate(tqdm(train_dl,total=num_iterations_per_epoch)):\n",
+    "                image0 = images[behav0[:,0,0].cpu().long()].float()\n",
+    "                image_iters[iter,s*batch_size:s*batch_size+batch_size] = image0\n",
+    "\n",
+    "                # if images are not fully preloaded, then can do this inefficient but more memory friendly approach\n",
+    "                # for ib,b in enumerate(behav0[:,0,0].cpu().long()):\n",
+    "                #     if ib==0:\n",
+    "                #         image0 = torch.Tensor(images[[b]])\n",
+    "                #     else:\n",
+    "                #         image0 = torch.vstack((image0, torch.Tensor(images[[b]])))\n",
+    "                # image_iters[iter,s*batch_size:s*batch_size+batch_size] = image0\n",
+    "                \n",
+    "                # get the corresponding raw voxel time series\n",
+    "                for ib,b in enumerate(behav0[:,0,5].cpu().long().numpy()):\n",
+    "                    tr = (nsddata_raw_stimuli[nsddata_raw_stimuli['global_trial'].isin([b.item()])]['global_TR_onsets'].values + TR_delay).astype(np.int32).item()\n",
+    "                    if ib==0:\n",
+    "                        voxels_raw = mindeye_funcs[tr-2:tr+2][None][None]\n",
+    "                    else:\n",
+    "                        voxels_raw = np.vstack((voxels_raw, mindeye_funcs[tr-2:tr+2][None][None]))\n",
+    "                voxels_raw = torch.Tensor(voxels_raw).clamp(0,1).to(device)\n",
+    "                \n",
+    "                ## Process it through pretrained MAE ##\n",
+    "                encoder_out = model(voxels_raw, encoder_mask=tube_mask)\n",
+    "                err\n",
+    "                # encoder_out = bn(encoder_out)\n",
+    "                \n",
+    "                voxel0 = encoder_out.flatten(1).unsqueeze(1)\n",
+    "                voxel0 = nn.functional.normalize(voxel0,dim=-1).cpu()\n",
+    "                \n",
+    "                assert len(voxel0) == batch_size\n",
+    "\n",
+    "                if epoch < int(mixup_pct * num_epochs):\n",
+    "                    voxel0, perm, betas, select = utils.mixco(voxel0)\n",
+    "                    perm_iters[f\"subj0{subj_list[s]}_iter{iter}\"] = perm\n",
+    "                    betas_iters[f\"subj0{subj_list[s]}_iter{iter}\"] = betas\n",
+    "                    select_iters[f\"subj0{subj_list[s]}_iter{iter}\"] = select\n",
+    "\n",
+    "                voxel_iters[f\"subj0{subj_list[s]}_iter{iter}\"] = voxel0\n",
+    "\n",
+    "                if iter >= num_iterations_per_epoch:\n",
+    "                    break\n",
+    "\n",
+    "    # you now have voxel_iters and image_iters with num_iterations_per_epoch batches each\n",
+    "    for train_i in range(num_iterations_per_epoch):\n",
+    "        with torch.cuda.amp.autocast(dtype=data_type):\n",
+    "            optimizer.zero_grad()\n",
+    "            loss=0.\n",
+    "\n",
+    "            voxel_list = [voxel_iters[f\"subj0{s}_iter{train_i}\"].detach().to(device) for s in subj_list]\n",
+    "            image = image_iters[train_i].detach()\n",
+    "            image = image.to(device)\n",
+    "\n",
+    "            clip_target = clip_img_embedder(image)\n",
+    "            assert not torch.any(torch.isnan(clip_target))\n",
+    "\n",
+    "            if epoch < int(mixup_pct * num_epochs):\n",
+    "                perm_list = [perm_iters[f\"subj0{s}_iter{train_i}\"].detach().to(device) for s in subj_list]\n",
+    "                perm = torch.cat(perm_list, dim=0)\n",
+    "                betas_list = [betas_iters[f\"subj0{s}_iter{train_i}\"].detach().to(device) for s in subj_list]\n",
+    "                betas = torch.cat(betas_list, dim=0)\n",
+    "                select_list = [select_iters[f\"subj0{s}_iter{train_i}\"].detach().to(device) for s in subj_list]\n",
+    "                select = torch.cat(select_list, dim=0)\n",
+    "\n",
+    "            voxel_ridge_list = [mindeye.ridge(voxel_list[si],si) for si,s in enumerate(subj_list)]\n",
+    "            clip_voxels = torch.cat(voxel_ridge_list, dim=0)\n",
+    "\n",
+    "#             backbone, clip_voxels = mindeye.backbone(voxel_ridge)\n",
+    "\n",
+    "            clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1)\n",
+    "            clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)\n",
+    "\n",
+    "            if epoch < int(mixup_pct * num_epochs):                \n",
+    "                loss_clip = utils.mixco_nce(\n",
+    "                    clip_voxels_norm,\n",
+    "                    clip_target_norm,\n",
+    "                    temp=.006,\n",
+    "                    perm=perm, betas=betas, select=select)\n",
+    "            else:\n",
+    "                epoch_temp = soft_loss_temps[epoch-int(mixup_pct*num_epochs)]\n",
+    "                loss_clip = utils.soft_clip_loss(\n",
+    "                    clip_voxels_norm,\n",
+    "                    clip_target_norm,\n",
+    "                    temp=epoch_temp)\n",
+    "\n",
+    "            loss_clip_total += loss_clip.item()\n",
+    "            loss += loss_clip\n",
+    "\n",
+    "            # forward and backward top 1 accuracy        \n",
+    "            labels = torch.arange(len(clip_voxels_norm)).to(clip_voxels_norm.device) \n",
+    "            fwd_percent_correct += utils.topk(utils.prenormed_batchwise_cosine_similarity(clip_voxels_norm, clip_target_norm), labels, k=1).item()\n",
+    "            bwd_percent_correct += utils.topk(utils.prenormed_batchwise_cosine_similarity(clip_target_norm, clip_voxels_norm), labels, k=1).item()\n",
+    "\n",
+    "            utils.check_loss(loss)\n",
+    "            accelerator.backward(loss)\n",
+    "            optimizer.step()\n",
+    "\n",
+    "            losses.append(loss.item())\n",
+    "            lrs.append(optimizer.param_groups[0]['lr'])\n",
+    "\n",
+    "            lr_scheduler.step()\n",
+    "\n",
+    "    mindeye.eval()\n",
+    "    if local_rank==0:\n",
+    "        with torch.no_grad(), torch.cuda.amp.autocast(dtype=data_type): \n",
+    "            for test_i, (behav, past_behav, future_behav, old_behav) in enumerate(test_dl):  \n",
+    "                loss=0.     \n",
+    "\n",
+    "                coco_idx = behav[:,0,0].cpu().long()\n",
+    "                _,test_indices = np.unique(coco_idx, return_index=True)\n",
+    "                test_indices = np.random.permutation(test_indices)[:num_test_eval]\n",
+    "                image = images[coco_idx[test_indices]].float().to(device)\n",
+    "                \n",
+    "                # get the corresponding raw voxel time series\n",
+    "                for ib,b in enumerate(behav[test_indices,0,5].cpu().long().numpy()):\n",
+    "                    tr = (nsddata_raw_stimuli[nsddata_raw_stimuli['global_trial'].isin([b.item()])]['global_TR_onsets'].values + TR_delay).astype(np.int32).item()\n",
+    "                    if ib==0:\n",
+    "                        voxels_raw = mindeye_funcs[tr-2:tr+2][None][None]\n",
+    "                    else:\n",
+    "                        voxels_raw = np.vstack((voxels_raw, mindeye_funcs[tr-2:tr+2][None][None]))\n",
+    "                voxels_raw = torch.Tensor(voxels_raw).clamp(0,1).to(device)\n",
+    "                \n",
+    "                ## Process it through pretrained MAE ##\n",
+    "                encoder_out = model(voxels_raw, encoder_mask=tube_mask)\n",
+    "                # encoder_out = bn(encoder_out)\n",
+    "                \n",
+    "                voxel = encoder_out.flatten(1).unsqueeze(1)\n",
+    "                voxel = nn.functional.normalize(voxel,dim=-1)\n",
+    "\n",
+    "                assert len(image) == num_test_eval\n",
+    "\n",
+    "                clip_target = clip_img_embedder(image.float())\n",
+    "                \n",
+    "                clip_voxels_norm = nn.functional.normalize(voxel.flatten(1), dim=-1)\n",
+    "\n",
+    "                clip_voxels = mindeye.ridge(voxel,0) # 0th index of subj_list\n",
+    "                \n",
+    "#                 backbone, clip_voxels = mindeye.backbone(voxel_ridge)\n",
+    "\n",
+    "                clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1)\n",
+    "                clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)\n",
+    "                \n",
+    "                loss_clip = utils.soft_clip_loss(\n",
+    "                    clip_voxels_norm,\n",
+    "                    clip_target_norm,\n",
+    "                    temp=.006)\n",
+    "\n",
+    "                test_loss_clip_total += loss_clip.item()\n",
+    "                loss += loss_clip\n",
+    "\n",
+    "                # forward and backward top 1 accuracy        \n",
+    "                labels = torch.arange(len(clip_voxels_norm)).to(clip_voxels_norm.device) \n",
+    "                test_fwd_percent_correct += utils.topk(utils.prenormed_batchwise_cosine_similarity(clip_voxels_norm, clip_target_norm), labels, k=1).item()\n",
+    "                test_bwd_percent_correct += utils.topk(utils.prenormed_batchwise_cosine_similarity(clip_target_norm, clip_voxels_norm), labels, k=1).item()\n",
+    "                \n",
+    "                utils.check_loss(loss)                \n",
+    "                test_losses.append(loss.item())\n",
+    "\n",
+    "            logs = {\"train/loss\": np.mean(losses[-(train_i+1):]),\n",
+    "                \"test/loss\": np.mean(test_losses[-(test_i+1):]),\n",
+    "                \"train/lr\": lrs[-1],\n",
+    "                \"train/num_steps\": len(losses),\n",
+    "                \"test/num_steps\": len(test_losses),\n",
+    "                \"train/fwd_pct_correct\": fwd_percent_correct / (train_i + 1),\n",
+    "                \"train/bwd_pct_correct\": bwd_percent_correct / (train_i + 1),\n",
+    "                \"test/test_fwd_pct_correct\": test_fwd_percent_correct / (test_i + 1),\n",
+    "                \"test/test_bwd_pct_correct\": test_bwd_percent_correct / (test_i + 1),\n",
+    "                \"train/loss_clip_total\": loss_clip_total / (train_i + 1),\n",
+    "                \"test/loss_clip_total\": test_loss_clip_total / (test_i + 1),\n",
+    "                }\n",
+    "\n",
+    "            progress_bar.set_postfix(**logs)\n",
+    "            if wandb_log: wandb.log(logs)\n",
+    "            \n",
+    "    # Save model checkpoint\n",
+    "    if (ckpt_saving) and (epoch % ckpt_interval == 0):\n",
+    "        save_ckpt()\n",
+    "\n",
+    "    # wait for other GPUs to catch up if needed\n",
+    "    accelerator.wait_for_everyone()\n",
+    "    torch.cuda.empty_cache()\n",
+    "    gc.collect()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6277d8b8-63ab-4741-942f-69bb0caff816",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# torch.save(image0,\"image0.pt\")\n",
+    "# torch.save(encoder_out,\"encoder_out0.pt\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a06050bc-86a2-4134-a661-ead4afe3e1e7",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "plt.figure(figsize=(8, 3))\n",
+    "plt.plot(losses)\n",
+    "plt.title(\"Training losses\")\n",
+    "plt.show()\n",
+    "\n",
+    "plt.figure(figsize=(8, 3))\n",
+    "plt.plot(test_losses)\n",
+    "plt.title(\"Test losses\")\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d38371ee-6e16-4db4-a096-f3f9ebdd73cf",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "voxel_list[0].shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0885cc21-c8b2-4b6f-950d-913d6317ea40",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "found",
+   "language": "python",
+   "name": "found"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

fMRI-foundation-model-old/fMRI-MAE/downstream.py

@@ -0,0 +1,731 @@
+# Import packages and setup gpu configuration.
+# This code block shouldnt need to be adjusted!
+import os
+import shutil
+import sys
+import json
+import yaml
+import numpy as np
+import pandas as pd
+import copy
+import math
+from einops import rearrange
+from einops.layers.torch import Rearrange
+import time
+import random
+import h5py
+import webdataset as wds
+import gc
+import matplotlib.pyplot as plt
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision import transforms
+import utils
+from models import *
+from mindeye_models import *
+import nibabel as nib
+from nilearn import plotting
+from functools import partial
+
+# tf32 data type is faster than standard float32
+torch.backends.cuda.matmul.allow_tf32 = True
+
+### Multi-GPU config ###
+local_rank = os.getenv('LOCAL_RANK')
+if local_rank is None: 
+    local_rank = 0
+else:
+    local_rank = int(local_rank)
+print(f"LOCAL RANK={local_rank}")
+
+num_devices = os.getenv('NUM_GPUS')
+if num_devices is None: 
+    num_devices = 1
+else:
+    num_devices = int(num_devices)
+print(f"NUM GPUS={num_devices}")
+distributed = True if num_devices>1 else False
+
+global_rank = os.getenv('RANK')
+if global_rank is None:
+    global_rank = 0
+else:
+    global_rank = int(global_rank)
+print(f"GLOBAL RANK={global_rank}")
+
+from tqdm import tqdm
+
+# Load parameters from yaml config
+config = yaml.load(open('config.yaml', 'r'), Loader=yaml.FullLoader)
+
+# create global variables from the config
+for attribute_name in config.keys():
+    globals()[attribute_name] = config[f'{attribute_name}']
+    
+# Load MindEye parameters from yaml config (will override any params with same name)
+mindeye_config = yaml.load(open('mindeye_config.yaml', 'r'), Loader=yaml.FullLoader)
+
+# create global variables from the config
+for attribute_name in mindeye_config.keys():
+    globals()[attribute_name] = mindeye_config[f'{attribute_name}']
+
+data_type = torch.float32 # change depending on your mixed_precision
+
+batch_size = global_batch_size // num_devices
+print("batch_size", batch_size)
+    
+# First use "accelerate config" in terminal and setup using deepspeed stage 2 with CPU offloading
+from accelerate import Accelerator
+accelerator = Accelerator(split_batches=False, mixed_precision="fp16")
+
+from accelerate.state import AcceleratorState
+try:
+    AcceleratorState().deepspeed_plugin.deepspeed_config['train_micro_batch_size_per_gpu'] = batch_size
+    print("deepspeed reconfigured, train_micro_batch_size_per_gpu = ", batch_size)
+except:
+    print("skipping deepspeed reconfiguration...")
+
+print("PID of this process =",os.getpid())
+device = accelerator.device
+print("device:",device)
+world_size = accelerator.state.num_processes
+num_workers = num_devices
+print(accelerator.state)
+
+print("distributed =",distributed, "num_devices =", num_devices, "local rank =", local_rank, "world size =", world_size, "data_type =", data_type)
+print = accelerator.print # only print if local_rank=0
+
+print("mae config\n\n",config)
+print("mindeye_config\n",mindeye_config)
+
+# if utils.is_interactive():
+#     ckpt_saving = False
+#     wandb_log = False
+
+# seed all random functions
+utils.seed_everything(seed)
+
+mae_ckpt_pth = os.path.abspath(f'../ckpts/{mae_model_name}/last.pth')
+print("mae_ckpt_pth", mae_ckpt_pth)
+
+outdir = os.path.abspath(f'../ckpts/{model_name}')
+os.makedirs(outdir,exist_ok=True)
+print("outdir", outdir)
+
+if type(patch_size) == int:
+    patch_size = [patch_size,patch_size,patch_size]
+patch_depth = patch_size[0]
+patch_height = patch_size[1]
+patch_width = patch_size[2]
+
+num_patches = int(
+    (img_size[0] / patch_depth)
+    * (img_size[1] / patch_height)
+    * (img_size[2] / patch_width)
+    * num_frames
+)
+print("num_patches", num_patches)
+
+vit_size = {
+    "encoder": encoder_model,
+    "decoder": decoder_model
+}
+    
+model = get_vit(
+    size=vit_size,
+    image_size=img_size,  # depth, height, width
+    image_patch_size=(patch_depth,patch_height,patch_width),  # depth, height, width patch size
+    frames=num_frames,
+    frame_patch_size=frame_patch_size,
+    channels=1,
+    use_rope_emb=use_rope_emb,
+    use_cls_token=use_cls_token,
+    use_decoder_same_emb_dim=use_decoder_same_emb_dim,
+    decoder_depth=decoder_depth
+)
+utils.count_params(model)
+
+
+def save_ckpt(tag):
+    ckpt_path = outdir+f'/{tag}.pth'
+    if accelerator.is_main_process:
+        unwrapped_model = accelerator.unwrap_model(model)
+        if tag == "last" and os.path.exists(ckpt_path):
+            shutil.copyfile(ckpt_path, os.path.join(outdir, f'{tag}_old.pth'))
+        torch.save({
+            'epoch': epoch,
+            'model_state_dict': unwrapped_model.state_dict(),
+            'optimizer_state_dict': optimizer.state_dict(),
+            'lr_scheduler': lr_scheduler.state_dict(),
+            'train_losses': losses,
+            'test_losses': test_losses,
+            'lrs': lrs,
+            }, ckpt_path)
+        if tag == "last" and os.path.exists(os.path.join(outdir, f'{tag}_old.pth')):
+            os.remove(os.path.join(outdir, f'{tag}_old.pth'))
+    print(f"\n---saved {outdir}/{tag} ckpt!---\n")
+
+def load_ckpt(tag,load_lr=True,load_optimizer=True,load_epoch=True,strict=True,outdir=outdir,multisubj_loading=False): 
+    print(f"\n---loading {outdir}/{tag}.pth ckpt---\n")
+    checkpoint = torch.load(outdir+'/last.pth', map_location='cpu')
+    state_dict = checkpoint['model_state_dict']
+    if multisubj_loading: # remove incompatible ridge layer that will otherwise error
+        state_dict.pop('ridge.linears.0.weight',None)
+    model.load_state_dict(state_dict, strict=strict)
+    if load_epoch:
+        globals()["epoch"] = checkpoint['epoch']
+        print("Epoch",epoch)
+    if load_optimizer:
+        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+    if load_lr:
+        lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
+    del checkpoint
+
+
+checkpoint = torch.load(mae_ckpt_pth, map_location=device)
+try:
+    model.load_state_dict(checkpoint['model_state_dict'], strict=False)
+except:
+    pass
+
+# set foundation model to evaluation
+model.eval()
+model.requires_grad_(False)
+model.to(device)
+pass
+
+
+nsddata_raw_stimuli = pd.read_csv(f"{nsd_raw_path}/nsddata_rawdata.csv")
+TR_delay = 3 # to account for bold hrf
+train_TRs = np.round(nsddata_raw_stimuli[nsddata_raw_stimuli['shared1000'] == False]['global_TR_onsets'].values + TR_delay).astype(np.int32)
+test_TRs = np.round(nsddata_raw_stimuli[nsddata_raw_stimuli['shared1000'] == True]['global_TR_onsets'].values + TR_delay).astype(np.int32)
+
+
+# Load 73k NSD images
+f = h5py.File(f'{nsd_image_path}/coco_images_224_float16.hdf5', 'r')
+images = f['images'][:] 
+images = torch.Tensor(images).to("cpu").to(data_type)
+print("Loaded all 73k possible NSD images!", images.shape)
+
+# Load MindEye hdf5
+f = h5py.File(f'{nsd_raw_path}/subj01_mnidata.h5', 'r') #subj01_rawdata_old.h5
+mindeye_global_trs = f['global_trs'][:]
+mindeye_funcs = f['funcs']
+
+
+# clip_img_embedder = FrozenOpenCLIPImageEmbedder(
+#     arch="ViT-bigG-14",
+#     version="laion2b_s39b_b160k",
+#     output_tokens=True,
+#     only_tokens=True,
+# )
+# clip_seq_dim = 256
+# clip_emb_dim = 1664
+# clip_img_embedder.to(device)
+
+clip_img_embedder = FrozenOpenCLIPImageEmbedder(
+    arch="ViT-bigG-14",
+    version="laion2b_s39b_b160k",
+    output_tokens=False,
+    only_tokens=False,
+    init_device=device,
+    device=device,
+    cache_dir="/weka/proj-fmri/shared/cache"
+)
+clip_seq_dim = 1
+clip_emb_dim = 1280
+clip_img_embedder.to(device)
+
+
+subj = s = 1
+subj_list = [subj]
+
+# if multi_subject:
+#     nsessions_allsubj=np.array([40, 40, 32, 30, 40, 32, 40, 30])
+#     num_samples_per_epoch = (750*40) // num_devices 
+# else:
+#     num_samples_per_epoch = (750*num_sessions) // num_devices 
+
+# num_samples_per_epoch = 1024 #(750*num_sessions) // num_devices
+
+print("dividing batch size by subj_list, which will then be concatenated across subj during training...") 
+batch_size = batch_size // len(subj_list)
+num_iterations_per_epoch = num_samples_per_epoch // (batch_size*len(subj_list))
+print("batch_size =", batch_size, "num_iterations_per_epoch =",num_iterations_per_epoch, "num_samples_per_epoch =",num_samples_per_epoch)
+
+train_data = {}
+train_dl = {}
+
+print(f"Training with {num_sessions} sessions")
+train_url = f"{nsd_wds_path}/subj0{s}/train/" + "{0.." + f"{num_sessions-1}" + "}.tar"
+print(train_url)
+    
+train_data[f'subj0{s}'] = wds.WebDataset(train_url,resampled=True,nodesplitter=utils.my_split_by_node)\
+                    .shuffle(750, initial=1500, rng=random.Random(42))\
+                    .decode("torch")\
+                    .rename(behav="behav.npy", past_behav="past_behav.npy", future_behav="future_behav.npy", olds_behav="olds_behav.npy")\
+                    .to_tuple(*["behav", "past_behav", "future_behav", "olds_behav"])
+# train_dl[f'subj0{s}'] = torch.utils.data.DataLoader(train_data[f'subj0{s}'], batch_size=batch_size, shuffle=False, drop_last=True, pin_memory=True)
+train_dl[f'subj0{s}'] = wds.WebLoader(
+    train_data[f'subj0{s}'].batched(batch_size), 
+    pin_memory=True,
+    shuffle=False,
+    batch_size=None,
+    num_workers=num_workers, 
+    persistent_workers=num_workers>0,
+).with_epoch(num_iterations_per_epoch)
+
+print("Loaded all subj train dls and betas!\n")
+if subj==3:
+    num_test=2371
+elif subj==4:
+    num_test=2188
+elif subj==6:
+    num_test=2371
+elif subj==8:
+    num_test=2188
+else:
+    num_test=300  # 3000
+test_url = f"{nsd_wds_path}/subj0{subj}/new_test/" + "0.tar"
+print(test_url)
+test_data = wds.WebDataset(test_url,resampled=True,nodesplitter=utils.my_split_by_node)\
+                    .decode("torch")\
+                    .rename(behav="behav.npy", past_behav="past_behav.npy", future_behav="future_behav.npy", olds_behav="olds_behav.npy")\
+                    .to_tuple(*["behav", "past_behav", "future_behav", "olds_behav"])
+# test_dl = torch.utils.data.DataLoader(test_data, batch_size=num_test, shuffle=False, drop_last=True, pin_memory=True)
+test_dl = wds.WebLoader(
+    test_data.batched(num_test),
+    pin_memory=True,
+    shuffle=False,
+    batch_size=None,
+    num_workers=num_workers,
+    persistent_workers=num_workers>0,
+).with_epoch(10)
+print(f"Loaded test dl for subj{subj}!\n")
+
+
+class MindEyeModule(nn.Module):
+    def __init__(self):
+        super(MindEyeModule, self).__init__()
+    def forward(self, x):
+        return x
+
+# class RidgeRegression(torch.nn.Module):
+#     # make sure to add weight_decay when initializing optimizer
+#     def __init__(self, input_sizes, out_features, seq_len=1): 
+#         super(RidgeRegression, self).__init__()
+#         self.seq_len = seq_len
+#         self.out_features = out_features
+#         self.linears = torch.nn.ModuleList([
+#                 torch.nn.Linear(input_size, out_features) for input_size in input_sizes
+#             ])
+#     def forward(self, x, subj_idx):
+#         out = torch.cat([self.linears[subj_idx](x[:,seq]).unsqueeze(1) for seq in range(self.seq_len)], dim=1)
+#         return out
+    
+class MLP(torch.nn.Module):
+    def __init__(self, input_sizes, out_features, seq_len=0): 
+        super(MLP, self).__init__()
+        self.input_sizes = input_sizes[0]
+        self.out_features = out_features
+        self.mlp = nn.Sequential(
+            nn.LayerNorm(self.input_sizes),
+            nn.GELU(),
+            nn.Linear(self.input_sizes, hidden_dim//2),
+            nn.LayerNorm(hidden_dim//2),
+            nn.GELU(),
+            nn.Dropout(drop),
+            nn.Linear(hidden_dim//2, hidden_dim),
+            nn.LayerNorm(hidden_dim),
+            nn.GELU(),
+            nn.Dropout(drop),
+            nn.Linear(hidden_dim, out_features)
+        )
+    def forward(self, x, z=None):
+        out = self.mlp(x[:,0])
+        return out
+
+class ReversibleBrainNetwork(nn.Module):
+    def __init__(self, out_dim=768, in_dim=15724, h=4096, n_blocks=4, norm_type='bn', act_first=True, 
+                 encoder_tokens=257, reverse=True, **kwargs):
+        super().__init__()
+        norm_func = partial(nn.BatchNorm1d, num_features=h) if norm_type == 'bn' else partial(nn.LayerNorm, normalized_shape=h)
+        act_fn = partial(nn.ReLU, inplace=True) if norm_type == 'bn' else nn.GELU
+        act_and_norm = (act_fn, norm_func) if act_first else (norm_func, act_fn)
+        # self.temp = nn.Parameter(torch.tensor(.006))
+        start_dim, end_dim = (out_dim*encoder_tokens, in_dim) if reverse else (in_dim, out_dim*encoder_tokens)
+        self.lin0 = nn.Sequential(
+            nn.Linear(start_dim, h),
+            *[item() for item in act_and_norm],
+            nn.Dropout(0.5),
+        )
+        
+        self.mlp = nn.ModuleList([
+            nn.Sequential(
+                nn.Linear(h, h),
+                *[item() for item in act_and_norm],
+                nn.Dropout(drop)
+            ) for _ in range(n_blocks)
+        ])
+        
+        self.lin1 = nn.Linear(h, end_dim, bias=True)
+        self.n_blocks = n_blocks
+        
+    def forward(self, x, *args, **kwargs):
+        x = self.lin0(x.flatten(1))  # bs, h
+        residual = x
+        for res_block in range(self.n_blocks):
+            x = self.mlp[res_block](x)
+            x += residual
+            residual = x
+        x = x.reshape(len(x), -1)
+        x = self.lin1(x)
+        return x
+
+
+mindeye = MindEyeModule()
+# mindeye.ridge = MLP(np.array([in_dim]), out_features=clip_emb_dim*clip_seq_dim)
+mindeye.ridge = ReversibleBrainNetwork(in_dim=in_dim, h=hidden_dim, out_dim=clip_emb_dim, 
+                                       encoder_tokens=1, norm_type='ln', act_first=False, reverse=False)
+
+# mindeye.backbone = BrainNetwork(h=hidden_dim, in_dim=hidden_dim, n_blocks=4, drop=drop,
+#                           clip_size=clip_emb_dim, out_dim=clip_emb_dim*clip_seq_dim, clip_scale=1)
+utils.count_params(mindeye.ridge)
+# utils.count_params(mindeye.backbone)
+utils.count_params(mindeye)
+
+
+no_decay = ["bias", "LayerNorm.bias", "LayerNorm.weight"]
+opt_grouped_parameters = [
+    {'params': [p for n, p in mindeye.ridge.named_parameters()], 'weight_decay': 1e-2},
+    # {'params': [p for n, p in mindeye.backbone.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': 1e-2},
+    # {'params': [p for n, p in mindeye.backbone.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0},
+]
+
+optimizer = torch.optim.AdamW(opt_grouped_parameters, lr=max_lr)
+
+total_steps = iters_scale_factor * num_epochs * num_iterations_per_epoch
+print("total_steps", total_steps)
+pct_start = 2/num_epochs if num_epochs>1 else 1.
+lr_scheduler = torch.optim.lr_scheduler.OneCycleLR(
+    optimizer,
+    max_lr=max_lr,
+    total_steps=total_steps,
+    pct_start=0.1
+)
+
+print("\nDone with model preparations!")
+num_params = utils.count_params(mindeye)
+
+
+if accelerator.is_main_process and wandb_log: # only use main process for wandb logging
+    import wandb
+    wandb_project = 'found_downstream'
+    print(f"wandb {wandb_project} run {model_name}")
+    # need to configure wandb beforehand in terminal with "wandb init"!
+    wandb_config = {
+      "model_name": model_name,
+      "mae_model_name": mae_model_name,
+      "global_batch_size": global_batch_size,
+      "batch_size": batch_size,
+      "num_epochs": num_epochs,
+      "num_sessions": num_sessions,
+      "num_samples_per_epoch": num_samples_per_epoch,
+      "in_dim": in_dim,
+      "hidden_dim": hidden_dim,
+      "mixup_pct": mixup_pct,
+      "num_params": num_params,
+      "max_lr": max_lr,
+      "ckpt_interval": ckpt_interval,
+      "ckpt_saving": ckpt_saving,
+      "seed": seed,
+      "distributed": distributed,
+      "num_devices": num_devices,
+      "world_size": world_size,
+      "train_url": train_url,
+      "test_url": test_url,
+    }
+    print("wandb_config:\n",wandb_config)
+    print("wandb_id:",model_name)
+    wandb.init(
+        id=model_name,
+        project=wandb_project,
+        name=model_name,
+        config=wandb_config,
+        resume="allow",
+    )
+else:
+    wandb_log = False
+
+
+epoch = 0
+losses, test_losses, lrs = [], [], []
+best_test_loss = 1e9
+torch.cuda.empty_cache()
+
+# if masking_strategy=="MNI":
+from einops.layers.torch import Rearrange
+
+MNI_brain = nib.load("/weka/proj-fmri/paulscotti/fMRI-foundation-model/dataset_creation/afni_conversion/tpl-MNI152NLin2009cAsym_res-02_T1w_brain.nii.gz").get_fdata()
+brain_pos_voxels = utils.crop_or_pad(torch.from_numpy(MNI_brain[6:94,8:112,10:82]), img_size)
+# brain_pos_voxels = MNI_brain[10:90,12:108,14:78]
+
+# brain_pos_voxels = brain_pos_voxels[:,30:31,:]
+
+brain_pos_pats = Rearrange(
+        "b c (f pf) (d pd) (h ph) (w pw) -> b f d h w (pd ph pw pf c)",
+        pd=patch_depth,
+        ph=patch_height,
+        pw=patch_width,
+        pf=1,
+    )(brain_pos_voxels[None,None,None])
+
+brain_pos_pats_vit = rearrange(brain_pos_pats, "b ... d -> b (...) d").mean(-1)[0]
+    
+tube_mask = torch.zeros(num_patches // num_frames).to(torch.bool)
+batch_positive_approx = (brain_pos_pats_vit > 0)
+mask_idx_candidates = torch.where(batch_positive_approx)[0]
+mask_idx_candidates = mask_idx_candidates[torch.randperm(len(mask_idx_candidates))]
+tube_idx = mask_idx_candidates[:int(num_patches / num_frames * (1 - tube_end_masking_ratio))]
+tube_mask[tube_idx] = True
+tube_mask = tube_mask.tile(num_frames//frame_patch_size)
+
+# load multisubject stage1 ckpt if set
+if multisubject_ckpt!="None" and not resume_from_ckpt:
+    load_ckpt("last",outdir=multisubject_ckpt,load_lr=False,load_optimizer=False,load_epoch=False,strict=False,multisubj_loading=True)
+    
+# load saved ckpt model weights into current model
+if resume_from_ckpt:
+    load_ckpt("last",load_lr=True,load_optimizer=True,load_epoch=True)
+elif wandb_log:
+    if wandb.run.resumed:
+        if os.path.exists(os.path.join(outdir, 'last.pth')) or os.path.exists(os.path.join(outdir, 'last_old.pth')):
+            if os.path.exists(os.path.join(outdir, 'last_old.pth')):
+                if os.path.exists(os.path.join(outdir, 'last.pth')):
+                    # this is corrupted
+                    os.remove(os.path.join(outdir, f'last.pth'))
+                # set last_old as last
+                shutil.move(os.path.join(outdir, f'last_old.pth'), os.path.join(outdir, f'last.pth'))
+            
+            # ckpt_path = os.path.join(outdir, 'last.pth')
+            # resume_from_ckpt = True
+            load_ckpt("last",load_lr=True,load_optimizer=True,load_epoch=True)
+
+train_dls = [train_dl[f'subj0{s}'] for s in subj_list]
+
+mindeye, optimizer, *train_dls, lr_scheduler = accelerator.prepare(mindeye, optimizer, *train_dls, lr_scheduler)
+# leaving out test_dl since we will only have local_rank 0 device do evals
+
+print(f"{model_name} starting with epoch {epoch} / {num_epochs}")
+progress_bar = tqdm(range(epoch, num_epochs), disable=not accelerator.is_main_process)
+mse = nn.MSELoss()
+soft_loss_temps = utils.cosine_anneal(0.004, 0.0075, num_epochs - int(mixup_pct * num_epochs))
+
+bn = nn.BatchNorm1d(512,affine=False).to(device)
+
+test_image=None
+num_test_eval=batch_size # should instead be average same-image 300 to mimic MindEye2 retrieval evaluation
+
+for epoch in progress_bar:
+    mindeye.train()
+
+    fwd_percent_correct = 0.
+    bwd_percent_correct = 0.
+    test_fwd_percent_correct = 0.
+    test_bwd_percent_correct = 0.
+    loss_clip_total = 0.
+    test_loss_clip_total = 0.
+
+    # pre-load all batches for this epoch (it's MUCH faster to pre-load in bulk than to separate loading per batch)
+    voxel_iters = {} # empty dict because diff subjects have differing # of voxels
+    image_iters = torch.zeros(num_iterations_per_epoch, batch_size*len(subj_list), 3, 224, 224).float()
+    annot_iters = {}
+    perm_iters, betas_iters, select_iters = {}, {}, {}
+    for s, train_dl in enumerate(train_dls):
+        with torch.cuda.amp.autocast(dtype=data_type):
+            for iter, (behav0, past_behav0, future_behav0, old_behav0) in enumerate(tqdm(train_dl,total=num_iterations_per_epoch)):
+                image0 = images[behav0[:,0,0].cpu().long()].float()
+                image_iters[iter,s*batch_size:s*batch_size+batch_size] = image0
+
+                # if images are not fully preloaded, then can do this inefficient but more memory friendly approach
+                # for ib,b in enumerate(behav0[:,0,0].cpu().long()):
+                #     if ib==0:
+                #         image0 = torch.Tensor(images[[b]])
+                #     else:
+                #         image0 = torch.vstack((image0, torch.Tensor(images[[b]])))
+                # image_iters[iter,s*batch_size:s*batch_size+batch_size] = image0
+                
+                # get the corresponding raw voxel time series
+                for ib,b in enumerate(behav0[:,0,5].cpu().long().numpy()):
+                    tr = (nsddata_raw_stimuli[nsddata_raw_stimuli['global_trial'].isin([b.item()])]['global_TR_onsets'].values + TR_delay).astype(np.int32).item()
+                    if ib==0:
+                        voxels_raw = mindeye_funcs[tr-2:tr+2][None][None]
+                    else:
+                        voxels_raw = np.vstack((voxels_raw, mindeye_funcs[tr-2:tr+2][None][None]))
+                voxels_raw = torch.Tensor(voxels_raw).clamp(0,1).to(device)
+                voxels_raw = utils.crop_or_pad(voxels_raw, img_size)
+                
+                ## Process it through pretrained MAE ##
+                encoder_out = model(voxels_raw, encoder_mask=tube_mask)
+                # encoder_out = bn(encoder_out)
+                
+                voxel0 = encoder_out.flatten(1).unsqueeze(1).cpu()
+                # voxel0 = nn.functional.normalize(voxel0,dim=-1).cpu()
+                
+                assert len(voxel0) == batch_size
+
+                if epoch < int(mixup_pct * num_epochs):
+                    voxel0, perm, betas, select = utils.mixco(voxel0)
+                    perm_iters[f"subj0{subj_list[s]}_iter{iter}"] = perm
+                    betas_iters[f"subj0{subj_list[s]}_iter{iter}"] = betas
+                    select_iters[f"subj0{subj_list[s]}_iter{iter}"] = select
+
+                voxel_iters[f"subj0{subj_list[s]}_iter{iter}"] = voxel0
+
+                if iter >= num_iterations_per_epoch:
+                    break
+
+    # you now have voxel_iters and image_iters with num_iterations_per_epoch batches each
+    for train_i in range(iters_scale_factor * num_iterations_per_epoch):
+        with torch.cuda.amp.autocast(dtype=data_type):
+            optimizer.zero_grad()
+            loss=0.
+
+            actual_idx = train_i % num_iterations_per_epoch
+
+            voxel_list = [voxel_iters[f"subj0{s}_iter{actual_idx}"].detach().to(device) for s in subj_list]
+            image = image_iters[actual_idx].detach()
+            image = image.to(device)
+
+            clip_target = clip_img_embedder(image)
+            assert not torch.any(torch.isnan(clip_target))
+
+            if epoch < int(mixup_pct * num_epochs):
+                perm_list = [perm_iters[f"subj0{s}_iter{actual_idx}"].detach().to(device) for s in subj_list]
+                perm = torch.cat(perm_list, dim=0)
+                betas_list = [betas_iters[f"subj0{s}_iter{actual_idx}"].detach().to(device) for s in subj_list]
+                betas = torch.cat(betas_list, dim=0)
+                select_list = [select_iters[f"subj0{s}_iter{actual_idx}"].detach().to(device) for s in subj_list]
+                select = torch.cat(select_list, dim=0)
+
+            voxel_ridge_list = [mindeye.ridge(voxel_list[si],si) for si,s in enumerate(subj_list)]
+            clip_voxels = torch.cat(voxel_ridge_list, dim=0)
+
+#             backbone, clip_voxels = mindeye.backbone(voxel_ridge)
+
+            clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1)
+            clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)
+
+            if epoch < int(mixup_pct * num_epochs):                
+                loss_clip = utils.mixco_nce(
+                    clip_voxels_norm,
+                    clip_target_norm,
+                    temp=.006,
+                    perm=perm, betas=betas, select=select)
+            else:
+                epoch_temp = soft_loss_temps[epoch-int(mixup_pct*num_epochs)]
+                loss_clip = utils.soft_clip_loss(
+                    clip_voxels_norm,
+                    clip_target_norm,
+                    temp=epoch_temp)
+
+            loss_clip_total += loss_clip.item()
+            loss += loss_clip
+
+            # forward and backward top 1 accuracy        
+            labels = torch.arange(len(clip_voxels_norm)).to(clip_voxels_norm.device) 
+            fwd_percent_correct += utils.topk(utils.prenormed_batchwise_cosine_similarity(clip_voxels_norm, clip_target_norm), labels, k=1).item()
+            bwd_percent_correct += utils.topk(utils.prenormed_batchwise_cosine_similarity(clip_target_norm, clip_voxels_norm), labels, k=1).item()
+
+            utils.check_loss(loss)
+            accelerator.backward(loss)
+            optimizer.step()
+
+            losses.append(loss.item())
+            lrs.append(optimizer.param_groups[0]['lr'])
+
+            lr_scheduler.step()
+
+    mindeye.eval()
+    if local_rank==0:
+        with torch.no_grad(), torch.cuda.amp.autocast(dtype=data_type): 
+            for test_i, (behav, past_behav, future_behav, old_behav) in enumerate(test_dl):  
+                loss=0.     
+
+                coco_idx = behav[:,0,0].cpu().long()
+                _,test_indices = np.unique(coco_idx, return_index=True)
+                test_indices = np.random.permutation(test_indices)[:num_test_eval]
+                image = images[coco_idx[test_indices]].float().to(device)
+                
+                # get the corresponding raw voxel time series
+                for ib,b in enumerate(behav[test_indices,0,5].cpu().long().numpy()):
+                    tr = (nsddata_raw_stimuli[nsddata_raw_stimuli['global_trial'].isin([b.item()])]['global_TR_onsets'].values + TR_delay).astype(np.int32).item()
+                    if ib==0:
+                        voxels_raw = mindeye_funcs[tr-2:tr+2][None][None]
+                    else:
+                        voxels_raw = np.vstack((voxels_raw, mindeye_funcs[tr-2:tr+2][None][None]))
+                voxels_raw = torch.Tensor(voxels_raw).clamp(0,1).to(device)
+                voxels_raw = utils.crop_or_pad(voxels_raw, img_size)
+                
+                ## Process it through pretrained MAE ##
+                encoder_out = model(voxels_raw, encoder_mask=tube_mask)
+                # encoder_out = bn(encoder_out)
+                
+                voxel = encoder_out.flatten(1).unsqueeze(1)
+                # voxel = nn.functional.normalize(voxel,dim=-1)
+
+                assert len(image) == num_test_eval
+
+                clip_target = clip_img_embedder(image.float())
+                
+                clip_voxels_norm = nn.functional.normalize(voxel.flatten(1), dim=-1)
+
+                clip_voxels = mindeye.ridge(voxel,0) # 0th index of subj_list
+                
+#                 backbone, clip_voxels = mindeye.backbone(voxel_ridge)
+
+                clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1)
+                clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)
+                
+                loss_clip = utils.soft_clip_loss(
+                    clip_voxels_norm,
+                    clip_target_norm,
+                    temp=.006)
+
+                test_loss_clip_total += loss_clip.item()
+                loss += loss_clip
+
+                # forward and backward top 1 accuracy        
+                labels = torch.arange(len(clip_voxels_norm)).to(clip_voxels_norm.device) 
+                test_fwd_percent_correct += utils.topk(utils.prenormed_batchwise_cosine_similarity(clip_voxels_norm, clip_target_norm), labels, k=1).item()
+                test_bwd_percent_correct += utils.topk(utils.prenormed_batchwise_cosine_similarity(clip_target_norm, clip_voxels_norm), labels, k=1).item()
+                
+                utils.check_loss(loss)                
+                test_losses.append(loss.item())
+
+            logs = {"train/loss": np.mean(losses[-(train_i+1):]),
+                "test/loss": np.mean(test_losses[-(test_i+1):]),
+                "train/lr": lrs[-1],
+                "train/num_steps": len(losses),
+                "test/num_steps": len(test_losses),
+                "train/fwd_pct_correct": fwd_percent_correct / (train_i + 1),
+                "train/bwd_pct_correct": bwd_percent_correct / (train_i + 1),
+                "test/test_fwd_pct_correct": test_fwd_percent_correct / (test_i + 1),
+                "test/test_bwd_pct_correct": test_bwd_percent_correct / (test_i + 1),
+                "train/loss_clip_total": loss_clip_total / (train_i + 1),
+                "test/loss_clip_total": test_loss_clip_total / (test_i + 1),
+                }
+
+            progress_bar.set_postfix(**logs)
+            if wandb_log: wandb.log(logs)
+            
+    # Save model checkpoint
+    if (ckpt_saving) and (epoch % ckpt_interval == 0):
+        save_ckpt('last')
+
+    # wait for other GPUs to catch up if needed
+    accelerator.wait_for_everyone()
+    torch.cuda.empty_cache()
+    gc.collect()
+

fMRI-foundation-model-old/fMRI-MAE/downstream.slurm

@@ -0,0 +1,29 @@
+#!/bin/bash
+#SBATCH --account=fmri
+#SBATCH --partition=a40
+#SBATCH --job-name=downMAE
+#SBATCH --nodes=1              
+#SBATCH --gres=gpu:1
+#SBATCH --time=20:00:00         # total run time limit (HH:MM:SS)
+#SBATCH -e slurms/%j.err        # first create a "slurms" folder in current directory to store logs
+#SBATCH -o slurms/%j.out
+#SBATCH --comment=medarc
+#SBATCH --no-requeue
+#SBATCH --exclusive
+
+export NUM_GPUS=1  # Set to equal gres=gpu:#!
+
+source ~/.bashrc
+cd /weka/proj-fmri/paulscotti/fMRI-foundation-model/fMRI-MAE/
+jupyter nbconvert downstream.ipynb --to python
+
+# Make sure another job doesnt use same port, here using random number
+export MASTER_PORT=$((RANDOM % (19000 - 11000 + 1) + 11000)) 
+export HOSTNAMES=$(scontrol show hostnames "$SLURM_JOB_NODELIST")
+export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
+export COUNT_NODE=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | wc -l)
+echo MASTER_ADDR=${MASTER_ADDR}
+echo MASTER_PORT=${MASTER_PORT}
+echo WORLD_SIZE=${COUNT_NODE}
+
+accelerate launch --num_processes=$(($NUM_GPUS * $COUNT_NODE)) --num_machines=$COUNT_NODE --main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT --mixed_precision=fp16 downstream.py

fMRI-foundation-model-old/fMRI-MAE/mindeye_models.py

@@ -0,0 +1,276 @@
+from typing import Optional
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class MindEyeModule(nn.Module):
+    def __init__(self):
+        super().__init__()
+    def forward(self, x):
+        return x
+class RidgeRegression(nn.Module):
+    # make sure to add weight_decay when initializing optimizer
+    def __init__(self, input_sizes, out_features, seq_len=1): 
+        super().__init__()
+        self.out_features = out_features
+        self.linears = torch.nn.ModuleList([
+                torch.nn.Linear(input_size, out_features) for input_size in input_sizes
+            ])
+        self.seq_len = seq_len
+    def forward(self, x, subj_idx):
+        out = torch.cat([self.linears[subj_idx](x[:,seq]).unsqueeze(1) for seq in range(self.seq_len)], dim=1)
+        return out
+class BrainNetwork(nn.Module):
+    def __init__(self, h=1024, in_dim=15724, out_dim=768, seq_len=1, n_blocks=4, drop=.15, clip_size=768, clip_scale=1):
+        super().__init__()
+        self.seq_len = seq_len
+        self.h = h
+        self.clip_size = clip_size
+        self.clip_scale = clip_scale
+        self.mixer_blocks1 = nn.ModuleList([
+            self.mixer_block1(h, drop) for _ in range(n_blocks)
+        ])
+        self.mixer_blocks2 = nn.ModuleList([
+            self.mixer_block2(seq_len, drop) for _ in range(n_blocks)
+        ])
+        
+        # Output linear layer
+        self.backbone_linear = nn.Linear(h * seq_len, out_dim, bias=True) 
+        self.clip_proj = self.projector(clip_size, clip_size, h=clip_size)
+            
+    def projector(self, in_dim, out_dim, h=2048):
+        return nn.Sequential(
+            nn.LayerNorm(in_dim),
+            nn.GELU(),
+            nn.Linear(in_dim, h),
+            nn.LayerNorm(h),
+            nn.GELU(),
+            nn.Linear(h, h),
+            nn.LayerNorm(h),
+            nn.GELU(),
+            nn.Linear(h, out_dim)
+        )
+    
+    def mlp(self, in_dim, out_dim, drop):
+        return nn.Sequential(
+            nn.Linear(in_dim, out_dim),
+            nn.GELU(),
+            nn.Dropout(drop),
+            nn.Linear(out_dim, out_dim),
+        )
+    
+    def mixer_block1(self, h, drop):
+        return nn.Sequential(
+            nn.LayerNorm(h),
+            self.mlp(h, h, drop),  # Token mixing
+        )
+
+    def mixer_block2(self, seq_len, drop):
+        return nn.Sequential(
+            nn.LayerNorm(seq_len),
+            self.mlp(seq_len, seq_len, drop)  # Channel mixing
+        )
+        
+    def forward(self, x):
+        # make empty tensors
+        c = torch.Tensor([0.])
+        
+        # Mixer blocks
+        residual1 = x
+        residual2 = x.permute(0,2,1)
+        for block1, block2 in zip(self.mixer_blocks1,self.mixer_blocks2):
+            x = block1(x) + residual1
+            residual1 = x
+            x = x.permute(0,2,1)
+            
+            x = block2(x) + residual2
+            residual2 = x
+            x = x.permute(0,2,1)
+            
+        x = x.reshape(x.size(0), -1)
+        backbone = self.backbone_linear(x).reshape(len(x), -1, self.clip_size)
+        if self.clip_scale>0:
+            c = self.clip_proj(backbone)
+        
+        return backbone, c
+    
+    
+import open_clip
+import kornia
+class FrozenOpenCLIPImageEmbedder(nn.Module):
+    """
+    Uses the OpenCLIP vision transformer encoder for images
+    """
+
+    def __init__(
+        self,
+        arch="ViT-H-14",
+        version="laion2b_s32b_b79k",
+        device="cuda",
+        init_device="cpu",
+        max_length=77,
+        freeze=True,
+        antialias=True,
+        ucg_rate=0.0,
+        unsqueeze_dim=False,
+        repeat_to_max_len=False,
+        num_image_crops=0,
+        output_tokens=False,
+        l2_norm_tokens=False,
+        only_tokens=False,
+        cache_dir: Optional[str] = None,
+    ):
+        super().__init__()
+        model, _, _ = open_clip.create_model_and_transforms(
+            arch,
+            device=torch.device(init_device),
+            pretrained=version,
+            cache_dir=cache_dir,
+        )
+        del model.transformer
+        self.model = model
+        self.max_crops = num_image_crops
+        self.pad_to_max_len = self.max_crops > 0
+        self.repeat_to_max_len = repeat_to_max_len and (not self.pad_to_max_len)
+        self.device = device
+        self.max_length = max_length
+        if freeze:
+            self.freeze()
+
+        self.antialias = antialias
+
+        self.register_buffer(
+            "mean", torch.Tensor([0.48145466, 0.4578275, 0.40821073]), persistent=False
+        )
+        self.register_buffer(
+            "std", torch.Tensor([0.26862954, 0.26130258, 0.27577711]), persistent=False
+        )
+        self.ucg_rate = ucg_rate
+        self.unsqueeze_dim = unsqueeze_dim
+        self.stored_batch = None
+        self.model.visual.output_tokens = output_tokens
+        self.output_tokens = output_tokens
+        if only_tokens:
+            assert output_tokens
+        self.only_tokens = only_tokens
+        self.l2_norm_tokens = l2_norm_tokens
+        if l2_norm_tokens:
+            assert output_tokens
+
+    def preprocess(self, x):
+        # normalize to [0,1]
+        x = kornia.geometry.resize(
+            x,
+            (224, 224),
+            interpolation="bicubic",
+            align_corners=True,
+            antialias=self.antialias,
+        )
+        x = (x + 1.0) / 2.0
+        # renormalize according to clip
+        x = kornia.enhance.normalize(x, self.mean, self.std)
+        return x
+
+    def freeze(self):
+        self.model = self.model.eval()
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def forward(self, image, no_dropout=False):
+        z = self.encode_with_vision_transformer(image)
+        tokens = None
+        if self.output_tokens:
+            z, tokens = z[0], z[1]
+        z = z.to(image.dtype)
+        if self.ucg_rate > 0.0 and not no_dropout and not (self.max_crops > 0):
+            z = (
+                torch.bernoulli(
+                    (1.0 - self.ucg_rate) * torch.ones(z.shape[0], device=z.device)
+                )[:, None]
+                * z
+            )
+            if tokens is not None:
+                tokens = (
+                    expand_dims_like(
+                        torch.bernoulli(
+                            (1.0 - self.ucg_rate)
+                            * torch.ones(tokens.shape[0], device=tokens.device)
+                        ),
+                        tokens,
+                    )
+                    * tokens
+                )
+        if self.unsqueeze_dim:
+            z = z[:, None, :]
+        if self.output_tokens:
+            assert not self.repeat_to_max_len
+            assert not self.pad_to_max_len
+            if self.only_tokens:
+                return tokens
+            return tokens, z
+        if self.repeat_to_max_len:
+            if z.dim() == 2:
+                z_ = z[:, None, :]
+            else:
+                z_ = z
+            return repeat(z_, "b 1 d -> b n d", n=self.max_length), z
+        elif self.pad_to_max_len:
+            assert z.dim() == 3
+            z_pad = torch.cat(
+                (
+                    z,
+                    torch.zeros(
+                        z.shape[0],
+                        self.max_length - z.shape[1],
+                        z.shape[2],
+                        device=z.device,
+                    ),
+                ),
+                1,
+            )
+            return z_pad, z_pad[:, 0, ...]
+        return z
+
+    def encode_with_vision_transformer(self, img):
+        # if self.max_crops > 0:
+        #    img = self.preprocess_by_cropping(img)
+        if img.dim() == 5:
+            assert self.max_crops == img.shape[1]
+            img = rearrange(img, "b n c h w -> (b n) c h w")
+        img = self.preprocess(img)
+        if not self.output_tokens:
+            assert not self.model.visual.output_tokens
+            x = self.model.visual(img)
+            tokens = None
+        else:
+            assert self.model.visual.output_tokens
+            x, tokens = self.model.visual(img)
+            if self.l2_norm_tokens:
+                token_shape = tokens.shape
+                tokens = tokens.flatten(1)
+                tokens = torch.nn.functional.normalize(tokens, dim=-1)
+                tokens = (tokens - .0002) / .0015
+                tokens = tokens.view(token_shape)
+                tokens = (tokens * 1.0957) + .1598
+        if self.max_crops > 0:
+            x = rearrange(x, "(b n) d -> b n d", n=self.max_crops)
+            # drop out between 0 and all along the sequence axis
+            x = (
+                torch.bernoulli(
+                    (1.0 - self.ucg_rate)
+                    * torch.ones(x.shape[0], x.shape[1], 1, device=x.device)
+                )
+                * x
+            )
+            if tokens is not None:
+                tokens = rearrange(tokens, "(b n) t d -> b t (n d)", n=self.max_crops)
+                logpy.warning(
+                    f"You are running very experimental token-concat in {self.__class__.__name__}. "
+                    f"Check what you are doing, and then remove this message."
+                )
+        if self.output_tokens:
+            return x, tokens
+        return x
+
+    def encode(self, text):
+        return self(text)

fMRI-foundation-model-old/fMRI-MAE/models.py

@@ -0,0 +1,423 @@
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from einops.layers.torch import Rearrange
+from rope import RotaryPositionalEmbeddings4D
+
+def posemb_sincos_4d(patches, temperature=10000, dtype=torch.float32):
+    _, f, d, h, w, dim, device, dtype = (*patches.shape, patches.device, patches.dtype)
+
+    z, y, x, t = torch.meshgrid(
+        torch.arange(f, device=device),
+        torch.arange(d, device=device),
+        torch.arange(h, device=device),
+        torch.arange(w, device=device),
+        indexing="ij",
+    )
+
+    fourier_dim = dim // 8
+
+    omega = torch.arange(fourier_dim, device=device) / (fourier_dim - 1)
+    omega = 1.0 / (temperature**omega)
+
+    z, y, x, t = [v.flatten()[:, None] * omega[None, :] for v in [z, y, x, t]]
+
+    pe = torch.cat(
+        (z.sin(), z.cos(), y.sin(), y.cos(), x.sin(), x.cos(), t.sin(), t.cos()), dim=1
+    )
+    pe = F.pad(pe, (0, dim - (fourier_dim * 8)))
+    return pe.type(dtype)
+
+
+class FeedForward(nn.Module):
+    def __init__(self, embed_dim, hidden_dim):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.LayerNorm(embed_dim),
+            nn.Linear(embed_dim, hidden_dim),
+            nn.GELU(),
+            nn.Linear(hidden_dim, embed_dim),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        embed_dim: int,
+        num_heads: int = 8,
+        dim_head: int = 64,
+        use_rope: bool = False,
+        cls_token: bool = False,
+    ):
+        super().__init__()
+        inner_dim = dim_head * num_heads
+        self.num_heads = num_heads
+        self.scale = dim_head**-0.5
+        self.use_rope = use_rope
+        self.cls_token = cls_token
+        self.norm = nn.LayerNorm(embed_dim)
+
+        self.attend = nn.Softmax(dim=-1)
+
+        self.to_qkv = nn.Linear(embed_dim, inner_dim * 3, bias=False)
+        self.to_out = nn.Linear(inner_dim, embed_dim, bias=False)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        pos_embed: Optional[nn.Module],
+        mask: Optional[torch.Tensor] = None,
+    ):
+        x = self.norm(x)
+
+        qkv = self.to_qkv(x).chunk(3, dim=-1)
+        q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=self.num_heads), qkv)
+        if self.use_rope:
+            if pos_embed is None:
+                raise ValueError(
+                    "For RoPE embeddings `pos_embed` should be \
+                    passed to the Attention forward."
+                )
+            # apply RoPE other than CLS token if it's included.
+            if self.cls_token:
+                q_cls = q[:, :, :1, :]
+                k_cls = k[:, :, :1, :]
+                q = q[:, :, 1:, :]
+                k = k[:, :, 1:, :]
+            q = pos_embed(q, mask=mask)
+            k = pos_embed(k, mask=mask)
+            if self.cls_token:
+                q = torch.cat([q_cls, q], dim=2)
+                k = torch.cat([k_cls, k], dim=2)
+
+        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
+
+        # print("q", q.shape) # B num_heads N D
+        # print("dots", dots.shape) # B num_heads N N
+        attn = self.attend(dots)
+
+        out = torch.matmul(attn, v)
+        out = rearrange(out, "b h n d -> b n (h d)")
+        return self.to_out(out)
+
+
+class Transformer(nn.Module):
+    def __init__(
+        self,
+        embed_dim: int,
+        depth: int,
+        num_heads: int,
+        dim_head: int,
+        mlp_dim: int,
+        use_rope: bool = False,
+        grid_height: Optional[int] = None,
+        grid_width: Optional[int] = None,
+        grid_depth: Optional[int] = None,
+        grid_time: Optional[int] = None,
+        cls_token: bool = False,
+        **args,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.mlp_dim = mlp_dim
+        self.dim_head = dim_head
+        self.norm = nn.LayerNorm(embed_dim)
+        self.layers = nn.ModuleList([])
+        for _ in range(depth):
+            self.layers.append(
+                nn.ModuleList(
+                    [
+                        Attention(
+                            embed_dim,
+                            num_heads=num_heads,
+                            dim_head=dim_head,
+                            use_rope=use_rope,
+                            cls_token=cls_token,
+                        ),
+                        FeedForward(embed_dim, mlp_dim),
+                    ]
+                )
+            )
+        # RoPE positional embeddings
+        self.use_rope = use_rope
+        if self.use_rope:
+            self.rope_pos_emb = RotaryPositionalEmbeddings4D(
+                d=dim_head,
+                grid_depth=grid_depth,
+                grid_height=grid_height,
+                grid_width=grid_width,
+                grid_time=grid_time,
+            )
+
+    def forward(self, x, mask: Optional[torch.Tensor] = None):
+        for attn, ff in self.layers:
+            x = (
+                attn(
+                    x, 
+                    pos_embed=self.rope_pos_emb if self.use_rope else None, 
+                    mask=mask,
+                )
+                + x
+            )
+            x = ff(x) + x
+        return self.norm(x)
+
+
+class VisionTransformerMAE(nn.Module):
+    def __init__(
+        self,
+        *,
+        encoder,
+        decoder,
+        image_size,
+        image_patch_size,
+        frames,
+        frame_patch_size,
+        channels=1,
+        use_rope_emb: bool = False,
+        use_cls_token: bool = False,
+        num_ids=1000,
+        **args,
+    ):
+        super().__init__()
+        image_depth, image_height, image_width = image_size
+        patch_depth, patch_height, patch_width = image_patch_size
+
+        self.encoder_transformer = encoder
+        self.decoder_transformer = decoder
+
+        # Check divisibility
+        assert (image_depth % patch_depth == 0 and image_height % patch_height == 0 and image_width % 
+                patch_width == 0), "Image dimensions must be divisible by the patch size."
+        assert (frames % frame_patch_size == 0), "Frames must be divisible by the frame patch size"
+
+        self.patch_dim = channels * patch_depth * patch_height * patch_width * frame_patch_size
+
+        self.num_patches = image_size[0]//image_patch_size[0] * image_size[1]//image_patch_size[1] * image_size[2]//image_patch_size[2] * frames
+
+        self.patchify = Rearrange(
+            "b c (f pf) (d pd) (h ph) (w pw) -> b f d h w (pd ph pw pf c)",
+            pd=patch_depth,
+            ph=patch_height,
+            pw=patch_width,
+            pf=frame_patch_size,
+        )
+
+        self.unpatchify = nn.Sequential(
+            Rearrange(
+                "b (f d h w) (pd ph pw pf c) -> b f d h w (pd ph pw pf c)",
+                c=channels,
+                d=image_depth,
+                h=image_height,
+                w=image_width,
+                pd=patch_depth,
+                ph=patch_height,
+                pw=patch_width,
+                pf=frame_patch_size,
+            )
+        )
+        self.encoder_embed_dim = self.encoder_transformer.embed_dim
+
+        self.patch_to_emb = nn.Sequential(
+            nn.LayerNorm(self.patch_dim),
+            nn.Linear(self.patch_dim, self.encoder_embed_dim),
+            nn.LayerNorm(self.encoder_embed_dim),
+        )
+
+        self.use_rope_emb = use_rope_emb
+        if not self.use_rope_emb:
+            self.posemb_sincos_4d = posemb_sincos_4d(
+                torch.zeros(
+                    1,
+                    frames // frame_patch_size,
+                    image_depth // patch_depth,
+                    image_height // patch_height,
+                    image_width // patch_width,
+                    self.encoder_embed_dim,
+                )
+            )
+
+        # cls token
+        self.use_cls_token = use_cls_token
+        if use_cls_token:
+            self.cls_token = nn.Parameter(torch.zeros(1, 1, self.encoder_embed_dim))
+
+        if not isinstance(self.decoder_transformer, nn.Identity):
+            self.decoder_embed_dim = self.decoder_transformer.embed_dim
+            if self.encoder_embed_dim != self.decoder_embed_dim:
+                self.encoder_to_decoder = nn.Linear(self.encoder_embed_dim, self.decoder_embed_dim, bias=False)
+            self.mask_token = nn.Parameter(torch.zeros(1, 1, self.decoder_embed_dim))
+
+            self.decoder_proj = nn.Sequential(
+                                    nn.LayerNorm(self.decoder_embed_dim), 
+                                    nn.GELU(), 
+                                    nn.Linear(self.decoder_embed_dim, self.patch_dim),
+                                )
+
+    def forward(self, x, encoder_mask=None, decoder_mask=None, device="cuda", verbose=False):
+        # ENCODER
+        if decoder_mask is None:            
+            if verbose: print(x.shape)
+            x = self.patchify(x)
+            if verbose: print("patched", x.shape)
+            x = rearrange(x, "b ... d -> b (...) d")
+            if verbose: print("reshaped", x.shape)
+            
+            x = x[:, encoder_mask]
+            if verbose: print("masked", x.shape)
+            
+            x = self.patch_to_emb(x.to(device))
+            if verbose: print("patched_emb", x.shape)
+
+            if not self.use_rope_emb:
+                if verbose: print("pe", self.posemb_sincos_4d.shape)
+                x = x + self.posemb_sincos_4d[encoder_mask].to(device)
+            if self.use_cls_token:
+                cls_tokens = self.cls_token.expand(len(x), -1, -1)
+                x = torch.cat((cls_tokens, x), dim=1)
+            if verbose: print("masked", x.shape)
+            x = self.encoder_transformer(x, mask=encoder_mask if self.use_rope_emb else None)
+            if verbose: print(x.shape)
+            
+            
+            # if verbose: print(x.shape)
+            # x = self.patchify(x)
+            # if verbose: print("patched", x.shape)
+            # x = self.patch_to_emb(x.to(device))
+            # if verbose: print("patched_emb", x.shape)
+            # x = rearrange(x, "b ... d -> b (...) d")
+            # if verbose: print("reshaped", x.shape)
+            # if not self.use_rope_emb:
+            #     if verbose: print("pe", self.posemb_sincos_4d.shape)
+            #     x = x + self.posemb_sincos_4d.to(x.device)
+            # if verbose: print("x", x.shape)
+            # x = x[:, encoder_mask]
+            # if self.use_cls_token:
+            #     cls_tokens = self.cls_token.expand(len(x), -1, -1)
+            #     x = torch.cat((cls_tokens, x), dim=1)
+            # if verbose: print("masked", x.shape)
+            # x = self.encoder_transformer(x, mask=encoder_mask if self.use_rope_emb else None)
+            # if verbose: print(x.shape)
+        else:  # DECODER
+            if verbose: print(x.shape)
+            if self.encoder_embed_dim != self.decoder_embed_dim:
+                x = self.encoder_to_decoder(x.to(device))
+                if verbose: print("Linear", x.shape)
+            B, _, _ = x.shape
+            N = decoder_mask.sum()
+            mask = None
+            if not self.use_rope_emb:
+                pos_embed = self.posemb_sincos_4d.to(x.device)
+                if verbose: print("pe", pos_embed.shape)
+                if self.encoder_embed_dim != self.decoder_embed_dim:
+                    pos_embed = self.encoder_to_decoder(pos_embed)
+                    if verbose: print("Linear pe", pos_embed.shape)
+                pos_emd_encoder = pos_embed[encoder_mask]
+                pos_emd_decoder = pos_embed[decoder_mask]
+                if verbose: print("pos_emd_encoder", pos_emd_encoder.shape)
+                if verbose: print("pos_emd_decoder", pos_emd_decoder.shape)
+                if self.use_cls_token:
+                    cls_tokens = x[:,:1,:]
+                    x = x[:,1:,:]
+                
+                x = torch.cat([x + pos_emd_encoder, 
+                               self.mask_token.repeat(B, N, 1) + pos_emd_decoder], 
+                              dim=1)
+                if self.use_cls_token:
+                    x = torch.cat([cls_tokens, x], dim=1)
+            else:
+                mask = torch.cat((torch.where(encoder_mask)[0], torch.where(decoder_mask)[0]))
+                # No abs positional embeddings for RoPE
+                x = torch.cat([x, 
+                               self.mask_token.repeat(B, N - 1 if self.use_cls_token else N, 1)],
+                              dim=1)
+            if verbose: print("x_concat", x.shape)
+            x = self.decoder_transformer(x, mask=mask)
+            if verbose: print(x.shape)
+            x = self.decoder_proj(x)
+            if verbose: print("proj", x.shape)
+        return x
+
+def transformer_mini(**args):
+    return Transformer(
+        embed_dim=48,
+        depth=6,
+        num_heads=2, 
+        mlp_dim=1024,
+        dim_head=64,
+        **args
+    )
+
+def transformer_small(**args):
+    return Transformer(
+        embed_dim=384,
+        depth=12,
+        num_heads=6,
+        mlp_dim=1536,
+        dim_head=64,
+        **args
+    )
+
+def transformer_base(**args):
+    return Transformer(
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_dim=3072,
+        dim_head=64,
+        **args
+    )    
+
+def transformer_large(**args):
+    return Transformer(
+        embed_dim=1024,
+        depth=24,
+        num_heads=16,
+        mlp_dim=4096,
+        dim_head=64,
+        **args
+    )
+
+def transformer_huge(**args):
+    return Transformer(
+        embed_dim=1280,
+        depth=32,
+        num_heads=16,
+        mlp_dim=5120,
+        dim_head=64,
+        **args
+    )
+
+transformer_mapping = {
+    "vit_mini": transformer_mini,
+    "vit_small": transformer_small,
+    "vit_base": transformer_base,
+    "vit_large": transformer_large,
+    "vit_huge": transformer_huge,
+    "none": nn.Identity
+}
+
+def get_vit(size, **args):
+    encoder = transformer_mapping[size["encoder"]](**args)
+    if args['use_decoder_same_emb_dim']:
+        depth = args.pop('decoder_depth', None)
+        assert depth is not None and depth > 0 , "for same emb dim decoder depth has to be provided and be >0"
+        decoder = Transformer(
+            embed_dim=encoder.embed_dim,
+            depth=depth,
+            num_heads=encoder.num_heads,
+            mlp_dim=encoder.mlp_dim,
+            dim_head=encoder.dim_head,
+            **args
+        )
+    else:
+        decoder = transformer_mapping[size["decoder"]](**args)
+    return VisionTransformerMAE(encoder=encoder, decoder=decoder, **args)
+    

fMRI-foundation-model-old/fMRI-MAE/rope.py

@@ -0,0 +1,111 @@
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+
+class RotaryPositionalEmbeddings4D(nn.Module):
+    def __init__(
+        self,
+        d: int,
+        grid_height: int,
+        grid_width: int,
+        grid_depth: int,
+        grid_time: int,
+        base: int = 10_000,
+    ):
+        super().__init__()
+        assert d % 4 == 0, f"{d} is not divisible by 4."
+        self.base = base
+        self.d = d
+        self.grid_height = grid_height
+        self.grid_width = grid_width
+        self.grid_depth = grid_depth
+        self.grid_time = grid_time
+        self.cos_cached = None
+        self.sin_cached = None
+
+    def _build_cache_1d(self, effective_d: int, seq_idx: torch.Tensor) -> torch.Tensor:
+        seq_idx = seq_idx.reshape(-1)  # List of positions
+        theta = 1.0 / (
+            self.base ** (torch.arange(0, effective_d, 2).float() / effective_d)
+        )
+
+        idx_theta = torch.einsum("n,d->nd", seq_idx, theta)
+        idx_theta2 = torch.cat([idx_theta, idx_theta], dim=1)
+        return idx_theta2
+
+    def _build_cache(self, x: torch.Tensor) -> None:
+        # x: batched tensor
+        # x.shape -> Batch, Seq Length, Embed Dim
+        if self.cos_cached is not None and x.shape[1] <= self.cos_cached.shape[1]:
+            # if cache is already built
+            return
+        # get the positions
+        grid_h = torch.arange(self.grid_height, dtype=torch.float32)
+        grid_w = torch.arange(self.grid_width, dtype=torch.float32)
+        grid_d = torch.arange(self.grid_depth, dtype=torch.float32)
+        grid_t = torch.arange(self.grid_time, dtype=torch.float32)
+
+        grid = torch.meshgrid(
+            grid_t, grid_d, grid_h, grid_w, indexing="xy"
+        )  # This order should match with i/p
+        grid = torch.stack(grid, axis=0)
+        grid = grid.reshape(
+            [4, 1, self.grid_time, self.grid_depth, self.grid_height, self.grid_width]
+        )
+
+        # Get the embedings
+        emb_t = self._build_cache_1d(self.d // 4, grid[0])  # 1/4 in-case of 4D
+        emb_d = self._build_cache_1d(self.d // 4, grid[1])
+        emb_h = self._build_cache_1d(self.d // 4, grid[2])  # (T*D*H*W, embedding_dim/4)
+        emb_w = self._build_cache_1d(self.d // 4, grid[3])
+        emb = torch.concatenate(
+            [emb_t, emb_d, emb_h, emb_w], axis=1
+        )  # (T*D*H*W, embedding_dim)
+        emb = emb.to(x.device)
+        # cache sin and cos
+        self.cos_cached = emb.cos()[
+            None, None, :, :
+        ]  # batch, Num Heads, Seq Len, Embed Dim
+        self.sin_cached = emb.sin()[None, None, :, :]
+
+    def _neg_half(self, x: torch.Tensor) -> torch.Tensor:
+        d_2 = self.d // 2
+        return torch.cat([-x[:, :, :, d_2:], x[:, :, :, :d_2]], dim=-1)
+
+    def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor]):
+        """
+        Args:
+        -----
+        x: query or key vector
+        mask: boolean vector of length sequence length.
+        True for the non-masked positions.
+        """
+        self._build_cache(x)
+        x_rope, x_pass = x[..., : self.d], x[..., self.d :]
+        neg_half_x = self._neg_half(x_rope)
+        x_rope = (
+            (x_rope * self.cos_cached[:, :, : x.shape[1], :])
+            + (neg_half_x * self.sin_cached[:, :, : x.shape[1], :])
+            if mask is None
+            else (
+                x_rope * self.cos_cached[:, :, mask, :]
+                + neg_half_x * self.sin_cached[:, :, mask, :]
+            )
+        )
+
+        return torch.cat((x_rope, x_pass), dim=-1)
+
+
+if __name__ == "__main__":
+    # img size 4, 64 64 48
+    # patch size 1, 8, 8, 8
+    # num of patches across time, depth, height, and width -> 4, 8, 8, 6
+    rot_embed = RotaryPositionalEmbeddings4D(
+        d=512, grid_depth=8, grid_height=8, grid_width=6, grid_time=4
+    )
+    query_ = torch.randn(
+        10, 1, 8 * 8 * 6 * 4, 512
+    )  # Batch, Heads, Num Tokens/Seq Length, Embedding Dims
+    print(rot_embed(query_).shape)

fMRI-foundation-model-old/fMRI-MAE/utils.py

@@ -0,0 +1,412 @@
+from io import BytesIO
+import os
+import random
+import numpy as np
+import torch
+from einops import rearrange
+from PIL import Image
+from skimage import filters
+from torchvision import transforms
+import nibabel as nib
+from nilearn import plotting
+import matplotlib.pyplot as plt
+import re
+import torch.nn.functional as F
+import torch.nn as nn
+
+def my_split_by_node(urls): return urls
+
+def is_interactive():
+    import __main__ as main
+
+    return not hasattr(main, "__file__")
+
+def my_split_by_node(urls): return urls
+
+def seed_everything(seed=0, cudnn_deterministic=True):
+    random.seed(seed)
+    os.environ['PYTHONHASHSEED'] = str(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    
+
+def grayscale_decoder(image_data):
+    return np.array(Image.open(BytesIO(image_data))).astype(np.float32) / 65535
+
+
+def numpy_decoder(npy_data):
+    return np.load(BytesIO(npy_data))
+
+
+def reshape_to_2d(tensor):
+    if tensor.ndim == 5:
+        tensor = tensor[0]
+    assert tensor.ndim == 4
+    return rearrange(tensor, "b h w c -> (b h) (c w)")
+
+
+def reshape_to_original(tensor_2d, h=64, w=64, c=48):
+    # print(tensor_2d.shape) # torch.Size([1, 256, 3072])
+    return rearrange(tensor_2d, "(tr h) (c w) -> tr h w c", h=h, w=w, c=c)
+
+
+def plot_numpy_nii(image):
+    while image.ndim > 3:
+        image = image[0]
+    nii = nib.Nifti1Image(image.astype(np.float32), np.eye(4))  # noqa
+    plotting.plot_epi(nii, cmap="gray")
+
+
+def threshold_based_masking(org_images):
+    thresholds = filters.threshold_multiotsu(org_images.numpy(), classes=3)
+    brain_segmentation = org_images > thresholds.min()
+    return brain_segmentation
+
+
+def get_brain_pos_patches(
+    func,
+    patch_depth=8,
+    patch_height=8,
+    patch_width=8,
+    frame_patch_size=1,
+    masking_strategy="conservative",
+):
+    _, _, depth = func.shape
+    if masking_strategy == "conservative":
+        func = func.sum(axis=(-1), keepdim=True).repeat(1, 1, depth)
+    else:
+        raise Exception("Not implemented other masking strategies than conservative.")
+
+    return func
+
+def crop_or_pad(tensor, new_shape):
+    # Ensure the tensor has at least three dimensions
+    if tensor.dim() < 3:
+        raise ValueError("Tensor must have at least 3 dimensions")
+
+    # Current dimensions of the last three axes
+    current_shape = tensor.shape[-3:]
+
+    # Compute padding and cropping needed for each dimension
+    padding_crop = [(ns - cs) for ns, cs in zip(new_shape, current_shape)]
+    if sum(padding_crop)==0:
+        return tensor
+
+    # Apply cropping if necessary
+    if any(pc < 0 for pc in padding_crop):
+        crop_slices = [slice(-pc//2, ns-pc//2) if pc < 0 else slice(None) for pc, ns in zip(padding_crop, new_shape)]
+        tensor = tensor[..., crop_slices[0], crop_slices[1], crop_slices[2]]
+
+    # Calculate padding to apply after cropping if necessary
+    pad_values = [(max(0, pc)//2, max(0, pc) - max(0, pc)//2) for pc in padding_crop]
+
+    # Apply padding
+    tensor = F.pad(tensor, pad_values[2] + pad_values[1] + pad_values[0])
+
+    return tensor
+
+
+class DataPrepper:
+    def __init__(
+        self,
+        num_frames=4,
+        masking_strategy="MNI",
+        patch_depth=8,
+        patch_height=8,
+        patch_width=8,
+        frame_patch_size=1,
+        image_size=[88, 104, 72]
+    ):
+        self.num_frames = num_frames
+        self.masking_strategy = masking_strategy
+        self.patch_depth = 8
+        self.patch_height = 8
+        self.patch_width = 8
+        self.frame_patch_size = 1
+        self.image_size=image_size
+
+    def __call__(self, func):
+        start_timepoint = np.random.choice(np.arange(func.shape[1] - self.num_frames))
+        timepoints = np.arange(start_timepoint, start_timepoint + self.num_frames)
+
+        func = func[:,timepoints]
+
+        # crop image_size acc to config
+        func = crop_or_pad(func, self.image_size)
+
+        if self.masking_strategy=="MNI" or self.masking_strategy=="None":
+            return func, None
+        
+        brain_segmentation = threshold_based_masking(func.mean(1))
+        pos_patches = None
+        for brain in brain_segmentation:
+            output = get_brain_pos_patches(
+                brain,
+                patch_depth=self.patch_depth,
+                patch_height=self.patch_height,
+                patch_width=self.patch_width,
+                frame_patch_size=self.frame_patch_size,
+                masking_strategy=self.masking_strategy,
+            )
+            if pos_patches is None:
+                pos_patches = output[None]
+            else:
+                pos_patches = torch.vstack((pos_patches, output[None]))
+        return func, pos_patches
+
+
+def plot_slices(unpatches):
+    if unpatches.ndim == 5:
+        unpatches = unpatches[0]
+    return transforms.ToPILImage()(reshape_to_2d(unpatches))
+
+
+def check_loss(loss):
+    if loss.isnan().any():
+        raise ValueError('NaN loss')
+        
+
+def count_params(model):
+    total = sum(p.numel() for p in model.parameters())
+    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print("param counts:\n{:,} total\n{:,} trainable".format(total, trainable))
+    return trainable
+
+
+def contrastive_loss(
+    cls_token1: torch.Tensor, cls_token2: torch.Tensor, temperature: torch.Tensor
+):
+    feat1 = cls_token1 / cls_token1.norm(dim=1, keepdim=True)
+    feat2 = cls_token2 / cls_token2.norm(dim=1, keepdim=True)
+
+    cosine_sim = feat1 @ feat2.T
+    logit_scale = temperature.exp()  # log scale, learned during training
+    feat1 = cosine_sim * logit_scale
+    feat2 = feat1.T
+
+    labels = torch.arange(feat1.shape[0]).to(feat1.device)
+    loss = (
+        torch.nn.functional.cross_entropy(feat1, labels)
+        + torch.nn.functional.cross_entropy(feat2, labels)
+    ) / 2
+    return loss
+
+### MindEye functions ###
+
+def soft_clip_loss(preds, targs, temp=0.006):
+    clip_clip = (targs @ targs.T)/temp
+    brain_clip = (preds @ targs.T)/temp
+    loss1 = -(brain_clip.log_softmax(-1) * clip_clip.softmax(-1)).sum(-1).mean()
+    loss2 = -(brain_clip.T.log_softmax(-1) * clip_clip.softmax(-1)).sum(-1).mean()
+    
+    loss = (loss1 + loss2)/2
+    return loss
+
+def cosine_anneal(start, end, steps):
+    return end + (start - end)/2 * (1 + torch.cos(torch.pi*torch.arange(steps)/(steps-1)))
+
+def mixco(voxels, beta=0.15, s_thresh=0.5, perm=None, betas=None, select=None):
+    if perm is None:
+        perm = torch.randperm(voxels.shape[0])
+    voxels_shuffle = voxels[perm].to(voxels.device,dtype=voxels.dtype)
+    if betas is None:
+        betas = torch.distributions.Beta(beta, beta).sample([voxels.shape[0]]).to(voxels.device,dtype=voxels.dtype)
+    if select is None:
+        select = (torch.rand(voxels.shape[0]) <= s_thresh).to(voxels.device)
+    betas_shape = [-1] + [1]*(len(voxels.shape)-1)
+    voxels[select] = voxels[select] * betas[select].reshape(*betas_shape) + \
+        voxels_shuffle[select] * (1 - betas[select]).reshape(*betas_shape)
+    betas[~select] = 1
+    return voxels, perm, betas, select
+
+def mixco_nce(preds, targs, temp=0.1, perm=None, betas=None, select=None, distributed=False, 
+              accelerator=None, local_rank=None, bidirectional=True):
+    brain_clip = (preds @ targs.T)/temp
+    
+    if perm is not None and betas is not None and select is not None:
+        probs = torch.diag(betas)
+        probs[torch.arange(preds.shape[0]).to(preds.device), perm] = 1 - betas
+
+        loss = -(brain_clip.log_softmax(-1) * probs).sum(-1).mean()
+        if bidirectional:
+            loss2 = -(brain_clip.T.log_softmax(-1) * probs.T).sum(-1).mean()
+            loss = (loss + loss2)/2
+        return loss
+    else:
+        loss =  F.cross_entropy(brain_clip, torch.arange(brain_clip.shape[0]).to(brain_clip.device))
+        if bidirectional:
+            loss2 = F.cross_entropy(brain_clip.T, torch.arange(brain_clip.shape[0]).to(brain_clip.device))
+            loss = (loss + loss2)/2
+        return loss
+    
+def topk(similarities,labels,k=5):
+    if k > similarities.shape[0]:
+        k = similarities.shape[0]
+    topsum=0
+    for i in range(k):
+        topsum += torch.sum(torch.argsort(similarities,axis=1)[:,-(i+1)] == labels)/len(labels)
+    return topsum
+
+def patchwise_cosine_similarity(latents1,latents2=None):
+    if latents2 is None:
+        latents_norm = latents1/latents1.norm(dim=-1, keepdim=True)
+        cos_sim = torch.bmm(latents_norm, latents_norm.permute(0,2,1))
+    else:
+        latents_norm1 = latents1/latents1.norm(dim=-1, keepdim=True)
+        latents_norm2 = latents2/latents2.norm(dim=-1, keepdim=True)
+        cos_sim = latents_norm1 @ latents_norm2.T
+    return cos_sim
+    
+def batchwise_cosine_similarity(Z,B):
+    Z = Z.flatten(1)
+    B = B.flatten(1).T
+    Z_norm = torch.linalg.norm(Z, dim=1, keepdim=True)  # Size (n, 1).
+    B_norm = torch.linalg.norm(B, dim=0, keepdim=True)  # Size (1, b).
+    cosine_similarity = ((Z @ B) / (Z_norm @ B_norm)).T
+    return cosine_similarity
+
+def prenormed_batchwise_cosine_similarity(Z,B):
+    return (Z @ B.T).T
+
+def torch_to_Image(x):
+    if x.ndim==4:
+        x=x[0]
+    return transforms.ToPILImage()(x)
+
+def get_masking_ratio(current_epoch, total_epochs, start_masking_ratio, end_masking_ratio):
+    """Returns the masking ratio for the current epochs. Linearly increase the masking ratio over the span of the training"""
+    return start_masking_ratio + (end_masking_ratio-start_masking_ratio) * ((current_epoch+1)/total_epochs)
+
+def view_brain(data,cut_coords=None):
+    if torch.is_tensor(data):
+        data = data.numpy()
+    if data.ndim==5:
+        new_nii = nib.Nifti1Image((data[0,0].astype(np.float32)-.5)*2, np.eye(4))
+    elif data.ndim==4:
+        new_nii = nib.Nifti1Image((data[0].astype(np.float32)-.5)*2, np.eye(4))
+    elif data.ndim==3:
+        new_nii = nib.Nifti1Image((data.astype(np.float32)-.5)*2, np.eye(4))
+    else:
+        raise Exception("Check dimensionality of your brain data")
+    return plotting.view_img(new_nii, bg_img=None, cut_coords=cut_coords, vmax=1, cmap=plt.cm.gray, threshold=None)
+
+def get_first_tar(train_urls):
+    if isinstance(train_urls, list):
+        # If train_urls is a list, get the first element
+        url = train_urls[0]
+    else:
+        # If train_urls is a string, treat it as the only element
+        url = train_urls
+
+    # Extract the first tar file using regular expression
+    match = re.search(r'\{(\d+)\.\.', url)
+    if match:
+        first_tar = match.group(1)
+        return f"/scratch/fmri_foundation_datasets/NSD_MNI_wds/{first_tar}.tar"
+    else:
+        return None
+
+
+class VICRegHandler(nn.Module):
+    def __init__(self, in_dim, num_layers=3, act=nn.GELU, h=1024, out_dim=4096):
+        super().__init__()
+        self.projector = nn.Sequential(
+            nn.Linear(in_dim, h),
+            nn.LayerNorm(h),
+            act(),
+            nn.Linear(h, h),
+            nn.LayerNorm(h),
+            act(),
+            nn.Linear(h, out_dim),
+        )
+
+    def forward(self, x):
+        return self.projector(x)
+
+    @staticmethod
+    def filter_global_to_local(l, enc_mask, dec_mask):
+        '''Get the subset of global tokens that correspond to encoder mask only'''
+        comb_mask = enc_mask | dec_mask
+        comb_indices = torch.where(comb_mask)[0]
+        enc_indices = torch.where(enc_mask)[0]
+        # enc_set = set(enc_indices.cpu().tolist()) 
+        
+        # new_mask = torch.zeros_like(comb_indices, dtype=bool)
+        # for i, idx in enumerate(comb_indices):
+        #     if idx in enc_set:
+        #         new_mask[i] = True
+        
+        new_mask = torch.isin(comb_indices, enc_indices)    
+        return l[:, new_mask]
+
+    @staticmethod
+    def vicreg_loss(l1, l2, gamma=1.0, lamda=25, mu=25, nu=1, rand_frac=0.2, use_vic_cls=True, eps=1e-4):
+        if use_vic_cls:
+            # always keep cls and pick a random set of tokens
+            rand_indices = torch.cat([torch.tensor([0]), 1+torch.randperm(l1.shape[1]-1)])[:int(rand_frac*l1.shape[1])]
+        else:
+            # drop cls tokens from loss calc
+            l1 = l1[:, 1:]
+            l2 = l2[:, 1:]
+            rand_indices = torch.randperm(l1.shape[1])[:int(rand_frac*l1.shape[1])]
+
+        std_l1 = torch.sqrt(l1.flatten(1).var(dim=0)+eps)  # nxd
+        std_l2 = torch.sqrt(l2.flatten(1).var(dim=0)+eps)  # nxd
+        var_loss = F.relu(gamma - std_l1).mean() + F.relu(gamma - std_l2).mean()
+        del std_l1, std_l2
+
+        sim_loss = F.mse_loss(l1, l2)
+
+        l1 = l1 - l1.mean(0, keepdim=True)  # b,n,d
+        l2 = l2 - l2.mean(0, keepdim=True)                
+        
+        l1_sub = l1[:, rand_indices]
+        del l1
+        cov_l1 = torch.bmm(l1_sub.permute(1,2,0), l1_sub.permute(1,0,2))/(l1_sub.shape[0]-1)  # 0.1*n,d,d
+        cov_loss = ((cov_l1**2).sum() - (torch.diagonal(cov_l1, dim1=1,dim2=2)**2).sum())/(l1_sub.shape[1]*l1_sub.shape[2])
+        del cov_l1, l1_sub
+
+        l2_sub = l2[:, rand_indices]
+        del l2
+        cov_l2 = torch.bmm(l2_sub.permute(1,2,0), l2_sub.permute(1,0,2))/(l2_sub.shape[0]-1)
+        cov_loss = cov_loss + ((cov_l2**2).sum() - (torch.diagonal(cov_l2, dim1=1,dim2=2)**2).sum())/(l2_sub.shape[1]*l2_sub.shape[2])  # div by nxd
+        del cov_l2, l2_sub
+
+        vic_loss = lamda * sim_loss + mu * var_loss + nu * cov_loss
+
+        return vic_loss
+
+
+class SimCLRHandler(nn.Module):
+    def __init__(self, in_dim, num_layers=2, act=nn.GELU, out_dim=1024):
+        super().__init__()
+        self.projector = nn.Sequential(
+            nn.Linear(in_dim, in_dim),
+            act(),
+            nn.Linear(in_dim, max(in_dim,out_dim)),
+        )
+
+    def forward(self, x):
+        return self.projector(x)
+
+    @staticmethod
+    def simclr_loss(lats, temp=0.006):
+        logits = (nn.functional.normalize(lats.flatten(1),dim=-1) @
+                    nn.functional.normalize(lats.flatten(1),dim=-1).T) / temp
+
+        labels = torch.diag_embed(
+            torch.ones(logits.shape[0] // 2), offset=logits.shape[0] // 2
+        ) + torch.diag_embed(torch.ones(logits.shape[0] // 2), offset=-logits.shape[0] // 2)
+        labels = labels.to(lats.device)
+        
+        mask = torch.ones_like(logits).bool()
+        torch.diagonal(mask).fill_(False)
+        
+        labels = labels[mask].reshape(logits.shape[0], logits.shape[0]-1)
+        logits = logits[mask].reshape(*labels.shape)
+
+        contr_loss = -(logits.log_softmax(-1) * labels).sum(-1).mean()
+
+        return contr_loss

fMRI-foundation-model-old/flat/.ipynb_checkpoints/main-checkpoint.py

@@ -0,0 +1,593 @@
+import os
+import shutil
+import sys
+import json
+import yaml
+import numpy as np
+import copy
+import math
+import time
+import random
+from tqdm import tqdm
+import webdataset as wds
+import matplotlib.pyplot as plt
+
+import torch
+import torch.nn as nn
+from torchvision import transforms
+import utils
+from flat_models import *
+
+# tf32 data type is faster than standard float32
+torch.backends.cuda.matmul.allow_tf32 = True
+# following fixes a Conv3D CUDNN_NOT_SUPPORTED error
+torch.backends.cudnn.benchmark = True
+
+### Multi-GPU config ###
+device_count = torch.cuda.device_count()
+print(f"Number of available CUDA devices: {device_count}")
+
+local_rank = os.getenv('LOCAL_RANK')
+if local_rank is None: 
+    local_rank = 0
+else:
+    local_rank = int(local_rank)
+print(f"LOCAL RANK={local_rank}")
+
+num_devices = os.getenv('NUM_GPUS')
+if num_devices is None: 
+    num_devices = 1
+else:
+    num_devices = int(num_devices)
+print(f"NUM GPUS={num_devices}")
+distributed = True if num_devices>1 else False
+if distributed: assert device_count==num_devices
+
+node = os.getenv('SLURM_NODEID')
+if node is None:
+    node = 0
+else:
+    node = int(node)
+print(f"NODE={node}")
+
+global_rank = os.getenv('RANK')
+if global_rank is None:
+    global_rank = 0
+else:
+    global_rank = int(global_rank)
+print(f"GLOBAL RANK={global_rank}")
+
+world_size = os.getenv('WORLD_SIZE')
+if world_size is None: 
+    world_size = 1
+else:
+    world_size = int(world_size)
+print(f"WORLD_SIZE={world_size}")
+
+# Load parameters from yaml config
+config = yaml.load(open('config.yaml', 'r'), Loader=yaml.FullLoader)
+
+print("\n__CONFIG__")
+for attribute_name in config.keys():
+    print(f"{attribute_name} = {config[attribute_name]}")
+    globals()[attribute_name] = config[f'{attribute_name}']
+print("\n")
+
+# Create outdir for ckpt and config.yaml
+outdir = os.path.abspath(f'checkpoints/{model_name}')
+print("outdir", outdir)
+
+# Load previous config.yaml if available
+if os.path.exists(f"{outdir}/config.yaml"):
+    config = yaml.load(open(f"{outdir}/config.yaml", 'r'), Loader=yaml.FullLoader)
+    print(f"Loaded config.yaml from ckpt folder {outdir}")
+
+    # create global variables from the config
+    print("\n__REPLACING_CONFIG__")
+    for attribute_name in config.keys():
+        print(f"{attribute_name} = {config[attribute_name]}")
+        globals()[attribute_name] = config[f'{attribute_name}']
+    print("\n")
+
+data_type = torch.float32 # change depending on your mixed_precision
+global_batch_size = batch_size * world_size
+
+# FSDP Setup
+if distributed:
+    import torch.distributed as dist
+    from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+    from torch.distributed.fsdp.api import BackwardPrefetch, CPUOffload, ShardingStrategy
+    import functools
+    from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy, transformer_auto_wrap_policy
+    print(f"setting device to cuda:{local_rank}")
+    torch.cuda.set_device(local_rank)
+    device = torch.device('cuda',local_rank)
+    dist.init_process_group("nccl", rank=global_rank, world_size=world_size)
+    print(f"\nSuccessfully set cuda:{local_rank} | global_rank{global_rank} | node{node}")
+    dist.barrier() 
+    print(f"global_rank{global_rank} passed barrier")
+else:
+    device = torch.device('cuda')
+
+print("PID of this process =",os.getpid())
+print("device =", device, "distributed =",distributed, "num_devices =", num_devices, "local rank =", local_rank, "world size =", world_size, "data_type =", data_type)
+
+# seed all random functions
+utils.seed_everything(seed + global_rank)
+
+from util.hcp_flat import load_hcp_flat_mask
+from util.hcp_flat import create_hcp_flat
+from util.losses import *
+import util.visualize as vis
+
+if utils.is_interactive(): # Use less samples per epoch for debugging
+    num_samples_per_epoch = 2000
+    test_num_samples_per_epoch = 2000
+
+model = mae_vit_small_fmri(
+    decoder_embed_dim=decoder_embed_dim,
+    t_patch_size=t_patch_size,
+    pred_t_dim=pred_t_dim,
+    decoder_depth=4,
+    cls_embed=cls_embed,
+    norm_pix_loss=norm_pix_loss,
+    no_qkv_bias=no_qkv_bias,
+    sep_pos_embed=sep_pos_embed,
+    trunc_init=trunc_init,
+)
+
+if use_contrastive_loss:
+    model.simclr_handler = SimCLRHandler(model.embed_dim).to(device)
+if use_vic_loss:
+    model.vicreg_handler = VICRegHandler(model.embed_dim).to(device)
+
+# state = torch.load("checkpoints/checkpoint-00099.pth", map_location="cpu")
+# model.load_state_dict(state["model"])
+
+num_batches = num_samples_per_epoch // (num_devices * batch_size)
+test_num_batches = test_num_samples_per_epoch // (num_devices * batch_size)
+print("num_batches", num_batches)
+print("test_num_batches", test_num_batches)
+
+## Train ##
+train_dataset = create_hcp_flat(root=hcp_flat_path, 
+                    training=True, frames=num_frames)
+train_dl = wds.WebLoader(
+    train_dataset.batched(batch_size, partial=False),
+    batch_size=None,
+    shuffle=False,
+    num_workers=num_workers,
+    pin_memory=True,
+)
+train_dl = train_dl.with_epoch(num_batches)
+
+## Test ##
+test_dataset = create_hcp_flat(root=hcp_flat_path,
+                    training=False, frames=num_frames)
+test_dl = wds.WebLoader(
+    test_dataset.batched(batch_size, partial=False),
+    batch_size=None,
+    shuffle=False,
+    num_workers=num_workers,
+    pin_memory=True,
+)
+test_dl = test_dl.with_epoch(test_num_batches)
+
+print(f"\nChecking distributed setup on global_rank {global_rank}...")
+from util.video_vit import Attention
+if distributed:
+    my_auto_wrap_policy = functools.partial(
+        transformer_auto_wrap_policy, 
+        transformer_layer_cls={
+            Attention, # <--- Your Transformer layer class
+        },
+    )
+    print(f"\nPrepping FSDP on {global_rank} {node}...\n")
+    model = model.to(device)
+    model = FSDP(
+        model,
+        sharding_strategy=ShardingStrategy.HYBRID_SHARD,
+        auto_wrap_policy=my_auto_wrap_policy,
+        use_orig_params=False,
+        cpu_offload=None, #CPUOffload(offload_params=True)
+        sync_module_states=True,
+        limit_all_gathers=True, # See https://github.com/pytorch/pytorch/issues/91165
+        device_id=device,
+    )
+    print(f"\nSuccessfully loaded FSDP model to device on global_rank {global_rank}\n")
+    dist.barrier()
+else:
+    print(f"\nNot training distributed! global_rank {global_rank}")
+    model.to(device)
+
+no_decay = ["bias", "LayerNorm.bias", "LayerNorm.weight"]
+opt_grouped_parameters = [
+    {'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': 0.05},
+    {'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0},
+]
+
+lr = base_lr * global_batch_size / 256
+print(f"multiply base lr {base_lr} by effective batch size {global_batch_size}")
+print(f"lr = {lr}")
+
+optimizer = torch.optim.AdamW(opt_grouped_parameters, lr=lr, betas=(0.9, 0.95))
+
+def adjust_learning_rate(optimizer, epoch, warmup_epochs=5, min_lr=0.0):
+    """Decay the learning rate with half-cycle cosine after warmup"""
+    if epoch < warmup_epochs:
+        lr_ = lr * epoch / warmup_epochs
+    else:
+        lr_ = min_lr + (lr - min_lr) * 0.5 * (
+            1.0
+            + math.cos(
+                math.pi
+                * (epoch - warmup_epochs)
+                / (num_epochs - warmup_epochs)
+            )
+        )
+    for param_group in optimizer.param_groups:
+        param_group["lr"] = lr_
+    return lr_
+
+print("\nDone with model preparations!")
+num_params = utils.count_params(model)
+
+def save_ckpt(model,tag="last"):
+    if distributed: dist.barrier()
+    model_states = model.state_dict()
+    if global_rank == 0:
+        os.makedirs(outdir,exist_ok=True)
+        ckpt_path = outdir+f'/{tag}.pth'
+        
+        if tag == "last" and os.path.exists(ckpt_path):
+            shutil.copyfile(os.path.join(outdir, f'{tag}.pth'), os.path.join(outdir, f'{tag}_old.pth'))
+        # print(f'saving {ckpt_path}',flush=True)
+        if tag=='last':
+            torch.save({
+                'epoch': epoch,
+                'model_state_dict': model_states,
+                'optimizer_state_dict': optimizer.state_dict(),
+                }, ckpt_path)
+        else:
+            torch.save({
+                'epoch': epoch,
+                'model_state_dict': model_states,
+                }, ckpt_path)
+            
+        if tag == "last" and os.path.exists(os.path.join(outdir, f'{tag}_old.pth')):
+            os.remove(os.path.join(outdir, f'{tag}_old.pth'))
+        print(f"\n---saved {ckpt_path}!---\n")
+        
+        # save the config.yaml
+        if not os.path.exists(f"{outdir}/config.yaml"):
+            with open(f"{outdir}/config.yaml", 'w') as file:
+                yaml.dump(config, file)
+            print(f"saved {outdir}/config.yaml!")
+
+# auto resume
+if os.path.exists(os.path.join(outdir, 'last.pth')) or os.path.exists(os.path.join(outdir, 'last_old.pth')):
+    if os.path.exists(os.path.join(outdir, 'last_old.pth')):
+        if os.path.exists(os.path.join(outdir, 'last.pth')):
+            # this is corrupted
+            os.remove(os.path.join(outdir, f'last.pth'))
+        # set last_old as last
+        shutil.move(os.path.join(outdir, f'last_old.pth'), os.path.join(outdir, f'last.pth'))
+    
+    ckpt_path = os.path.join(outdir, 'last.pth')
+    resume_from_ckpt = True
+
+epoch = 0
+if resume_from_ckpt:
+    print("\n---resuming from ckpt_path---\n", ckpt_path)
+    checkpoint = torch.load(ckpt_path, map_location=device)
+    epoch = checkpoint['epoch']+1
+    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])        
+    model.load_state_dict(checkpoint['model_state_dict'])
+    # total_steps_done = epoch*num_iterations_per_epoch
+    # for _ in range(total_steps_done):
+    #     lr_scheduler.step()
+    del checkpoint
+    torch.cuda.empty_cache()
+
+if utils.is_interactive():
+    wandb_log = False
+    ckpt_saving = False
+if local_rank==0 and wandb_log: # only use main process for wandb logging
+    import wandb
+    wandb_project = 'fMRI-foundation-model'
+    print(f"wandb {wandb_project} run {model_name}")
+    # need to configure wandb beforehand in terminal with "wandb init"!
+    wandb_config = {
+      "model_name": model_name,
+      "global_batch_size": global_batch_size,
+      "batch_size": batch_size,
+      "num_epochs": num_epochs,
+      "cls_embed": cls_embed,
+      "norm_pix_loss": norm_pix_loss,
+      "t_patch_size": t_patch_size,
+      "pred_t_dim": pred_t_dim,
+      "mask_ratio": mask_ratio,
+      "num_frames": num_frames,
+      "sep_pos_embed": sep_pos_embed,
+      "decoder_embed_dim": decoder_embed_dim,
+      "use_contrastive_loss": use_contrastive_loss,
+      "num_params": num_params,
+      "base_lr": base_lr,
+      "lr": lr,
+      "num_samples_per_epoch": num_samples_per_epoch,
+      "test_num_samples_per_epoch": test_num_samples_per_epoch,
+      "num_epochs": num_epochs,
+      "grad_clip": grad_clip,
+      "ckpt_interval": ckpt_interval,
+      "ckpt_saving": ckpt_saving,
+      "print_interval": print_interval,
+      "seed": seed,
+      "distributed": distributed,
+      "num_devices": num_devices,
+      "world_size": world_size,
+    }
+    print("wandb_config:\n",wandb_config)
+    print("wandb_id:",model_name)
+    wandb.init(
+        id=model_name,
+        project=wandb_project,
+        name=model_name,
+        config=wandb_config,
+        resume="allow",
+    )
+else:
+    wandb_log = False
+
+# lrs, train_losses, test_losses = [], [], []
+# train_losses1, train_losses2, train_losses3, train_losses4 = [], [], [], []
+
+epoch = 0
+lrs, train_losses, recon_losses, contrastive_losses, vic_losses = [], [], [], [], []
+cos_sim_encoder_output, cos_sim_decoder_output, cos_sim_encoder_output_patchwise = [], [], []
+probe_losses, probe_accs, test_losses, test_accs = [], [], [], []
+cos_sim_encoder_output_patchwise_test, cos_sim_encoder_output_test = [], []
+
+mse = nn.MSELoss()
+l1 = nn.L1Loss()
+crossentropy = nn.CrossEntropyLoss()
+
+if use_contrastive_loss:
+    contrastive_temps = utils.cosine_anneal(0.004, 0.0075, num_epochs)
+
+grad_scaler = torch.cuda.amp.GradScaler(enabled=True)
+progress_bar = tqdm(range(epoch, num_epochs), disable=local_rank!=0, desc="Overall")
+for epoch in progress_bar:
+    model.train()
+    for train_i, batch in enumerate(train_dl):
+        optimizer.zero_grad()
+        adjust_learning_rate(optimizer, train_i / num_batches + epoch)
+
+        input_func = batch[1]
+
+        input_func = input_func.to(dtype=data_type, device=device, non_blocking=True)
+        if len(input_func.shape) == 6:
+            b, r, c, t, h, w = input_func.shape
+            input_func = input_func.reshape(b * r, c, t, h, w)
+
+        ids_shuffle, ids_restore = get_ids_shuffle(input_func.shape[0], input_func.device, model)
+
+        # adjust number to keep relative to image mask
+        if model.img_mask is not None:
+            len_keep = int(model.patch_embed.t_grid_size * model.n_mask_patches * (1 - mask_ratio))
+        else:
+            len_keep = int(ids_restore.shape[1] * (1 - mask_ratio))
+
+        # keep the first subset
+        ids_keep = ids_shuffle[:, :len_keep]
+        lats = model.forward_encoder_with_mask(input_func, ids_keep)
+
+        if use_decoder:
+            mask = torch.ones_like(ids_restore)
+            mask[:, :len_keep] = 0
+            # unshuffle to get the binary mask
+            mask = torch.gather(mask, dim=1, index=ids_restore)
+            
+            pred = model.forward_decoder(lats, ids_restore)
+            recon_loss = model.forward_loss(input_func, pred, mask)
+            
+            recon_losses.append(recon_loss.item())
+            loss = recon_loss
+        else:
+            recon_loss = torch.nan
+            recon_losses.append(recon_loss)
+            loss = 0
+        
+        if use_contrastive_loss and not use_vic_loss:
+            ids_keep2 = ids_shuffle[:, :2*len_keep]
+            lats2 = model.forward_encoder_with_mask(input_func, ids_keep2)
+            
+            enc_cls_token = lats[:, :1]
+            enc_cls_token2 = lats2[:,:1]
+            
+            temp = contrastive_temps[epoch]
+
+            all_cls = torch.cat([enc_cls_token, enc_cls_token2], dim=0)
+            all_cls_proj = model.simclr_handler(all_cls)
+
+            contr_loss = SimCLRHandler.simclr_loss(all_cls_proj, temp)
+            
+            contrastive_losses.append(contr_loss.item())
+            loss += (contr_loss * contrastive_loss_weight)
+        # new loss
+        elif use_vic_loss:
+            ids_keep2 = ids_shuffle[:, :2*len_keep]
+            lats2 = model.forward_encoder_with_mask(input_func, ids_keep2)
+
+            enc_cls_token = lats[:, :1]
+            enc_cls_token2 = lats2[:,:1]
+
+            l1 = lats
+            l2 = lats2[:, :lats.shape[1]]
+
+            l1_proj = model.vicreg_handler(l1)
+            l2_proj = model.vicreg_handler(l2)
+
+            vic_loss = VICRegHandler.vicreg_loss(l1_proj, l2_proj, gamma=gamma, lamda=lamda, mu=mu, nu=nu, 
+                                                        rand_frac=rand_frac, use_vic_cls=use_vic_cls)
+
+            if use_contrastive_loss:
+                temp = contrastive_temps[epoch]
+                all_cls = torch.cat([enc_cls_token, enc_cls_token2], dim=0)
+                all_cls_proj = model.simclr_handler(all_cls)
+
+                contr_loss = SimCLRHandler.simclr_loss(all_cls_proj, temp)
+                contrastive_losses.append(contr_loss.item())
+            else:
+                contr_loss = 0
+                contrastive_losses.append(0)
+            
+            vic_losses.append(vic_loss.item())
+            loss += (contr_loss * contrastive_loss_weight + vic_loss * vic_loss_weight)
+        else:
+            vic_losses.append(0)
+            contrastive_losses.append(0)
+
+        grad_scaler.scale(loss).backward()
+        grad_scaler.unscale_(optimizer)
+        if grad_clip:
+            nn.utils.clip_grad_norm_(model.parameters(), max_norm=grad_clip)
+        grad_scaler.step(optimizer)
+        grad_scaler.update()
+        torch.cuda.synchronize()
+
+        cos_sim_encoder_output_patchwise.append(utils.patchwise_cosine_similarity(lats)[~torch.eye(lats.shape[1], dtype=bool)[None].expand(lats.shape[0],-1,-1)].mean().item())
+        cos_sim_encoder_output.append(utils.batchwise_cosine_similarity(lats.flatten(1)/1e3,lats.flatten(1)/1e3)[~torch.eye(len(lats),dtype=torch.bool)].mean().item())
+        if use_decoder:
+            cos_sim_decoder_output.append(utils.batchwise_cosine_similarity(pred, pred)[~torch.eye(len(pred),dtype=torch.bool)].mean().item())
+
+        lrs.append(optimizer.param_groups[0]["lr"])
+        train_losses.append(loss.item())
+        
+        if ((train_i%print_interval)==0 or (train_i==num_batches-1)) and train_i>0:
+            print(f"Ep. {epoch} | loss {np.mean(train_losses[-print_interval:]):.3f} | lr {optimizer.param_groups[0]['lr']} | {train_i}/{num_batches}")
+            if wandb_log: 
+                logs = {"train/loss": np.mean(train_losses[-print_interval:])}
+                # epoch_1000x as x-axis calibrates different curves when batch size changes
+                epoch_1000x = int((train_i / num_batches + epoch) * 1000)
+                wandb.log(logs, step=epoch_1000x)
+
+    if utils.is_interactive() or wandb_log:
+        print(f"Ep. {epoch} | loss {np.mean(train_losses[-print_interval:]):.3f} | lr {optimizer.param_groups[0]['lr']} | {train_i}/{num_batches}")
+        with torch.no_grad():
+            if norm_pix_loss:
+                normed_input_func, patch_info = model.patchify(input_func, alter_patch_info=False, return_patch_info=True)
+                target_mean = normed_input_func.mean(dim=-1, keepdim=True)
+                target_var = normed_input_func.var(dim=-1, keepdim=True)
+                normed_input_func = (normed_input_func - target_mean) / (target_var + 1.0e-6) ** 0.5
+                normed_input_func = model.unpatchify(normed_input_func, patch_info=patch_info)
+
+                vis_out = vis.plot_mask_pred(
+                    model, normed_input_func, pred, mask, 
+                    mean=0.5, std=0.2, 
+                )
+            else:
+                vis_out = vis.plot_mask_pred(
+                    model, input_func, pred, mask, 
+                    mean=0.5, std=0.2, 
+                )
+
+            if wandb_log:
+                logs = {"train/mask_pred": wandb.Image(vis_out)}
+                epoch_1000x = int((train_i / num_batches + epoch) * 1000)
+                wandb.log(logs, step=epoch_1000x)
+            else:
+                display(vis_out)
+        
+    logs = {
+        "train/loss": np.mean(train_losses[-(train_i + 1) :]),
+        "train/recon_losses": np.mean(recon_losses[-(train_i + 1) :]),
+        "train/contrastive_losses": np.mean(contrastive_losses[-(train_i + 1) :]),
+        "train/vic_losses": np.mean(vic_losses[-(train_i + 1) :]),
+        "train/num_steps": len(recon_losses),
+        "train/cos_sim_encoder_output": np.mean(cos_sim_encoder_output[-(train_i + 1) :]),
+        "train/cos_sim_decoder_output": np.mean(cos_sim_decoder_output[-(train_i + 1) :]) if use_decoder else np.nan,
+        "train/cos_sim_encoder_output_patchwise": np.mean(cos_sim_encoder_output_patchwise[-(train_i + 1) :]),
+        "lr": np.mean(lrs[-(train_i + 1) :]),
+        "epoch": epoch,
+    }
+        
+    # Evaluate performance on held-out test dataset
+    model.eval()
+    with torch.no_grad():
+        for test_i, batch in enumerate(test_dl):
+            input_func = batch[1]
+
+            input_func = input_func.to(dtype=data_type, device=device, non_blocking=True)
+            if len(input_func.shape) == 6:
+                b, r, c, t, h, w = input_func.shape
+                input_func = input_func.reshape(b * r, c, t, h, w)
+
+            if not use_contrastive_loss:
+                loss, pred, mask, latent = model(input_func, mask_ratio=mask_ratio, use_contrastive_loss=use_contrastive_loss)
+            else:
+                loss1, loss2, loss3, pred1, pred2, mask1, mask2, true_mask, latent1, latent2 = model(input_func, mask_ratio=mask_ratio, use_contrastive_loss=use_contrastive_loss)
+                pred, mask, latent = pred1, mask1, latent1
+                
+                # contrastive loss
+                temp = contrastive_temps[epoch]
+                logits = (nn.functional.normalize(latent1[:,:1].flatten(1),dim=-1) @
+                            nn.functional.normalize(latent2[:,:1].flatten(1),dim=-1).T) / temp
+                labels = torch.arange(len(logits)).long().to(device)
+                contr_loss = (crossentropy(logits, labels) + crossentropy(logits.T, labels)) / 2
+                
+                loss = loss1 + loss2 + loss3 + contr_loss
+
+            test_losses.append(loss.item())
+
+            cos_sim_encoder_output_patchwise_test.append(utils.patchwise_cosine_similarity(latent)[~torch.eye(latent.shape[1], dtype=bool)[None].expand(latent.shape[0],-1,-1)].mean().item())
+            cos_sim_encoder_output_test.append(utils.batchwise_cosine_similarity(latent.flatten(1)/1e3, latent.flatten(1)/1e3)[~torch.eye(len(latent),dtype=torch.bool)].mean().item())
+
+            if test_i%print_interval==0 and test_i>0:
+                print(f"Test | loss {np.mean(test_losses[-print_interval:]):.3f} | {test_i}/{test_num_batches}")
+
+    print(f"Test | iter {test_i} | loss {np.mean(test_losses[-test_i:]):.3f}")
+    if wandb_log: 
+        logs.update({
+            "test/loss": np.mean(test_losses[-test_i:]),
+            "test/cos_sim_encoder_output": np.mean(cos_sim_encoder_output_test[-(test_i + 1) :]),
+            "test/cos_sim_encoder_output_patchwise": np.mean(cos_sim_encoder_output_patchwise_test[-(test_i + 1) :]),
+        })
+        wandb.log(logs)
+
+    # Plot progress (first sample in batch)
+    if utils.is_interactive() or wandb_log:
+        with torch.no_grad():
+            if norm_pix_loss:
+                normed_input_func, patch_info = model.patchify(input_func, alter_patch_info=False, return_patch_info=True)
+                target_mean = normed_input_func.mean(dim=-1, keepdim=True)
+                target_var = normed_input_func.var(dim=-1, keepdim=True)
+                normed_input_func = (normed_input_func - target_mean) / (target_var + 1.0e-6) ** 0.5
+                normed_input_func = model.unpatchify(normed_input_func, patch_info=patch_info)
+
+                vis_out = vis.plot_mask_pred(
+                    model, normed_input_func, pred, mask, 
+                    mean=0.5, std=0.2, 
+                )
+            else:
+                vis_out = vis.plot_mask_pred(
+                    model, input_func, pred, mask, 
+                    mean=0.5, std=0.2, 
+                )
+
+            if wandb_log:
+                logs = {"test/mask_pred": wandb.Image(vis_out)}
+                wandb.log(logs)
+            else:
+                display(vis_out)
+
+    # Save model checkpoint
+    if ckpt_saving and epoch>0 and ((epoch % ckpt_interval == 0) or (epoch==num_epochs-1)):
+        save_ckpt(model,f"epoch{epoch}")
+
+    # wait for other GPUs to catch up if needed
+    if distributed: dist.barrier()
+    
+    # close any open plots
+    plt.close()
+    
+if distributed: dist.destroy_process_group()

fMRI-foundation-model-old/flat/.ipynb_checkpoints/main2-checkpoint.ipynb

@@ -0,0 +1,6 @@
+{
+ "cells": [],
+ "metadata": {},
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

fMRI-foundation-model-old/flat/accel.slurm

@@ -0,0 +1,44 @@
+#!/bin/bash
+#SBATCH --account=fmri
+#SBATCH --partition=p5
+#SBATCH --job-name=found
+#SBATCH --ntasks-per-node=1
+#SBATCH --nodes=1
+#SBATCH --gres=gpu:1
+#SBATCH --gpus-per-task=1       # Set to equal gres=gpu:#!
+#SBATCH --cpus-per-task=176     # 40 / 80 / 176 distributed across node
+#SBATCH --time=23:00:00         # total run time limit (HH:MM:SS)
+#SBATCH -e slurms/%j.err        # first create a "slurms" folder in current directory to store logs
+#SBATCH -o slurms/%j.out
+#SBATCH --comment=medarc
+#SBATCH --no-requeue
+
+export NUM_GPUS=1 # Set to equal gres=gpu:#!
+echo NUM_GPUS=$NUM_GPUS
+
+source ~/.bashrc
+
+cd /weka/proj-fmri/paulscotti/fMRI-foundation-model/flat/
+jupyter nbconvert main.ipynb --to python
+if [ $? -ne 0 ]; then
+  echo "Error: Conversion of ipynb to Python failed. Exiting."
+  exit 1
+fi
+
+# Make sure another job doesnt use same port, here using random number
+export MASTER_PORT=$((RANDOM % (19000 - 11000 + 1) + 11000))
+export HOSTNAMES=$(scontrol show hostnames "$SLURM_JOB_NODELIST")
+export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
+echo MASTER_ADDR=${MASTER_ADDR}
+echo MASTER_PORT=${MASTER_PORT}
+echo WORLD_SIZE=$((${SLURM_NNODES} * ${NUM_GPUS}))
+
+srun torchrun \
+--nproc_per_node=$NUM_GPUS \
+--nnodes=$SLURM_NNODES \
+main.py
+
+if [ $? -ne 0 ]; then
+  echo "Error: srun command failed. Please check the logs for more details."
+  exit 1
+fi

fMRI-foundation-model-old/flat/config.yaml

@@ -0,0 +1,58 @@
+model_name: "flat_NORM_contrastive_"
+hcp_flat_path: "/weka/proj-medarc/shared/hcp_flat"
+
+use_contrastive_loss: True
+use_decoder: True
+patch_size: 16
+
+decoder_embed_dim: 512
+num_frames: 16
+mask_ratio: .75
+pred_t_dim: 8
+t_patch_size: 2
+cls_embed: True
+no_qkv_bias: False
+sep_pos_embed: True
+trunc_init: False
+norm_pix_loss: True
+contrastive_loss_weight: 1.0  # 2.0
+
+# VICReg params
+use_vic_loss: True
+vic_loss_weight: 1.0  # 0.05  # 1.0
+gamma: 0.5  # 1.0  # 0.5
+# mu->v, lamda->i, nu->c
+mu: 25
+lamda: 25
+nu: 1
+# fraction of tokens to use for cov loss (all tokens don't fit in mem)
+rand_frac: 0.1  # 0.2
+# whether to apply vic to cls tok (if cont is also on, 2 losses will apply to cls token)
+use_vic_cls: True
+norm_pix_loss: False
+pct_masks_to_decode: 1
+
+# Training Configs
+batch_size: 32 
+num_workers: 10
+num_epochs: 100
+seed: 42
+base_lr: 3.0e-4 # Keep the x.0 else will be converted to string
+num_samples_per_epoch: 200000
+test_num_samples_per_epoch: 50000
+grad_clip: 1.0  # set 0 for no clip
+grad_accumulation_steps: 1
+test_set: False
+
+# Downstream probe config
+probe_num_samples_per_epoch: 100000
+probe_num_epochs: 30
+probe_batch_size: 8
+probe_base_lr: 3.0e-4
+
+# Saving progress
+ckpt_saving: True
+ckpt_interval: 5 # in epochs
+print_interval: 20 # in steps
+resume_from_ckpt: False
+wandb_log: False

fMRI-foundation-model-old/flat/flat_models.py

@@ -0,0 +1,784 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# References:
+# timm: https://github.com/rwightman/pytorch-image-models/tree/master/timm
+# DeiT: https://github.com/facebookresearch/deit
+# MAE: https://github.com/facebookresearch/mae
+# MAE-ST: https://github.com/facebookresearch/mae_st
+# --------------------------------------------------------
+
+from functools import partial
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+from util import video_vit
+from util.logging import master_print as print
+from util.hcp_flat import load_hcp_flat_mask
+import copy
+
+
+class MaskedAutoencoderViT(nn.Module):
+    """Masked Autoencoder with VisionTransformer backbone"""
+
+    def __init__(
+        self,
+        img_size=224,
+        patch_size=16,
+        in_chans=3,
+        embed_dim=1024,
+        depth=24,
+        num_heads=16,
+        decoder_embed_dim=512,
+        decoder_depth=8,
+        decoder_num_heads=16,
+        mlp_ratio=4.0,
+        norm_layer=nn.LayerNorm,
+        norm_pix_loss=False,
+        num_frames=16,
+        t_patch_size=2,
+        patch_embed=video_vit.PatchEmbed,
+        no_qkv_bias=False,
+        sep_pos_embed=True,
+        trunc_init=False,
+        cls_embed=True,
+        pred_t_dim=8,
+        img_mask=None,
+        pct_masks_to_decode=1,
+        **kwargs,
+    ):
+        super().__init__()
+        self.trunc_init = trunc_init
+        self.sep_pos_embed = sep_pos_embed
+        self.cls_embed = cls_embed
+        self.pred_t_dim = pred_t_dim
+        self.t_pred_patch_size = t_patch_size * pred_t_dim // num_frames
+        self.embed_dim = embed_dim
+
+        self.pct_masks_to_decode = pct_masks_to_decode
+
+        self.patch_embed = patch_embed(
+            img_size,
+            patch_size,
+            in_chans,
+            embed_dim,
+            num_frames,
+            t_patch_size,
+        )
+        num_patches = self.patch_embed.num_patches
+        input_size = self.patch_embed.input_size
+        self.input_size = input_size
+
+        if self.cls_embed:
+            self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+            self.decoder_cls_token = nn.Parameter(torch.zeros(1, 1, decoder_embed_dim))
+
+        if sep_pos_embed:
+            self.pos_embed_spatial = nn.Parameter(
+                torch.zeros(1, input_size[1] * input_size[2], embed_dim)
+            )
+            self.pos_embed_temporal = nn.Parameter(
+                torch.zeros(1, input_size[0], embed_dim)
+            )
+            if self.cls_embed:
+                self.pos_embed_class = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        else:
+            if self.cls_embed:
+                _num_patches = num_patches + 1
+            else:
+                _num_patches = num_patches
+
+            self.pos_embed = nn.Parameter(
+                torch.zeros(1, _num_patches, embed_dim),
+            )
+
+        self.blocks = nn.ModuleList(
+            [
+                video_vit.Block(
+                    embed_dim,
+                    num_heads,
+                    mlp_ratio,
+                    qkv_bias=not no_qkv_bias,
+                    qk_scale=None,
+                    norm_layer=norm_layer,
+                )
+                for i in range(depth)
+            ]
+        )
+        self.norm = norm_layer(embed_dim)
+
+        self.decoder_embed = nn.Linear(embed_dim, decoder_embed_dim, bias=True)
+
+        self.mask_token = nn.Parameter(torch.zeros(1, 1, decoder_embed_dim))
+
+        if sep_pos_embed:
+            self.decoder_pos_embed_spatial = nn.Parameter(
+                torch.zeros(1, input_size[1] * input_size[2], decoder_embed_dim)
+            )
+            self.decoder_pos_embed_temporal = nn.Parameter(
+                torch.zeros(1, input_size[0], decoder_embed_dim)
+            )
+            if self.cls_embed:
+                self.decoder_pos_embed_class = nn.Parameter(
+                    torch.zeros(1, 1, decoder_embed_dim)
+                )
+        else:
+            if self.cls_embed:
+                _num_patches = num_patches + 1
+            else:
+                _num_patches = num_patches
+
+            self.decoder_pos_embed = nn.Parameter(
+                torch.zeros(1, _num_patches, decoder_embed_dim),
+            )
+
+        self.decoder_blocks = nn.ModuleList(
+            [
+                video_vit.Block(
+                    decoder_embed_dim,
+                    decoder_num_heads,
+                    mlp_ratio,
+                    qkv_bias=not no_qkv_bias,
+                    qk_scale=None,
+                    norm_layer=norm_layer,
+                )
+                for i in range(decoder_depth)
+            ]
+        )
+
+        self.decoder_norm = norm_layer(decoder_embed_dim)
+        self.decoder_pred = nn.Linear(
+            decoder_embed_dim,
+            self.t_pred_patch_size * patch_size**2 * in_chans,
+            bias=True,
+        )
+
+        self.norm_pix_loss = norm_pix_loss
+
+        self.initialize_mask(img_mask)
+        self.initialize_weights()
+
+        print("model initialized")
+
+    def initialize_weights(self):
+        if self.cls_embed:
+            torch.nn.init.trunc_normal_(self.cls_token, std=0.02)
+        if self.sep_pos_embed:
+            torch.nn.init.trunc_normal_(self.pos_embed_spatial, std=0.02)
+            torch.nn.init.trunc_normal_(self.pos_embed_temporal, std=0.02)
+
+            torch.nn.init.trunc_normal_(self.decoder_pos_embed_spatial, std=0.02)
+            torch.nn.init.trunc_normal_(self.decoder_pos_embed_temporal, std=0.02)
+
+            if self.cls_embed:
+                torch.nn.init.trunc_normal_(self.pos_embed_class, std=0.02)
+                torch.nn.init.trunc_normal_(self.decoder_pos_embed_class, std=0.02)
+        else:
+            torch.nn.init.trunc_normal_(self.pos_embed, std=0.02)
+            torch.nn.init.trunc_normal_(self.decoder_pos_embed, std=0.02)
+        w = self.patch_embed.proj.weight.data
+        if self.trunc_init:
+            torch.nn.init.trunc_normal_(w)
+            torch.nn.init.trunc_normal_(self.mask_token, std=0.02)
+        else:
+            torch.nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
+            torch.nn.init.normal_(self.mask_token, std=0.02)
+
+        # initialize nn.Linear and nn.LayerNorm
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            # we use xavier_uniform following official JAX ViT:
+            if self.trunc_init:
+                nn.init.trunc_normal_(m.weight, std=0.02)
+            else:
+                torch.nn.init.xavier_uniform_(m.weight)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def initialize_mask(self, img_mask):
+        if img_mask is not None:
+            img_mask = torch.as_tensor(img_mask > 0).float()
+
+            H, W = img_mask.shape
+            img_mask_patches = self.patchify(
+                img_mask
+                .view(1, 1, 1, H, W)
+                .repeat(1, self.patch_embed.in_chans, self.pred_t_dim, 1, 1)
+            )
+
+            patch_mask = rearrange(
+                img_mask,
+                "(h ph) (w pw) -> (h w) (ph pw)",
+                ph=self.patch_embed.patch_size[0],
+                pw=self.patch_embed.patch_size[1],
+            ).any(dim=1).float()
+            patch_mask_indices, = patch_mask.nonzero(as_tuple=True)
+
+            self.register_buffer("img_mask", img_mask)
+            self.register_buffer("img_mask_patches", img_mask_patches)
+            self.register_buffer("patch_mask", patch_mask)
+            self.register_buffer("patch_mask_indices", patch_mask_indices)
+            self.n_mask_patches = int(len(patch_mask_indices) * self.pct_masks_to_decode)
+        else:
+            self.register_buffer("img_mask", None)
+            self.register_buffer("img_mask_patches", None)
+            self.register_buffer("patch_mask", None)
+            self.register_buffer("patch_mask_indices", None)
+            self.n_mask_patches = None
+
+    def patchify(self, imgs):
+        """
+        imgs: (N, C, T, H, W)
+        x: (N, L, patch_size**2 *C)
+        """
+        N, C, T, H, W = imgs.shape
+        ph, pw = self.patch_embed.patch_size
+        u = self.t_pred_patch_size
+        assert H % ph == 0 and W % pw == 0 and T % u == 0
+        h = H // ph
+        w = W // pw
+        t = T // u
+
+        x = imgs.reshape(shape=(N, C, t, u, h, ph, w, pw))
+        x = torch.einsum("nctuhpwq->nthwupqc", x)
+        x = x.reshape(shape=(N, t * h * w, u * ph * pw * C))
+        self.patch_info = (N, C, T, H, W, ph, pw, u, t, h, w)
+        return x
+
+    def unpatchify(self, x):
+        """
+        x: (N, L, patch_size**2 *C)
+        imgs: (N, C, H, W)
+        """
+        N, C, T, H, W, ph, pw, u, t, h, w = self.patch_info
+
+        x = x.reshape(shape=(N, t, h, w, u, ph, pw, C))
+
+        x = torch.einsum("nthwupqc->nctuhpwq", x)
+        imgs = x.reshape(shape=(N, C, T, H, W))
+        return imgs
+
+    def random_masking(self, x, mask_ratio, use_contrastive_loss=False):
+        """
+        Perform per-sample random masking by per-sample shuffling.
+        Per-sample shuffling is done by argsort random noise.
+        x: [N, L, D], sequence
+        """
+        N, L, D = x.shape  # batch, length, dim
+        T = self.patch_embed.t_grid_size
+        H, W = self.patch_embed.grid_size
+        assert L == T * H * W
+
+        # adjust number to keep relative to image mask
+        if self.img_mask is not None:
+            len_keep = int(T * self.n_mask_patches * (1 - mask_ratio))
+        else:
+            len_keep = int(L * (1 - mask_ratio))
+
+        noise = torch.rand(N, L, device=x.device)  # noise in [0, 1]
+
+        # shift missing patches to not be selected
+        if self.img_mask is not None:
+            noise = noise.view(N, T, H * W)
+            noise = noise + (1.0 - self.patch_mask)
+            noise = noise.view(N, L)
+
+        # sort noise for each sample
+        ids_shuffle = torch.argsort(
+            noise, dim=1
+        )  # ascend: small is keep, large is remove
+        ids_restore = torch.argsort(ids_shuffle, dim=1)
+
+        # keep the first subset
+        ids_keep = ids_shuffle[:, :len_keep]
+        if not use_contrastive_loss:
+            x_masked = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D))
+        else:
+            x_masked1 = torch.gather(x, dim=1, index=ids_keep[:,:len_keep//2].unsqueeze(-1).repeat(1, 1, D))
+            x_masked2 = torch.gather(x, dim=1, index=ids_keep[:,len_keep//2:len_keep].unsqueeze(-1).repeat(1, 1, D))
+
+        if not use_contrastive_loss:
+            # generate the binary mask: 0 is keep, 1 is remove
+            mask = torch.ones([N, L], device=x.device)
+            mask[:, :len_keep] = 0
+            # unshuffle to get the binary mask
+            mask = torch.gather(mask, dim=1, index=ids_restore)
+        else:
+            # generate the binary mask: 0 is keep, 1 is remove
+            mask1 = torch.ones([N, L], device=x.device)
+            mask2 = torch.ones([N, L], device=x.device)
+            mask1[:, :len_keep//2] = 0
+            mask2[:, len_keep//2:len_keep] = 0
+            # unshuffle to get the binary mask
+            mask1 = torch.gather(mask1, dim=1, index=ids_restore)
+            mask2 = torch.gather(mask2, dim=1, index=ids_restore)
+
+        if not use_contrastive_loss:
+            return x_masked, mask, ids_restore, ids_keep
+        else:
+            return [x_masked1,x_masked2], [mask1,mask2], ids_restore, ids_keep
+        
+    def forward_encoder(self, x, mask_ratio, use_contrastive_loss=False):
+        x = self.patch_embed(x)
+        
+        N, T, L, C = x.shape
+
+        x = x.reshape(N, T * L, C)
+
+        # masking: length -> length * mask_ratio
+        if not use_contrastive_loss:
+            x, mask, ids_restore, ids_keep = self.random_masking(x, mask_ratio)
+            x = x.view(N, -1, C)
+        else:
+            [x1,x2], [mask1,mask2], ids_restore, ids_keep = self.random_masking(x, mask_ratio, use_contrastive_loss=use_contrastive_loss)
+            x1 = x1.view(len(x1), -1, C)
+            x2 = x2.view(len(x2), -1, C)
+        # append cls token
+        if self.cls_embed:
+            cls_token = self.cls_token
+            cls_tokens = cls_token.expand(x.shape[0], -1, -1)
+            if not use_contrastive_loss:
+                x = torch.cat((cls_tokens, x), dim=1)
+            else:
+                x1 = torch.cat((cls_tokens, x1), dim=1)
+                x2 = torch.cat((cls_tokens, x2), dim=1)
+
+        # add pos embed w/o cls token
+        if self.sep_pos_embed:
+            pos_embed = self.pos_embed_spatial.repeat(
+                1, self.input_size[0], 1
+            ) + torch.repeat_interleave(
+                self.pos_embed_temporal,
+                self.input_size[1] * self.input_size[2],
+                dim=1,
+            )
+            pos_embed = pos_embed.expand(x.shape[0], -1, -1)
+            pos_embed = torch.gather(
+                pos_embed,
+                dim=1,
+                index=ids_keep.unsqueeze(-1).repeat(1, 1, pos_embed.shape[2]),
+            )
+            if self.cls_embed:
+                pos_embed = torch.cat(
+                    [
+                        self.pos_embed_class.expand(pos_embed.shape[0], -1, -1),
+                        pos_embed,
+                    ],
+                    1,
+                )
+        else:
+            if self.cls_embed:
+                cls_ind = 1
+            else:
+                cls_ind = 0
+            pos_embed = self.pos_embed[:, cls_ind:, :].expand(x.shape[0], -1, -1)
+            pos_embed = torch.gather(
+                pos_embed,
+                dim=1,
+                index=ids_keep.unsqueeze(-1).repeat(1, 1, pos_embed.shape[2]),
+            )
+            if self.cls_embed:
+                pos_embed = torch.cat(
+                    [
+                        self.pos_embed[:, :1, :].expand(x.shape[0], -1, -1),
+                        pos_embed,
+                    ],
+                    1,
+                )
+        if not use_contrastive_loss:
+            x = x.view([N, -1, C]) + pos_embed
+        else:
+            x1 = x1.view([len(x1), -1, C]) + pos_embed[:,:x1.shape[1]]
+            x2 = x2.view([len(x2), -1, C]) + torch.cat((pos_embed[:,:1], pos_embed[:,x1.shape[1]:]),dim=1)
+
+        if not use_contrastive_loss:
+            # apply Transformer blocks
+            for blk in self.blocks:
+                x = blk(x)
+            x = self.norm(x)
+        else:
+            # apply Transformer blocks
+            for blk in self.blocks:
+                x1 = blk(x1)
+                x2 = blk(x2)
+            x1 = self.norm(x1)
+            x2 = self.norm(x2)
+
+        if not use_contrastive_loss:
+            if self.cls_embed:
+                # remove cls token
+                x = x[:, 1:, :]
+
+            return x, mask, ids_restore
+        else:
+            if self.cls_embed:
+                # remove cls token
+                x1 = x1[:, 1:, :]
+                x2 = x2[:, 1:, :]
+
+            return [x1,x2], [mask1,mask2], ids_restore
+
+    def forward_encoder_with_mask(self, x, ids_keep):
+        # embed patches
+        x = self.patch_embed(x)
+        N, T, L, C = x.shape
+
+        x = x.reshape(N, T * L, C)
+        # mask out tokens
+        x = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, C))
+        
+        # append cls token
+        if self.cls_embed:
+            cls_token = self.cls_token
+            cls_tokens = cls_token.expand(x.shape[0], -1, -1)
+            x = torch.cat((cls_tokens, x), dim=1)
+        
+        # add pos embed w/o cls token
+        if self.sep_pos_embed:
+            pos_embed = self.pos_embed_spatial.repeat(
+                1, self.input_size[0], 1
+            ) + torch.repeat_interleave(
+                self.pos_embed_temporal,
+                self.input_size[1] * self.input_size[2],
+                dim=1,
+            )
+            pos_embed = pos_embed.expand(x.shape[0], -1, -1)
+            pos_embed = torch.gather(
+                pos_embed,
+                dim=1,
+                index=ids_keep.unsqueeze(-1).repeat(1, 1, pos_embed.shape[2]),
+            )
+            if self.cls_embed:
+                pos_embed = torch.cat(
+                    [
+                        self.pos_embed_class.expand(pos_embed.shape[0], -1, -1),
+                        pos_embed,
+                    ],
+                    1,
+                )
+        else:
+            if self.cls_embed:
+                cls_ind = 1
+            else:
+                cls_ind = 0
+            pos_embed = self.pos_embed[:, cls_ind:, :].expand(x.shape[0], -1, -1)
+            pos_embed = torch.gather(
+                pos_embed,
+                dim=1,
+                index=ids_keep.unsqueeze(-1).repeat(1, 1, pos_embed.shape[2]),
+            )
+            if self.cls_embed:
+                pos_embed = torch.cat(
+                    [
+                        self.pos_embed[:, :1, :].expand(x.shape[0], -1, -1),
+                        pos_embed,
+                    ],
+                    1,
+                )
+
+        x = x.view([N, -1, C]) + pos_embed
+
+        for blk in self.blocks:
+            x = blk(x)
+        x = self.norm(x)
+        return x
+
+    def forward_decoder(self, x, ids_restore, use_contrastive_loss=False):
+        N = x.shape[0]
+        T = self.patch_embed.t_grid_size
+        H, W = self.patch_embed.grid_size
+
+        # embed tokens
+        x = self.decoder_embed(x)
+        C = x.shape[-1]
+        
+        # append mask tokens to sequence
+        mask_tokens = self.mask_token.repeat(N, T * H * W + 0 - x.shape[1], 1)
+        x_ = torch.cat([x[:, :, :], mask_tokens], dim=1)  # no cls token
+        x_ = x_.view([N, T * H * W, C])
+        x_ = torch.gather(
+            x_, dim=1, index=ids_restore.unsqueeze(-1).repeat(1, 1, x_.shape[2])
+        )  # unshuffle
+        x = x_.view([N, T * H * W, C])
+        # append cls token
+        if self.cls_embed:
+            decoder_cls_token = self.decoder_cls_token
+            decoder_cls_tokens = decoder_cls_token.expand(x.shape[0], -1, -1)
+            x = torch.cat((decoder_cls_tokens, x), dim=1)
+
+        if self.sep_pos_embed:
+            decoder_pos_embed = self.decoder_pos_embed_spatial.repeat(
+                1, self.input_size[0], 1
+            ) + torch.repeat_interleave(
+                self.decoder_pos_embed_temporal,
+                self.input_size[1] * self.input_size[2],
+                dim=1,
+            )
+            if self.cls_embed:
+                decoder_pos_embed = torch.cat(
+                    [
+                        self.decoder_pos_embed_class.expand(
+                            decoder_pos_embed.shape[0], -1, -1
+                        ),
+                        decoder_pos_embed,
+                    ],
+                    1,
+                )
+        else:
+            decoder_pos_embed = self.decoder_pos_embed[:, :, :]
+
+        # add pos embed
+        x = x + decoder_pos_embed
+
+        attn = self.decoder_blocks[0].attn
+
+        # drop patches outside image mask (and then only keep a subset a la VideoMAE2)
+        if self.img_mask is not None:
+            if self.cls_embed:
+                decoder_cls_tokens, x = x[:, :1, :], x[:, 1:, :]
+            x = x.view([N, T, H * W, C])
+            # x = x[:, :, self.patch_mask_indices]
+            
+            # drop patches randomly to preserve memory (VideoMAE2 approach)
+            included_patches = self.patch_mask_indices
+            num_to_select = int(self.pct_masks_to_decode * len(included_patches))
+            selected_idx = torch.randperm(len(included_patches))[:num_to_select]
+            included_patches = included_patches[selected_idx]
+            x = x[:, :, included_patches]
+            
+            x = x.view([N, T * self.n_mask_patches, C])
+            if self.cls_embed:
+                x = torch.cat((decoder_cls_tokens, x), dim=1)
+
+        # apply Transformer blocks
+        for blk in self.decoder_blocks:
+            x = blk(x)
+        x = self.decoder_norm(x)
+
+        # predictor projection
+        x = self.decoder_pred(x)
+
+        if self.cls_embed:
+            # remove cls token
+            x = x[:, 1:, :]
+
+        # fill outside mask with zeros
+        if self.img_mask is not None:
+            C = x.shape[-1]
+            x = x.view([N, T, self.n_mask_patches, C])
+            x_ = torch.zeros([N, T, H * W, C], dtype=x.dtype, device=x.device)
+            x = x_.scatter(
+                2, included_patches.view(1, 1, -1, 1).expand(N, T, self.n_mask_patches, C), x,
+            )
+            x = x.view([N, T * H * W, C])
+
+        return x
+
+    def forward_loss(self, imgs, pred, mask):
+        """
+        imgs: [N, C, T, H, W]
+        pred: [N, t*h*w, u*p*p*C]
+        mask: [N, t*h*w], 0 is keep, 1 is remove,
+        """
+        _imgs = torch.index_select(
+            imgs,
+            2,
+            torch.linspace(
+                0,
+                imgs.shape[2] - 1,
+                self.pred_t_dim,
+            )
+            .long()
+            .to(imgs.device),
+        )
+        target = self.patchify(_imgs)
+        if self.norm_pix_loss:
+            mean = target.mean(dim=-1, keepdim=True)
+            var = target.var(dim=-1, keepdim=True)
+            target = (target - mean) / (var + 1.0e-6) ** 0.5
+
+        loss = (pred - target) ** 2
+        if self.img_mask is not None:
+            # exclude missing pixels from loss
+            mask = mask.unsqueeze(-1) * self.img_mask_patches
+        else:
+            loss = loss.mean(dim=-1)  # [N, L], mean loss per patch
+
+        loss = (loss * mask).sum() / mask.sum()  # mean loss on removed patches
+        return loss
+    
+    def forward_cyclic_loss(self, pred1, pred2, mask):
+        """
+        mask1 and mask2 encoder outputs should be the same since they are predicting the same held-out true mask
+        """
+        loss = (pred1 - pred2) ** 2
+        if self.img_mask is not None:
+            # exclude missing pixels from loss
+            mask = mask.unsqueeze(-1) * self.img_mask_patches
+        else:
+            loss = loss.mean(dim=-1)  # [N, L], mean loss per patch
+
+        loss = (loss * mask).sum() / mask.sum()  # mean loss on removed patches
+        return loss
+
+    def forward(self, imgs, mask_ratio=0.75, use_contrastive_loss=False):
+        latent, mask, ids_restore = self.forward_encoder(imgs, mask_ratio, use_contrastive_loss=use_contrastive_loss)
+        if not use_contrastive_loss:
+            pred = self.forward_decoder(latent, ids_restore, use_contrastive_loss=use_contrastive_loss)  # [N, L, p*p*C]
+            loss = self.forward_loss(imgs, pred, mask)
+            return loss, pred, mask, latent
+        else:
+            latent1, latent2 = latent
+            mask1, mask2 = mask
+            true_mask = copy.deepcopy(mask1)
+            true_mask[mask2==0]=0 # dont try to predict the masks that were fed to the other encoder
+            pred1 = self.forward_decoder(latent1, ids_restore, use_contrastive_loss=use_contrastive_loss)  # [N, L, p*p*C]
+            pred2 = self.forward_decoder(latent2, ids_restore, use_contrastive_loss=use_contrastive_loss)  # [N, L, p*p*C]
+            loss1 = self.forward_loss(imgs, pred1, true_mask)
+            loss2 = self.forward_loss(imgs, pred2, true_mask)
+            loss3 = self.forward_cyclic_loss(pred1, pred2, true_mask)
+            return loss1, loss2, loss3, pred1, pred2, mask1, mask2, true_mask, latent1, latent2
+    
+    def forward_features(self, x, global_pool=True):
+        # embed patches
+        x = self.patch_embed(x)
+        N, T, L, C = x.shape  # T: temporal; L: spatial
+
+        x = x.view([N, T * L, C])
+
+        # append cls token
+        if self.cls_embed:
+            cls_token = self.cls_token
+            cls_tokens = cls_token.expand(x.shape[0], -1, -1)
+            x = torch.cat((cls_tokens, x), dim=1)
+
+        if self.sep_pos_embed:
+            pos_embed = self.pos_embed_spatial.repeat(
+                1, self.input_size[0], 1
+            ) + torch.repeat_interleave(
+                self.pos_embed_temporal,
+                self.input_size[1] * self.input_size[2],
+                dim=1,
+            )
+            if self.cls_embed:
+                pos_embed = torch.cat(
+                    [
+                        self.pos_embed_class.expand(pos_embed.shape[0], -1, -1),
+                        pos_embed,
+                    ],
+                    1,
+                )
+        else:
+            pos_embed = self.pos_embed[:, :, :]
+        x = x + pos_embed
+
+        # drop patches outside image mask
+        if self.img_mask is not None:
+            if self.cls_embed:
+                cls_tokens, x = x[:, :1, :], x[:, 1:, :]
+            x = x.view([N, T, L, C])
+            x = x[:, :, self.patch_mask_indices]
+            x = x.view([N, T * self.n_mask_patches, C])
+            if self.cls_embed:
+                x = torch.cat((cls_tokens, x), dim=1)
+
+        # apply Transformer blocks
+        for blk in self.blocks:
+            x = blk(x)
+
+        if global_pool:
+            if self.cls_embed:
+                x = x[:, 1:, :]
+            x = x.mean(dim=1)
+        return x
+
+    def forward_head(self, x):
+        # classifier
+        x = self.norm(x)
+        # x = self.fc_norm(x)
+        x = self.dropout(x)
+        x = self.head(x)
+
+        return x
+    
+    def mask_fill(self, x):
+        N, L, C = x.shape
+        T = self.patch_embed.t_grid_size
+        H, W = self.patch_embed.grid_size
+        assert L == T * self.n_mask_patches
+
+        x = x.view(N, T, -1, C)
+        x_ = torch.zeros([N, T, H * W, C], dtype=x.dtype, device=x.device)
+        x = x_.scatter(
+            2, self.patch_mask_indices.view(1, 1, -1, 1).expand(N, T, -1, C), x,
+        )
+        return x
+
+
+def mae_vit_small_fmri(num_heads=6,**kwargs):
+    model = MaskedAutoencoderViT(
+        img_size=(144, 320),
+        in_chans=1,
+        embed_dim=384, 
+        depth=12, 
+        num_heads=num_heads, 
+        mlp_ratio=4,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),
+        img_mask=load_hcp_flat_mask(),
+        **kwargs,
+    )
+    return model
+
+
+def mae_vit_base_fmri(**kwargs):
+    model = MaskedAutoencoderViT(
+        img_size=(144, 320),
+        in_chans=1,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),
+        img_mask=load_hcp_flat_mask(),
+        **kwargs,
+    )
+    return model
+
+
+def mae_vit_large_fmri(**kwargs):
+    model = MaskedAutoencoderViT(
+        img_size=(144, 320),
+        in_chans=1,
+        embed_dim=1024,
+        depth=24,
+        num_heads=16,
+        mlp_ratio=4,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),
+        img_mask=load_hcp_flat_mask(),
+        **kwargs,
+    )
+    return model
+
+
+def mae_vit_huge_fmri(**kwargs):
+    model = MaskedAutoencoderViT(
+        img_size=(144, 320),
+        in_chans=1,
+        embed_dim=1280,
+        depth=32,
+        num_heads=16,
+        mlp_ratio=4,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),
+        img_mask=load_hcp_flat_mask(),
+        **kwargs,
+    )
+    return model

fMRI-foundation-model-old/flat/main.py

@@ -0,0 +1,593 @@
+import os
+import shutil
+import sys
+import json
+import yaml
+import numpy as np
+import copy
+import math
+import time
+import random
+from tqdm import tqdm
+import webdataset as wds
+import matplotlib.pyplot as plt
+
+import torch
+import torch.nn as nn
+from torchvision import transforms
+import utils
+from flat_models import *
+
+# tf32 data type is faster than standard float32
+torch.backends.cuda.matmul.allow_tf32 = True
+# following fixes a Conv3D CUDNN_NOT_SUPPORTED error
+torch.backends.cudnn.benchmark = True
+
+### Multi-GPU config ###
+device_count = torch.cuda.device_count()
+print(f"Number of available CUDA devices: {device_count}")
+
+local_rank = os.getenv('LOCAL_RANK')
+if local_rank is None: 
+    local_rank = 0
+else:
+    local_rank = int(local_rank)
+print(f"LOCAL RANK={local_rank}")
+
+num_devices = os.getenv('NUM_GPUS')
+if num_devices is None: 
+    num_devices = 1
+else:
+    num_devices = int(num_devices)
+print(f"NUM GPUS={num_devices}")
+distributed = True if num_devices>1 else False
+if distributed: assert device_count==num_devices
+
+node = os.getenv('SLURM_NODEID')
+if node is None:
+    node = 0
+else:
+    node = int(node)
+print(f"NODE={node}")
+
+global_rank = os.getenv('RANK')
+if global_rank is None:
+    global_rank = 0
+else:
+    global_rank = int(global_rank)
+print(f"GLOBAL RANK={global_rank}")
+
+world_size = os.getenv('WORLD_SIZE')
+if world_size is None: 
+    world_size = 1
+else:
+    world_size = int(world_size)
+print(f"WORLD_SIZE={world_size}")
+
+# Load parameters from yaml config
+config = yaml.load(open('config.yaml', 'r'), Loader=yaml.FullLoader)
+
+print("\n__CONFIG__")
+for attribute_name in config.keys():
+    print(f"{attribute_name} = {config[attribute_name]}")
+    globals()[attribute_name] = config[f'{attribute_name}']
+print("\n")
+
+# Create outdir for ckpt and config.yaml
+outdir = os.path.abspath(f'checkpoints/{model_name}')
+print("outdir", outdir)
+
+# Load previous config.yaml if available
+if os.path.exists(f"{outdir}/config.yaml"):
+    config = yaml.load(open(f"{outdir}/config.yaml", 'r'), Loader=yaml.FullLoader)
+    print(f"Loaded config.yaml from ckpt folder {outdir}")
+
+    # create global variables from the config
+    print("\n__REPLACING_CONFIG__")
+    for attribute_name in config.keys():
+        print(f"{attribute_name} = {config[attribute_name]}")
+        globals()[attribute_name] = config[f'{attribute_name}']
+    print("\n")
+
+data_type = torch.float32 # change depending on your mixed_precision
+global_batch_size = batch_size * world_size
+
+# FSDP Setup
+if distributed:
+    import torch.distributed as dist
+    from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+    from torch.distributed.fsdp.api import BackwardPrefetch, CPUOffload, ShardingStrategy
+    import functools
+    from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy, transformer_auto_wrap_policy
+    print(f"setting device to cuda:{local_rank}")
+    torch.cuda.set_device(local_rank)
+    device = torch.device('cuda',local_rank)
+    dist.init_process_group("nccl", rank=global_rank, world_size=world_size)
+    print(f"\nSuccessfully set cuda:{local_rank} | global_rank{global_rank} | node{node}")
+    dist.barrier() 
+    print(f"global_rank{global_rank} passed barrier")
+else:
+    device = torch.device('cuda')
+
+print("PID of this process =",os.getpid())
+print("device =", device, "distributed =",distributed, "num_devices =", num_devices, "local rank =", local_rank, "world size =", world_size, "data_type =", data_type)
+
+# seed all random functions
+utils.seed_everything(seed + global_rank)
+
+from util.hcp_flat import load_hcp_flat_mask
+from util.hcp_flat import create_hcp_flat
+from util.losses import *
+import util.visualize as vis
+
+if utils.is_interactive(): # Use less samples per epoch for debugging
+    num_samples_per_epoch = 2000
+    test_num_samples_per_epoch = 2000
+
+model = mae_vit_small_fmri(
+    decoder_embed_dim=decoder_embed_dim,
+    t_patch_size=t_patch_size,
+    pred_t_dim=pred_t_dim,
+    decoder_depth=4,
+    cls_embed=cls_embed,
+    norm_pix_loss=norm_pix_loss,
+    no_qkv_bias=no_qkv_bias,
+    sep_pos_embed=sep_pos_embed,
+    trunc_init=trunc_init,
+)
+
+if use_contrastive_loss:
+    model.simclr_handler = SimCLRHandler(model.embed_dim).to(device)
+if use_vic_loss:
+    model.vicreg_handler = VICRegHandler(model.embed_dim).to(device)
+
+# state = torch.load("checkpoints/checkpoint-00099.pth", map_location="cpu")
+# model.load_state_dict(state["model"])
+
+num_batches = num_samples_per_epoch // (num_devices * batch_size)
+test_num_batches = test_num_samples_per_epoch // (num_devices * batch_size)
+print("num_batches", num_batches)
+print("test_num_batches", test_num_batches)
+
+## Train ##
+train_dataset = create_hcp_flat(root=hcp_flat_path, 
+                    training=True, frames=num_frames)
+train_dl = wds.WebLoader(
+    train_dataset.batched(batch_size, partial=False),
+    batch_size=None,
+    shuffle=False,
+    num_workers=num_workers,
+    pin_memory=True,
+)
+train_dl = train_dl.with_epoch(num_batches)
+
+## Test ##
+test_dataset = create_hcp_flat(root=hcp_flat_path,
+                    training=False, frames=num_frames)
+test_dl = wds.WebLoader(
+    test_dataset.batched(batch_size, partial=False),
+    batch_size=None,
+    shuffle=False,
+    num_workers=num_workers,
+    pin_memory=True,
+)
+test_dl = test_dl.with_epoch(test_num_batches)
+
+print(f"\nChecking distributed setup on global_rank {global_rank}...")
+from util.video_vit import Attention
+if distributed:
+    my_auto_wrap_policy = functools.partial(
+        transformer_auto_wrap_policy, 
+        transformer_layer_cls={
+            Attention, # <--- Your Transformer layer class
+        },
+    )
+    print(f"\nPrepping FSDP on {global_rank} {node}...\n")
+    model = model.to(device)
+    model = FSDP(
+        model,
+        sharding_strategy=ShardingStrategy.HYBRID_SHARD,
+        auto_wrap_policy=my_auto_wrap_policy,
+        use_orig_params=False,
+        cpu_offload=None, #CPUOffload(offload_params=True)
+        sync_module_states=True,
+        limit_all_gathers=True, # See https://github.com/pytorch/pytorch/issues/91165
+        device_id=device,
+    )
+    print(f"\nSuccessfully loaded FSDP model to device on global_rank {global_rank}\n")
+    dist.barrier()
+else:
+    print(f"\nNot training distributed! global_rank {global_rank}")
+    model.to(device)
+
+no_decay = ["bias", "LayerNorm.bias", "LayerNorm.weight"]
+opt_grouped_parameters = [
+    {'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': 0.05},
+    {'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0},
+]
+
+lr = base_lr * global_batch_size / 256
+print(f"multiply base lr {base_lr} by effective batch size {global_batch_size}")
+print(f"lr = {lr}")
+
+optimizer = torch.optim.AdamW(opt_grouped_parameters, lr=lr, betas=(0.9, 0.95))
+
+def adjust_learning_rate(optimizer, epoch, warmup_epochs=5, min_lr=0.0):
+    """Decay the learning rate with half-cycle cosine after warmup"""
+    if epoch < warmup_epochs:
+        lr_ = lr * epoch / warmup_epochs
+    else:
+        lr_ = min_lr + (lr - min_lr) * 0.5 * (
+            1.0
+            + math.cos(
+                math.pi
+                * (epoch - warmup_epochs)
+                / (num_epochs - warmup_epochs)
+            )
+        )
+    for param_group in optimizer.param_groups:
+        param_group["lr"] = lr_
+    return lr_
+
+print("\nDone with model preparations!")
+num_params = utils.count_params(model)
+
+def save_ckpt(model,tag="last"):
+    if distributed: dist.barrier()
+    model_states = model.state_dict()
+    if global_rank == 0:
+        os.makedirs(outdir,exist_ok=True)
+        ckpt_path = outdir+f'/{tag}.pth'
+        
+        if tag == "last" and os.path.exists(ckpt_path):
+            shutil.copyfile(os.path.join(outdir, f'{tag}.pth'), os.path.join(outdir, f'{tag}_old.pth'))
+        # print(f'saving {ckpt_path}',flush=True)
+        if tag=='last':
+            torch.save({
+                'epoch': epoch,
+                'model_state_dict': model_states,
+                'optimizer_state_dict': optimizer.state_dict(),
+                }, ckpt_path)
+        else:
+            torch.save({
+                'epoch': epoch,
+                'model_state_dict': model_states,
+                }, ckpt_path)
+            
+        if tag == "last" and os.path.exists(os.path.join(outdir, f'{tag}_old.pth')):
+            os.remove(os.path.join(outdir, f'{tag}_old.pth'))
+        print(f"\n---saved {ckpt_path}!---\n")
+        
+        # save the config.yaml
+        if not os.path.exists(f"{outdir}/config.yaml"):
+            with open(f"{outdir}/config.yaml", 'w') as file:
+                yaml.dump(config, file)
+            print(f"saved {outdir}/config.yaml!")
+
+# auto resume
+if os.path.exists(os.path.join(outdir, 'last.pth')) or os.path.exists(os.path.join(outdir, 'last_old.pth')):
+    if os.path.exists(os.path.join(outdir, 'last_old.pth')):
+        if os.path.exists(os.path.join(outdir, 'last.pth')):
+            # this is corrupted
+            os.remove(os.path.join(outdir, f'last.pth'))
+        # set last_old as last
+        shutil.move(os.path.join(outdir, f'last_old.pth'), os.path.join(outdir, f'last.pth'))
+    
+    ckpt_path = os.path.join(outdir, 'last.pth')
+    resume_from_ckpt = True
+
+epoch = 0
+if resume_from_ckpt:
+    print("\n---resuming from ckpt_path---\n", ckpt_path)
+    checkpoint = torch.load(ckpt_path, map_location=device)
+    epoch = checkpoint['epoch']+1
+    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])        
+    model.load_state_dict(checkpoint['model_state_dict'])
+    # total_steps_done = epoch*num_iterations_per_epoch
+    # for _ in range(total_steps_done):
+    #     lr_scheduler.step()
+    del checkpoint
+    torch.cuda.empty_cache()
+
+if utils.is_interactive():
+    wandb_log = False
+    ckpt_saving = False
+if local_rank==0 and wandb_log: # only use main process for wandb logging
+    import wandb
+    wandb_project = 'fMRI-foundation-model'
+    print(f"wandb {wandb_project} run {model_name}")
+    # need to configure wandb beforehand in terminal with "wandb init"!
+    wandb_config = {
+      "model_name": model_name,
+      "global_batch_size": global_batch_size,
+      "batch_size": batch_size,
+      "num_epochs": num_epochs,
+      "cls_embed": cls_embed,
+      "norm_pix_loss": norm_pix_loss,
+      "t_patch_size": t_patch_size,
+      "pred_t_dim": pred_t_dim,
+      "mask_ratio": mask_ratio,
+      "num_frames": num_frames,
+      "sep_pos_embed": sep_pos_embed,
+      "decoder_embed_dim": decoder_embed_dim,
+      "use_contrastive_loss": use_contrastive_loss,
+      "num_params": num_params,
+      "base_lr": base_lr,
+      "lr": lr,
+      "num_samples_per_epoch": num_samples_per_epoch,
+      "test_num_samples_per_epoch": test_num_samples_per_epoch,
+      "num_epochs": num_epochs,
+      "grad_clip": grad_clip,
+      "ckpt_interval": ckpt_interval,
+      "ckpt_saving": ckpt_saving,
+      "print_interval": print_interval,
+      "seed": seed,
+      "distributed": distributed,
+      "num_devices": num_devices,
+      "world_size": world_size,
+    }
+    print("wandb_config:\n",wandb_config)
+    print("wandb_id:",model_name)
+    wandb.init(
+        id=model_name,
+        project=wandb_project,
+        name=model_name,
+        config=wandb_config,
+        resume="allow",
+    )
+else:
+    wandb_log = False
+
+# lrs, train_losses, test_losses = [], [], []
+# train_losses1, train_losses2, train_losses3, train_losses4 = [], [], [], []
+
+epoch = 0
+lrs, train_losses, recon_losses, contrastive_losses, vic_losses = [], [], [], [], []
+cos_sim_encoder_output, cos_sim_decoder_output, cos_sim_encoder_output_patchwise = [], [], []
+probe_losses, probe_accs, test_losses, test_accs = [], [], [], []
+cos_sim_encoder_output_patchwise_test, cos_sim_encoder_output_test = [], []
+
+mse = nn.MSELoss()
+l1 = nn.L1Loss()
+crossentropy = nn.CrossEntropyLoss()
+
+if use_contrastive_loss:
+    contrastive_temps = utils.cosine_anneal(0.004, 0.0075, num_epochs)
+
+grad_scaler = torch.cuda.amp.GradScaler(enabled=True)
+progress_bar = tqdm(range(epoch, num_epochs), disable=local_rank!=0, desc="Overall")
+for epoch in progress_bar:
+    model.train()
+    for train_i, batch in enumerate(train_dl):
+        optimizer.zero_grad()
+        adjust_learning_rate(optimizer, train_i / num_batches + epoch)
+
+        input_func = batch[1]
+
+        input_func = input_func.to(dtype=data_type, device=device, non_blocking=True)
+        if len(input_func.shape) == 6:
+            b, r, c, t, h, w = input_func.shape
+            input_func = input_func.reshape(b * r, c, t, h, w)
+
+        ids_shuffle, ids_restore = get_ids_shuffle(input_func.shape[0], input_func.device, model)
+
+        # adjust number to keep relative to image mask
+        if model.img_mask is not None:
+            len_keep = int(model.patch_embed.t_grid_size * model.n_mask_patches * (1 - mask_ratio))
+        else:
+            len_keep = int(ids_restore.shape[1] * (1 - mask_ratio))
+
+        # keep the first subset
+        ids_keep = ids_shuffle[:, :len_keep]
+        lats = model.forward_encoder_with_mask(input_func, ids_keep)
+
+        if use_decoder:
+            mask = torch.ones_like(ids_restore)
+            mask[:, :len_keep] = 0
+            # unshuffle to get the binary mask
+            mask = torch.gather(mask, dim=1, index=ids_restore)
+            
+            pred = model.forward_decoder(lats, ids_restore)
+            recon_loss = model.forward_loss(input_func, pred, mask)
+            
+            recon_losses.append(recon_loss.item())
+            loss = recon_loss
+        else:
+            recon_loss = torch.nan
+            recon_losses.append(recon_loss)
+            loss = 0
+        
+        if use_contrastive_loss and not use_vic_loss:
+            ids_keep2 = ids_shuffle[:, :2*len_keep]
+            lats2 = model.forward_encoder_with_mask(input_func, ids_keep2)
+            
+            enc_cls_token = lats[:, :1]
+            enc_cls_token2 = lats2[:,:1]
+            
+            temp = contrastive_temps[epoch]
+
+            all_cls = torch.cat([enc_cls_token, enc_cls_token2], dim=0)
+            all_cls_proj = model.simclr_handler(all_cls)
+
+            contr_loss = SimCLRHandler.simclr_loss(all_cls_proj, temp)
+            
+            contrastive_losses.append(contr_loss.item())
+            loss += (contr_loss * contrastive_loss_weight)
+        # new loss
+        elif use_vic_loss:
+            ids_keep2 = ids_shuffle[:, :2*len_keep]
+            lats2 = model.forward_encoder_with_mask(input_func, ids_keep2)
+
+            enc_cls_token = lats[:, :1]
+            enc_cls_token2 = lats2[:,:1]
+
+            l1 = lats
+            l2 = lats2[:, :lats.shape[1]]
+
+            l1_proj = model.vicreg_handler(l1)
+            l2_proj = model.vicreg_handler(l2)
+
+            vic_loss = VICRegHandler.vicreg_loss(l1_proj, l2_proj, gamma=gamma, lamda=lamda, mu=mu, nu=nu, 
+                                                        rand_frac=rand_frac, use_vic_cls=use_vic_cls)
+
+            if use_contrastive_loss:
+                temp = contrastive_temps[epoch]
+                all_cls = torch.cat([enc_cls_token, enc_cls_token2], dim=0)
+                all_cls_proj = model.simclr_handler(all_cls)
+
+                contr_loss = SimCLRHandler.simclr_loss(all_cls_proj, temp)
+                contrastive_losses.append(contr_loss.item())
+            else:
+                contr_loss = 0
+                contrastive_losses.append(0)
+            
+            vic_losses.append(vic_loss.item())
+            loss += (contr_loss * contrastive_loss_weight + vic_loss * vic_loss_weight)
+        else:
+            vic_losses.append(0)
+            contrastive_losses.append(0)
+
+        grad_scaler.scale(loss).backward()
+        grad_scaler.unscale_(optimizer)
+        if grad_clip:
+            nn.utils.clip_grad_norm_(model.parameters(), max_norm=grad_clip)
+        grad_scaler.step(optimizer)
+        grad_scaler.update()
+        torch.cuda.synchronize()
+
+        cos_sim_encoder_output_patchwise.append(utils.patchwise_cosine_similarity(lats)[~torch.eye(lats.shape[1], dtype=bool)[None].expand(lats.shape[0],-1,-1)].mean().item())
+        cos_sim_encoder_output.append(utils.batchwise_cosine_similarity(lats.flatten(1)/1e3,lats.flatten(1)/1e3)[~torch.eye(len(lats),dtype=torch.bool)].mean().item())
+        if use_decoder:
+            cos_sim_decoder_output.append(utils.batchwise_cosine_similarity(pred, pred)[~torch.eye(len(pred),dtype=torch.bool)].mean().item())
+
+        lrs.append(optimizer.param_groups[0]["lr"])
+        train_losses.append(loss.item())
+        
+        if ((train_i%print_interval)==0 or (train_i==num_batches-1)) and train_i>0:
+            print(f"Ep. {epoch} | loss {np.mean(train_losses[-print_interval:]):.3f} | lr {optimizer.param_groups[0]['lr']} | {train_i}/{num_batches}")
+            if wandb_log: 
+                logs = {"train/loss": np.mean(train_losses[-print_interval:])}
+                # epoch_1000x as x-axis calibrates different curves when batch size changes
+                epoch_1000x = int((train_i / num_batches + epoch) * 1000)
+                wandb.log(logs, step=epoch_1000x)
+
+    if utils.is_interactive() or wandb_log:
+        print(f"Ep. {epoch} | loss {np.mean(train_losses[-print_interval:]):.3f} | lr {optimizer.param_groups[0]['lr']} | {train_i}/{num_batches}")
+        with torch.no_grad():
+            if norm_pix_loss:
+                normed_input_func, patch_info = model.patchify(input_func, alter_patch_info=False, return_patch_info=True)
+                target_mean = normed_input_func.mean(dim=-1, keepdim=True)
+                target_var = normed_input_func.var(dim=-1, keepdim=True)
+                normed_input_func = (normed_input_func - target_mean) / (target_var + 1.0e-6) ** 0.5
+                normed_input_func = model.unpatchify(normed_input_func, patch_info=patch_info)
+
+                vis_out = vis.plot_mask_pred(
+                    model, normed_input_func, pred, mask, 
+                    mean=0.5, std=0.2, 
+                )
+            else:
+                vis_out = vis.plot_mask_pred(
+                    model, input_func, pred, mask, 
+                    mean=0.5, std=0.2, 
+                )
+
+            if wandb_log:
+                logs = {"train/mask_pred": wandb.Image(vis_out)}
+                epoch_1000x = int((train_i / num_batches + epoch) * 1000)
+                wandb.log(logs, step=epoch_1000x)
+            else:
+                display(vis_out)
+        
+    logs = {
+        "train/loss": np.mean(train_losses[-(train_i + 1) :]),
+        "train/recon_losses": np.mean(recon_losses[-(train_i + 1) :]),
+        "train/contrastive_losses": np.mean(contrastive_losses[-(train_i + 1) :]),
+        "train/vic_losses": np.mean(vic_losses[-(train_i + 1) :]),
+        "train/num_steps": len(recon_losses),
+        "train/cos_sim_encoder_output": np.mean(cos_sim_encoder_output[-(train_i + 1) :]),
+        "train/cos_sim_decoder_output": np.mean(cos_sim_decoder_output[-(train_i + 1) :]) if use_decoder else np.nan,
+        "train/cos_sim_encoder_output_patchwise": np.mean(cos_sim_encoder_output_patchwise[-(train_i + 1) :]),
+        "lr": np.mean(lrs[-(train_i + 1) :]),
+        "epoch": epoch,
+    }
+        
+    # Evaluate performance on held-out test dataset
+    model.eval()
+    with torch.no_grad():
+        for test_i, batch in enumerate(test_dl):
+            input_func = batch[1]
+
+            input_func = input_func.to(dtype=data_type, device=device, non_blocking=True)
+            if len(input_func.shape) == 6:
+                b, r, c, t, h, w = input_func.shape
+                input_func = input_func.reshape(b * r, c, t, h, w)
+
+            if not use_contrastive_loss:
+                loss, pred, mask, latent = model(input_func, mask_ratio=mask_ratio, use_contrastive_loss=use_contrastive_loss)
+            else:
+                loss1, loss2, loss3, pred1, pred2, mask1, mask2, true_mask, latent1, latent2 = model(input_func, mask_ratio=mask_ratio, use_contrastive_loss=use_contrastive_loss)
+                pred, mask, latent = pred1, mask1, latent1
+                
+                # contrastive loss
+                temp = contrastive_temps[epoch]
+                logits = (nn.functional.normalize(latent1[:,:1].flatten(1),dim=-1) @
+                            nn.functional.normalize(latent2[:,:1].flatten(1),dim=-1).T) / temp
+                labels = torch.arange(len(logits)).long().to(device)
+                contr_loss = (crossentropy(logits, labels) + crossentropy(logits.T, labels)) / 2
+                
+                loss = loss1 + loss2 + loss3 + contr_loss
+
+            test_losses.append(loss.item())
+
+            cos_sim_encoder_output_patchwise_test.append(utils.patchwise_cosine_similarity(latent)[~torch.eye(latent.shape[1], dtype=bool)[None].expand(latent.shape[0],-1,-1)].mean().item())
+            cos_sim_encoder_output_test.append(utils.batchwise_cosine_similarity(latent.flatten(1)/1e3, latent.flatten(1)/1e3)[~torch.eye(len(latent),dtype=torch.bool)].mean().item())
+
+            if test_i%print_interval==0 and test_i>0:
+                print(f"Test | loss {np.mean(test_losses[-print_interval:]):.3f} | {test_i}/{test_num_batches}")
+
+    print(f"Test | iter {test_i} | loss {np.mean(test_losses[-test_i:]):.3f}")
+    if wandb_log: 
+        logs.update({
+            "test/loss": np.mean(test_losses[-test_i:]),
+            "test/cos_sim_encoder_output": np.mean(cos_sim_encoder_output_test[-(test_i + 1) :]),
+            "test/cos_sim_encoder_output_patchwise": np.mean(cos_sim_encoder_output_patchwise_test[-(test_i + 1) :]),
+        })
+        wandb.log(logs)
+
+    # Plot progress (first sample in batch)
+    if utils.is_interactive() or wandb_log:
+        with torch.no_grad():
+            if norm_pix_loss:
+                normed_input_func, patch_info = model.patchify(input_func, alter_patch_info=False, return_patch_info=True)
+                target_mean = normed_input_func.mean(dim=-1, keepdim=True)
+                target_var = normed_input_func.var(dim=-1, keepdim=True)
+                normed_input_func = (normed_input_func - target_mean) / (target_var + 1.0e-6) ** 0.5
+                normed_input_func = model.unpatchify(normed_input_func, patch_info=patch_info)
+
+                vis_out = vis.plot_mask_pred(
+                    model, normed_input_func, pred, mask, 
+                    mean=0.5, std=0.2, 
+                )
+            else:
+                vis_out = vis.plot_mask_pred(
+                    model, input_func, pred, mask, 
+                    mean=0.5, std=0.2, 
+                )
+
+            if wandb_log:
+                logs = {"test/mask_pred": wandb.Image(vis_out)}
+                wandb.log(logs)
+            else:
+                display(vis_out)
+
+    # Save model checkpoint
+    if ckpt_saving and epoch>0 and ((epoch % ckpt_interval == 0) or (epoch==num_epochs-1)):
+        save_ckpt(model,f"epoch{epoch}")
+
+    # wait for other GPUs to catch up if needed
+    if distributed: dist.barrier()
+    
+    # close any open plots
+    plt.close()
+    
+if distributed: dist.destroy_process_group()

fMRI-foundation-model-old/flat/main2.ipynb

@@ -0,0 +1,450 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "a0834886-8336-4648-bd48-804b5e472b3d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import shutil\n",
+    "import sys\n",
+    "import json\n",
+    "import yaml\n",
+    "import numpy as np\n",
+    "import copy\n",
+    "import math\n",
+    "import time\n",
+    "import random\n",
+    "from tqdm import tqdm\n",
+    "import webdataset as wds\n",
+    "import matplotlib.pyplot as plt\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "from torchvision import transforms\n",
+    "import utils\n",
+    "from flat_models import *\n",
+    "\n",
+    "# tf32 data type is faster than standard float32\n",
+    "torch.backends.cuda.matmul.allow_tf32 = True\n",
+    "# following fixes a Conv3D CUDNN_NOT_SUPPORTED error\n",
+    "torch.backends.cudnn.benchmark = True\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "7a8e5e12-b9ac-4baf-91eb-b4c25740c578",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Number of available CUDA devices: 1\n",
+      "LOCAL RANK=0\n",
+      "NUM GPUS=1\n",
+      "NODE=0\n",
+      "GLOBAL RANK=0\n",
+      "WORLD_SIZE=1\n"
+     ]
+    }
+   ],
+   "source": [
+    "\n",
+    "### Multi-GPU config ###\n",
+    "device_count = torch.cuda.device_count()\n",
+    "print(f\"Number of available CUDA devices: {device_count}\")\n",
+    "\n",
+    "local_rank = os.getenv('LOCAL_RANK')\n",
+    "if local_rank is None: \n",
+    "    local_rank = 0\n",
+    "else:\n",
+    "    local_rank = int(local_rank)\n",
+    "print(f\"LOCAL RANK={local_rank}\")\n",
+    "\n",
+    "num_devices = os.getenv('NUM_GPUS')\n",
+    "if num_devices is None: \n",
+    "    num_devices = 1\n",
+    "else:\n",
+    "    num_devices = int(num_devices)\n",
+    "print(f\"NUM GPUS={num_devices}\")\n",
+    "distributed = True if num_devices>1 else False\n",
+    "if distributed: assert device_count==num_devices\n",
+    "\n",
+    "node = os.getenv('SLURM_NODEID')\n",
+    "if node is None:\n",
+    "    node = 0\n",
+    "else:\n",
+    "    node = int(node)\n",
+    "print(f\"NODE={node}\")\n",
+    "\n",
+    "global_rank = os.getenv('RANK')\n",
+    "if global_rank is None:\n",
+    "    global_rank = 0\n",
+    "else:\n",
+    "    global_rank = int(global_rank)\n",
+    "print(f\"GLOBAL RANK={global_rank}\")\n",
+    "\n",
+    "world_size = os.getenv('WORLD_SIZE')\n",
+    "if world_size is None: \n",
+    "    world_size = 1\n",
+    "else:\n",
+    "    world_size = int(world_size)\n",
+    "print(f\"WORLD_SIZE={world_size}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "65ff5246-4ff5-457d-8509-646b319093c6",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "__CONFIG__\n",
+      "model_name = flat_NORM_contrastive_\n",
+      "hcp_flat_path = /weka/proj-medarc/shared/hcp_flat\n",
+      "use_contrastive_loss = True\n",
+      "use_decoder = True\n",
+      "patch_size = 16\n",
+      "decoder_embed_dim = 512\n",
+      "num_frames = 16\n",
+      "mask_ratio = 0.75\n",
+      "pred_t_dim = 8\n",
+      "t_patch_size = 2\n",
+      "cls_embed = True\n",
+      "no_qkv_bias = False\n",
+      "sep_pos_embed = True\n",
+      "trunc_init = False\n",
+      "norm_pix_loss = False\n",
+      "contrastive_loss_weight = 1.0\n",
+      "use_vic_loss = True\n",
+      "vic_loss_weight = 1.0\n",
+      "gamma = 0.5\n",
+      "mu = 25\n",
+      "lamda = 25\n",
+      "nu = 1\n",
+      "rand_frac = 0.1\n",
+      "use_vic_cls = True\n",
+      "pct_masks_to_decode = 1\n",
+      "batch_size = 32\n",
+      "num_workers = 10\n",
+      "num_epochs = 100\n",
+      "seed = 42\n",
+      "base_lr = 0.0003\n",
+      "num_samples_per_epoch = 200000\n",
+      "test_num_samples_per_epoch = 50000\n",
+      "grad_clip = 1.0\n",
+      "grad_accumulation_steps = 1\n",
+      "test_set = False\n",
+      "probe_num_samples_per_epoch = 100000\n",
+      "probe_num_epochs = 30\n",
+      "probe_batch_size = 8\n",
+      "probe_base_lr = 0.0003\n",
+      "ckpt_saving = True\n",
+      "ckpt_interval = 5\n",
+      "print_interval = 20\n",
+      "resume_from_ckpt = False\n",
+      "wandb_log = False\n",
+      "\n",
+      "\n",
+      "outdir /weka/proj-fmri/ckadirt/fMRI-foundation-model/flat/checkpoints/flat_NORM_contrastive_\n"
+     ]
+    }
+   ],
+   "source": [
+    "\n",
+    "# Load parameters from yaml config\n",
+    "config = yaml.load(open('config.yaml', 'r'), Loader=yaml.FullLoader)\n",
+    "\n",
+    "print(\"\\n__CONFIG__\")\n",
+    "for attribute_name in config.keys():\n",
+    "    print(f\"{attribute_name} = {config[attribute_name]}\")\n",
+    "    globals()[attribute_name] = config[f'{attribute_name}']\n",
+    "print(\"\\n\")\n",
+    "\n",
+    "# Create outdir for ckpt and config.yaml\n",
+    "outdir = os.path.abspath(f'checkpoints/{model_name}')\n",
+    "print(\"outdir\", outdir)\n",
+    "\n",
+    "# Load previous config.yaml if available\n",
+    "if os.path.exists(f\"{outdir}/config.yaml\"):\n",
+    "    config = yaml.load(open(f\"{outdir}/config.yaml\", 'r'), Loader=yaml.FullLoader)\n",
+    "    print(f\"Loaded config.yaml from ckpt folder {outdir}\")\n",
+    "\n",
+    "    # create global variables from the config\n",
+    "    print(\"\\n__REPLACING_CONFIG__\")\n",
+    "    for attribute_name in config.keys():\n",
+    "        print(f\"{attribute_name} = {config[attribute_name]}\")\n",
+    "        globals()[attribute_name] = config[f'{attribute_name}']\n",
+    "    print(\"\\n\")\n",
+    "\n",
+    "data_type = torch.float32 # change depending on your mixed_precision\n",
+    "global_batch_size = batch_size * world_size\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "b5b30c27-04e9-4566-ac33-cc2d69e27711",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "PID of this process = 925118\n",
+      "device = cuda distributed = False num_devices = 1 local rank = 0 world size = 1 data_type = torch.float32\n"
+     ]
+    }
+   ],
+   "source": [
+    "\n",
+    "# FSDP Setup\n",
+    "if distributed:\n",
+    "    import torch.distributed as dist\n",
+    "    from torch.distributed.fsdp import FullyShardedDataParallel as FSDP\n",
+    "    from torch.distributed.fsdp.api import BackwardPrefetch, CPUOffload, ShardingStrategy\n",
+    "    import functools\n",
+    "    from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy, transformer_auto_wrap_policy\n",
+    "    print(f\"setting device to cuda:{local_rank}\")\n",
+    "    torch.cuda.set_device(local_rank)\n",
+    "    device = torch.device('cuda',local_rank)\n",
+    "    dist.init_process_group(\"nccl\", rank=global_rank, world_size=world_size)\n",
+    "    print(f\"\\nSuccessfully set cuda:{local_rank} | global_rank{global_rank} | node{node}\")\n",
+    "    dist.barrier() \n",
+    "    print(f\"global_rank{global_rank} passed barrier\")\n",
+    "else:\n",
+    "    device = torch.device('cuda')\n",
+    "\n",
+    "print(\"PID of this process =\",os.getpid())\n",
+    "print(\"device =\", device, \"distributed =\",distributed, \"num_devices =\", num_devices, \"local rank =\", local_rank, \"world size =\", world_size, \"data_type =\", data_type)\n",
+    "\n",
+    "# seed all random functions\n",
+    "utils.seed_everything(seed + global_rank)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "ad692bcb-a166-46d5-97f7-292c91cd4443",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "img_size (144, 320) patch_size (16, 16) frames 16 t_patch_size 2\n",
+      "model initialized\n",
+      "num_batches 62\n",
+      "test_num_batches 62\n"
+     ]
+    }
+   ],
+   "source": [
+    "\n",
+    "from util.hcp_flat import load_hcp_flat_mask\n",
+    "from util.hcp_flat import create_hcp_flat\n",
+    "from util.losses import *\n",
+    "import util.visualize as vis\n",
+    "\n",
+    "if utils.is_interactive(): # Use less samples per epoch for debugging\n",
+    "    num_samples_per_epoch = 2000\n",
+    "    test_num_samples_per_epoch = 2000\n",
+    "\n",
+    "model = mae_vit_small_fmri(\n",
+    "    decoder_embed_dim=decoder_embed_dim,\n",
+    "    t_patch_size=t_patch_size,\n",
+    "    pred_t_dim=pred_t_dim,\n",
+    "    decoder_depth=4,\n",
+    "    cls_embed=cls_embed,\n",
+    "    norm_pix_loss=norm_pix_loss,\n",
+    "    no_qkv_bias=no_qkv_bias,\n",
+    "    sep_pos_embed=sep_pos_embed,\n",
+    "    trunc_init=trunc_init,\n",
+    ")\n",
+    "\n",
+    "if use_contrastive_loss:\n",
+    "    model.simclr_handler = SimCLRHandler(model.embed_dim).to(device)\n",
+    "if use_vic_loss:\n",
+    "    model.vicreg_handler = VICRegHandler(model.embed_dim).to(device)\n",
+    "\n",
+    "# state = torch.load(\"checkpoints/checkpoint-00099.pth\", map_location=\"cpu\")\n",
+    "# model.load_state_dict(state[\"model\"])\n",
+    "\n",
+    "num_batches = num_samples_per_epoch // (num_devices * batch_size)\n",
+    "test_num_batches = test_num_samples_per_epoch // (num_devices * batch_size)\n",
+    "print(\"num_batches\", num_batches)\n",
+    "print(\"test_num_batches\", test_num_batches)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "63da3ad3-116f-498d-ae01-d75585b010f2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "## Train ##\n",
+    "train_dataset = create_hcp_flat(root=hcp_flat_path, \n",
+    "                    training=True, frames=num_frames)\n",
+    "train_dl = wds.WebLoader(\n",
+    "    train_dataset.batched(batch_size, partial=False),\n",
+    "    batch_size=None,\n",
+    "    shuffle=False,\n",
+    "    num_workers=num_workers,\n",
+    "    pin_memory=True,\n",
+    ")\n",
+    "train_dl = train_dl.with_epoch(num_batches)\n",
+    "\n",
+    "## Test ##\n",
+    "test_dataset = create_hcp_flat(root=hcp_flat_path,\n",
+    "                    training=False, frames=num_frames)\n",
+    "test_dl = wds.WebLoader(\n",
+    "    test_dataset.batched(batch_size, partial=False),\n",
+    "    batch_size=None,\n",
+    "    shuffle=False,\n",
+    "    num_workers=num_workers,\n",
+    "    pin_memory=True,\n",
+    ")\n",
+    "test_dl = test_dl.with_epoch(test_num_batches)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "9ad20649-e51e-4e5c-9fb4-f78d33edaea0",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Checking distributed setup on global_rank 0...\n",
+      "\n",
+      "Not training distributed! global_rank 0\n",
+      "multiply base lr 0.0003 by effective batch size 32\n",
+      "lr = 3.75e-05\n"
+     ]
+    }
+   ],
+   "source": [
+    "\n",
+    "print(f\"\\nChecking distributed setup on global_rank {global_rank}...\")\n",
+    "from util.video_vit import Attention\n",
+    "if distributed:\n",
+    "    my_auto_wrap_policy = functools.partial(\n",
+    "        transformer_auto_wrap_policy, \n",
+    "        transformer_layer_cls={\n",
+    "            Attention, # <--- Your Transformer layer class\n",
+    "        },\n",
+    "    )\n",
+    "    print(f\"\\nPrepping FSDP on {global_rank} {node}...\\n\")\n",
+    "    model = model.to(device)\n",
+    "    model = FSDP(\n",
+    "        model,\n",
+    "        sharding_strategy=ShardingStrategy.HYBRID_SHARD,\n",
+    "        auto_wrap_policy=my_auto_wrap_policy,\n",
+    "        use_orig_params=False,\n",
+    "        cpu_offload=None, #CPUOffload(offload_params=True)\n",
+    "        sync_module_states=True,\n",
+    "        limit_all_gathers=True, # See https://github.com/pytorch/pytorch/issues/91165\n",
+    "        device_id=device,\n",
+    "    )\n",
+    "    print(f\"\\nSuccessfully loaded FSDP model to device on global_rank {global_rank}\\n\")\n",
+    "    dist.barrier()\n",
+    "else:\n",
+    "    print(f\"\\nNot training distributed! global_rank {global_rank}\")\n",
+    "    model.to(device)\n",
+    "\n",
+    "no_decay = [\"bias\", \"LayerNorm.bias\", \"LayerNorm.weight\"]\n",
+    "opt_grouped_parameters = [\n",
+    "    {'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': 0.05},\n",
+    "    {'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0},\n",
+    "]\n",
+    "\n",
+    "lr = base_lr * global_batch_size / 256\n",
+    "print(f\"multiply base lr {base_lr} by effective batch size {global_batch_size}\")\n",
+    "print(f\"lr = {lr}\")\n",
+    "\n",
+    "optimizer = torch.optim.AdamW(opt_grouped_parameters, lr=lr, betas=(0.9, 0.95))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "cef70e06-d181-46c0-a36c-12aa944a8441",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Done with model preparations!\n",
+      "param counts:\n",
+      "40,789,504 total\n",
+      "40,789,504 trainable\n"
+     ]
+    }
+   ],
+   "source": [
+    "\n",
+    "def adjust_learning_rate(optimizer, epoch, warmup_epochs=5, min_lr=0.0):\n",
+    "    \"\"\"Decay the learning rate with half-cycle cosine after warmup\"\"\"\n",
+    "    if epoch < warmup_epochs:\n",
+    "        lr_ = lr * epoch / warmup_epochs\n",
+    "    else:\n",
+    "        lr_ = min_lr + (lr - min_lr) * 0.5 * (\n",
+    "            1.0\n",
+    "            + math.cos(\n",
+    "                math.pi\n",
+    "                * (epoch - warmup_epochs)\n",
+    "                / (num_epochs - warmup_epochs)\n",
+    "            )\n",
+    "        )\n",
+    "    for param_group in optimizer.param_groups:\n",
+    "        param_group[\"lr\"] = lr_\n",
+    "    return lr_\n",
+    "\n",
+    "print(\"\\nDone with model preparations!\")\n",
+    "num_params = utils.count_params(model)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1deff37d-c9c6-48e8-b1c4-1cec41f23fca",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.14"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

fMRI-foundation-model-old/flat/parquet_probe.ipynb

@@ -0,0 +1,407 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "b8e236f1-385a-4d93-bb39-bea3ee384d76",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "outdir /weka/proj-fmri/paulscotti/fMRI-foundation-model/flat/checkpoints/clane\n",
+      "Loaded config.yaml from ckpt folder /weka/proj-fmri/paulscotti/fMRI-foundation-model/flat/checkpoints/clane\n",
+      "\n",
+      "__CONFIG__\n",
+      "base_lr = 0.001\n",
+      "batch_size = 32\n",
+      "ckpt_interval = 5\n",
+      "ckpt_saving = True\n",
+      "cls_embed = True\n",
+      "decoder_embed_dim = 512\n",
+      "grad_clip = 1.0\n",
+      "hcp_flat_path = /weka/proj-fmri/shared/hcp_flat\n",
+      "mask_ratio = 0.75\n",
+      "model_name = clane\n",
+      "no_qkv_bias = False\n",
+      "norm_pix_loss = False\n",
+      "num_epochs = 50\n",
+      "num_frames = 16\n",
+      "num_samples_per_epoch = 200000\n",
+      "num_workers = 10\n",
+      "pred_t_dim = 8\n",
+      "print_interval = 20\n",
+      "probe_num_samples_per_epoch = 100000\n",
+      "probe_num_epochs = 30\n",
+      "probe_batch_size = 8\n",
+      "probe_base_lr = 0.0003\n",
+      "resume_from_ckpt = True\n",
+      "seed = 42\n",
+      "sep_pos_embed = True\n",
+      "t_patch_size = 2\n",
+      "test_num_samples_per_epoch = 50000\n",
+      "trunc_init = False\n",
+      "use_contrastive_loss = False\n",
+      "wandb_log = True\n",
+      "\n",
+      "\n",
+      "Number of available CUDA devices: 1\n",
+      "LOCAL RANK=0\n",
+      "NUM GPUS=1\n",
+      "NODE=0\n",
+      "GLOBAL RANK=0\n",
+      "WORLD_SIZE=1\n",
+      "PID of this process = 448416\n",
+      "device = cuda distributed = False num_devices = 1 local rank = 0 world size = 1 data_type = torch.float32\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Import packages and setup gpu configuration.\n",
+    "# This code block shouldnt need to be adjusted!\n",
+    "import os\n",
+    "import sys\n",
+    "import json\n",
+    "import yaml\n",
+    "import numpy as np\n",
+    "import copy\n",
+    "import math\n",
+    "import time\n",
+    "import random\n",
+    "from tqdm import tqdm\n",
+    "import webdataset as wds\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "from torchvision import transforms\n",
+    "import utils\n",
+    "from flat_models import *\n",
+    "\n",
+    "# tf32 data type is faster than standard float32\n",
+    "torch.backends.cuda.matmul.allow_tf32 = True\n",
+    "# following fixes a Conv3D CUDNN_NOT_SUPPORTED error\n",
+    "torch.backends.cudnn.benchmark = True\n",
+    "\n",
+    "## MODEL TO LOAD ##\n",
+    "model_name = \"clane\"\n",
+    "parquet_folder = \"epoch99\"\n",
+    "outdir = os.path.abspath(f'checkpoints/{model_name}')\n",
+    "print(\"outdir\", outdir)\n",
+    "# Load previous config.yaml if available\n",
+    "if os.path.exists(f\"{outdir}/config.yaml\"):\n",
+    "    config = yaml.load(open(f\"{outdir}/config.yaml\", 'r'), Loader=yaml.FullLoader)\n",
+    "    print(f\"Loaded config.yaml from ckpt folder {outdir}\")\n",
+    "    # create global variables from the config\n",
+    "    print(\"\\n__CONFIG__\")\n",
+    "    for attribute_name in config.keys():\n",
+    "        print(f\"{attribute_name} = {config[attribute_name]}\")\n",
+    "        globals()[attribute_name] = config[f'{attribute_name}']\n",
+    "    print(\"\\n\")\n",
+    "\n",
+    "### Multi-GPU config ###\n",
+    "device_count = torch.cuda.device_count()\n",
+    "print(f\"Number of available CUDA devices: {device_count}\")\n",
+    "\n",
+    "local_rank = os.getenv('LOCAL_RANK')\n",
+    "if local_rank is None: \n",
+    "    local_rank = 0\n",
+    "else:\n",
+    "    local_rank = int(local_rank)\n",
+    "print(f\"LOCAL RANK={local_rank}\")\n",
+    "\n",
+    "num_devices = os.getenv('NUM_GPUS')\n",
+    "if num_devices is None: \n",
+    "    num_devices = 1\n",
+    "else:\n",
+    "    num_devices = int(num_devices)\n",
+    "print(f\"NUM GPUS={num_devices}\")\n",
+    "distributed = True if num_devices>1 else False\n",
+    "if distributed: assert device_count==num_devices\n",
+    "\n",
+    "node = os.getenv('SLURM_NODEID')\n",
+    "if node is None:\n",
+    "    node = 0\n",
+    "else:\n",
+    "    node = int(node)\n",
+    "print(f\"NODE={node}\")\n",
+    "\n",
+    "global_rank = os.getenv('RANK')\n",
+    "if global_rank is None:\n",
+    "    global_rank = 0\n",
+    "else:\n",
+    "    global_rank = int(global_rank)\n",
+    "print(f\"GLOBAL RANK={global_rank}\")\n",
+    "\n",
+    "world_size = os.getenv('WORLD_SIZE')\n",
+    "if world_size is None: \n",
+    "    world_size = 1\n",
+    "else:\n",
+    "    world_size = int(world_size)\n",
+    "print(f\"WORLD_SIZE={world_size}\")\n",
+    "\n",
+    "if utils.is_interactive():\n",
+    "    # Following allows you to change functions in models.py or utils.py and \n",
+    "    # have this notebook automatically update with your revisions\n",
+    "    %load_ext autoreload\n",
+    "    %autoreload 2\n",
+    "\n",
+    "batch_size = probe_batch_size\n",
+    "num_epochs = probe_num_epochs\n",
+    "\n",
+    "data_type = torch.float32 # change depending on your mixed_precision\n",
+    "global_batch_size = batch_size * world_size\n",
+    "\n",
+    "# FSDP Setup\n",
+    "if distributed:\n",
+    "    import torch.distributed as dist\n",
+    "    import torch.multiprocessing as mp\n",
+    "    from torch.distributed.fsdp import FullyShardedDataParallel as FSDP\n",
+    "    from torch.distributed.fsdp.api import BackwardPrefetch, CPUOffload, ShardingStrategy\n",
+    "    import functools\n",
+    "    from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy, transformer_auto_wrap_policy\n",
+    "    print(f\"setting device to cuda:{local_rank}\")\n",
+    "    torch.cuda.set_device(local_rank)\n",
+    "    device = torch.device('cuda',local_rank)\n",
+    "    dist.init_process_group(\"nccl\", rank=global_rank, world_size=world_size)\n",
+    "    print(f\"\\nSuccessfully set cuda:{local_rank} | global_rank{global_rank} | node{node}\")\n",
+    "    dist.barrier()\n",
+    "    print(f\"global_rank{global_rank} passed barrier\")\n",
+    "else:\n",
+    "    device = torch.device('cuda')\n",
+    "\n",
+    "print(\"PID of this process =\",os.getpid())\n",
+    "print(\"device =\", device, \"distributed =\",distributed, \"num_devices =\", num_devices, \"local rank =\", local_rank, \"world size =\", world_size, \"data_type =\", data_type)\n",
+    "\n",
+    "# seed all random functions\n",
+    "utils.seed_everything(seed + global_rank)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ab15aca0-148e-435f-b8f2-7a708b61a6d9",
+   "metadata": {},
+   "source": [
+    "# hcp_flat"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "de1a5b87-fa69-44e1-bdb9-bd9e257485c6",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "img_size (144, 320) patch_size (16, 16) frames 16 t_patch_size 2\n",
+      "model initialized\n",
+      "\n",
+      "Loaded checkpoint epoch99.pth from /weka/proj-fmri/paulscotti/fMRI-foundation-model/flat/checkpoints/clane\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "from util.hcp_flat import load_hcp_flat_mask\n",
+    "from util.hcp_flat import create_hcp_flat\n",
+    "import util.visualize as vis\n",
+    "\n",
+    "if utils.is_interactive(): # Use less samples per epoch for debugging\n",
+    "    probe_num_samples_per_epoch = 100000\n",
+    "    test_num_samples_per_epoch = 100000\n",
+    "    num_epochs = 10\n",
+    "\n",
+    "model = mae_vit_small_fmri(\n",
+    "    patch_size=16,\n",
+    "    decoder_embed_dim=decoder_embed_dim,\n",
+    "    t_patch_size=t_patch_size,\n",
+    "    pred_t_dim=pred_t_dim,\n",
+    "    decoder_depth=4,\n",
+    "    cls_embed=cls_embed,\n",
+    "    norm_pix_loss=norm_pix_loss,\n",
+    "    no_qkv_bias=no_qkv_bias,\n",
+    "    sep_pos_embed=sep_pos_embed,\n",
+    "    trunc_init=trunc_init,\n",
+    ")\n",
+    "\n",
+    "# Load ckpt\n",
+    "if not os.path.exists(outdir) or not os.path.isdir(outdir):\n",
+    "    print(f\"\\nCheckpoint folder {outdir} does not exist.\\n\")\n",
+    "else:\n",
+    "    checkpoint_files = [f for f in os.listdir(outdir) if f.endswith('.pth')]\n",
+    "\n",
+    "    # Find the latest ckpt to load\n",
+    "    epoch_numbers = []\n",
+    "    for file in checkpoint_files:\n",
+    "        try:\n",
+    "            epoch_number = int(file.split('epoch')[-1].split('.')[0])\n",
+    "            epoch_numbers.append(epoch_number)\n",
+    "        except ValueError:\n",
+    "            continue\n",
+    "    latest_epoch = max(epoch_numbers)\n",
+    "    checkpoint_name = f\"epoch{latest_epoch}.pth\"\n",
+    "    \n",
+    "    ### Or provide the specific checkpoint you want to load\n",
+    "    # checkpoint_name = \"epoch10.pth\" #\"epoch15.pth\"\n",
+    "\n",
+    "    # Load the checkpoint\n",
+    "    checkpoint_path = os.path.join(outdir, checkpoint_name)\n",
+    "    state = torch.load(checkpoint_path)\n",
+    "    model.load_state_dict(state[\"model\"], strict=True) #model_state_dict\n",
+    "\n",
+    "    print(f\"\\nLoaded checkpoint {checkpoint_name} from {outdir}\\n\")\n",
+    "\n",
+    "model.eval()\n",
+    "model.requires_grad_(False)\n",
+    "model.to(device)\n",
+    "pass"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "c3461199-e805-4e9c-8c91-894e83cf8bc3",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Target: trial_type\n",
+      "train: (118656, 9), test: (12816, 9)\n",
+      "X_train: (118656, 384), X_test: (12816, 384)\n",
+      "classes (21): ['0bk_body' '0bk_faces' '0bk_places' '0bk_tools' '2bk_body' '2bk_faces'\n",
+      " '2bk_places' '2bk_tools' 'fear' 'lf' 'lh' 'match' 'math' 'mental' 'neut'\n",
+      " 'relation' 'rf' 'rh' 'rnd' 'story' 't']\n",
+      "\n",
+      "y_train: (118656,) [14  8 14  8 14 13 13 13 13 18 18 13 15 11 11 15 11 15 15 11]\n",
+      "y_test: (12816,) [ 7  7  0  0  5  5  3  3  4  4  6  6 17  9 20 16 10 20  9 17]\n",
+      "\n",
+      "train_ind: 106790 [ 71756  44749 102949   6302  76306  43438  88189  46933  66308 102792]\n",
+      "val_ind: 11866 [73083 36571 52192 51200  2619   399  2375  2327 41809 84134]\n",
+      "Fitting PCA projection\n",
+      "Fitting logistic regression\n",
+      "Done:\n",
+      "{\"target\": \"trial_type\", \"train_acc\": 0.9191310047757281, \"val_acc\": 0.9014832293949099, \"test_acc\": 0.899812734082397}\n"
+     ]
+    }
+   ],
+   "source": [
+    "import argparse\n",
+    "import json\n",
+    "import os\n",
+    "import pickle\n",
+    "from pathlib import Path\n",
+    "\n",
+    "import pandas as pd\n",
+    "import numpy as np\n",
+    "from sklearn.decomposition import PCA\n",
+    "from sklearn.linear_model import LogisticRegressionCV\n",
+    "from sklearn.model_selection import train_test_split\n",
+    "from sklearn.preprocessing import LabelEncoder\n",
+    "\n",
+    "target = \"trial_type\"\n",
+    "print(f\"Target: {target}\")\n",
+    "\n",
+    "train_features = pd.read_parquet(f\"{outdir}/{parquet_folder}/train.parquet\")\n",
+    "test_features = pd.read_parquet(f\"{outdir}/{parquet_folder}/test.parquet\")\n",
+    "print(f\"train: {train_features.shape}, test: {test_features.shape}\")\n",
+    "\n",
+    "X_train = np.stack(train_features[\"feature\"])\n",
+    "X_test = np.stack(test_features[\"feature\"])\n",
+    "print(f\"X_train: {X_train.shape}, X_test: {X_test.shape}\")\n",
+    "\n",
+    "if target == \"task\":\n",
+    "    labels_train = train_features[\"task\"].str.rstrip(\"1234\").values\n",
+    "    labels_test = test_features[\"task\"].str.rstrip(\"1234\").values\n",
+    "elif target == \"trial_type\":\n",
+    "    labels_train = train_features[\"trial_type\"].values\n",
+    "    labels_test = test_features[\"trial_type\"].values\n",
+    "\n",
+    "label_enc = LabelEncoder()\n",
+    "y_train = label_enc.fit_transform(labels_train)\n",
+    "y_test = label_enc.transform(labels_test)\n",
+    "\n",
+    "print(f\"classes ({len(label_enc.classes_)}): {label_enc.classes_}\")\n",
+    "print(\n",
+    "    f\"\\ny_train: {y_train.shape} {y_train[:20]}\\n\"\n",
+    "    f\"y_test: {y_test.shape} {y_test[:20]}\"\n",
+    ")\n",
+    "del train_features, test_features\n",
+    "\n",
+    "train_ind, val_ind = train_test_split(\n",
+    "    np.arange(len(X_train)), train_size=0.9, random_state=42\n",
+    ")\n",
+    "print(\n",
+    "    f\"\\ntrain_ind: {len(train_ind)} {train_ind[:10]}\\n\"\n",
+    "    f\"val_ind: {len(val_ind)} {val_ind[:10]}\"\n",
+    ")\n",
+    "X_train, X_val = X_train[train_ind], X_train[val_ind]\n",
+    "y_train, y_val = y_train[train_ind], y_train[val_ind]\n",
+    "\n",
+    "print(\"Fitting PCA projection\")\n",
+    "pca = PCA(n_components=384, whiten=True, svd_solver=\"randomized\")\n",
+    "pca.fit(X_train)\n",
+    "\n",
+    "X_train = pca.transform(X_train)\n",
+    "X_val = pca.transform(X_val)\n",
+    "X_test = pca.transform(X_test)\n",
+    "\n",
+    "print(\"Fitting logistic regression\")\n",
+    "clf = LogisticRegressionCV()\n",
+    "clf.fit(X_train, y_train)\n",
+    "\n",
+    "train_acc = clf.score(X_train, y_train)\n",
+    "val_acc = clf.score(X_val, y_val)\n",
+    "test_acc = clf.score(X_test, y_test)\n",
+    "\n",
+    "result = {\n",
+    "    \"target\": target,\n",
+    "    \"train_acc\": train_acc,\n",
+    "    \"val_acc\": val_acc,\n",
+    "    \"test_acc\": test_acc,\n",
+    "}\n",
+    "\n",
+    "print(f\"Done:\\n{json.dumps(result)}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f9d84788-0ba5-42b9-9e04-b03790997504",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "found2",
+   "language": "python",
+   "name": "found2"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.14"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

fMRI-foundation-model-old/flat/probe_prepper.ipynb

@@ -0,0 +1,587 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "b8e236f1-385a-4d93-bb39-bea3ee384d76",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "outdir /weka/proj-fmri/paulscotti/fMRI-foundation-model/flat/checkpoints/ps16_mask9_3losses_bs32\n",
+      "Loaded config.yaml from ckpt folder /weka/proj-fmri/paulscotti/fMRI-foundation-model/flat/checkpoints/ps16_mask9_3losses_bs32\n",
+      "\n",
+      "__CONFIG__\n",
+      "base_lr = 0.001\n",
+      "batch_size = 32\n",
+      "ckpt_interval = 5\n",
+      "ckpt_saving = True\n",
+      "cls_embed = True\n",
+      "decoder_embed_dim = 512\n",
+      "grad_clip = 1.0\n",
+      "hcp_flat_path = /weka/proj-medarc/shared/hcp_flat\n",
+      "mask_ratio = 0.9\n",
+      "model_name = ps16_mask9_3losses_bs32\n",
+      "no_qkv_bias = False\n",
+      "norm_pix_loss = False\n",
+      "num_epochs = 100\n",
+      "num_frames = 16\n",
+      "num_samples_per_epoch = 200000\n",
+      "num_workers = 10\n",
+      "patch_size = 16\n",
+      "pred_t_dim = 8\n",
+      "print_interval = 20\n",
+      "probe_base_lr = 0.0003\n",
+      "probe_batch_size = 8\n",
+      "probe_num_epochs = 30\n",
+      "probe_num_samples_per_epoch = 100000\n",
+      "resume_from_ckpt = True\n",
+      "seed = 42\n",
+      "sep_pos_embed = True\n",
+      "t_patch_size = 2\n",
+      "test_num_samples_per_epoch = 50000\n",
+      "trunc_init = False\n",
+      "use_contrastive_loss = True\n",
+      "wandb_log = True\n",
+      "\n",
+      "\n",
+      "Number of available CUDA devices: 1\n",
+      "LOCAL RANK=0\n",
+      "NUM GPUS=1\n",
+      "NODE=0\n",
+      "GLOBAL RANK=0\n",
+      "WORLD_SIZE=1\n",
+      "PID of this process = 2361819\n",
+      "device = cuda distributed = False num_devices = 1 local rank = 0 world size = 1 data_type = torch.float32\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Import packages and setup gpu configuration.\n",
+    "# This code block shouldnt need to be adjusted!\n",
+    "import os\n",
+    "import sys\n",
+    "import json\n",
+    "import yaml\n",
+    "import numpy as np\n",
+    "import copy\n",
+    "import math\n",
+    "import time\n",
+    "import datetime\n",
+    "import random\n",
+    "from tqdm import tqdm\n",
+    "import webdataset as wds\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "from torchvision import transforms\n",
+    "import utils\n",
+    "from flat_models import *\n",
+    "\n",
+    "from elbow.sinks import BufferedParquetWriter\n",
+    "\n",
+    "# tf32 data type is faster than standard float32\n",
+    "torch.backends.cuda.matmul.allow_tf32 = True\n",
+    "# following fixes a Conv3D CUDNN_NOT_SUPPORTED error\n",
+    "torch.backends.cudnn.benchmark = True\n",
+    "\n",
+    "## MODEL TO LOAD ##\n",
+    "model_name = \"ps16_mask9_3losses_bs32\"\n",
+    "outdir = os.path.abspath(f'checkpoints/{model_name}')\n",
+    "print(\"outdir\", outdir)\n",
+    "# Load previous config.yaml if available\n",
+    "if os.path.exists(f\"{outdir}/config.yaml\"):\n",
+    "    config = yaml.load(open(f\"{outdir}/config.yaml\", 'r'), Loader=yaml.FullLoader)\n",
+    "    print(f\"Loaded config.yaml from ckpt folder {outdir}\")\n",
+    "    # create global variables from the config\n",
+    "    print(\"\\n__CONFIG__\")\n",
+    "    for attribute_name in config.keys():\n",
+    "        print(f\"{attribute_name} = {config[attribute_name]}\")\n",
+    "        globals()[attribute_name] = config[f'{attribute_name}']\n",
+    "    print(\"\\n\")\n",
+    "\n",
+    "### Multi-GPU config ###\n",
+    "device_count = torch.cuda.device_count()\n",
+    "print(f\"Number of available CUDA devices: {device_count}\")\n",
+    "\n",
+    "local_rank = os.getenv('LOCAL_RANK')\n",
+    "if local_rank is None: \n",
+    "    local_rank = 0\n",
+    "else:\n",
+    "    local_rank = int(local_rank)\n",
+    "print(f\"LOCAL RANK={local_rank}\")\n",
+    "\n",
+    "num_devices = os.getenv('NUM_GPUS')\n",
+    "if num_devices is None: \n",
+    "    num_devices = 1\n",
+    "else:\n",
+    "    num_devices = int(num_devices)\n",
+    "print(f\"NUM GPUS={num_devices}\")\n",
+    "distributed = True if num_devices>1 else False\n",
+    "if distributed: assert device_count==num_devices\n",
+    "\n",
+    "node = os.getenv('SLURM_NODEID')\n",
+    "if node is None:\n",
+    "    node = 0\n",
+    "else:\n",
+    "    node = int(node)\n",
+    "print(f\"NODE={node}\")\n",
+    "\n",
+    "global_rank = os.getenv('RANK')\n",
+    "if global_rank is None:\n",
+    "    global_rank = 0\n",
+    "else:\n",
+    "    global_rank = int(global_rank)\n",
+    "print(f\"GLOBAL RANK={global_rank}\")\n",
+    "\n",
+    "world_size = os.getenv('WORLD_SIZE')\n",
+    "if world_size is None: \n",
+    "    world_size = 1\n",
+    "else:\n",
+    "    world_size = int(world_size)\n",
+    "print(f\"WORLD_SIZE={world_size}\")\n",
+    "\n",
+    "if utils.is_interactive():\n",
+    "    # Following allows you to change functions in models.py or utils.py and \n",
+    "    # have this notebook automatically update with your revisions\n",
+    "    %load_ext autoreload\n",
+    "    %autoreload 2\n",
+    "\n",
+    "batch_size = probe_batch_size\n",
+    "num_epochs = probe_num_epochs\n",
+    "hcp_flat_path = \"/weka/proj-medarc/shared/hcp_flat\"\n",
+    "\n",
+    "data_type = torch.float32 # change depending on your mixed_precision\n",
+    "global_batch_size = batch_size * world_size\n",
+    "\n",
+    "device = torch.device('cuda')\n",
+    "\n",
+    "print(\"PID of this process =\",os.getpid())\n",
+    "print(\"device =\", device, \"distributed =\",distributed, \"num_devices =\", num_devices, \"local rank =\", local_rank, \"world size =\", world_size, \"data_type =\", data_type)\n",
+    "\n",
+    "# seed all random functions\n",
+    "utils.seed_everything(seed + global_rank)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ab15aca0-148e-435f-b8f2-7a708b61a6d9",
+   "metadata": {},
+   "source": [
+    "# hcp_flat"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "de1a5b87-fa69-44e1-bdb9-bd9e257485c6",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "img_size (144, 320) patch_size (16, 16) frames 16 t_patch_size 2\n",
+      "model initialized\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "MaskedAutoencoderViT(\n",
+       "  (patch_embed): PatchEmbed(\n",
+       "    (proj): Conv3d(1, 384, kernel_size=(2, 16, 16), stride=(2, 16, 16))\n",
+       "  )\n",
+       "  (blocks): ModuleList(\n",
+       "    (0-11): 12 x Block(\n",
+       "      (norm1): LayerNorm((384,), eps=1e-06, elementwise_affine=True)\n",
+       "      (attn): Attention(\n",
+       "        (q): Linear(in_features=384, out_features=384, bias=True)\n",
+       "        (k): Linear(in_features=384, out_features=384, bias=True)\n",
+       "        (v): Linear(in_features=384, out_features=384, bias=True)\n",
+       "        (proj): Linear(in_features=384, out_features=384, bias=True)\n",
+       "        (proj_drop): Dropout(p=0.0, inplace=False)\n",
+       "      )\n",
+       "      (drop_path): Identity()\n",
+       "      (norm2): LayerNorm((384,), eps=1e-06, elementwise_affine=True)\n",
+       "      (mlp): Mlp(\n",
+       "        (fc1): Linear(in_features=384, out_features=1536, bias=True)\n",
+       "        (act): GELU(approximate='none')\n",
+       "        (drop1): Dropout(p=0.0, inplace=False)\n",
+       "        (norm): Identity()\n",
+       "        (fc2): Linear(in_features=1536, out_features=384, bias=True)\n",
+       "        (drop2): Dropout(p=0.0, inplace=False)\n",
+       "      )\n",
+       "    )\n",
+       "  )\n",
+       "  (norm): LayerNorm((384,), eps=1e-06, elementwise_affine=True)\n",
+       "  (decoder_embed): Linear(in_features=384, out_features=512, bias=True)\n",
+       "  (decoder_blocks): ModuleList(\n",
+       "    (0-3): 4 x Block(\n",
+       "      (norm1): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
+       "      (attn): Attention(\n",
+       "        (q): Linear(in_features=512, out_features=512, bias=True)\n",
+       "        (k): Linear(in_features=512, out_features=512, bias=True)\n",
+       "        (v): Linear(in_features=512, out_features=512, bias=True)\n",
+       "        (proj): Linear(in_features=512, out_features=512, bias=True)\n",
+       "        (proj_drop): Dropout(p=0.0, inplace=False)\n",
+       "      )\n",
+       "      (drop_path): Identity()\n",
+       "      (norm2): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
+       "      (mlp): Mlp(\n",
+       "        (fc1): Linear(in_features=512, out_features=2048, bias=True)\n",
+       "        (act): GELU(approximate='none')\n",
+       "        (drop1): Dropout(p=0.0, inplace=False)\n",
+       "        (norm): Identity()\n",
+       "        (fc2): Linear(in_features=2048, out_features=512, bias=True)\n",
+       "        (drop2): Dropout(p=0.0, inplace=False)\n",
+       "      )\n",
+       "    )\n",
+       "  )\n",
+       "  (decoder_norm): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
+       "  (decoder_pred): Linear(in_features=512, out_features=256, bias=True)\n",
+       ")"
+      ]
+     },
+     "execution_count": 2,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from util.hcp_flat import load_hcp_flat_mask\n",
+    "from util.hcp_flat import create_hcp_flat\n",
+    "import util.visualize as vis\n",
+    "\n",
+    "model = mae_vit_small_fmri(\n",
+    "    patch_size=patch_size,\n",
+    "    decoder_embed_dim=decoder_embed_dim,\n",
+    "    t_patch_size=t_patch_size,\n",
+    "    pred_t_dim=pred_t_dim,\n",
+    "    decoder_depth=4,\n",
+    "    cls_embed=cls_embed,\n",
+    "    norm_pix_loss=norm_pix_loss,\n",
+    "    no_qkv_bias=no_qkv_bias,\n",
+    "    sep_pos_embed=sep_pos_embed,\n",
+    "    trunc_init=trunc_init,\n",
+    ")\n",
+    "\n",
+    "model.eval()\n",
+    "model.requires_grad_(False)\n",
+    "model.to(device)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3dd51ddf-fb71-48f4-bdbd-88753b44d2aa",
+   "metadata": {},
+   "source": [
+    "## Create dataset and data loaders"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "3e330cd2-8f4f-4ca7-ae68-16698e90060f",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "changed batch_size to 64\n"
+     ]
+    }
+   ],
+   "source": [
+    "batch_size = 64 #12\n",
+    "print(f\"changed batch_size to {batch_size}\")\n",
+    "\n",
+    "## Train ##\n",
+    "train_dataset = create_hcp_flat(root=hcp_flat_path, \n",
+    "                    split=\"train\", frames=num_frames, \n",
+    "                    clip_mode=\"event\", shuffle=False)\n",
+    "train_dl = wds.WebLoader(\n",
+    "    train_dataset.batched(batch_size, partial=False),\n",
+    "    batch_size=None,\n",
+    "    shuffle=False,\n",
+    "    num_workers=num_workers,\n",
+    "    pin_memory=True,\n",
+    ")\n",
+    "\n",
+    "## Test ##\n",
+    "test_dataset = create_hcp_flat(root=hcp_flat_path, \n",
+    "                    split=\"test\", frames=num_frames, \n",
+    "                    clip_mode=\"event\", shuffle=False)\n",
+    "test_dl = wds.WebLoader(\n",
+    "    test_dataset.batched(batch_size, partial=False),\n",
+    "    batch_size=None,\n",
+    "    shuffle=False,\n",
+    "    num_workers=num_workers,\n",
+    "    pin_memory=True,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2b8e6baa-4b1c-4f38-b078-70b2b092d14d",
+   "metadata": {},
+   "source": [
+    "# Load checkpoint"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "4da73c08-ca61-48ef-9e63-b70db6f07a59",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Loaded checkpoint epoch99.pth from /weka/proj-fmri/paulscotti/fMRI-foundation-model/flat/checkpoints/ps16_mask9_3losses_bs32\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "if not os.path.exists(outdir) or not os.path.isdir(outdir):\n",
+    "    print(f\"\\nCheckpoint folder {outdir} does not exist.\\n\")\n",
+    "    err\n",
+    "else:\n",
+    "    checkpoint_files = [f for f in os.listdir(outdir) if f.endswith('.pth')]\n",
+    "\n",
+    "    # Extract epoch numbers and find the largest\n",
+    "    epoch_numbers = []\n",
+    "    for file in checkpoint_files:\n",
+    "        try:\n",
+    "            epoch_number = int(file.split('epoch')[-1].split('.')[0])\n",
+    "            epoch_numbers.append(epoch_number)\n",
+    "        except ValueError:\n",
+    "            continue\n",
+    "    latest_epoch = max(epoch_numbers)\n",
+    "    latest_checkpoint = f\"epoch{latest_epoch}.pth\"\n",
+    "    \n",
+    "    # # Or specify epoch number \n",
+    "    # latest_checkpoint = \"epoch15.pth\"\n",
+    "\n",
+    "    # Load the checkpoint\n",
+    "    checkpoint_path = os.path.join(outdir, latest_checkpoint)\n",
+    "\n",
+    "    state = torch.load(checkpoint_path)\n",
+    "    model.load_state_dict(state[\"model_state_dict\"], strict=True)\n",
+    "\n",
+    "    print(f\"\\nLoaded checkpoint {latest_checkpoint} from {outdir}\\n\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c43a5055-8afd-468a-93bf-32f94bd1d042",
+   "metadata": {},
+   "source": [
+    "# Start extraction"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "bf1e2df3-6c55-4218-89ea-348c5aa4f8b8",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "/weka/proj-fmri/paulscotti/fMRI-foundation-model/flat/checkpoints/ps16_mask9_3losses_bs32/epoch99\n"
+     ]
+    }
+   ],
+   "source": [
+    "outdir_parquet = f'{outdir}/{latest_checkpoint[:-4]}'\n",
+    "print(outdir_parquet)\n",
+    "\n",
+    "os.makedirs(outdir_parquet,exist_ok=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "1e1bbe95-5fbf-4f8d-bc7a-17d5983a94d3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# for i, (samples, samples_meta) in enumerate(tqdm(train_dl)):\n",
+    "#     samples = samples.to(device, non_blocking=True)\n",
+    "## using batch_size=12: 49284it [14:03, 58.44it/s]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "6507067f-1ece-46ac-ad8c-24ef7a9b3a58",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "@torch.no_grad()\n",
+    "def extract_features(dl):\n",
+    "    for samples, samples_meta in tqdm(dl,total=49284):\n",
+    "        samples = samples.to(device, non_blocking=True)\n",
+    "\n",
+    "        features = model.forward_features(samples)\n",
+    "\n",
+    "        features = features.cpu().numpy()\n",
+    "\n",
+    "        for feat, meta in zip(features, samples_meta):\n",
+    "            yield {\"feature\": feat, **meta}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "c3461199-e805-4e9c-8c91-894e83cf8bc3",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Start extract\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  0%|▏                                    | 197/49284 [00:29<2:03:32,  6.62it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Extract time 0:00:29\n",
+      "389945344\n"
+     ]
+    }
+   ],
+   "source": [
+    "utils.seed_everything(seed)\n",
+    "\n",
+    "print(\"Start extract\")\n",
+    "start_time = time.time()\n",
+    "\n",
+    "with BufferedParquetWriter(f\"{outdir_parquet}/test.parquet\", blocking=True) as writer:\n",
+    "    for sample in extract_features(test_dl):\n",
+    "        writer.write(sample)\n",
+    "\n",
+    "total_time = time.time() - start_time\n",
+    "total_time_str = str(datetime.timedelta(seconds=int(total_time)))\n",
+    "print(\"Extract time {}\".format(total_time_str))\n",
+    "print(torch.cuda.memory_allocated())"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "a898429c-b013-4989-86ef-fbf75ca15971",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Start extract\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  4%|█▎                                  | 1849/49284 [03:59<1:42:32,  7.71it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Extract time 0:03:59\n",
+      "389945344\n"
+     ]
+    }
+   ],
+   "source": [
+    "utils.seed_everything(seed)\n",
+    "\n",
+    "print(\"Start extract\")\n",
+    "start_time = time.time()\n",
+    "\n",
+    "with BufferedParquetWriter(f\"{outdir_parquet}/train.parquet\", blocking=True) as writer:\n",
+    "    for sample in extract_features(train_dl):\n",
+    "        writer.write(sample)\n",
+    "\n",
+    "total_time = time.time() - start_time\n",
+    "total_time_str = str(datetime.timedelta(seconds=int(total_time)))\n",
+    "print(\"Extract time {}\".format(total_time_str))\n",
+    "print(torch.cuda.memory_allocated())"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6ffb97e6-c61c-4026-82fb-b1729f51c40a",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2404f355-090f-48b7-8b04-5c76e7bb7e04",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "found",
+   "language": "python",
+   "name": "found"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.14"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

fMRI-foundation-model-old/flat/unique_tasks.npy

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

fMRI-foundation-model-old/flat/util/.ipynb_checkpoints/hcp_flat-checkpoint.py

@@ -0,0 +1,120 @@
+import os
+import random
+from functools import partial
+from pathlib import Path
+from typing import Any, Dict, Iterable, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import webdataset as wds
+from torch.utils.data import IterableDataset
+
+HCP_FLAT_ROOT = "https://huggingface.co/datasets/bold-ai/HCP-Flat/resolve/main"
+NUM_SHARDS = {"train": 1629, "test": 174}
+
+
+def create_hcp_flat(
+    root: Optional[str] = None,
+    training: bool = True,
+    shards: Optional[Union[int, Iterable[int]]] = None,
+    frames: int = 16,
+    cache_dir: Optional[str] = None,
+) -> wds.WebDataset:
+    """
+    Create HCP-Flat dataset. Yields samples of (key, images) where key is the webdataset
+    sample key and images is shape (C, T, H, W).
+
+    References:
+        https://github.com/webdataset/webdataset/issues/250#issuecomment-1454094496
+        https://github.com/tmbdev-archive/webdataset-imagenet-2/blob/main/imagenet.py
+        https://github.com/huggingface/pytorch-image-models/blob/main/timm/data/readers/reader_wds.py
+    """
+    root = root or os.environ.get("HCP_FLAT_ROOT") or HCP_FLAT_ROOT
+    split = "train" if training else "test"
+
+    shards = shards or NUM_SHARDS[split]
+    if isinstance(shards, int):
+        shards = range(shards)
+    assert (
+        min(shards) >= 0 and max(shards) < NUM_SHARDS[split]
+    ), f"Invalid shards {shards}; expected in [0, {NUM_SHARDS[split]})"
+
+    urls = [f"{root}/{split}/hcp-flat_{split}_{shard:06d}.tar" for shard in shards]
+
+    # In training, we resample shards with replacement independently in every worker and
+    # yield batches up to the target number of samples. In test, we iterate over the
+    # shards in order, with workers getting interleaving shards, and yield batches up to
+    # the target samples. In a distributed setting with variable size shards, setting a
+    # fixed number of samples is the easiest way to get balanced batches per worker. In
+    # training we will still see all data eventually. But in test, it means we cut off
+    # some data.
+
+    # Note that in training this does not do deterministic shuffling, which we would
+    # need for exact reproducibility. They get determistic shuffling in timm, but it's
+    # more complicated.
+
+    # Note that we are splitting the long timeseries into clips after shuffling, which
+    # means clips from the same series will appear consecutively in the batch(es).
+    # I think this is not too bad. It is basically equivalent to training on longer
+    # sequences. Clipping before shuffling might be preferred, but it results in a bad
+    # system memory leak (https://github.com/webdataset/webdataset/issues/354).
+    dataset = (
+        wds.WebDataset(
+            urls,
+            resampled=training,
+            shardshuffle=1000 if training else False,
+            nodesplitter=wds.split_by_node,
+            select_files=select_extensions(("bold.npy","meta.json")),
+            cache_dir=cache_dir,
+        )
+        .shuffle(1000 if training else 0)
+        .decode()
+        .map(partial(extract_images, mask=load_hcp_flat_mask()))
+        .compose(to_clips(frames))
+    )
+    return dataset
+
+
+def select_extensions(extensions: Tuple[str, ...]):
+    extensions_set = set(extensions)
+    def select_files(fname: str):
+        suffix = ".".join(fname.split(".")[1:])
+        return suffix in extensions_set
+    return select_files
+
+
+def extract_images(sample: Dict[str, Any], mask: torch.Tensor):
+    key = sample["__key__"]
+    images = sample["bold.npy"]
+    meta = sample["meta.json"]['task']
+    
+    images = torch.from_numpy(images) / 255.0
+    images = (images - 0.5) / 0.2
+    images = unmask(images, mask)
+    # (C, T, H, W,)
+    images = images.unsqueeze(0)
+    
+    return key, images, meta
+
+
+def unmask(images: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+    images_unmasked = torch.zeros(
+        (images.shape[0], *mask.shape), dtype=images.dtype, device=images.device
+    )
+    images_unmasked[:, mask] = images
+    return images_unmasked
+
+
+def to_clips(frames: int = 16):
+    def _filter(src: IterableDataset[Tuple[str, torch.Tensor, str]]):
+        for key, images, meta in src:
+            offset = random.randint(0, frames)
+            for start in range(offset, images.shape[1] - frames, frames):
+                yield key, images[:, start : start + frames], meta
+    return _filter
+
+
+def load_hcp_flat_mask() -> torch.Tensor:
+    mask = np.load(Path(__file__).parents[1] / "hcp-flat_mask.npy")
+    mask = torch.as_tensor(mask)
+    return mask

fMRI-foundation-model-old/flat/util/.ipynb_checkpoints/misc-checkpoint.py

@@ -0,0 +1,508 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# References:
+# DeiT: https://github.com/facebookresearch/deit
+# BEiT: https://github.com/microsoft/unilm/tree/master/beit
+# --------------------------------------------------------
+
+import builtins
+import datetime
+import math
+import os
+import time
+from collections import defaultdict, deque, OrderedDict
+
+import util.logging as logging
+import psutil
+import torch
+import torch.distributed as dist
+from iopath.common.file_io import g_pathmgr as pathmgr
+from util.logging import master_print as print
+from torch import inf
+
+
+logger = logging.get_logger(__name__)
+
+
+class SmoothedValue:
+    """Track a series of values and provide access to smoothed values over a
+    window or the global series average.
+    """
+
+    def __init__(self, window_size=20, fmt=None):
+        if fmt is None:
+            fmt = "{median:.4f} ({global_avg:.4f})"
+        self.deque = deque(maxlen=window_size)
+        self.total = 0.0
+        self.count = 0
+        self.fmt = fmt
+
+    def update(self, value, n=1):
+        self.deque.append(value)
+        self.count += n
+        self.total += value * n
+
+    def synchronize_between_processes(self):
+        """
+        Warning: does not synchronize the deque!
+        """
+        if not is_dist_avail_and_initialized():
+            return
+        t = torch.tensor([self.count, self.total], dtype=torch.float64, device="cuda")
+        dist.barrier()
+        dist.all_reduce(t)
+        t = t.tolist()
+        self.count = int(t[0])
+        self.total = t[1]
+
+    @property
+    def median(self):
+        d = torch.tensor(list(self.deque))
+        return d.median().item()
+
+    @property
+    def avg(self):
+        d = torch.tensor(list(self.deque), dtype=torch.float32)
+        return d.mean().item()
+
+    @property
+    def global_avg(self):
+        return self.total / self.count
+
+    @property
+    def max(self):
+        return max(self.deque)
+
+    @property
+    def value(self):
+        return self.deque[-1]
+
+    def __str__(self):
+        return self.fmt.format(
+            median=self.median,
+            avg=self.avg,
+            global_avg=self.global_avg,
+            max=self.max,
+            value=self.value,
+        )
+
+
+class MetricLogger:
+    def __init__(self, delimiter="\t"):
+        self.meters = defaultdict(SmoothedValue)
+        self.delimiter = delimiter
+
+    def update(self, **kwargs):
+        for k, v in kwargs.items():
+            if v is None:
+                continue
+            if isinstance(v, torch.Tensor):
+                v = v.item()
+            assert isinstance(v, (float, int))
+            self.meters[k].update(v)
+
+    def __getattr__(self, attr):
+        if attr in self.meters:
+            return self.meters[attr]
+        if attr in self.__dict__:
+            return self.__dict__[attr]
+        raise AttributeError(
+            "'{}' object has no attribute '{}'".format(type(self).__name__, attr)
+        )
+
+    def __str__(self):
+        loss_str = []
+        for name, meter in self.meters.items():
+            loss_str.append("{}: {}".format(name, str(meter)))
+        return self.delimiter.join(loss_str)
+
+    def synchronize_between_processes(self):
+        for meter in self.meters.values():
+            meter.synchronize_between_processes()
+
+    def add_meter(self, name, meter):
+        self.meters[name] = meter
+
+    def log_every(self, iterable, print_freq, header=None, total_steps=None):
+        i = 0
+        total_steps = total_steps or len(iterable)
+        if not header:
+            header = ""
+        start_time = time.time()
+        end = time.time()
+        iter_time = SmoothedValue(fmt="{avg:.4f}")
+        data_time = SmoothedValue(fmt="{avg:.4f}")
+        space_fmt = ":" + str(len(str(total_steps))) + "d"
+        log_msg = [
+            header,
+            "[{0" + space_fmt + "}/{1}]",
+            "eta: {eta}",
+            "{meters}",
+            "time: {time}",
+            "data: {data}",
+        ]
+        if torch.cuda.is_available():
+            log_msg.append("max mem: {memory:.0f}")
+        log_msg = self.delimiter.join(log_msg)
+        MB = 1024.0 * 1024.0
+        for obj in iterable:
+            data_time.update(time.time() - end)
+            yield obj
+            iter_time.update(time.time() - end)
+            if i % print_freq == 0 or i == total_steps - 1:
+                eta_seconds = iter_time.global_avg * (total_steps - i)
+                eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
+                if torch.cuda.is_available():
+                    print(
+                        log_msg.format(
+                            i,
+                            total_steps,
+                            eta=eta_string,
+                            meters=str(self),
+                            time=str(iter_time),
+                            data=str(data_time),
+                            memory=torch.cuda.max_memory_allocated() / MB,
+                        )
+                    )
+
+                else:
+                    print(
+                        log_msg.format(
+                            i,
+                            total_steps,
+                            eta=eta_string,
+                            meters=str(self),
+                            time=str(iter_time),
+                            data=str(data_time),
+                        )
+                    )
+            i += 1
+            end = time.time()
+        total_time = time.time() - start_time
+        total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+        print(
+            "{} Total time: {} ({:.4f} s / it)".format(
+                header, total_time_str, total_time / total_steps
+            )
+        )
+
+
+def setup_for_distributed(is_master):
+    """
+    This function disables printing when not in master process
+    """
+    builtin_print = builtins.print
+
+    def print(*args, **kwargs):
+        force = kwargs.pop("force", False)
+        force = force or (get_world_size() > 8)
+        if is_master or force:
+            now = datetime.datetime.now().time()
+            builtin_print("[{}] ".format(now), end="")  # print with time stamp
+            builtin_print(*args, **kwargs)
+
+    builtins.print = print
+
+
+def is_dist_avail_and_initialized():
+    if not dist.is_available():
+        return False
+    if not dist.is_initialized():
+        return False
+    return True
+
+
+def get_world_size():
+    if not is_dist_avail_and_initialized():
+        return 1
+    return dist.get_world_size()
+
+
+def get_rank():
+    if not is_dist_avail_and_initialized():
+        return 0
+    return dist.get_rank()
+
+
+def is_main_process():
+    return get_rank() == 0
+
+
+def save_on_master(state, path):
+    if is_main_process():
+        print(f"save path {path}")
+        with pathmgr.open(path, "wb") as f:
+            torch.save(state, f)
+
+
+def init_distributed_mode(args):
+    if args.no_env:
+        pass
+    elif args.dist_on_itp:
+        args.rank = int(os.environ["OMPI_COMM_WORLD_RANK"])
+        args.world_size = int(os.environ["OMPI_COMM_WORLD_SIZE"])
+        args.gpu = int(os.environ["OMPI_COMM_WORLD_LOCAL_RANK"])
+        args.dist_url = "tcp://%s:%s" % (
+            os.environ["MASTER_ADDR"],
+            os.environ["MASTER_PORT"],
+        )
+        os.environ["LOCAL_RANK"] = str(args.gpu)
+        os.environ["RANK"] = str(args.rank)
+        os.environ["WORLD_SIZE"] = str(args.world_size)
+        # ["RANK", "WORLD_SIZE", "MASTER_ADDR", "MASTER_PORT", "LOCAL_RANK"]
+    elif "RANK" in os.environ and "WORLD_SIZE" in os.environ:
+        args.rank = int(os.environ["RANK"])
+        args.world_size = int(os.environ["WORLD_SIZE"])
+        args.gpu = int(os.environ["LOCAL_RANK"])
+    elif "SLURM_PROCID" in os.environ:
+        args.rank = int(os.environ["SLURM_PROCID"])
+        args.gpu = args.rank % torch.cuda.device_count()
+    else:
+        print("Not using distributed mode")
+        setup_for_distributed(is_master=True)  # hack
+        args.distributed = False
+        return
+
+    args.distributed = True
+
+    torch.cuda.set_device(args.gpu)
+    args.dist_backend = "nccl"
+    print(
+        "| distributed init (rank {}): {}, gpu {}".format(
+            args.rank, args.dist_url, args.gpu
+        ),
+        # flush=True,
+    )
+    torch.distributed.init_process_group(
+        backend=args.dist_backend,
+        world_size=args.world_size,
+        rank=args.rank,
+    )
+    # init_method=args.dist_url,
+    torch.distributed.barrier()
+    setup_for_distributed(args.rank == 0)
+
+
+class NativeScalerWithGradNormCount:
+    state_dict_key = "amp_scaler"
+
+    def __init__(self, fp32=False):
+        self._scaler = torch.cuda.amp.GradScaler(enabled=not fp32)
+
+    def __call__(
+        self,
+        loss,
+        optimizer,
+        clip_grad=None,
+        parameters=None,
+        create_graph=False,
+        update_grad=True,
+    ):
+        self._scaler.scale(loss).backward(create_graph=create_graph)
+        if update_grad:
+            if clip_grad is not None:
+                assert parameters is not None
+                self._scaler.unscale_(
+                    optimizer
+                )  # unscale the gradients of optimizer's assigned params in-place
+                norm = torch.nn.utils.clip_grad_norm_(parameters, clip_grad)
+            else:
+                self._scaler.unscale_(optimizer)
+                norm = get_grad_norm_(parameters)
+            self._scaler.step(optimizer)
+            self._scaler.update()
+        else:
+            norm = None
+        return norm
+
+    def state_dict(self):
+        return self._scaler.state_dict()
+
+    def load_state_dict(self, state_dict):
+        self._scaler.load_state_dict(state_dict)
+
+
+def get_grad_norm_(parameters, norm_type: float = 2.0) -> torch.Tensor:
+    if isinstance(parameters, torch.Tensor):
+        parameters = [parameters]
+    parameters = [p for p in parameters if p.grad is not None]
+    norm_type = float(norm_type)
+    if len(parameters) == 0:
+        return torch.tensor(0.0)
+    device = parameters[0].grad.device
+    if norm_type == inf:
+        total_norm = max(p.grad.detach().abs().max().to(device) for p in parameters)
+    else:
+        total_norm = torch.norm(
+            torch.stack(
+                [torch.norm(p.grad.detach(), norm_type).to(device) for p in parameters]
+            ),
+            norm_type,
+        )
+    return total_norm
+
+
+def save_model(args, epoch, model, model_without_ddp, optimizer, loss_scaler):
+    checkpoint_path = "{}/checkpoint-{:05d}.pth".format(args.output_dir, epoch)
+    to_save = {
+        "model": model_without_ddp.state_dict(),
+        "optimizer": optimizer.state_dict(),
+        "epoch": epoch,
+        "scaler": loss_scaler.state_dict(),
+        "args": args,
+    }
+
+    save_on_master(to_save, checkpoint_path)
+    return checkpoint_path
+
+
+def get_last_checkpoint(args):
+    """
+    Get the last checkpoint from the checkpointing folder.
+    Args:
+        path_to_job (string): the path to the folder of the current job.
+    """
+    d = args.output_dir
+    names = pathmgr.ls(d) if pathmgr.exists(d) else []
+    names = [f for f in names if "checkpoint" in f]
+    if len(names) == 0:
+        print("No checkpoints found in '{}'.".format(d))
+        return None
+    else:
+        # Sort the checkpoints by epoch.
+        name = sorted(names)[-1]
+        return os.path.join(d, name)
+
+
+def load_model(args, model_without_ddp, optimizer, loss_scaler):
+    if not args.resume:
+        args.resume = get_last_checkpoint(args)
+    if args.resume:
+        if args.resume.startswith("https"):
+            checkpoint = torch.hub.load_state_dict_from_url(
+                args.resume, map_location="cpu", check_hash=True
+            )
+        else:
+            with pathmgr.open(args.resume, "rb") as f:
+                checkpoint = torch.load(f, map_location="cpu")
+        model_without_ddp.load_state_dict(checkpoint["model"])
+        print("Resume checkpoint %s" % args.resume)
+        if (
+            "optimizer" in checkpoint
+            and "epoch" in checkpoint
+            and not (hasattr(args, "eval") and args.eval)
+        ):
+            optimizer.load_state_dict(checkpoint["optimizer"])
+            args.start_epoch = checkpoint["epoch"] + 1
+            if "scaler" in checkpoint:
+                loss_scaler.load_state_dict(checkpoint["scaler"])
+            print("With optim & sched!")
+
+
+def all_reduce_mean(x):
+    world_size = get_world_size()
+    if world_size > 1:
+        x_reduce = torch.tensor(x).cuda()
+        dist.all_reduce(x_reduce)
+        x_reduce /= world_size
+        return x_reduce.item()
+    else:
+        return x
+
+
+def gpu_mem_usage():
+    """
+    Compute the GPU memory usage for the current device (GB).
+    """
+    if torch.cuda.is_available():
+        mem_usage_bytes = torch.cuda.max_memory_allocated()
+    else:
+        mem_usage_bytes = 0
+    return mem_usage_bytes / 1024**3
+
+
+def cpu_mem_usage():
+    """
+    Compute the system memory (RAM) usage for the current device (GB).
+    Returns:
+        usage (float): used memory (GB).
+        total (float): total memory (GB).
+    """
+    vram = psutil.virtual_memory()
+    usage = (vram.total - vram.available) / 1024**3
+    total = vram.total / 1024**3
+
+    return usage, total
+
+
+def all_gather(tensors):
+    """
+    All gathers the provided tensors from all processes across machines.
+    Args:
+        tensors (list): tensors to perform all gather across all processes in
+        all machines.
+    """
+
+    gather_list = []
+    output_tensor = []
+    world_size = dist.get_world_size()
+    for tensor in tensors:
+        tensor_placeholder = [torch.ones_like(tensor) for _ in range(world_size)]
+        dist.all_gather(tensor_placeholder, tensor, async_op=False)
+        gather_list.append(tensor_placeholder)
+    for gathered_tensor in gather_list:
+        output_tensor.append(torch.cat(gathered_tensor, dim=0))
+    return output_tensor
+
+
+def add_weight_decay(model, weight_decay=1e-5, skip_list=(), bias_wd=False):
+    decay = []
+    no_decay = []
+    for name, param in model.named_parameters():
+        if not param.requires_grad:
+            continue  # frozen weights
+        if (
+            (not bias_wd)
+            and len(param.shape) == 1
+            or name.endswith(".bias")
+            or name in skip_list
+        ):
+            no_decay.append(param)
+        else:
+            decay.append(param)
+    return [
+        {"params": no_decay, "weight_decay": 0.0},
+        {"params": decay, "weight_decay": weight_decay},
+    ]
+
+
+def inflate(model_2d, model_3d):
+    state_dict_inflated = OrderedDict()
+    for k, v2d in model_2d.items():
+        if "patch_embed.proj.weight" in k:
+            v3d = model_3d[k]
+            v3d = v2d.unsqueeze(2).repeat(1, 1, v3d.shape[2], 1, 1) / v3d.shape[2]
+            state_dict_inflated[k] = v3d.clone()
+        elif "pos_embed" in k:
+            pos_embed_cls, pos_embed_spatial = torch.split(v2d, [1, 196], dim=1)
+            state_dict_inflated["pos_embed_cls"] = pos_embed_cls.clone()
+            state_dict_inflated["pos_embed"] = pos_embed_spatial.clone()
+        else:
+            state_dict_inflated[k] = v2d.clone()
+    return state_dict_inflated
+
+
+def convert_checkpoint(model_2d):
+    state_dict_inflated = OrderedDict()
+    for k, v2d in model_2d.items():
+        if "head.projection.weight" in k:
+            state_dict_inflated["head.weight"] = v2d.clone()
+        elif "head.projection.bias" in k:
+            state_dict_inflated["head.bias"] = v2d.clone()
+        else:
+            state_dict_inflated[k] = v2d.clone()
+    return state_dict_inflated

fMRI-foundation-model-old/flat/util/.ipynb_checkpoints/visualize-checkpoint.py

@@ -0,0 +1,123 @@
+import io
+from typing import Any, Optional
+
+import torch
+from matplotlib import pyplot as plt
+from matplotlib.figure import Figure
+from PIL import Image
+
+plt.rcParams["figure.dpi"] = 150
+
+
+def plot_mask_pred(
+    model: torch.nn.Module,
+    imgs: torch.Tensor,
+    pred: torch.Tensor,
+    mask: torch.Tensor,
+    mean: Optional[Any] = None,
+    std: Optional[Any] = None,
+    nrow: int = 8,
+    norm_pix_loss: bool = False,
+):
+    # imgs: [N, C, T, H, W]
+    # pred: [N, t*h*w, u*p*p*C]
+    # mask: [N, t*h*w], 0 is keep, 1 is remove,
+    target = torch.index_select(
+        imgs,
+        2,
+        torch.linspace(
+            0,
+            imgs.shape[2] - 1,
+            model.pred_t_dim,
+        )
+        .long()
+        .to(imgs.device),
+    )
+    if norm_pix_loss:
+        target_mean = target.mean(dim=-1, keepdim=True)
+        target_var = target.var(dim=-1, keepdim=True)
+    target = torch.einsum("ncthw->nthwc", target)
+    target = target.flatten(0, 1)[:nrow].cpu()
+
+    pred = pred.detach()
+    if norm_pix_loss:
+        pred = (pred * target_var) + target_mean
+    pred = model.unpatchify(pred)
+    pred = torch.einsum("ncthw->nthwc", pred).cpu()
+    pred = pred.flatten(0, 1)[:nrow].cpu()
+
+    mask = mask.unsqueeze(-1).repeat(
+        1, 1, model.patch_embed.patch_size[0]**2 * imgs.shape[1]
+    )  # (N, T*H*W, p*p*c)
+    mask = model.unpatchify(mask)  # 1 is removing, 0 is keeping
+    mask = torch.einsum("ncthw->nthwc", mask).cpu()
+    mask = mask.flatten(0, 1)[:nrow].cpu()
+
+    # masked image
+    im_masked = target * (1 - mask)
+
+    # MAE reconstruction pasted with visible patches
+    im_paste = target * (1 - mask) + pred * mask
+
+    if model.img_mask is not None:
+        img_mask = model.img_mask.cpu()
+    else:
+        img_mask = None
+
+    H, W = target.shape[1:3]
+    ploth = 2.0
+    plotw = (W / H) * ploth
+    nrow = len(target)
+    ncol = 3
+    fig, axs = plt.subplots(
+        nrow, ncol, figsize=(plotw * ncol, ploth * nrow), squeeze=False
+    )
+
+    for ii in range(nrow):
+        plt.sca(axs[ii, 0])
+        imshow(im_masked[ii], mean=mean, std=std, mask=img_mask)
+
+        plt.sca(axs[ii, 1])
+        imshow(im_paste[ii], mean=mean, std=std, mask=img_mask)
+
+        plt.sca(axs[ii, 2])
+        imshow(target[ii], mean=mean, std=std, mask=img_mask)
+
+    plt.tight_layout(pad=0.25)
+    return fig
+
+
+def imshow(
+    image: torch.Tensor,
+    mean: Optional[Any] = None,
+    std: Optional[Any] = None,
+    mask: Optional[torch.Tensor] = None,
+    **kwargs,
+):
+    # image: (H, W, C)
+    assert image.shape[2] in (1, 3)
+    if image.shape[2] == 1:
+        kwargs = {
+            "cmap": "gray",
+            "vmin": 0.0,
+            "vmax": 1.0,
+            "interpolation": "nearest",
+            **kwargs,
+        }
+    if mean is not None:
+        mean = torch.as_tensor(mean)
+        std = torch.as_tensor(std)
+        image = torch.clip(image * std + mean, 0.0, 1.0)
+    if mask is not None:
+        image = mask.unsqueeze(-1) * image
+    plt.imshow(image, **kwargs)
+    plt.axis("off")
+
+
+def fig2pil(fig: Figure, format: str = "png") -> Image.Image:
+    with io.BytesIO() as f:
+        fig.savefig(f, format=format)
+        f.seek(0)
+        img = Image.open(f)
+        img.load()
+    return img

fMRI-foundation-model-old/flat/util/hcp_flat.py

@@ -0,0 +1,121 @@
+import os
+import random
+from functools import partial
+from pathlib import Path
+from typing import Any, Dict, Iterable, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import webdataset as wds
+from torch.utils.data import IterableDataset
+
+HCP_FLAT_ROOT = "https://huggingface.co/datasets/bold-ai/HCP-Flat/resolve/main"
+NUM_SHARDS = {"train": 1629, "test": 174}
+
+
+def create_hcp_flat(
+    root: Optional[str] = None,
+    training: bool = True,
+    shards: Optional[Union[int, Iterable[int]]] = None,
+    frames: int = 16,
+    cache_dir: Optional[str] = None,
+) -> wds.WebDataset:
+    """
+    Create HCP-Flat dataset. Yields samples of (key, images) where key is the webdataset
+    sample key and images is shape (C, T, H, W).
+
+    References:
+        https://github.com/webdataset/webdataset/issues/250#issuecomment-1454094496
+        https://github.com/tmbdev-archive/webdataset-imagenet-2/blob/main/imagenet.py
+        https://github.com/huggingface/pytorch-image-models/blob/main/timm/data/readers/reader_wds.py
+    """
+    root = root or os.environ.get("HCP_FLAT_ROOT") or HCP_FLAT_ROOT
+    split = "train" if training else "test"
+
+    shards = shards or NUM_SHARDS[split]
+    if isinstance(shards, int):
+        shards = range(shards)
+    assert (
+        min(shards) >= 0 and max(shards) < NUM_SHARDS[split]
+    ), f"Invalid shards {shards}; expected in [0, {NUM_SHARDS[split]})"
+
+    urls = [f"{root}/{split}/hcp-flat_{split}_{shard:06d}.tar" for shard in shards]
+
+    # In training, we resample shards with replacement independently in every worker and
+    # yield batches up to the target number of samples. In test, we iterate over the
+    # shards in order, with workers getting interleaving shards, and yield batches up to
+    # the target samples. In a distributed setting with variable size shards, setting a
+    # fixed number of samples is the easiest way to get balanced batches per worker. In
+    # training we will still see all data eventually. But in test, it means we cut off
+    # some data.
+
+    # Note that in training this does not do deterministic shuffling, which we would
+    # need for exact reproducibility. They get determistic shuffling in timm, but it's
+    # more complicated.
+
+    # Note that we are splitting the long timeseries into clips after shuffling, which
+    # means clips from the same series will appear consecutively in the batch(es).
+    # I think this is not too bad. It is basically equivalent to training on longer
+    # sequences. Clipping before shuffling might be preferred, but it results in a bad
+    # system memory leak (https://github.com/webdataset/webdataset/issues/354).
+    dataset = (
+        wds.WebDataset(
+            urls,
+            resampled=training,
+            shardshuffle=1000 if training else False,
+            nodesplitter=wds.split_by_node,
+            select_files=select_extensions(("bold.npy","meta.json")),
+            cache_dir=cache_dir,
+        )
+        .decode()
+        .map(partial(extract_images, mask=load_hcp_flat_mask()))
+        .compose(to_clips(frames))
+        .shuffle(2500 if training else 0)
+    )
+    return dataset
+
+
+def select_extensions(extensions: Tuple[str, ...]):
+    extensions_set = set(extensions)
+    def select_files(fname: str):
+        suffix = ".".join(fname.split(".")[1:])
+        return suffix in extensions_set
+    return select_files
+
+
+def extract_images(sample: Dict[str, Any], mask: torch.Tensor):
+    key = sample["__key__"]
+    images = sample["bold.npy"]
+    meta = sample["meta.json"]['task']
+    
+    images = torch.from_numpy(images) / 255.0
+    images = (images - 0.5) / 0.2
+    images = unmask(images, mask)
+    # (C, T, H, W,)
+    images = images.unsqueeze(0)
+    
+    return key, images, meta
+
+
+def unmask(images: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+    images_unmasked = torch.zeros(
+        (images.shape[0], *mask.shape), dtype=images.dtype, device=images.device
+    )
+    images_unmasked[:, mask] = images
+    return images_unmasked
+
+
+def to_clips(frames: int = 16):
+    def _filter(src: IterableDataset[Tuple[str, torch.Tensor, str]]):
+        for key, images, meta in src:
+            offset = random.randint(0, frames)
+            for start in range(offset, images.shape[1] - frames, frames):
+                # yield key, images[:, start : start + frames].copy(), meta
+                yield key, images[:, start : start + frames].clone(), meta
+    return _filter
+
+
+def load_hcp_flat_mask() -> torch.Tensor:
+    mask = np.load(Path(__file__).parents[1] / "hcp-flat_mask.npy")
+    mask = torch.as_tensor(mask)
+    return mask

fMRI-foundation-model-old/flat/util/logging.py

@@ -0,0 +1,117 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+
+"""Logging."""
+
+import atexit
+import builtins
+import decimal
+import functools
+import logging
+import os
+import sys
+
+import simplejson
+import torch
+import torch.distributed as dist
+from iopath.common.file_io import g_pathmgr as pathmgr
+
+
+def is_master_proc(multinode=False):
+    """
+    Determines if the current process is the master process.
+    """
+    if dist.is_initialized():
+        if multinode:
+            return dist.get_rank() % dist.get_world_size() == 0
+        else:
+            return dist.get_rank() % torch.cuda.device_count() == 0
+    else:
+        return True
+
+
+def _suppress_print():
+    """
+    Suppresses printing from the current process.
+    """
+
+    def print_pass(*objects, sep=" ", end="\n", file=sys.stdout, flush=False):
+        pass
+
+    builtins.print = print_pass
+
+
+@functools.lru_cache(maxsize=None)
+def _cached_log_stream(filename):
+    # Use 1K buffer if writing to cloud storage.
+    io = pathmgr.open(filename, "a", buffering=1024 if "://" in filename else -1)
+    atexit.register(io.close)
+    return io
+
+
+def setup_logging(output_dir=None):
+    """
+    Sets up the logging for multiple processes. Only enable the logging for the
+    master process, and suppress logging for the non-master processes.
+    """
+    # Set up logging format.
+    if is_master_proc():
+        # Enable logging for the master process.
+        logging.root.handlers = []
+    else:
+        # Suppress logging for non-master processes.
+        _suppress_print()
+
+    logger = logging.getLogger()
+    logger.setLevel(logging.DEBUG)
+    logger.propagate = False
+    plain_formatter = logging.Formatter(
+        "[%(asctime)s][%(levelname)s] %(filename)s: %(lineno)3d: %(message)s",
+        datefmt="%m/%d %H:%M:%S",
+    )
+
+    if is_master_proc():
+        ch = logging.StreamHandler(stream=sys.stdout)
+        ch.setLevel(logging.DEBUG)
+        ch.setFormatter(plain_formatter)
+        logger.addHandler(ch)
+
+    if output_dir is not None and is_master_proc(multinode=True):
+        filename = os.path.join(output_dir, "stdout.log")
+        fh = logging.StreamHandler(_cached_log_stream(filename))
+        fh.setLevel(logging.DEBUG)
+        fh.setFormatter(plain_formatter)
+        logger.addHandler(fh)
+
+
+def get_logger(name):
+    """
+    Retrieve the logger with the specified name or, if name is None, return a
+    logger which is the root logger of the hierarchy.
+    Args:
+        name (string): name of the logger.
+    """
+    return logging.getLogger(name)
+
+
+def log_json_stats(stats):
+    """
+    Logs json stats.
+    Args:
+        stats (dict): a dictionary of statistical information to log.
+    """
+    stats = {
+        k: decimal.Decimal("{:.5f}".format(v)) if isinstance(v, float) else v
+        for k, v in stats.items()
+    }
+    json_stats = simplejson.dumps(stats, sort_keys=True, use_decimal=True)
+    logger = get_logger(__name__)
+    print("json_stats: {:s}".format(json_stats))
+
+
+def master_print(*args, **kwargs):
+    if is_master_proc():
+        print(*args, **kwargs)
+    else:
+        pass

fMRI-foundation-model-old/flat/util/losses.py

@@ -0,0 +1,132 @@
+import os
+import random
+import numpy as np
+import torch
+from einops import rearrange
+import re
+import torch.nn.functional as F
+import torch.nn as nn
+
+def get_ids_shuffle(batch_size, device, model, mask_ratio=0.75):
+    N = batch_size
+    T = model.patch_embed.t_grid_size
+    H, W = model.patch_embed.grid_size
+    L = T * H * W
+
+    noise = torch.rand(N, L, device=device)  # noise in [0, 1]
+
+    # shift missing patches to not be selected
+    if model.img_mask is not None:
+        noise = noise.view(N, T, H * W)
+        noise = noise + (1.0 - model.patch_mask)
+        noise = noise.view(N, L)
+
+    # sort noise for each sample
+    ids_shuffle = torch.argsort(
+        noise, dim=1
+    )  # ascend: small is keep, large is remove
+    ids_restore = torch.argsort(ids_shuffle, dim=1)
+
+    return ids_shuffle, ids_restore
+
+class VICRegHandler(nn.Module):
+    def __init__(self, in_dim, num_layers=3, act=nn.GELU, h=1024, out_dim=4096):
+        super().__init__()
+        self.projector = nn.Sequential(
+            nn.Linear(in_dim, h),
+            nn.LayerNorm(h),
+            act(),
+            nn.Linear(h, h),
+            nn.LayerNorm(h),
+            act(),
+            nn.Linear(h, out_dim),
+        )
+
+    def forward(self, x):
+        return self.projector(x)
+
+    @staticmethod
+    def filter_global_to_local(l, enc_mask, dec_mask):
+        '''Get the subset of global tokens that correspond to encoder mask only'''
+        comb_mask = enc_mask | dec_mask
+        comb_indices = torch.where(comb_mask)[0]
+        enc_indices = torch.where(enc_mask)[0]
+        # enc_set = set(enc_indices.cpu().tolist()) 
+        
+        # new_mask = torch.zeros_like(comb_indices, dtype=bool)
+        # for i, idx in enumerate(comb_indices):
+        #     if idx in enc_set:
+        #         new_mask[i] = True
+        
+        new_mask = torch.isin(comb_indices, enc_indices)    
+        return l[:, new_mask]
+
+    @staticmethod
+    def vicreg_loss(l1, l2, gamma=1.0, lamda=25, mu=25, nu=1, rand_frac=0.2, use_vic_cls=True, eps=1e-4):
+        if use_vic_cls:
+            # always keep cls and pick a random set of tokens
+            rand_indices = torch.cat([torch.tensor([0]), 1+torch.randperm(l1.shape[1]-1)])[:int(rand_frac*l1.shape[1])]
+        else:
+            # drop cls tokens from loss calc
+            l1 = l1[:, 1:]
+            l2 = l2[:, 1:]
+            rand_indices = torch.randperm(l1.shape[1])[:int(rand_frac*l1.shape[1])]
+
+        std_l1 = torch.sqrt(l1.flatten(1).var(dim=0)+eps)  # nxd
+        std_l2 = torch.sqrt(l2.flatten(1).var(dim=0)+eps)  # nxd
+        var_loss = F.relu(gamma - std_l1).mean() + F.relu(gamma - std_l2).mean()
+        del std_l1, std_l2
+
+        sim_loss = F.mse_loss(l1, l2)
+
+        l1 = l1 - l1.mean(0, keepdim=True)  # b,n,d
+        l2 = l2 - l2.mean(0, keepdim=True)                
+        
+        l1_sub = l1[:, rand_indices]
+        del l1
+        cov_l1 = torch.bmm(l1_sub.permute(1,2,0), l1_sub.permute(1,0,2))/(l1_sub.shape[0]-1)  # 0.1*n,d,d
+        cov_loss = ((cov_l1**2).sum() - (torch.diagonal(cov_l1, dim1=1,dim2=2)**2).sum())/(l1_sub.shape[1]*l1_sub.shape[2])
+        del cov_l1, l1_sub
+
+        l2_sub = l2[:, rand_indices]
+        del l2
+        cov_l2 = torch.bmm(l2_sub.permute(1,2,0), l2_sub.permute(1,0,2))/(l2_sub.shape[0]-1)
+        cov_loss = cov_loss + ((cov_l2**2).sum() - (torch.diagonal(cov_l2, dim1=1,dim2=2)**2).sum())/(l2_sub.shape[1]*l2_sub.shape[2])  # div by nxd
+        del cov_l2, l2_sub
+
+        vic_loss = lamda * sim_loss + mu * var_loss + nu * cov_loss
+
+        return vic_loss
+
+
+class SimCLRHandler(nn.Module):
+    def __init__(self, in_dim, num_layers=2, act=nn.GELU, out_dim=1024):
+        super().__init__()
+        self.projector = nn.Sequential(
+            nn.Linear(in_dim, in_dim),
+            act(),
+            nn.Linear(in_dim, max(in_dim,out_dim)),
+        )
+
+    def forward(self, x):
+        return self.projector(x)
+
+    @staticmethod
+    def simclr_loss(lats, temp=0.006):
+        logits = (nn.functional.normalize(lats.flatten(1),dim=-1) @
+                    nn.functional.normalize(lats.flatten(1),dim=-1).T) / temp
+
+        labels = torch.diag_embed(
+            torch.ones(logits.shape[0] // 2), offset=logits.shape[0] // 2
+        ) + torch.diag_embed(torch.ones(logits.shape[0] // 2), offset=-logits.shape[0] // 2)
+        labels = labels.to(lats.device)
+        
+        mask = torch.ones_like(logits).bool()
+        torch.diagonal(mask).fill_(False)
+        
+        labels = labels[mask].reshape(logits.shape[0], logits.shape[0]-1)
+        logits = logits[mask].reshape(*labels.shape)
+
+        contr_loss = -(logits.log_softmax(-1) * labels).sum(-1).mean()
+
+        return contr_loss

fMRI-foundation-model-old/flat/util/lr_decay.py

@@ -0,0 +1,83 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# References:
+# ELECTRA https://github.com/google-research/electra
+# BEiT: https://github.com/microsoft/unilm/tree/master/beit
+# --------------------------------------------------------
+
+import json
+
+
+def param_groups_lrd(
+    model, weight_decay=0.05, no_weight_decay_list=[], layer_decay=0.75
+):
+    """
+    Parameter groups for layer-wise lr decay
+    Following BEiT: https://github.com/microsoft/unilm/blob/master/beit/optim_factory.py#L58
+    """
+    param_group_names = {}
+    param_groups = {}
+
+    num_layers = len(model.blocks) + 1
+
+    layer_scales = list(layer_decay ** (num_layers - i) for i in range(num_layers + 1))
+
+    for n, p in model.named_parameters():
+        if not p.requires_grad:
+            continue
+
+        # no decay: all 1D parameters and model specific ones
+        if p.ndim == 1 or n in no_weight_decay_list:
+            g_decay = "no_decay"
+            this_decay = 0.0
+        else:
+            g_decay = "decay"
+            this_decay = weight_decay
+
+        layer_id = get_layer_id_for_vit(n, num_layers)
+        group_name = "layer_%d_%s" % (layer_id, g_decay)
+
+        if group_name not in param_group_names:
+            this_scale = layer_scales[layer_id]
+
+            param_group_names[group_name] = {
+                "lr_scale": this_scale,
+                "weight_decay": this_decay,
+                "params": [],
+            }
+            param_groups[group_name] = {
+                "lr_scale": this_scale,
+                "weight_decay": this_decay,
+                "params": [],
+            }
+
+        param_group_names[group_name]["params"].append(n)
+        param_groups[group_name]["params"].append(p)
+
+    print("parameter groups: \n%s" % json.dumps(param_group_names, indent=2))
+
+    return list(param_groups.values())
+
+
+def get_layer_id_for_vit(name, num_layers):
+    """
+    Assign a parameter with its layer id
+    Following BEiT: https://github.com/microsoft/unilm/blob/master/beit/optim_factory.py#L33
+    """
+    if name in [
+        "cls_token",
+        "mask_token",
+    ]:
+        return 0
+    elif name.startswith("patch_embed"):
+        return 0
+    elif name.startswith("pos_embed"):
+        return 0
+    elif name.startswith("blocks"):
+        return int(name.split(".")[1]) + 1
+    else:
+        return num_layers

fMRI-foundation-model-old/flat/util/lr_sched.py

@@ -0,0 +1,28 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+
+
+def adjust_learning_rate(optimizer, epoch, args):
+    """Decay the learning rate with half-cycle cosine after warmup"""
+    if epoch < args.warmup_epochs:
+        lr = args.lr * epoch / args.warmup_epochs
+    else:
+        lr = args.min_lr + (args.lr - args.min_lr) * 0.5 * (
+            1.0
+            + math.cos(
+                math.pi
+                * (epoch - args.warmup_epochs)
+                / (args.epochs - args.warmup_epochs)
+            )
+        )
+    for param_group in optimizer.param_groups:
+        if "lr_scale" in param_group:
+            param_group["lr"] = lr * param_group["lr_scale"]
+        else:
+            param_group["lr"] = lr
+    return lr

fMRI-foundation-model-old/flat/util/meters.py

@@ -0,0 +1,234 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+
+import numpy as np
+import torch
+from sklearn.metrics import average_precision_score
+
+
+def topks_correct(preds, labels, ks):
+    """
+    Given the predictions, labels, and a list of top-k values, compute the
+    number of correct predictions for each top-k value.
+
+    Args:
+        preds (array): array of predictions. Dimension is batchsize
+            N x ClassNum.
+        labels (array): array of labels. Dimension is batchsize N.
+        ks (list): list of top-k values. For example, ks = [1, 5] correspods
+            to top-1 and top-5.
+
+    Returns:
+        topks_correct (list): list of numbers, where the `i`-th entry
+            corresponds to the number of top-`ks[i]` correct predictions.
+    """
+    assert preds.size(0) == labels.size(
+        0
+    ), "Batch dim of predictions and labels must match"
+    # Find the top max_k predictions for each sample
+    _top_max_k_vals, top_max_k_inds = torch.topk(
+        preds, max(ks), dim=1, largest=True, sorted=True
+    )
+    # (batch_size, max_k) -> (max_k, batch_size).
+    top_max_k_inds = top_max_k_inds.t()
+    # (batch_size, ) -> (max_k, batch_size).
+    rep_max_k_labels = labels.view(1, -1).expand_as(top_max_k_inds)
+    # (i, j) = 1 if top i-th prediction for the j-th sample is correct.
+    top_max_k_correct = top_max_k_inds.eq(rep_max_k_labels)
+    # Compute the number of topk correct predictions for each k.
+    topks_correct = [top_max_k_correct[:k, :].float().sum() for k in ks]
+    return topks_correct
+
+
+def topk_errors(preds, labels, ks):
+    """
+    Computes the top-k error for each k.
+    Args:
+        preds (array): array of predictions. Dimension is N.
+        labels (array): array of labels. Dimension is N.
+        ks (list): list of ks to calculate the top accuracies.
+    """
+    num_topks_correct = topks_correct(preds, labels, ks)
+    return [(1.0 - x / preds.size(0)) * 100.0 for x in num_topks_correct]
+
+
+def topk_accuracies(preds, labels, ks):
+    """
+    Computes the top-k accuracy for each k.
+    Args:
+        preds (array): array of predictions. Dimension is N.
+        labels (array): array of labels. Dimension is N.
+        ks (list): list of ks to calculate the top accuracies.
+    """
+    num_topks_correct = topks_correct(preds, labels, ks)
+    return [(x / preds.size(0)) * 100.0 for x in num_topks_correct]
+
+
+def get_map(preds, labels):
+    """
+    Compute mAP for multi-label case.
+    Args:
+        preds (numpy tensor): num_examples x num_classes.
+        labels (numpy tensor): num_examples x num_classes.
+    Returns:
+        mean_ap (int): final mAP score.
+    """
+
+    print("Getting mAP for {} examples".format(preds.shape[0]))
+
+    preds = preds[:, ~(np.all(labels == 0, axis=0))]
+    labels = labels[:, ~(np.all(labels == 0, axis=0))]
+    aps = [0]
+    try:
+        aps = average_precision_score(labels, preds, average=None)
+    except ValueError:
+        print(
+            "Average precision requires a sufficient number of samples \
+            in a batch which are missing in this sample."
+        )
+
+    mean_ap = np.mean(aps)
+    return mean_ap
+
+
+class TestMeter:
+    """
+    Perform the multi-view ensemble for testing: each video with an unique index
+    will be sampled with multiple clips, and the predictions of the clips will
+    be aggregated to produce the final prediction for the video.
+    The accuracy is calculated with the given ground truth labels.
+    """
+
+    def __init__(
+        self,
+        num_videos,
+        num_clips,
+        num_cls,
+        overall_iters,
+        multi_label=False,
+        ensemble_method="sum",
+    ):
+        """
+        Construct tensors to store the predictions and labels. Expect to get
+        num_clips predictions from each video, and calculate the metrics on
+        num_videos videos.
+        Args:
+            num_videos (int): number of videos to test.
+            num_clips (int): number of clips sampled from each video for
+                aggregating the final prediction for the video.
+            num_cls (int): number of classes for each prediction.
+            overall_iters (int): overall iterations for testing.
+            multi_label (bool): if True, use map as the metric.
+            ensemble_method (str): method to perform the ensemble, options
+                include "sum", and "max".
+        """
+
+        self.num_clips = num_clips
+        self.overall_iters = overall_iters
+        self.multi_label = multi_label
+        self.ensemble_method = ensemble_method
+        # Initialize tensors.
+        self.video_preds = torch.zeros((num_videos, num_cls))
+        if multi_label:
+            self.video_preds -= 1e10
+
+        self.video_labels = (
+            torch.zeros((num_videos, num_cls))
+            if multi_label
+            else torch.zeros((num_videos)).long()
+        )
+        self.clip_count = torch.zeros((num_videos)).long()
+        self.topk_accs = []
+        self.stats = {}
+
+        # Reset metric.
+        self.reset()
+
+    def reset(self):
+        """
+        Reset the metric.
+        """
+        self.clip_count.zero_()
+        self.video_preds.zero_()
+        if self.multi_label:
+            self.video_preds -= 1e10
+        self.video_labels.zero_()
+
+    def update_stats(self, preds, labels, clip_ids):
+        """
+        Collect the predictions from the current batch and perform on-the-flight
+        summation as ensemble.
+        Args:
+            preds (tensor): predictions from the current batch. Dimension is
+                N x C where N is the batch size and C is the channel size
+                (num_cls).
+            labels (tensor): the corresponding labels of the current batch.
+                Dimension is N.
+            clip_ids (tensor): clip indexes of the current batch, dimension is
+                N.
+        """
+        for ind in range(preds.shape[0]):
+            vid_id = int(clip_ids[ind]) // self.num_clips
+            if self.video_labels[vid_id].sum() > 0:
+                assert torch.equal(
+                    self.video_labels[vid_id].type(torch.FloatTensor),
+                    labels[ind].type(torch.FloatTensor),
+                )
+            self.video_labels[vid_id] = labels[ind]
+            if self.ensemble_method == "sum":
+                self.video_preds[vid_id] += preds[ind]
+            elif self.ensemble_method == "max":
+                self.video_preds[vid_id] = torch.max(
+                    self.video_preds[vid_id], preds[ind]
+                )
+            else:
+                raise NotImplementedError(
+                    "Ensemble Method {} is not supported".format(self.ensemble_method)
+                )
+            self.clip_count[vid_id] += 1
+
+    def log_iter_stats(self, cur_iter):
+        """
+        Log the stats.
+        Args:
+            cur_iter (int): the current iteration of testing.
+        """
+        stats = {
+            "split": "test_iter",
+            "cur_iter": "{}".format(cur_iter + 1),
+        }
+        print(stats)
+
+    def finalize_metrics(self, ks=(1, 5)):
+        """
+        Calculate and log the final ensembled metrics.
+        ks (tuple): list of top-k values for topk_accuracies. For example,
+            ks = (1, 5) correspods to top-1 and top-5 accuracy.
+        """
+        if not all(self.clip_count == self.num_clips):
+            print(
+                "clip count {} ~= num clips {}".format(
+                    ", ".join(
+                        [
+                            "{}: {}".format(i, k)
+                            for i, k in enumerate(self.clip_count.tolist())
+                        ]
+                    ),
+                    self.num_clips,
+                )
+            )
+
+        self.stats = {"split": "test_final"}
+        if self.multi_label:
+            map = get_map(
+                self.video_preds.cpu().numpy(), self.video_labels.cpu().numpy()
+            )
+            self.stats["map"] = map
+        else:
+            num_topks_correct = topks_correct(self.video_preds, self.video_labels, ks)
+            topks = [(x / self.video_preds.size(0)) * 100.0 for x in num_topks_correct]
+            assert len({len(ks), len(topks)}) == 1
+            for k, topk in zip(ks, topks):
+                self.stats["top{}_acc".format(k)] = "{:.{prec}f}".format(topk, prec=2)
+        print(self.stats)

fMRI-foundation-model-old/flat/util/misc.py

@@ -0,0 +1,508 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# References:
+# DeiT: https://github.com/facebookresearch/deit
+# BEiT: https://github.com/microsoft/unilm/tree/master/beit
+# --------------------------------------------------------
+
+import builtins
+import datetime
+import math
+import os
+import time
+from collections import defaultdict, deque, OrderedDict
+
+import util.logging as logging
+import psutil
+import torch
+import torch.distributed as dist
+from iopath.common.file_io import g_pathmgr as pathmgr
+from util.logging import master_print as print
+from torch import inf
+
+
+logger = logging.get_logger(__name__)
+
+
+class SmoothedValue:
+    """Track a series of values and provide access to smoothed values over a
+    window or the global series average.
+    """
+
+    def __init__(self, window_size=20, fmt=None):
+        if fmt is None:
+            fmt = "{median:.4f} ({global_avg:.4f})"
+        self.deque = deque(maxlen=window_size)
+        self.total = 0.0
+        self.count = 0
+        self.fmt = fmt
+
+    def update(self, value, n=1):
+        self.deque.append(value)
+        self.count += n
+        self.total += value * n
+
+    def synchronize_between_processes(self):
+        """
+        Warning: does not synchronize the deque!
+        """
+        if not is_dist_avail_and_initialized():
+            return
+        t = torch.tensor([self.count, self.total], dtype=torch.float64, device="cuda")
+        dist.barrier()
+        dist.all_reduce(t)
+        t = t.tolist()
+        self.count = int(t[0])
+        self.total = t[1]
+
+    @property
+    def median(self):
+        d = torch.tensor(list(self.deque))
+        return d.median().item()
+
+    @property
+    def avg(self):
+        d = torch.tensor(list(self.deque), dtype=torch.float32)
+        return d.mean().item()
+
+    @property
+    def global_avg(self):
+        return self.total / self.count
+
+    @property
+    def max(self):
+        return max(self.deque)
+
+    @property
+    def value(self):
+        return self.deque[-1]
+
+    def __str__(self):
+        return self.fmt.format(
+            median=self.median,
+            avg=self.avg,
+            global_avg=self.global_avg,
+            max=self.max,
+            value=self.value,
+        )
+
+
+class MetricLogger:
+    def __init__(self, delimiter="\t"):
+        self.meters = defaultdict(SmoothedValue)
+        self.delimiter = delimiter
+
+    def update(self, **kwargs):
+        for k, v in kwargs.items():
+            if v is None:
+                continue
+            if isinstance(v, torch.Tensor):
+                v = v.item()
+            assert isinstance(v, (float, int))
+            self.meters[k].update(v)
+
+    def __getattr__(self, attr):
+        if attr in self.meters:
+            return self.meters[attr]
+        if attr in self.__dict__:
+            return self.__dict__[attr]
+        raise AttributeError(
+            "'{}' object has no attribute '{}'".format(type(self).__name__, attr)
+        )
+
+    def __str__(self):
+        loss_str = []
+        for name, meter in self.meters.items():
+            loss_str.append("{}: {}".format(name, str(meter)))
+        return self.delimiter.join(loss_str)
+
+    def synchronize_between_processes(self):
+        for meter in self.meters.values():
+            meter.synchronize_between_processes()
+
+    def add_meter(self, name, meter):
+        self.meters[name] = meter
+
+    def log_every(self, iterable, print_freq, header=None, total_steps=None):
+        i = 0
+        total_steps = total_steps or len(iterable)
+        if not header:
+            header = ""
+        start_time = time.time()
+        end = time.time()
+        iter_time = SmoothedValue(fmt="{avg:.4f}")
+        data_time = SmoothedValue(fmt="{avg:.4f}")
+        space_fmt = ":" + str(len(str(total_steps))) + "d"
+        log_msg = [
+            header,
+            "[{0" + space_fmt + "}/{1}]",
+            "eta: {eta}",
+            "{meters}",
+            "time: {time}",
+            "data: {data}",
+        ]
+        if torch.cuda.is_available():
+            log_msg.append("max mem: {memory:.0f}")
+        log_msg = self.delimiter.join(log_msg)
+        MB = 1024.0 * 1024.0
+        for obj in iterable:
+            data_time.update(time.time() - end)
+            yield obj
+            iter_time.update(time.time() - end)
+            if i % print_freq == 0 or i == total_steps - 1:
+                eta_seconds = iter_time.global_avg * (total_steps - i)
+                eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
+                if torch.cuda.is_available():
+                    print(
+                        log_msg.format(
+                            i,
+                            total_steps,
+                            eta=eta_string,
+                            meters=str(self),
+                            time=str(iter_time),
+                            data=str(data_time),
+                            memory=torch.cuda.max_memory_allocated() / MB,
+                        )
+                    )
+
+                else:
+                    print(
+                        log_msg.format(
+                            i,
+                            total_steps,
+                            eta=eta_string,
+                            meters=str(self),
+                            time=str(iter_time),
+                            data=str(data_time),
+                        )
+                    )
+            i += 1
+            end = time.time()
+        total_time = time.time() - start_time
+        total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+        print(
+            "{} Total time: {} ({:.4f} s / it)".format(
+                header, total_time_str, total_time / total_steps
+            )
+        )
+
+
+def setup_for_distributed(is_master):
+    """
+    This function disables printing when not in master process
+    """
+    builtin_print = builtins.print
+
+    def print(*args, **kwargs):
+        force = kwargs.pop("force", False)
+        force = force or (get_world_size() > 8)
+        if is_master or force:
+            now = datetime.datetime.now().time()
+            builtin_print("[{}] ".format(now), end="")  # print with time stamp
+            builtin_print(*args, **kwargs)
+
+    builtins.print = print
+
+
+def is_dist_avail_and_initialized():
+    if not dist.is_available():
+        return False
+    if not dist.is_initialized():
+        return False
+    return True
+
+
+def get_world_size():
+    if not is_dist_avail_and_initialized():
+        return 1
+    return dist.get_world_size()
+
+
+def get_rank():
+    if not is_dist_avail_and_initialized():
+        return 0
+    return dist.get_rank()
+
+
+def is_main_process():
+    return get_rank() == 0
+
+
+def save_on_master(state, path):
+    if is_main_process():
+        print(f"save path {path}")
+        with pathmgr.open(path, "wb") as f:
+            torch.save(state, f)
+
+
+def init_distributed_mode(args):
+    if args.no_env:
+        pass
+    elif args.dist_on_itp:
+        args.rank = int(os.environ["OMPI_COMM_WORLD_RANK"])
+        args.world_size = int(os.environ["OMPI_COMM_WORLD_SIZE"])
+        args.gpu = int(os.environ["OMPI_COMM_WORLD_LOCAL_RANK"])
+        args.dist_url = "tcp://%s:%s" % (
+            os.environ["MASTER_ADDR"],
+            os.environ["MASTER_PORT"],
+        )
+        os.environ["LOCAL_RANK"] = str(args.gpu)
+        os.environ["RANK"] = str(args.rank)
+        os.environ["WORLD_SIZE"] = str(args.world_size)
+        # ["RANK", "WORLD_SIZE", "MASTER_ADDR", "MASTER_PORT", "LOCAL_RANK"]
+    elif "RANK" in os.environ and "WORLD_SIZE" in os.environ:
+        args.rank = int(os.environ["RANK"])
+        args.world_size = int(os.environ["WORLD_SIZE"])
+        args.gpu = int(os.environ["LOCAL_RANK"])
+    elif "SLURM_PROCID" in os.environ:
+        args.rank = int(os.environ["SLURM_PROCID"])
+        args.gpu = args.rank % torch.cuda.device_count()
+    else:
+        print("Not using distributed mode")
+        setup_for_distributed(is_master=True)  # hack
+        args.distributed = False
+        return
+
+    args.distributed = True
+
+    torch.cuda.set_device(args.gpu)
+    args.dist_backend = "nccl"
+    print(
+        "| distributed init (rank {}): {}, gpu {}".format(
+            args.rank, args.dist_url, args.gpu
+        ),
+        # flush=True,
+    )
+    torch.distributed.init_process_group(
+        backend=args.dist_backend,
+        world_size=args.world_size,
+        rank=args.rank,
+    )
+    # init_method=args.dist_url,
+    torch.distributed.barrier()
+    setup_for_distributed(args.rank == 0)
+
+
+class NativeScalerWithGradNormCount:
+    state_dict_key = "amp_scaler"
+
+    def __init__(self, fp32=False):
+        self._scaler = torch.cuda.amp.GradScaler(enabled=not fp32)
+
+    def __call__(
+        self,
+        loss,
+        optimizer,
+        clip_grad=None,
+        parameters=None,
+        create_graph=False,
+        update_grad=True,
+    ):
+        self._scaler.scale(loss).backward(create_graph=create_graph)
+        if update_grad:
+            if clip_grad is not None:
+                assert parameters is not None
+                self._scaler.unscale_(
+                    optimizer
+                )  # unscale the gradients of optimizer's assigned params in-place
+                norm = torch.nn.utils.clip_grad_norm_(parameters, clip_grad)
+            else:
+                self._scaler.unscale_(optimizer)
+                norm = get_grad_norm_(parameters)
+            self._scaler.step(optimizer)
+            self._scaler.update()
+        else:
+            norm = None
+        return norm
+
+    def state_dict(self):
+        return self._scaler.state_dict()
+
+    def load_state_dict(self, state_dict):
+        self._scaler.load_state_dict(state_dict)
+
+
+def get_grad_norm_(parameters, norm_type: float = 2.0) -> torch.Tensor:
+    if isinstance(parameters, torch.Tensor):
+        parameters = [parameters]
+    parameters = [p for p in parameters if p.grad is not None]
+    norm_type = float(norm_type)
+    if len(parameters) == 0:
+        return torch.tensor(0.0)
+    device = parameters[0].grad.device
+    if norm_type == inf:
+        total_norm = max(p.grad.detach().abs().max().to(device) for p in parameters)
+    else:
+        total_norm = torch.norm(
+            torch.stack(
+                [torch.norm(p.grad.detach(), norm_type).to(device) for p in parameters]
+            ),
+            norm_type,
+        )
+    return total_norm
+
+
+def save_model(args, epoch, model, model_without_ddp, optimizer, loss_scaler):
+    checkpoint_path = "{}/checkpoint-{:05d}.pth".format(args.output_dir, epoch)
+    to_save = {
+        "model": model_without_ddp.state_dict(),
+        "optimizer": optimizer.state_dict(),
+        "epoch": epoch,
+        "scaler": loss_scaler.state_dict(),
+        "args": args,
+    }
+
+    save_on_master(to_save, checkpoint_path)
+    return checkpoint_path
+
+
+def get_last_checkpoint(args):
+    """
+    Get the last checkpoint from the checkpointing folder.
+    Args:
+        path_to_job (string): the path to the folder of the current job.
+    """
+    d = args.output_dir
+    names = pathmgr.ls(d) if pathmgr.exists(d) else []
+    names = [f for f in names if "checkpoint" in f]
+    if len(names) == 0:
+        print("No checkpoints found in '{}'.".format(d))
+        return None
+    else:
+        # Sort the checkpoints by epoch.
+        name = sorted(names)[-1]
+        return os.path.join(d, name)
+
+
+def load_model(args, model_without_ddp, optimizer, loss_scaler):
+    if not args.resume:
+        args.resume = get_last_checkpoint(args)
+    if args.resume:
+        if args.resume.startswith("https"):
+            checkpoint = torch.hub.load_state_dict_from_url(
+                args.resume, map_location="cpu", check_hash=True
+            )
+        else:
+            with pathmgr.open(args.resume, "rb") as f:
+                checkpoint = torch.load(f, map_location="cpu")
+        model_without_ddp.load_state_dict(checkpoint["model"])
+        print("Resume checkpoint %s" % args.resume)
+        if (
+            "optimizer" in checkpoint
+            and "epoch" in checkpoint
+            and not (hasattr(args, "eval") and args.eval)
+        ):
+            optimizer.load_state_dict(checkpoint["optimizer"])
+            args.start_epoch = checkpoint["epoch"] + 1
+            if "scaler" in checkpoint:
+                loss_scaler.load_state_dict(checkpoint["scaler"])
+            print("With optim & sched!")
+
+
+def all_reduce_mean(x):
+    world_size = get_world_size()
+    if world_size > 1:
+        x_reduce = torch.tensor(x).cuda()
+        dist.all_reduce(x_reduce)
+        x_reduce /= world_size
+        return x_reduce.item()
+    else:
+        return x
+
+
+def gpu_mem_usage():
+    """
+    Compute the GPU memory usage for the current device (GB).
+    """
+    if torch.cuda.is_available():
+        mem_usage_bytes = torch.cuda.max_memory_allocated()
+    else:
+        mem_usage_bytes = 0
+    return mem_usage_bytes / 1024**3
+
+
+def cpu_mem_usage():
+    """
+    Compute the system memory (RAM) usage for the current device (GB).
+    Returns:
+        usage (float): used memory (GB).
+        total (float): total memory (GB).
+    """
+    vram = psutil.virtual_memory()
+    usage = (vram.total - vram.available) / 1024**3
+    total = vram.total / 1024**3
+
+    return usage, total
+
+
+def all_gather(tensors):
+    """
+    All gathers the provided tensors from all processes across machines.
+    Args:
+        tensors (list): tensors to perform all gather across all processes in
+        all machines.
+    """
+
+    gather_list = []
+    output_tensor = []
+    world_size = dist.get_world_size()
+    for tensor in tensors:
+        tensor_placeholder = [torch.ones_like(tensor) for _ in range(world_size)]
+        dist.all_gather(tensor_placeholder, tensor, async_op=False)
+        gather_list.append(tensor_placeholder)
+    for gathered_tensor in gather_list:
+        output_tensor.append(torch.cat(gathered_tensor, dim=0))
+    return output_tensor
+
+
+def add_weight_decay(model, weight_decay=1e-5, skip_list=(), bias_wd=False):
+    decay = []
+    no_decay = []
+    for name, param in model.named_parameters():
+        if not param.requires_grad:
+            continue  # frozen weights
+        if (
+            (not bias_wd)
+            and len(param.shape) == 1
+            or name.endswith(".bias")
+            or name in skip_list
+        ):
+            no_decay.append(param)
+        else:
+            decay.append(param)
+    return [
+        {"params": no_decay, "weight_decay": 0.0},
+        {"params": decay, "weight_decay": weight_decay},
+    ]
+
+
+def inflate(model_2d, model_3d):
+    state_dict_inflated = OrderedDict()
+    for k, v2d in model_2d.items():
+        if "patch_embed.proj.weight" in k:
+            v3d = model_3d[k]
+            v3d = v2d.unsqueeze(2).repeat(1, 1, v3d.shape[2], 1, 1) / v3d.shape[2]
+            state_dict_inflated[k] = v3d.clone()
+        elif "pos_embed" in k:
+            pos_embed_cls, pos_embed_spatial = torch.split(v2d, [1, 196], dim=1)
+            state_dict_inflated["pos_embed_cls"] = pos_embed_cls.clone()
+            state_dict_inflated["pos_embed"] = pos_embed_spatial.clone()
+        else:
+            state_dict_inflated[k] = v2d.clone()
+    return state_dict_inflated
+
+
+def convert_checkpoint(model_2d):
+    state_dict_inflated = OrderedDict()
+    for k, v2d in model_2d.items():
+        if "head.projection.weight" in k:
+            state_dict_inflated["head.weight"] = v2d.clone()
+        elif "head.projection.bias" in k:
+            state_dict_inflated["head.bias"] = v2d.clone()
+        else:
+            state_dict_inflated[k] = v2d.clone()
+    return state_dict_inflated

fMRI-foundation-model-old/flat/util/pos_embed.py

@@ -0,0 +1,53 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# Position embedding utils
+# --------------------------------------------------------
+
+import util.logging as logging
+import numpy as np
+import torch
+
+
+logger = logging.get_logger(__name__)
+
+
+# --------------------------------------------------------
+# 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

fMRI-foundation-model-old/flat/util/video_vit.py

@@ -0,0 +1,173 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+
+import util.logging as logging
+import torch
+import torch.nn as nn
+from timm.layers import to_2tuple
+from timm.models.vision_transformer import DropPath, Mlp
+
+
+logger = logging.get_logger(__name__)
+
+
+class PatchEmbed(nn.Module):
+    """Image to Patch Embedding"""
+
+    def __init__(
+        self,
+        img_size=224,
+        patch_size=16,
+        in_chans=3,
+        embed_dim=768,
+        # temporal related:
+        frames=32,
+        t_patch_size=4,
+    ):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        assert img_size[1] % patch_size[1] == 0
+        assert img_size[0] % patch_size[0] == 0
+        assert frames % t_patch_size == 0
+        num_patches = (
+            (img_size[1] // patch_size[1])
+            * (img_size[0] // patch_size[0])
+            * (frames // t_patch_size)
+        )
+        self.input_size = (
+            frames // t_patch_size,
+            img_size[0] // patch_size[0],
+            img_size[1] // patch_size[1],
+        )
+        print(
+            f"img_size {img_size} patch_size {patch_size} frames {frames} t_patch_size {t_patch_size}"
+        )
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.in_chans = in_chans
+
+        self.frames = frames
+        self.t_patch_size = t_patch_size
+
+        self.num_patches = num_patches
+
+        self.grid_size = img_size[0] // patch_size[0], img_size[1] // patch_size[1]
+        self.t_grid_size = frames // t_patch_size
+
+        kernel_size = [t_patch_size] + list(patch_size)
+        self.proj = nn.Conv3d(
+            in_chans, embed_dim, kernel_size=kernel_size, stride=kernel_size
+        )
+
+    def forward(self, x):
+        B, C, T, H, W = x.shape
+        assert (
+            H == self.img_size[0] and W == self.img_size[1]
+        ), f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+        assert T == self.frames
+        x = self.proj(x).flatten(3)
+        x = torch.einsum("ncts->ntsc", x)  # [N, T, H*W, C]
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        dim,
+        num_heads=8,
+        qkv_bias=False,
+        qk_scale=None,
+        attn_drop=0.0,
+        proj_drop=0.0,
+        input_size=(4, 14, 14),
+    ):
+        super().__init__()
+        assert dim % num_heads == 0, "dim should be divisible by num_heads"
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim**-0.5
+
+        self.q = nn.Linear(dim, dim, bias=qkv_bias)
+        self.k = nn.Linear(dim, dim, bias=qkv_bias)
+        self.v = nn.Linear(dim, dim, bias=qkv_bias)
+        assert attn_drop == 0.0  # do not use
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+        self.input_size = input_size
+        assert input_size[1] == input_size[2]
+
+    def forward(self, x):
+        B, N, C = x.shape
+        q = (
+            self.q(x)
+            .reshape(B, N, self.num_heads, C // self.num_heads)
+            .permute(0, 2, 1, 3)
+        )
+        k = (
+            self.k(x)
+            .reshape(B, N, self.num_heads, C // self.num_heads)
+            .permute(0, 2, 1, 3)
+        )
+        v = (
+            self.v(x)
+            .reshape(B, N, self.num_heads, C // self.num_heads)
+            .permute(0, 2, 1, 3)
+        )
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+
+        attn = attn.softmax(dim=-1)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        x = x.view(B, -1, C)
+        return x
+
+
+class Block(nn.Module):
+    """
+    Transformer Block with specified Attention function
+    """
+
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        mlp_ratio=4.0,
+        qkv_bias=False,
+        qk_scale=None,
+        drop=0.0,
+        attn_drop=0.0,
+        drop_path=0.0,
+        act_layer=nn.GELU,
+        norm_layer=nn.LayerNorm,
+        attn_func=Attention,
+    ):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = attn_func(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            qk_scale=qk_scale,
+            attn_drop=attn_drop,
+            proj_drop=drop,
+        )
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            act_layer=act_layer,
+            drop=drop,
+        )
+
+    def forward(self, x):
+        x = x + self.drop_path(self.attn(self.norm1(x)))
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x