Upload model

cls_token.py

@@ -0,0 +1,47 @@
+import torch
+from torch import nn
+
+
+class ClsToken(nn.Module):
+    def __init__(self, ndim: int,
+                 num_tokens: int = 1,
+                 enabled: bool = True,
+                 register_multiple: int = 0,
+    ):
+        super().__init__()
+
+        self.ndim = ndim
+        self.enabled = enabled
+        self.num_registers = 0
+        self.num_tokens = num_tokens
+        if enabled:
+            if register_multiple > 0:
+                self.num_registers = register_multiple - (num_tokens % register_multiple)
+
+            scale = ndim ** -0.5
+            self.token = nn.Parameter(torch.randn(num_tokens + self.num_registers, ndim) * scale)
+        else:
+            self.token = None
+
+        self.num_patches = self.num_tokens + self.num_registers
+
+    def disable(self):
+        self.token = None
+        self.enabled = False
+
+    def forward(self, x: torch.Tensor):
+        if self.token is None:
+            return x
+
+        token = self.token.unsqueeze(0).expand(x.shape[0], -1, -1)
+        x = torch.cat([
+            token,
+            x,
+        ], dim=1)
+
+        return x
+
+    def no_weight_decay(self):
+        return [
+            'token',
+        ]

config.json

@@ -0,0 +1,235 @@
+{
+  "architectures": [
+    "RADIOModel"
+  ],
+  "args": {
+    "aa": null,
+    "amp": true,
+    "amp_dtype": "bfloat16",
+    "amp_impl": "native",
+    "aug_repeats": 0,
+    "aug_splits": 0,
+    "auto_loss_balance_mode": "adaloss",
+    "batch_size": 32,
+    "bn_eps": null,
+    "bn_momentum": null,
+    "cache_dir": null,
+    "channels_last": false,
+    "checkpoint_hist": 10,
+    "class_map": "",
+    "clip_grad": null,
+    "clip_mode": "norm",
+    "cls_token_per_teacher": true,
+    "coco_annotations_file": "/datasets/coco2017-adlsa/annotations/captions_val2017.json",
+    "coco_image_dir": "/datasets/coco2017-adlsa/val2017",
+    "color_jitter": 0.4,
+    "cooldown_epochs": 0,
+    "cpe_max_size": 1050,
+    "crd_loss": false,
+    "crd_loss_weight": 0.8,
+    "crop_pct": null,
+    "cutmix": 0.0,
+    "cutmix_minmax": null,
+    "data_dir": "/lustre/fsw/portfolios/llmservice/projects/llmservice_nlp_fm/datasets/captioning/datacomp/dc1b/stage2",
+    "dataset": "nvgpt4",
+    "dataset_download": false,
+    "debug_full_knn": false,
+    "decay_epochs": 90,
+    "decay_milestones": [
+      90,
+      180,
+      270
+    ],
+    "decay_rate": 0.1,
+    "device": "cuda:0",
+    "dist_bn": "reduce",
+    "distributed": true,
+    "drop": 0.0,
+    "drop_block": null,
+    "drop_connect": null,
+    "drop_path": null,
+    "epoch_repeats": 0.0,
+    "epochs": 300,
+    "eval": false,
+    "eval_metric": "knn_top1",
+    "eval_teacher": false,
+    "eval_teacher_only": false,
+    "eval_throughput": false,
+    "experiment": "checkpoints",
+    "fast_norm": false,
+    "feature_summarizer": "cls_token",
+    "feature_upscale_factor": null,
+    "fuser": "",
+    "gp": "avg",
+    "grad_accum_steps": 1,
+    "grad_checkpointing": false,
+    "head_init_bias": null,
+    "head_init_scale": null,
+    "hflip": 0.5,
+    "img_size": null,
+    "in_chans": null,
+    "initial_checkpoint": "",
+    "input_size": null,
+    "interpolation": "",
+    "layer_decay": null,
+    "local_rank": 0,
+    "log_interval": 50,
+    "log_mlflow": false,
+    "log_wandb": true,
+    "loss": "cosine",
+    "loss_auto_balance": false,
+    "lr": 0.001,
+    "lr_base": 0.1,
+    "lr_base_scale": "",
+    "lr_base_size": 256,
+    "lr_cycle_decay": 0.5,
+    "lr_cycle_limit": 1,
+    "lr_cycle_mul": 1.0,
+    "lr_k_decay": 1.0,
+    "lr_noise": null,
+    "lr_noise_pct": 0.67,
+    "lr_noise_std": 1.0,
+    "mean": null,
+    "min_lr": 0,
+    "mixup": 0.0,
+    "mixup_mode": "batch",
+    "mixup_off_epoch": 0,
+    "mixup_prob": 1.0,
+    "mixup_switch_prob": 0.5,
+    "mlp_hidden_size": 1520,
+    "mlp_num_inner": 3,
+    "mlp_version": "v2",
+    "model": "vit_huge_patch14_224",
+    "model_ema": false,
+    "model_ema_decay": 0.9998,
+    "model_ema_force_cpu": false,
+    "model_kwargs": {},
+    "momentum": 0.9,
+    "no_aug": false,
+    "no_ddp_bb": false,
+    "no_prefetcher": false,
+    "no_resume_opt": false,
+    "num_classes": null,
+    "opt": "fusedlamb",
+    "opt_betas": null,
+    "opt_eps": null,
+    "opt_kwargs": {},
+    "output": "/lustre/fs6/portfolios/llmservice/users/mranzinger/output/evfm/dfn_oai/11-29-23_vit-h-14-cpe_dfn-oai-dino_maxres",
+    "patience_epochs": 10,
+    "pin_mem": false,
+    "prefetcher": true,
+    "pretrained": false,
+    "rank": 0,
+    "ratio": [
+      0.75,
+      1.3333333333333333
+    ],
+    "recount": 1,
+    "recovery_interval": 0,
+    "register_multiple": 8,
+    "remode": "pixel",
+    "reprob": 0.0,
+    "resplit": false,
+    "resume": "/lustre/fs6/portfolios/llmservice/users/mranzinger/output/evfm/dfn_oai/11-29-23_vit-h-14-cpe_dfn-oai-dino_maxres/checkpoints/last.pth.tar",
+    "save_images": false,
+    "scale": [
+      0.5,
+      1.0
+    ],
+    "sched": "cosine",
+    "sched_on_updates": true,
+    "seed": 42,
+    "smoothing": 0.1,
+    "split_bn": false,
+    "start_epoch": null,
+    "std": null,
+    "steps_per_epoch": 2000,
+    "sync_bn": false,
+    "synchronize_step": false,
+    "teachers": [
+      {
+        "amp": true,
+        "amp_dtype": "bfloat16",
+        "batch_size": 16,
+        "fd_loss_weight": 1.0,
+        "fd_normalize": false,
+        "feature_distillation": true,
+        "input_size": 378,
+        "model": "ViT-H-14-378-quickgelu",
+        "name": "clip",
+        "pretrained": "dfn5b",
+        "sample_rate": 16,
+        "summary_loss_weight": 1.0,
+        "type": "open_clip",
+        "vitdet_prob": 0.05,
+        "vitdet_window_sizes": [
+          3,
+          9,
+          9,
+          9
+        ]
+      },
+      {
+        "amp": false,
+        "amp_dtype": "bfloat16",
+        "batch_size": 16,
+        "fd_loss_weight": 0.8,
+        "fd_normalize": false,
+        "feature_distillation": true,
+        "input_size": 336,
+        "model": "ViT-L/14@336px",
+        "name": "openai_clip",
+        "pretrained": "openai",
+        "sample_rate": 16,
+        "summary_loss_weight": 0.8,
+        "type": "openai_clip",
+        "use_summary": false
+      },
+      {
+        "amp": true,
+        "amp_dtype": "bfloat16",
+        "batch_size": 16,
+        "fd_loss_weight": 1.0,
+        "fd_normalize": false,
+        "feature_distillation": true,
+        "input_size": 224,
+        "model": "dinov2_vitg14",
+        "name": "dino_v2",
+        "sample_rate": 16,
+        "summary_loss_weight": 1.0,
+        "type": "dino_v2"
+      }
+    ],
+    "torchcompile": null,
+    "torchscript": false,
+    "train_interpolation": "random",
+    "train_split": "train",
+    "tta": 0,
+    "use_coco": false,
+    "use_multi_epochs_loader": false,
+    "val_data_dir": "/lustre/fsw/portfolios/llmservice/projects/llmservice_nlp_fm/datasets/classification/imagenet-1k/webdataset",
+    "val_img_size": 378,
+    "val_split": "val",
+    "validation_batch_size": 128,
+    "vflip": 0.0,
+    "wandb_entity": "",
+    "wandb_group": "dfn_oai",
+    "wandb_job_type": "",
+    "wandb_name": "",
+    "wandb_project": "",
+    "warmup_epochs": 2.5,
+    "warmup_lr": 1e-05,
+    "warmup_prefix": false,
+    "weight_decay": 2e-05,
+    "worker_seeding": "all",
+    "workers": 4,
+    "world_size": 64
+  },
+  "auto_map": {
+    "AutoConfig": "hf_model.RADIOConfig",
+    "AutoModel": "hf_model.RADIOModel"
+  },
+  "torch_dtype": "float32",
+  "transformers_version": "4.29.0",
+  "version": "v1"
+}

enable_cpe_support.py

@@ -0,0 +1,59 @@
+from typing import Union, Tuple
+from types import MethodType
+
+import torch
+from torch import nn
+
+from timm.models import VisionTransformer, checkpoint_seq
+
+from .vit_patch_generator import ViTPatchGenerator
+
+
+def _forward_cpe(self: VisionTransformer, x: torch.Tensor) -> torch.Tensor:
+    x = self.patch_generator(x)
+    if self.grad_checkpointing and not torch.jit.is_scripting():
+        x = checkpoint_seq(self.blocks, x)
+    else:
+        x = self.blocks(x)
+    x = self.norm(x)
+    return x
+
+
+def enable_cpe(model: nn.Module,
+               max_img_size: Union[int, Tuple[int, int]] = 1024,
+               num_cls_tokens: int = 1,
+               pos_dropout: float = 0.1,
+               register_multiple: int = 0,
+):
+    if not isinstance(model, VisionTransformer):
+        raise ValueError("CPE only support for VisionTransformer models!")
+
+    patch_size = model.patch_embed.patch_size[0]
+    embed_dim = model.embed_dim
+    input_dims = model.patch_embed.img_size
+    normalize_patches = not isinstance(model.patch_embed.norm, nn.Identity)
+    cls_token = model.cls_token is not None
+
+    max_img_size = int(round(max_img_size / patch_size) * patch_size)
+
+    patch_generator = ViTPatchGenerator(
+        patch_size=patch_size,
+        embed_dim=embed_dim,
+        input_dims=input_dims,
+        normalize_patches=normalize_patches,
+        cls_token=cls_token,
+        max_input_dims=max_img_size,
+        pos_dropout=pos_dropout,
+        num_cls_tokens=num_cls_tokens,
+        register_multiple=register_multiple,
+    )
+
+    model.patch_generator = patch_generator
+    model.patch_embed = None
+    model.cls_token = None
+    model.pos_embed = None
+    model.pos_drop = None
+    model.num_cls_tokens = num_cls_tokens
+    model.num_registers = patch_generator.num_registers
+
+    model.forward_features = MethodType(_forward_cpe, model)

hf_model.py

@@ -0,0 +1,84 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from collections import namedtuple
+from typing import Optional
+
+from timm.models import VisionTransformer
+import torch
+from transformers import PretrainedConfig, PreTrainedModel
+
+
+from .model import create_model_from_args
+from .input_conditioner import get_default_conditioner, InputConditioner
+
+
+resource_map = {
+    'radio_v1': 'https://huggingface.co/nvidia/RADIO/raw/main/radio_v1.pth.tar'
+}
+
+
+class RADIOConfig(PretrainedConfig):
+    """Pretrained Hugging Face configuration for RADIO models."""
+
+    def __init__(
+        self,
+        args: Optional[dict] = None,
+        version: Optional[str]="v1",
+        **kwargs,
+    ):
+        self.args = args
+        self.version = version
+        super().__init__(**kwargs)
+
+
+class RADIOModel(PreTrainedModel):
+    """Pretrained Hugging Face model for RADIO."""
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        RADIOArgs = namedtuple("RADIOArgs", config.args.keys())
+        args = RADIOArgs(**config.args)
+        self.model = create_model_from_args(args)
+
+        self.input_conditioner: InputConditioner = get_default_conditioner()
+
+        #return RADIOModel(mod, conditioner, return_summary=return_summary, return_spatial_features=return_spatial_features)
+
+    def forward(self, x: torch.Tensor):
+        x = self.input_conditioner(x)
+
+        y = self.model.forward_features(x)
+
+        if isinstance(y, (list, tuple)):
+            summary, all_feat = y
+        elif isinstance(self.model, VisionTransformer):
+            patch_gen = getattr(self.model, 'patch_generator', None)
+            if patch_gen is not None:
+                summary = y[:, :patch_gen.num_cls_tokens].flatten(1)
+                all_feat = y[:, patch_gen.num_skip:]
+            elif self.model.global_pool == 'avg':
+                summary = y[:, self.model.num_prefix_tokens:].mean(dim=1)
+                all_feat = y
+            else:
+                summary = y[:, 0]
+                all_feat = y[:, 1:]
+        else:
+            raise ValueError("Unsupported model type")
+
+        if self.return_summary and self.return_spatial_features:
+            return summary, all_feat
+        elif self.return_summary:
+            return summary
+        return all_feat

input_conditioner.py

@@ -0,0 +1,41 @@
+from typing import Union, Tuple
+
+import torch
+from torch import nn
+
+
+norm_t = Union[Tuple[float, float, float], torch.Tensor]
+
+class InputConditioner(nn.Module):
+    def __init__(self,
+                 input_scale: float,
+                 norm_mean: norm_t,
+                 norm_std: norm_t,
+                 dtype: torch.dtype = torch.float32,
+    ):
+        super().__init__()
+
+        self.dtype = dtype
+
+        # self.input_scale = input_scale
+        self.register_buffer("norm_mean", _to_tensor(norm_mean) / input_scale)
+        self.register_buffer("norm_std", _to_tensor(norm_std) / input_scale)
+
+    def forward(self, x: torch.Tensor):
+        # x = x * self.input_scale
+        y = (x - self.norm_mean) / self.norm_std
+        return y.to(self.dtype)
+
+
+def get_default_conditioner():
+    from timm.data.constants import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD
+
+    return InputConditioner(
+        input_scale=1.0,
+        norm_mean=OPENAI_CLIP_MEAN,
+        norm_std=OPENAI_CLIP_STD,
+    )
+
+
+def _to_tensor(v: norm_t):
+    return torch.as_tensor(v, dtype=torch.float32).view(-1, 1, 1)

model.py

@@ -0,0 +1,40 @@
+from torch import nn
+
+from timm.models import create_model
+
+from .enable_cpe_support import enable_cpe
+
+
+def create_model_from_args(args) -> nn.Module:
+    in_chans = 3
+    if args.in_chans is not None:
+        in_chans = args.in_chans
+    elif args.input_size is not None:
+        in_chans = args.input_size[0]
+
+    model = create_model(
+        args.model,
+        pretrained=args.pretrained,
+        in_chans=in_chans,
+        num_classes=args.num_classes,
+        drop_rate=args.drop,
+        drop_path_rate=args.drop_path,
+        drop_block_rate=args.drop_block,
+        global_pool=args.gp,
+        bn_momentum=args.bn_momentum,
+        bn_eps=args.bn_eps,
+        scriptable=args.torchscript,
+        checkpoint_path=args.initial_checkpoint,
+        **args.model_kwargs,
+    )
+
+    assert not args.cls_token_per_teacher or args.cpe_max_size is not None, "CPE must be enabled for multiple CLS tokens!"
+
+    if args.cpe_max_size is not None:
+        enable_cpe(model,
+                   args.cpe_max_size,
+                   num_cls_tokens=len(args.teachers) if args.cls_token_per_teacher else 1,
+                   register_multiple=args.register_multiple,
+        )
+
+    return model

pytorch_model.bin

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

vit_patch_generator.py

@@ -0,0 +1,291 @@
+import math
+from typing import Union, Tuple, Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from einops import rearrange
+
+from .cls_token import ClsToken
+
+input_dim_t = Union[int, Tuple[int, int]]
+
+try:
+    # raise ImportError()
+    from indirect_grid_sample import indirect_grid_sample
+except ImportError:
+    indirect_grid_sample = None
+
+class ViTPatchGenerator(nn.Module):
+    def __init__(self,
+                 patch_size: int,
+                 embed_dim: int,
+                 input_dims: input_dim_t,
+                 abs_pos: bool = True,
+                 normalize_patches: bool = False,
+                 cls_token: bool = False,
+                 max_input_dims: Optional[input_dim_t] = None,
+                 pos_dropout: float = 0.0,
+                 return_pos_enc: bool = False,
+                 num_cls_tokens: int = 1,
+                 register_multiple: int = 0,
+                 device=None, dtype=None,
+    ):
+        super().__init__()
+
+        if isinstance(input_dims, int):
+            input_dims = (input_dims, input_dims)
+
+        if max_input_dims is None:
+            max_input_dims = input_dims
+        if isinstance(max_input_dims, int):
+            max_input_dims = (max_input_dims, max_input_dims)
+
+        max_input_dims = tuple(
+            int(math.ceil(d / patch_size) * patch_size)
+            for d in max_input_dims
+        )
+
+        self.cpe_mode = max_input_dims != input_dims
+        self.pos_dropout = pos_dropout
+        self.return_pos_enc = return_pos_enc
+
+        factory = dict(device=device, dtype=dtype)
+
+        self.patch_size = patch_size
+        self.abs_pos = abs_pos
+        self.embed_dim = embed_dim
+
+        self.num_rows = max_input_dims[0] // patch_size
+        self.num_cols = max_input_dims[1] // patch_size
+        self.input_dims = tuple(d // patch_size for d in input_dims)
+        self.num_patches = self.num_rows * self.num_cols
+        self.max_input_dims = max_input_dims
+
+        self.im_to_patches = Im2Patches(patch_size)
+        self.embedder = ViTPatchLinear(patch_size, embed_dim, **factory)
+
+        if abs_pos:
+            scale = embed_dim ** -0.5
+            self.pos_embed = nn.Parameter(torch.randn(1, self.num_patches, embed_dim, **factory) * scale)
+
+        self.cls_token = ClsToken(
+            embed_dim,
+            num_tokens=num_cls_tokens,
+            enabled=cls_token,
+            register_multiple=register_multiple,
+        )
+
+        self.patch_normalizer = nn.LayerNorm(embed_dim) if normalize_patches else nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        patches = self.embed_patches(x)
+        patches, pos_enc = self.apply_pos_enc(patches, input_size=x.shape[2:])
+        patches = self.cls_token(patches)
+        patches = self.patch_normalizer(patches)
+        if self.return_pos_enc:
+            return patches, pos_enc
+        return patches
+
+    @property
+    def apply_cls_token(self):
+        return self.cls_token.enabled
+
+    @property
+    def num_cls_tokens(self):
+        return self.cls_token.num_tokens
+
+    @property
+    def num_registers(self):
+        return self.cls_token.num_registers
+
+    @property
+    def num_skip(self):
+        return self.num_cls_tokens + self.num_registers
+
+    def no_weight_decay(self):
+        return [
+            'pos_embed',
+        ]
+
+    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
+        if self.abs_pos:
+            self._load_embed(state_dict[f'{prefix}pos_embed'], self.pos_embed)
+
+    def _load_embed(self, src_embed: torch.Tensor, targ_embed: nn.Parameter):
+        if src_embed.shape != targ_embed.shape:
+            src_size = int(math.sqrt(src_embed.shape[1]))
+
+            assert src_size ** 2 == src_embed.shape[1], 'Unable to interpolate non-square embedding'
+
+            src_embed = rearrange(src_embed, 'b (h w) c -> b c h w', h=src_size, w=src_size)
+            src_embed = F.interpolate(src_embed, size=(self.num_rows, self.num_cols), mode='bicubic', align_corners=True, antialias=False)
+            src_embed = rearrange(src_embed, 'b c h w -> b (h w) c')
+        targ_embed.data.copy_(src_embed)
+
+    def _load_projection(self, src_proj_weight: torch.Tensor, targ_proj_weight: torch.Tensor):
+        if src_proj_weight.shape != targ_proj_weight.shape:
+            src_patch_size = int(math.sqrt(src_proj_weight.shape[1] // 3))
+
+            assert (src_patch_size ** 2) * 3 == src_proj_weight.shape[1], 'Unable to interpolate non-square patch size'
+
+            src_proj_weight = rearrange(src_proj_weight, 'b (c h w) -> b c h w', c=3, h=src_patch_size, w=src_patch_size)
+            src_proj_weight = F.interpolate(src_proj_weight, size=(self.patch_size, self.patch_size), mode='bicubic', align_corners=True, antialias=False)
+            src_proj_weight = rearrange(src_proj_weight, 'b c h w -> b (c h w)')
+        targ_proj_weight.data.copy_(src_proj_weight)
+
+    def embed_patches(self, x: torch.Tensor) -> torch.Tensor:
+        patches = self.im_to_patches(x)
+        patches = self.embedder(patches)
+        return patches
+
+    def apply_pos_enc(self,
+                      patches: torch.Tensor,
+                      patch_idxs: Optional[torch.Tensor] = None,
+                      input_size: Optional[Tuple[int, int]] = None,
+    ) -> torch.Tensor:
+        if not self.abs_pos:
+            return patches
+
+        pos_enc = self.get_pos_enc(patches.shape[0], patch_idxs, input_size)
+
+        if self.training and self.pos_dropout > 0:
+            keeps = torch.rand(patches.shape[0], 1, 1, dtype=pos_enc.dtype, device=pos_enc.device) > self.pos_dropout
+            pos_enc_drop = torch.where(keeps, pos_enc, 0)
+        else:
+            pos_enc_drop = pos_enc
+
+        return patches + pos_enc_drop, pos_enc
+
+    def get_pos_enc(self,
+                    batch_size: int,
+                    patch_idxs: Optional[torch.Tensor] = None,
+                    input_size: Optional[Tuple[int, int]] = None,
+    ) -> torch.Tensor:
+        if input_size is None:
+            input_dims = self.input_dims
+        else:
+            input_dims = tuple(d // self.patch_size for d in input_size)
+
+        pos_embed = self._get_pos_embeddings(batch_size, input_dims)
+
+        if patch_idxs is None:
+            return pos_embed
+
+        exp_patch_idxs = patch_idxs.unsqueeze(-1).expand(-1, -1, pos_embed.shape[-1])
+
+        pos_embed = torch.gather(pos_embed.expand(patch_idxs.shape[0], -1, -1), dim=1, index=exp_patch_idxs)
+        return pos_embed
+
+
+    def _get_pos_embeddings(self, batch_size: int, input_dims: Tuple[int, int]):
+        if (self.num_rows, self.num_cols) == input_dims:
+            return self.pos_embed
+
+        pos_embed = self.pos_embed.reshape(1, self.num_rows, self.num_cols, -1).permute(0, 3, 1, 2)
+
+        def window_select(pos_embed):
+            if input_dims[0] < pos_embed.shape[-2]:
+                pos_embed = pos_embed[..., :input_dims[0], :]
+            if input_dims[1] < pos_embed.shape[-1]:
+                pos_embed = pos_embed[..., :, :input_dims[1]]
+            return pos_embed
+
+        if self.cpe_mode:
+            if self.training:
+                min_scale = math.sqrt(0.1)
+                scale = torch.rand(batch_size, 1, 1, device=pos_embed.device) * (1 - min_scale) + min_scale
+                aspect_min = math.log(3 / 4)
+                aspect_max = -aspect_min
+                aspect = torch.exp(torch.rand(batch_size, 1, 1, device=pos_embed.device) * (aspect_max - aspect_min) + aspect_min)
+
+                scale_x = scale * aspect
+                scale_y = scale * (1 / aspect)
+                scale_xy = torch.stack([scale_x, scale_y], dim=-1).clamp_(0, 1)
+
+                pos_xy = torch.rand(batch_size, 1, 1, 2, device=pos_embed.device) * (1 - scale_xy)
+
+                lin_x = torch.linspace(0, 1, steps=input_dims[1], device=pos_embed.device)[None, None].expand(batch_size, input_dims[0], -1)
+                lin_y = torch.linspace(0, 1, steps=input_dims[0], device=pos_embed.device)[None, :, None].expand(batch_size, -1, input_dims[1])
+
+                lin_xy = torch.stack([lin_x, lin_y], dim=-1)
+
+                grid_xy = lin_xy * scale_xy + pos_xy
+
+                # Convert to [-1, 1] range
+                grid_xy.mul_(2).sub_(1)
+
+                pos_embed = F.grid_sample(
+                    pos_embed.expand(batch_size, -1, -1, -1),
+                    grid=grid_xy,
+                    mode='bilinear',
+                    padding_mode='zeros',
+                    align_corners=True,
+                )
+            else:
+                # i_rows, i_cols = input_dims
+                # p_rows, p_cols = pos_embed.shape[2:]
+                # if i_rows <= p_rows and i_cols <= p_cols:
+                #     left = (p_cols - i_cols) // 2
+                #     top = (p_rows - i_rows) // 2
+                #     pos_embed = pos_embed[..., top:top+i_rows, left:left+i_cols]
+                # else:
+                max_dim = max(input_dims)
+                pos_embed = F.interpolate(pos_embed, size=(max_dim, max_dim), align_corners=True, mode='bilinear')
+
+                pos_embed = window_select(pos_embed)
+        else:
+            pos_embed = window_select(pos_embed)
+
+        if pos_embed.shape[-2:] != input_dims:
+            pos_embed = F.interpolate(pos_embed, size=input_dims, align_corners=True, mode='bilinear')
+
+        pos_embed = pos_embed.flatten(2).permute(0, 2, 1)
+
+        return pos_embed
+
+
+class Im2Patches(nn.Module):
+    def __init__(self, patch_size: int):
+        super().__init__()
+        self.patch_size = patch_size
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.patch_size == 1:
+            patches = x.flatten(2)
+            patches = patches.permute(0, 2, 1)
+            return patches
+
+        py = x.shape[-2] // self.patch_size
+        px = x.shape[-1] // self.patch_size
+        patches = rearrange(x, 'b c (py yy) (px xx) -> b (py px) (c yy xx)',
+                            py=py, yy=self.patch_size,
+                            px=px, xx=self.patch_size,
+        )
+        return patches
+
+
+class ViTPatchLinear(nn.Linear):
+    def __init__(self, patch_size: int, embed_dim: int, **factory):
+        super().__init__(
+            3 * (patch_size ** 2),
+            embed_dim,
+            bias=False,
+            **factory
+        )
+        self.patch_size = patch_size
+
+    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
+        if self.bias is not None:
+            self.bias.data.copy_(state_dict[f'{prefix}bias'])
+
+        chk_weight = state_dict[f'{prefix}weight']
+        if chk_weight.shape != self.weight.shape:
+            src_patch_size = int(math.sqrt(chk_weight.shape[1] // 3))
+
+            assert (src_patch_size ** 2) * 3 == chk_weight.shape[1], 'Unable to interpolate non-square patch size'
+
+            chk_weight = rearrange(chk_weight, 'b (c h w) -> b c h w', c=3, h=src_patch_size, w=src_patch_size)
+            chk_weight = F.interpolate(chk_weight, size=(self.patch_size, self.patch_size), mode='bicubic', align_corners=True, antialias=False)
+            chk_weight = rearrange(chk_weight, 'b c h w -> b (c h w)')
+        self.weight.data.copy_(chk_weight)