🐣 born

.gitignore

@@ -0,0 +1,167 @@
+# MacOS
+**/.DS_Store
+
+**/*.pth
+models/clip/_clip/configs/*
+models/clip/_clip/weights/*
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/

app.py

@@ -0,0 +1,180 @@
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+import numpy as np
+from PIL import Image
+import json, os, random
+import gradio as gr
+import torchvision.transforms.functional as TF
+from safetensors.torch import load_file  # Import the load_file function from safetensors
+from matplotlib import cm
+from huggingface_hub import hf_hub_download
+
+from typing import Tuple
+
+from models import get_model
+
+
+def resize_density_map(x: Tensor, size: Tuple[int, int]) -> Tensor:
+    x_sum = torch.sum(x, dim=(-1, -2))
+    x = F.interpolate(x, size=size, mode="bilinear")
+    scale_factor = torch.nan_to_num(torch.sum(x, dim=(-1, -2)) / x_sum, nan=0.0, posinf=0.0, neginf=0.0)
+    return x * scale_factor
+
+
+def init_seeds(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+
+
+mean = (0.485, 0.456, 0.406)
+std = (0.229, 0.224, 0.225)
+alpha = 0.8
+init_seeds(42)
+
+# -----------------------------
+# Define the model architecture
+# -----------------------------
+truncation = 4
+reduction = 8
+granularity = "fine"
+anchor_points = "average"
+
+model_name = "clip_vit_l_14"
+input_size = 224
+
+# Comment the lines below to test non-CLIP models.
+prompt_type = "word"
+num_vpt = 32
+vpt_drop = 0.
+deep_vpt = True
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+if truncation is None:  # regression, no truncation.
+    bins, anchor_points = None, None
+else:
+    with open(os.path.join("configs", f"reduction_{reduction}.json"), "r") as f:
+        config = json.load(f)[str(truncation)]["nwpu"]
+    bins = config["bins"][granularity]
+    anchor_points = config["anchor_points"][granularity]["average"] if anchor_points == "average" else config["anchor_points"][granularity]["middle"]
+    bins = [(float(b[0]), float(b[1])) for b in bins]
+    anchor_points = [float(p) for p in anchor_points]
+
+
+model = get_model(
+    backbone=model_name,
+    input_size=input_size, 
+    reduction=reduction,
+    bins=bins,
+    anchor_points=anchor_points,
+    # CLIP parameters
+    prompt_type=prompt_type,
+    num_vpt=num_vpt,
+    vpt_drop=vpt_drop,
+    deep_vpt=deep_vpt
+)
+
+repo_id = "Yiming-M/CLIP-EBC"
+filename = "nwpu_weights/CLIP_EBC_ViT_L_14/model.safetensors"
+weights_path = hf_hub_download(repo_id, filename)
+# weights_path = os.path.join("CLIP_EBC_ViT_L_14", "model.safetensors")
+state_dict = load_file(weights_path)
+new_state_dict = {}
+for k, v in state_dict.items():
+    new_state_dict[k.replace("model.", "")] = v
+model.load_state_dict(new_state_dict)
+model.to(device)
+model.eval()
+
+
+# -----------------------------
+# Preprocessing function
+# -----------------------------
+# Adjust the image transforms to match what your model expects.
+def transform(image: Image.Image):
+    assert isinstance(image, Image.Image), "Input must be a PIL Image"
+    image_tensor = TF.to_tensor(image)
+
+    image_height, image_width = image_tensor.shape[-2:]
+    if image_height < input_size or image_width < input_size:
+        # Find the ratio to resize the image while maintaining the aspect ratio
+        ratio = max(input_size / image_height, input_size / image_width)
+        new_height = int(image_height * ratio) + 1
+        new_width = int(image_width * ratio) + 1
+        image_tensor = TF.resize(image_tensor, (new_height, new_width), interpolation=TF.InterpolationMode.BICUBIC, antialias=True)
+
+    image_tensor = TF.normalize(image_tensor, mean=mean, std=std)
+    return image_tensor.unsqueeze(0)  # Add batch dimension
+
+
+
+# -----------------------------
+# Inference function
+# -----------------------------
+def predict(image: Image.Image):
+    """
+    Given an input image, preprocess it, run the model to obtain a density map,
+    compute the total crowd count, and prepare the density map for display.
+    """
+    # Preprocess the image
+    input_width, input_height = image.size
+    input_tensor = transform(image).to(device)  # shape: (1, 3, H, W)
+    
+    with torch.no_grad():
+        density_map = model(input_tensor)  # expected shape: (1, 1, H, W)
+        total_count = density_map.sum().item()
+        resized_density_map = resize_density_map(density_map, (input_height, input_width)).cpu().squeeze().numpy()
+    
+    # Normalize the density map for display purposes
+    eps = 1e-8
+    density_map_norm = (resized_density_map - resized_density_map.min()) / (resized_density_map.max() - resized_density_map.min() + eps)
+    
+    # Apply a colormap (e.g., 'jet') to get an RGBA image
+    colormap = cm.get_cmap("jet")
+    # The colormap returns values in [0,1]. Scale to [0,255] and convert to uint8.
+    density_map_color = (colormap(density_map_norm) * 255).astype(np.uint8)
+    density_map_color_img = Image.fromarray(density_map_color).convert("RGBA")
+    
+    # Ensure the original image is in RGBA format.
+    image_rgba = image.convert("RGBA")
+    overlayed_image = Image.blend(image_rgba, density_map_color_img, alpha=alpha)
+    
+    return image, overlayed_image, f"Predicted Count: {total_count:.2f}"
+
+
+# -----------------------------
+# Build Gradio Interface using Blocks for a two-column layout
+# -----------------------------
+with gr.Blocks() as demo:
+    gr.Markdown("# Crowd Counting Demo")
+    gr.Markdown("Upload an image or select an example below to see the predicted crowd density map and total count.")
+    
+    with gr.Row():
+        with gr.Column():
+            input_img = gr.Image(
+                label="Input Image",
+                sources=["upload", "clipboard"],
+                type="pil",
+            )
+            submit_btn = gr.Button("Predict")
+        with gr.Column():
+            output_img = gr.Image(label="Predicted Density Map", type="pil")
+            output_text = gr.Textbox(label="Total Count")
+    
+    submit_btn.click(fn=predict, inputs=input_img, outputs=[input_img, output_img, output_text])
+    
+    # Optional: add example images. Ensure these files are in your repo.
+    gr.Examples(
+        examples=[
+            ["example1.jpg"],
+            ["example2.jpg"]
+        ],
+        inputs=input_img,
+        label="Try an example"
+    )
+
+# Launch the app
+demo.launch()

configs/reduction_16.json

@@ -0,0 +1,33 @@
+{
+    "8":{
+        "qnrf": {
+            "bins": {
+                "fine":[
+                    [0, 0], [1, 1], [2, 2], [3, 3], [4, 4],
+                    [5, 5], [6, 6], [7, 7], [8, "inf"]
+                ],
+                "dynamic": [
+                    [0, 0], [1, 1], [2, 2], [3, 3],
+                    [4, 5], [6, 7], [8, "inf"]
+                ],
+                "coarse": [
+                    [0, 0], [1, 2], [3, 4], [5, 6], [7, "inf"]
+                ]
+            },
+            "anchor_points": {
+                "fine": {
+                    "middle": [0, 1, 2, 3, 4, 5, 6, 7, 8],
+                    "average": [0, 1, 2, 3, 4, 5, 6, 7, 9.23349]
+                },
+                "dynamic": {
+                    "middle": [0, 1, 2, 3, 4.5, 6.5, 8],
+                    "average": [0, 1, 2, 3, 4.29278, 6.31441, 9.23349]
+                },
+                "coarse": {
+                    "middle": [0, 1.5, 3.5, 5.5, 7],
+                    "average": [0, 1.14978, 3.27641, 5.30609, 8.11466]
+                }
+            }
+        }
+    }
+}

configs/reduction_32.json

@@ -0,0 +1,56 @@
+{
+    "19": {
+        "qnrf": {
+            "bins": {
+                "fine": [
+                    [0, 0], [1, 1], [2, 2], [3, 3], [4, 4],
+                    [5, 5], [6, 6], [7, 7], [8, 8], [9, 9],
+                    [10, 10], [11, 11], [12, 12], [13, 13], [14, 14],
+                    [15, 15], [16, 16], [17, 17], [18, 18], [19, "inf"]
+                ],
+                "dynamic": [
+                    [0, 0], [1, 1], [2, 2], [3, 3], [4, 4],
+                    [5, 5], [6, 6], [7, 7], [8, 8], [9, 9],
+                    [10, 11], [12, 13], [14, 15], [16, 17], [18, "inf"]
+                ],
+                "coarse": [
+                    [0, 0], [1, 2], [3, 4], [5, 6], [7, 8],
+                    [9, 10], [11, 12], [13, 14], [15, 16], [17, 18],
+                    [19, "inf"]
+                ]
+            },
+            "anchor_points": {
+                "fine": {
+                    "middle": [
+                        0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10,
+                        11, 12, 13, 14, 15, 16, 17, 18, 19
+                    ],
+                    "average": [
+                        0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10,
+                        11, 12, 13, 14, 15, 16, 17, 18, 23.01897
+                    ]
+                },
+                "dynamic": {
+                    "middle": [
+                        0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10.5,
+                        12.5, 14.5, 16.5, 18
+                    ],
+                    "average": [
+                        0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10.42903,
+                        12.43320, 14.43341, 16.43521, 21.93548
+                    ]
+                },
+                "coarse": {
+                    "middle": [
+                        0, 1.5, 3.5, 5.5, 7.5, 9.5,
+                        11.5, 13.5, 15.5, 17.5, 19
+                    ],
+                    "average": [
+                        0, 1.23498, 3.36108, 5.40298, 7.41406, 9.42356,
+                        11.43094, 13.43244, 15.43697, 17.43759, 23.01897
+                    ]
+                }
+            }
+        }
+    }
+}

configs/reduction_8.json

@@ -0,0 +1,129 @@
+{   
+    "2": {
+        "sha": {
+            "bins": {
+                "fine": [[0, 0], [1, 1], [2, "inf"]]
+            },
+            "anchor_points": {
+                "fine": {
+                    "middle": [0, 1, 2],
+                    "average": [0, 1, 2.24479]
+                }
+            }
+        },
+        "shb": {
+            "bins": {
+                "fine": [[0, 0], [1, 1], [2, "inf"]]
+            },
+            "anchor_points": {
+                "fine": {
+                    "middle": [0, 1, 2],
+                    "average": [0, 1, 2.15171]
+                }
+            }
+        },
+        "nwpu": {
+            "bins": {
+                "fine": [[0, 0], [1, 1], [2, "inf"]]
+            },
+            "anchor_points": {
+                "fine": {
+                    "middle": [0, 1, 2],
+                    "average": [0, 1, 2.10737]
+                }
+            }
+        },
+        "qnrf": {
+            "bins": {
+                "fine": [[0, 0], [1, 1], [2, "inf"]]
+            },
+            "anchor_points": {
+                "fine": {
+                    "middle": [0, 1, 2],
+                    "average": [0, 1, 2.09296]
+                }
+            }
+        },
+        "jhu": {
+            "bins": {
+                "fine": [[0, 0], [1, 1], [2, "inf"]]
+            },
+            "anchor_points": {
+                "fine": {
+                    "middle": [0, 1, 2],
+                    "average": [0, 1, 2.18589]
+                }
+            }
+        }
+    },
+    "4": {
+        "sha": {
+            "bins": {
+                "fine": [[0, 0], [1, 1], [2, 2], [3, 3], [4, "inf"]]
+            },
+            "anchor_points": {
+                "fine": {
+                    "middle": [0, 1, 2, 3, 4],
+                    "average": [0, 1, 2, 3, 4.29992]
+                }
+            }
+        },
+        "shb": {
+            "bins": {
+                "fine": [[0, 0], [1, 1], [2, 2], [3, 3], [4, "inf"]]
+            },
+            "anchor_points": {
+                "fine": {
+                    "middle": [0, 1, 2, 3, 4],
+                    "average": [0, 1, 2, 3, 4.41009]
+                }
+            }
+        },
+        "nwpu": {
+            "bins": {
+                "fine": [[0, 0], [1, 1], [2, 2], [3, 3], [4, "inf"]]
+            },
+            "anchor_points": {
+                "fine": {
+                    "middle": [0, 1, 2, 3, 4],
+                    "average": [0, 1, 2, 3, 4.21931]
+                }
+            }
+        },
+        "qnrf": {
+            "bins": {
+                "fine": [[0, 0], [1, 1], [2, 2], [3, 3], [4, "inf"]]
+            },
+            "anchor_points": {
+                "fine": {
+                    "middle": [0, 1, 2, 3, 4],
+                    "average": [0, 1, 2, 3, 4.21937]
+                }
+            }
+        },
+        "jhu": {
+            "bins": {
+                "fine": [[0, 0], [1, 1], [2, 2], [3, 3], [4, "inf"]]
+            },
+            "anchor_points": {
+                "fine": {
+                    "middle": [0, 1, 2, 3, 4],
+                    "average": [0, 1, 2, 3, 4.24058]
+                }
+            }
+        }
+    },
+    "11": {
+        "qnrf": {
+            "bins": {
+                "fine": [[0, 0], [1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6], [7, 7], [8, 8], [9, 9], [10, 10], [11, "inf"]]
+            },
+            "anchor_points": {
+                "fine": {
+                    "middle": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11],
+                    "average": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
+                }
+            }
+        }
+    }
+}

models/__init__.py

@@ -0,0 +1,49 @@
+from typing import List, Tuple, Optional, Any, Union
+
+from .model import _classifier, _regressor, Classifier, Regressor
+from .clip import _clip_ebc, CLIP_EBC
+
+
+clip_names = ["resnet50", "resnet50x4", "resnet50x16", "resnet50x64", "resnet101", "vit_b_16", "vit_b_32", "vit_l_14"]
+
+
+def get_model(
+    backbone: str,
+    input_size: int,
+    reduction: int,
+    bins: Optional[List[Tuple[float, float]]] = None,
+    anchor_points: Optional[List[float]] = None,
+    **kwargs: Any,
+) -> Union[Regressor, Classifier, CLIP_EBC]:
+    backbone = backbone.lower()
+    if "clip" in backbone:
+        backbone = backbone[5:]
+        assert backbone in clip_names, f"Expected backbone to be in {clip_names}, got {backbone}"
+        return _clip_ebc(
+            backbone=backbone,
+            input_size=input_size,
+            reduction=reduction,
+            bins=bins,
+            anchor_points=anchor_points,
+            **kwargs
+        )
+    elif bins is None and anchor_points is None:
+        return _regressor(
+            backbone=backbone,
+            input_size=input_size,
+            reduction=reduction,
+        )
+    else:
+        assert bins is not None and anchor_points is not None, f"Expected bins and anchor_points to be both None or not None, got {bins} and {anchor_points}"
+        return _classifier(
+            backbone=backbone,
+            input_size=input_size,
+            reduction=reduction,
+            bins=bins,
+            anchor_points=anchor_points,
+        )
+
+
+__all__ = [
+    "get_model",
+]

models/clip/__init__.py

@@ -0,0 +1,7 @@
+from .model import CLIP_EBC, _clip_ebc
+
+
+__all__ = [
+    "CLIP_EBC",
+    "_clip_ebc",
+]

models/clip/_clip/__init__.py

@@ -0,0 +1,273 @@
+import torch
+import os
+from typing import Tuple, Optional, Any, Union
+import json
+
+from .utils import tokenize, transform
+from .prepare import prepare
+from .text_encoder import CLIPTextEncoder
+from .image_encoder import ModifiedResNet, VisionTransformer
+from .model import CLIP
+
+
+curr_dir = os.path.dirname(os.path.abspath(__file__))
+
+clip_model_names = [
+    "clip_resnet50",
+    "clip_resnet101",
+    "clip_resnet50x4",
+    "clip_resnet50x16",
+    "clip_resnet50x64",
+    "clip_vit_b_32",
+    "clip_vit_b_16",
+    "clip_vit_l_14",
+    "clip_vit_l_14_336px",
+]
+
+clip_image_encoder_names = [f"clip_image_encoder_{name[5:]}" for name in clip_model_names]
+clip_text_encoder_names = [f"clip_text_encoder_{name[5:]}" for name in clip_model_names]
+
+
+for name in clip_model_names + clip_image_encoder_names + clip_text_encoder_names:
+    model_weights_path = os.path.join(curr_dir, "weights", f"{name}.pth")
+    model_config_path = os.path.join(curr_dir, "configs", f"{name}.json")
+    if not os.path.exists(os.path.join(curr_dir, "weights", f"{name}.pth")) or not os.path.exists(os.path.join(curr_dir, "configs", f"{name}.json")):
+        prepare()
+        break
+
+
+for name in clip_model_names + clip_image_encoder_names + clip_text_encoder_names:
+    assert os.path.exists(os.path.join(curr_dir, "weights", f"{name}.pth")), f"Missing {name}.pth in weights folder. Please run models/clip/prepare.py to download the weights."
+    assert os.path.exists(os.path.join(curr_dir, "configs", f"{name}.json")), f"Missing {name}.json in configs folder. Please run models/clip/prepare.py to download the configs."
+
+
+def _clip(name: str, input_size: Optional[Union[int, Tuple[int, int]]] = None) -> CLIP:
+    with open(os.path.join(curr_dir, "configs", f"clip_{name}.json"), "r") as f:
+        config = json.load(f)
+
+    model = CLIP(
+        embed_dim=config["embed_dim"],
+        # vision
+        image_resolution=config["image_resolution"],
+        vision_layers=config["vision_layers"],
+        vision_width=config["vision_width"],
+        vision_patch_size=config["vision_patch_size"],
+        # text
+        context_length=config["context_length"],
+        vocab_size=config["vocab_size"],
+        transformer_width=config["transformer_width"],
+        transformer_heads=config["transformer_heads"],
+        transformer_layers=config["transformer_layers"]
+    )
+    state_dict = torch.load(os.path.join(curr_dir, "weights", f"clip_{name}.pth"), map_location="cpu")
+    model.load_state_dict(state_dict, strict=True)
+
+    if input_size is not None:
+        input_size = (input_size, input_size) if isinstance(input_size, int) else input_size
+        if name.startswith("vit"):
+            model.visual.adjust_pos_embed(*input_size)
+
+    return model
+
+
+def _resnet(
+    name: str,
+    reduction: int = 32,
+    features_only: bool = False,
+    out_indices: Optional[Tuple[int, ...]] = None,
+    **kwargs: Any
+) -> ModifiedResNet:
+    with open(os.path.join(curr_dir, "configs", f"clip_image_encoder_{name}.json"), "r") as f:
+        config = json.load(f)
+    model = ModifiedResNet(
+        layers=config["vision_layers"],
+        output_dim=config["embed_dim"],
+        input_resolution=config["image_resolution"],
+        width=config["vision_width"],
+        heads=config["vision_heads"],
+        features_only=features_only,
+        out_indices=out_indices,
+        reduction=reduction
+    )
+    state_dict = torch.load(os.path.join(curr_dir, "weights", f"clip_image_encoder_{name}.pth"), map_location="cpu")
+    missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
+    if len(missing_keys) > 0 or len(unexpected_keys) > 0:
+        print(f"Missing keys: {missing_keys}")
+        print(f"Unexpected keys: {unexpected_keys}")
+    else:
+        print(f"All keys matched successfully.")
+
+    return model
+
+
+def _vit(name: str, features_only: bool = False, input_size: Optional[Union[int, Tuple[int, int]]] = None, **kwargs: Any) -> VisionTransformer:
+    with open(os.path.join(curr_dir, "configs", f"clip_image_encoder_{name}.json"), "r") as f:
+        config = json.load(f)
+    model = VisionTransformer(
+        input_resolution=config["image_resolution"],
+        patch_size=config["vision_patch_size"],
+        output_dim=config["embed_dim"],
+        width=config["vision_width"],
+        layers=config["vision_layers"],
+        heads=config["vision_heads"],
+        features_only=features_only
+    )
+    state_dict = torch.load(os.path.join(curr_dir, "weights", f"clip_image_encoder_{name}.pth"), map_location="cpu")
+    missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
+    if len(missing_keys) > 0 or len(unexpected_keys) > 0:
+        print(f"Missing keys: {missing_keys}")
+        print(f"Unexpected keys: {unexpected_keys}")
+    else:
+        print(f"All keys matched successfully.")
+
+    if input_size is not None:
+        input_size = (input_size, input_size) if isinstance(input_size, int) else input_size
+        model.adjust_pos_embed(*input_size)
+    return model
+
+
+def _text_encoder(name: str) -> CLIPTextEncoder:
+    with open(os.path.join(curr_dir, "configs", f"clip_text_encoder_{name}.json"), "r") as f:
+        config = json.load(f)
+    model = CLIPTextEncoder(
+        embed_dim=config["embed_dim"],
+        context_length=config["context_length"],
+        vocab_size=config["vocab_size"],
+        transformer_width=config["transformer_width"],
+        transformer_heads=config["transformer_heads"],
+        transformer_layers=config["transformer_layers"]
+    )
+    state_dict = torch.load(os.path.join(curr_dir, "weights", f"clip_text_encoder_{name}.pth"), map_location="cpu")
+    missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
+    if len(missing_keys) > 0 or len(unexpected_keys) > 0:
+        print(f"Missing keys: {missing_keys}")
+        print(f"Unexpected keys: {unexpected_keys}")
+    else:
+        print(f"All keys matched successfully.")
+
+    return model
+
+
+
+# CLIP models
+def resnet50_clip(input_size: Optional[Union[int, Tuple[int, int]]] = None) -> CLIP:
+    return _clip("resnet50", input_size)
+
+def resnet101_clip(input_size: Optional[Union[int, Tuple[int, int]]] = None) -> CLIP:
+    return _clip("resnet101", input_size)
+
+def resnet50x4_clip(input_size: Optional[Union[int, Tuple[int, int]]] = None) -> CLIP:
+    return _clip("resnet50x4", input_size)
+
+def resnet50x16_clip(input_size: Optional[Union[int, Tuple[int, int]]] = None) -> CLIP:
+    return _clip("resnet50x16", input_size)
+
+def resnet50x64_clip(input_size: Optional[Union[int, Tuple[int, int]]] = None) -> CLIP:
+    return _clip("resnet50x64", input_size)
+
+def vit_b_32_clip(input_size: Optional[Union[int, Tuple[int, int]]] = None) -> CLIP:
+    return _clip("vit_b_32", input_size)
+
+def vit_b_16_clip(input_size: Optional[Union[int, Tuple[int, int]]] = None) -> CLIP:
+    return _clip("vit_b_16", input_size)
+
+def vit_l_14_clip(input_size: Optional[Union[int, Tuple[int, int]]] = None) -> CLIP:
+    return _clip("vit_l_14", input_size)
+
+def vit_l_14_336px_clip(input_size: Optional[Union[int, Tuple[int, int]]] = None) -> CLIP:
+    return _clip("vit_l_14_336px", input_size)
+
+
+# CLIP image encoders
+def resnet50_img(features_only: bool = False, out_indices: Optional[Tuple[int, ...]] = None, **kwargs: Any) -> ModifiedResNet:
+    return _resnet("resnet50", features_only=features_only, out_indices=out_indices, **kwargs)
+
+def resnet101_img(features_only: bool = False, out_indices: Optional[Tuple[int, ...]] = None, **kwargs: Any) -> ModifiedResNet:
+    return _resnet("resnet101", features_only=features_only, out_indices=out_indices, **kwargs)
+
+def resnet50x4_img(features_only: bool = False, out_indices: Optional[Tuple[int, ...]] = None, **kwargs: Any) -> ModifiedResNet:
+    return _resnet("resnet50x4", features_only=features_only, out_indices=out_indices, **kwargs)
+
+def resnet50x16_img(features_only: bool = False, out_indices: Optional[Tuple[int, ...]] = None, **kwargs: Any) -> ModifiedResNet:
+    return _resnet("resnet50x16", features_only=features_only, out_indices=out_indices, **kwargs)
+
+def resnet50x64_img(features_only: bool = False, out_indices: Optional[Tuple[int, ...]] = None, **kwargs: Any) -> ModifiedResNet:
+    return _resnet("resnet50x64", features_only=features_only, out_indices=out_indices, **kwargs)
+
+def vit_b_32_img(features_only: bool = False, input_size: Optional[Union[int, Tuple[int, int]]] = None, **kwargs: Any) -> VisionTransformer:
+    return _vit("vit_b_32", features_only=features_only, input_size=input_size, **kwargs)
+
+def vit_b_16_img(features_only: bool = False, input_size: Optional[Union[int, Tuple[int, int]]] = None, **kwargs: Any) -> VisionTransformer:
+    return _vit("vit_b_16", features_only=features_only, input_size=input_size, **kwargs)
+
+def vit_l_14_img(features_only: bool = False, input_size: Optional[Union[int, Tuple[int, int]]] = None, **kwargs: Any) -> VisionTransformer:
+    return _vit("vit_l_14", features_only=features_only, input_size=input_size, **kwargs)
+
+def vit_l_14_336px_img(features_only: bool = False, input_size: Optional[Union[int, Tuple[int, int]]] = None, **kwargs: Any) -> VisionTransformer:
+    return _vit("vit_l_14_336px", features_only=features_only, input_size=input_size, **kwargs)
+
+
+# CLIP text encoders
+def resnet50_txt() -> CLIPTextEncoder:
+    return _text_encoder("resnet50")
+
+def resnet101_txt() -> CLIPTextEncoder:
+    return _text_encoder("resnet101")
+
+def resnet50x4_txt() -> CLIPTextEncoder:
+    return _text_encoder("resnet50x4")
+
+def resnet50x16_txt() -> CLIPTextEncoder:
+    return _text_encoder("resnet50x16")
+
+def resnet50x64_txt() -> CLIPTextEncoder:
+    return _text_encoder("resnet50x64")
+
+def vit_b_32_txt() -> CLIPTextEncoder:
+    return _text_encoder("vit_b_32")
+
+def vit_b_16_txt() -> CLIPTextEncoder:
+    return _text_encoder("vit_b_16")
+
+def vit_l_14_txt() -> CLIPTextEncoder:
+    return _text_encoder("vit_l_14")
+
+def vit_l_14_336px_txt() -> CLIPTextEncoder:
+    return _text_encoder("vit_l_14_336px")
+
+
+__all__ = [
+    # utils
+    "tokenize",
+    "transform",
+    # clip models
+    "resnet50_clip",
+    "resnet101_clip",
+    "resnet50x4_clip",
+    "resnet50x16_clip",
+    "resnet50x64_clip",
+    "vit_b_32_clip",
+    "vit_b_16_clip",
+    "vit_l_14_clip",
+    "vit_l_14_336px_clip",
+    # clip image encoders
+    "resnet50_img",
+    "resnet101_img",
+    "resnet50x4_img",
+    "resnet50x16_img",
+    "resnet50x64_img",
+    "vit_b_32_img",
+    "vit_b_16_img",
+    "vit_l_14_img",
+    "vit_l_14_336px_img",
+    # clip text encoders
+    "resnet50_txt",
+    "resnet101_txt",
+    "resnet50x4_txt",
+    "resnet50x16_txt",
+    "resnet50x64_txt",
+    "vit_b_32_txt",
+    "vit_b_16_txt",
+    "vit_l_14_txt",
+    "vit_l_14_336px_txt",
+]

models/clip/_clip/blocks.py

@@ -0,0 +1,137 @@
+import torch
+from torch import nn, Tensor
+import torch.nn.functional as F
+from collections import OrderedDict
+from typing import Optional, Iterable
+
+
+class LayerNorm(nn.LayerNorm):
+    """Subclass torch's LayerNorm to handle fp16."""
+
+    def forward(self, x: Tensor):
+        orig_type = x.dtype
+        ret = super().forward(x.type(torch.float32))
+        return ret.type(orig_type)
+
+
+class QuickGELU(nn.Module):
+    def forward(self, x: Tensor):
+        return x * torch.sigmoid(1.702 * x)
+
+
+class ResidualAttentionBlock(nn.Module):
+    def __init__(self, d_model: int, n_head: int, attn_mask: Tensor = None):
+        super().__init__()
+        self.attn = nn.MultiheadAttention(d_model, n_head)
+        self.ln_1 = LayerNorm(d_model)
+        self.mlp = nn.Sequential(OrderedDict([
+            ("c_fc", nn.Linear(d_model, d_model * 4)),
+            ("gelu", QuickGELU()),
+            ("c_proj", nn.Linear(d_model * 4, d_model))
+        ]))
+        self.ln_2 = LayerNorm(d_model)
+        self.attn_mask = attn_mask
+
+    def attention(self, x: Tensor):
+        self.attn_mask = self.attn_mask.to(dtype=x.dtype, device=x.device) if self.attn_mask is not None else None
+        return self.attn(x, x, x, need_weights=False, attn_mask=self.attn_mask)[0]
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = x + self.attention(self.ln_1(x))
+        x = x + self.mlp(self.ln_2(x))
+        return x
+
+
+class Transformer(nn.Module):
+    def __init__(self, width: int, layers: int, heads: int, attn_mask: Tensor = None):
+        super().__init__()
+        self.width = width
+        self.layers = layers
+        self.resblocks = nn.Sequential(*[ResidualAttentionBlock(width, heads, attn_mask) for _ in range(layers)])
+
+    def forward(self, x: Tensor):
+        return self.resblocks(x)
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1):
+        super().__init__()
+
+        # all conv layers have stride 1. an avgpool is performed after the second convolution when stride > 1
+        self.conv1 = nn.Conv2d(inplanes, planes, 1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.relu1 = nn.ReLU(inplace=True)
+
+        self.conv2 = nn.Conv2d(planes, planes, 3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.relu2 = nn.ReLU(inplace=True)
+
+        self.avgpool = nn.AvgPool2d(stride) if stride > 1 else nn.Identity()
+
+        self.conv3 = nn.Conv2d(planes, planes * self.expansion, 1, bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+        self.relu3 = nn.ReLU(inplace=True)
+
+        self.downsample = None
+        self.stride = stride
+
+        if stride > 1 or inplanes != planes * Bottleneck.expansion:
+            # downsampling layer is prepended with an avgpool, and the subsequent convolution has stride 1
+            self.downsample = nn.Sequential(OrderedDict([
+                ("-1", nn.AvgPool2d(stride)),
+                ("0", nn.Conv2d(inplanes, planes * self.expansion, 1, stride=1, bias=False)),
+                ("1", nn.BatchNorm2d(planes * self.expansion))
+            ]))
+
+    def forward(self, x: Tensor):
+        identity = x
+
+        out = self.relu1(self.bn1(self.conv1(x)))
+        out = self.relu2(self.bn2(self.conv2(out)))
+        out = self.avgpool(out)
+        out = self.bn3(self.conv3(out))
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu3(out)
+        return out
+
+
+class AttentionPool2d(nn.Module):
+    def __init__(self, spacial_dim: int, embed_dim: int, num_heads: int, output_dim: int = None):
+        super().__init__()
+        self.positional_embedding = nn.Parameter(torch.randn(spacial_dim + 1, embed_dim) / embed_dim ** 0.5)
+        self.k_proj = nn.Linear(embed_dim, embed_dim)
+        self.q_proj = nn.Linear(embed_dim, embed_dim)
+        self.v_proj = nn.Linear(embed_dim, embed_dim)
+        self.c_proj = nn.Linear(embed_dim, output_dim or embed_dim)
+        self.num_heads = num_heads
+
+    def forward(self, x):
+        x = x.flatten(start_dim=2).permute(2, 0, 1)  # NCHW -> (HW)NC
+        x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0)  # (HW+1)NC
+        x = x + self.positional_embedding[:, None, :].to(x.dtype)  # (HW+1)NC
+        x, _ = F.multi_head_attention_forward(
+            query=x[:1], key=x, value=x,
+            embed_dim_to_check=x.shape[-1],
+            num_heads=self.num_heads,
+            q_proj_weight=self.q_proj.weight,
+            k_proj_weight=self.k_proj.weight,
+            v_proj_weight=self.v_proj.weight,
+            in_proj_weight=None,
+            in_proj_bias=torch.cat([self.q_proj.bias, self.k_proj.bias, self.v_proj.bias]),
+            bias_k=None,
+            bias_v=None,
+            add_zero_attn=False,
+            dropout_p=0,
+            out_proj_weight=self.c_proj.weight,
+            out_proj_bias=self.c_proj.bias,
+            use_separate_proj_weight=True,
+            training=self.training,
+            need_weights=False
+        )
+        return x.squeeze(0)

models/clip/_clip/bpe_simple_vocab_16e6.txt.gz

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

models/clip/_clip/image_encoder.py

@@ -0,0 +1,225 @@
+import torch
+from torch import nn, Tensor
+import torch.nn.functional as F
+from einops import rearrange
+from typing import Tuple, Union, Any, List, Iterable, Optional
+
+from .blocks import LayerNorm, Transformer, Bottleneck, AttentionPool2d
+
+
+class ModifiedResNet(nn.Module):
+    """
+    A ResNet class that is similar to torchvision's but contains the following changes:
+    - There are now 3 "stem" convolutions as opposed to 1, with an average pool instead of a max pool.
+    - Performs anti-aliasing strided convolutions, where an avgpool is prepended to convolutions with stride > 1
+    - The final pooling layer is a QKV attention instead of an average pool
+    """
+    def __init__(
+        self,
+        layers: Tuple[int, int, int, int],
+        output_dim: int,
+        input_resolution: int = 224,
+        width: int = 64,
+        heads: int = 8,
+        features_only: bool = False,
+        out_indices: Optional[Iterable[int]] = None,
+        reduction: int = 32,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__()
+        input_resolution = (input_resolution, input_resolution) if isinstance(input_resolution, int) else input_resolution
+        assert isinstance(input_resolution, tuple) and len(input_resolution) == 2, f"input_resolution should be a tuple of length 2, but got {input_resolution}"
+        self.input_resolution = input_resolution
+        self.downsampling_rate = 32  # the rate at which the input is downsampled by the network
+
+        # the 3-layer stem
+        self.conv1 = nn.Conv2d(3, width // 2, kernel_size=3, stride=2, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(width // 2)
+        self.relu1 = nn.ReLU(inplace=True)
+        self.conv2 = nn.Conv2d(width // 2, width // 2, kernel_size=3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(width // 2)
+        self.relu2 = nn.ReLU(inplace=True)
+        self.conv3 = nn.Conv2d(width // 2, width, kernel_size=3, padding=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(width)
+        self.relu3 = nn.ReLU(inplace=True)
+        self.avgpool = nn.AvgPool2d(2)
+
+        # residual layers
+        self._inplanes = width  # this is a *mutable* variable used during construction
+        self.layer1 = self._make_layer(width, layers[0])
+        self.layer2 = self._make_layer(width * 2, layers[1], stride=2)
+        self.layer3 = self._make_layer(width * 4, layers[2], stride=2)
+        self.layer4 = self._make_layer(width * 8, layers[3], stride=1 if reduction <= 16 else 2)
+
+        self.features_only = features_only
+        if features_only:
+            self.out_indices = out_indices if out_indices is not None else range(5)
+            self.out_indices = [idx + 5 if idx < 0 else idx for idx in self.out_indices]  # map negative indices to positive indices
+            self.out_indices = sorted(set(self.out_indices))  # remove duplicates and sort
+            assert min(self.out_indices) >= 0 and max(self.out_indices) <= 4, f"out_indices={self.out_indices} is invalid for a ResNet with 5 stages"
+            self.channels = width * 32  # the ResNet feature dimension
+        else:
+            self.out_indices = None
+            embed_dim = width * 32  # the ResNet feature dimension
+            self.attnpool = AttentionPool2d((input_resolution[0] // 32) * (input_resolution[1] // 32), embed_dim, heads, output_dim)
+            self.channels = output_dim
+
+        self.reduction = self.downsampling_rate // 2 if reduction <= 16 else self.downsampling_rate
+        self.clip_embed_dim = output_dim
+
+    def _make_layer(self, planes, blocks, stride=1):
+        layers = [Bottleneck(self._inplanes, planes, stride)]
+
+        self._inplanes = planes * Bottleneck.expansion
+        for _ in range(1, blocks):
+            layers.append(Bottleneck(self._inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def _stem(self, x: Tensor) -> Tensor:
+        x = self.relu1(self.bn1(self.conv1(x)))
+        x = self.relu2(self.bn2(self.conv2(x)))
+        x = self.relu3(self.bn3(self.conv3(x)))
+        x = self.avgpool(x)
+        return x
+
+    def forward(self, x: Tensor) -> Union[Tensor, List[Tensor]]:
+        x = x.type(self.conv1.weight.dtype)
+        x = self._stem(x)
+
+        feats = [x] if self.features_only and 0 in self.out_indices else []
+
+        x = self.layer1(x)
+        if self.features_only and 1 in self.out_indices:
+            feats.append(x)
+
+        x = self.layer2(x)
+        if self.features_only and 2 in self.out_indices:
+            feats.append(x)
+
+        x = self.layer3(x)
+        if self.features_only and 3 in self.out_indices:
+            feats.append(x)
+
+        x = self.layer4(x)
+        if self.features_only and 4 in self.out_indices:
+            feats.append(x)
+
+        if self.features_only:
+            if len(self.out_indices) == 1:
+                return feats[0]
+            else:
+                return feats
+        else:
+            x = self.attnpool(x)
+            return x
+
+
+class VisionTransformer(nn.Module):
+    def __init__(
+        self,
+        input_resolution: Union[int, Tuple[int, int]],
+        patch_size: Union[int, Tuple[int, int]],
+        output_dim: int,
+        width: int,
+        layers: int,
+        heads: int,
+        features_only: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__()
+        input_resolution = (input_resolution, input_resolution) if isinstance(input_resolution, int) else input_resolution
+        patch_size = (patch_size, patch_size) if isinstance(patch_size, int) else patch_size
+        assert isinstance(input_resolution, tuple) and len(input_resolution) == 2, f"input_resolution should be a tuple of length 2, but got {input_resolution}"
+        assert isinstance(patch_size, tuple) and len(patch_size) == 2, f"patch_size should be a tuple of length 2, but got {patch_size}"
+        assert patch_size[0] == patch_size[1], f"ViT only supports square patches, patch_size={patch_size} is invalid."
+        assert input_resolution[0] % patch_size[0] == 0 and input_resolution[1] % patch_size[1] == 0, f"input_resolution {input_resolution} should be divisible by patch_size {patch_size}"
+        self.input_resolution = input_resolution
+        self.patch_size = patch_size
+        self.downsampling_rate = patch_size[0]
+
+        self.conv1 = nn.Conv2d(in_channels=3, out_channels=width, kernel_size=patch_size, stride=patch_size, bias=False)
+
+        scale = width ** -0.5
+        self.class_embedding = nn.Parameter(scale * torch.randn(width))
+        self.num_patches_h = int(input_resolution[0] // patch_size[0])
+        self.num_patches_w = int(input_resolution[1] // patch_size[1])
+        self.positional_embedding = nn.Parameter(scale * torch.randn(self.num_patches_h * self.num_patches_w + 1, width))
+        self.ln_pre = LayerNorm(width)
+
+        self.transformer = Transformer(width, layers, heads)
+        self.ln_post = LayerNorm(width)
+
+        self.features_only = features_only  # if True, return the final patches instead of the CLS token
+        if features_only:
+            self.channels = width
+        else:
+            self.proj = nn.Parameter(scale * torch.randn(width, output_dim))
+            self.channels = output_dim
+
+        self.reduction = patch_size[0]
+        self.clip_embed_dim = output_dim
+
+    def adjust_pos_embed(self, h: int, w: int) -> None:
+        """
+        Permanently adjust the size of the positional embedding matrix.
+
+        Args:
+            h: the height of the original input image.
+            w: the width of the original input image.
+        """
+        assert h % self.patch_size[0] == 0 and w % self.patch_size[1] == 0, f"input_resolution {h, w} should be divisible by patch_size {self.patch_size}"
+        if self.input_resolution[0] != h or self.input_resolution[1] != w:
+            new_num_patches_h = int(h // self.patch_size[0])
+            new_num_patches_w = int(w // self.patch_size[1])
+            positional_embedding = rearrange(self.positional_embedding[1:, :], "(h w) c -> c h w", h=self.num_patches_h, w=self.num_patches_w).unsqueeze(0)  # add batch dimension
+            positional_embedding = F.interpolate(positional_embedding, size=(new_num_patches_h, new_num_patches_w), mode="bicubic", ).squeeze(0)  # remove batch dimension
+            positional_embedding = rearrange(positional_embedding, "c h w -> (h w) c")
+            self.positional_embedding = nn.Parameter(torch.cat([self.positional_embedding[:1, :], positional_embedding], dim=0))
+            self.input_resolution = (h, w)
+            self.num_patches_h = new_num_patches_h
+            self.num_patches_w = new_num_patches_w
+
+    def _interpolate_pos_embed(self, h: int, w: int) -> Tensor:
+        """
+        Interpolate the positional embedding matrix to match the size of the input image.
+
+        Args:
+            h: the required number of patches along the height dimension.
+            w: the required number of patches along the width dimension.
+        """
+        if h == self.num_patches_h and w == self.num_patches_w:
+            return self.positional_embedding
+        else:
+            positional_embedding = rearrange(self.positional_embedding[1:, :], "(h w) c -> c h w", h=self.num_patches_h, w=self.num_patches_w).unsqueeze(0)  # add batch dimension
+            positional_embedding = F.interpolate(positional_embedding, size=(h, w), mode="bicubic").squeeze(0)  # remove batch dimension
+            positional_embedding = rearrange(positional_embedding, "c h w -> (h w) c")
+            positional_embedding = torch.cat([self.positional_embedding[:1, :], positional_embedding], dim=0)
+            return positional_embedding
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.conv1(x) # shape = [*, width, grid, grid]
+        num_patches_h, num_patches_w = x.shape[-2:]
+
+        positional_embedding = self._interpolate_pos_embed(num_patches_h, num_patches_w).to(x.dtype)
+        x = x.reshape(x.shape[0], x.shape[1], -1)  # shape = [*, width, grid ** 2]
+        x = x.permute(0, 2, 1)  # shape = [*, grid ** 2, width]
+        x = torch.cat([
+                self.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device), 
+                x
+            ], dim=1)
+        x = x + positional_embedding
+        x = self.ln_pre(x)
+
+        x = x.permute(1, 0, 2)  # NLD -> LND. N: batch size, L: sequence length, D: feature dimension
+        x = self.transformer(x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        x = self.ln_post(x)
+
+        if self.features_only:
+            x = x[:, 1:, :]  # remove the CLS token
+            x = rearrange(x, "n (h w) c -> n c h w", h=num_patches_h, w=num_patches_w)
+        else:
+            x = x[:, 0, :]
+            x = x @ self.proj
+        return x

models/clip/_clip/model.py

@@ -0,0 +1,214 @@
+import torch
+from torch import nn
+import numpy as np
+
+from typing import Tuple, Union
+
+from .image_encoder import ModifiedResNet, VisionTransformer
+from .text_encoder import LayerNorm, Transformer
+
+
+class CLIP(nn.Module):
+    def __init__(
+        self,
+        embed_dim: int,
+        # vision
+        image_resolution: int,
+        vision_layers: Union[Tuple[int, int, int, int], int],
+        vision_width: int,
+        vision_patch_size: int,
+        # text
+        context_length: int,
+        vocab_size: int,
+        transformer_width: int,
+        transformer_heads: int,
+        transformer_layers: int
+        ) -> None:
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.image_resolution = image_resolution
+        self.vision_layers = vision_layers
+        self.vision_width = vision_width
+        self.vision_patch_size = vision_patch_size
+        self.context_length = context_length
+        self.vocab_size = vocab_size
+        self.transformer_width = transformer_width
+        self.transformer_heads = transformer_heads
+        self.transformer_layers = transformer_layers
+
+        if isinstance(vision_layers, (tuple, list)):
+            vision_heads = vision_width * 32 // 64
+            self.visual = ModifiedResNet(
+                layers=vision_layers,
+                output_dim=embed_dim,
+                heads=vision_heads,
+                input_resolution=image_resolution,
+                width=vision_width,
+                features_only=False,
+            )
+        else:
+            vision_heads = vision_width // 64
+            self.visual = VisionTransformer(
+                input_resolution=image_resolution,
+                patch_size=vision_patch_size,
+                width=vision_width,
+                layers=vision_layers,
+                heads=vision_heads,
+                output_dim=embed_dim,
+                features_only=False,
+            )
+        self.vision_heads = vision_heads
+        self.transformer = Transformer(
+            width=transformer_width,
+            layers=transformer_layers,
+            heads=transformer_heads,
+            attn_mask=self.build_attention_mask()
+        )
+
+        self.token_embedding = nn.Embedding(vocab_size, transformer_width)
+        self.positional_embedding = nn.Parameter(torch.empty(self.context_length, transformer_width))
+        self.ln_final = LayerNorm(transformer_width)
+
+        self.text_projection = nn.Parameter(torch.empty(transformer_width, embed_dim))
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+        self.initialize_parameters()
+
+    def initialize_parameters(self):
+        nn.init.normal_(self.token_embedding.weight, std=0.02)
+        nn.init.normal_(self.positional_embedding, std=0.01)
+
+        if isinstance(self.visual, ModifiedResNet):
+            if self.visual.attnpool is not None:
+                std = self.visual.attnpool.c_proj.in_features ** -0.5
+                nn.init.normal_(self.visual.attnpool.q_proj.weight, std=std)
+                nn.init.normal_(self.visual.attnpool.k_proj.weight, std=std)
+                nn.init.normal_(self.visual.attnpool.v_proj.weight, std=std)
+                nn.init.normal_(self.visual.attnpool.c_proj.weight, std=std)
+
+            for resnet_block in [self.visual.layer1, self.visual.layer2, self.visual.layer3, self.visual.layer4]:
+                for name, param in resnet_block.named_parameters():
+                    if name.endswith("bn3.weight"):
+                        nn.init.zeros_(param)
+
+        proj_std = (self.transformer.width ** -0.5) * ((2 * self.transformer.layers) ** -0.5)
+        attn_std = self.transformer.width ** -0.5
+        fc_std = (2 * self.transformer.width) ** -0.5
+        for block in self.transformer.resblocks:
+            nn.init.normal_(block.attn.in_proj_weight, std=attn_std)
+            nn.init.normal_(block.attn.out_proj.weight, std=proj_std)
+            nn.init.normal_(block.mlp.c_fc.weight, std=fc_std)
+            nn.init.normal_(block.mlp.c_proj.weight, std=proj_std)
+
+        if self.text_projection is not None:
+            nn.init.normal_(self.text_projection, std=self.transformer.width ** -0.5)
+
+    def build_attention_mask(self):
+        # lazily create causal attention mask, with full attention between the vision tokens
+        # pytorch uses additive attention mask; fill with -inf
+        mask = torch.empty(self.context_length, self.context_length)
+        mask.fill_(float("-inf"))
+        mask.triu_(1)  # zero out the lower diagonal
+        return mask
+
+    @property
+    def dtype(self):
+        return self.visual.conv1.weight.dtype
+
+    def encode_image(self, image):
+        return self.visual(image.type(self.dtype))
+
+    def encode_text(self, text):
+        x = self.token_embedding(text).type(self.dtype)  # [batch_size, n_ctx, d_model]
+
+        x = x + self.positional_embedding.type(self.dtype)
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.transformer(x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        x = self.ln_final(x).type(self.dtype)
+
+        # x.shape = [batch_size, n_ctx, transformer.width]
+        # take features from the eot embedding (eot_token is the highest number in each sequence)
+        x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
+
+        return x
+
+    def forward(self, image, text):
+        image_features = self.encode_image(image)
+        text_features = self.encode_text(text)
+
+        # normalized features
+        image_features = image_features / image_features.norm(dim=1, keepdim=True)
+        text_features = text_features / text_features.norm(dim=1, keepdim=True)
+
+        # cosine similarity as logits
+        logit_scale = self.logit_scale.exp()
+        logits_per_image = logit_scale * image_features @ text_features.t()
+        logits_per_text = logits_per_image.t()
+
+        # shape = [global_batch_size, global_batch_size]
+        return logits_per_image, logits_per_text
+
+
+def convert_weights(model: nn.Module):
+    """Convert applicable model parameters to fp16"""
+
+    def _convert_weights_to_fp16(l):
+        if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)):
+            l.weight.data = l.weight.data.half()
+            if l.bias is not None:
+                l.bias.data = l.bias.data.half()
+
+        if isinstance(l, nn.MultiheadAttention):
+            for attr in [*[f"{s}_proj_weight" for s in ["in", "q", "k", "v"]], "in_proj_bias", "bias_k", "bias_v"]:
+                tensor = getattr(l, attr)
+                if tensor is not None:
+                    tensor.data = tensor.data.half()
+
+        for name in ["text_projection", "proj"]:
+            if hasattr(l, name):
+                attr = getattr(l, name)
+                if attr is not None:
+                    attr.data = attr.data.half()
+
+    model.apply(_convert_weights_to_fp16)
+
+
+def build_model(state_dict: dict):
+    vit = "visual.proj" in state_dict
+
+    if vit:
+        vision_width = state_dict["visual.conv1.weight"].shape[0]
+        vision_layers = len([k for k in state_dict.keys() if k.startswith("visual.") and k.endswith(".attn.in_proj_weight")])
+        vision_patch_size = state_dict["visual.conv1.weight"].shape[-1]
+        grid_size = round((state_dict["visual.positional_embedding"].shape[0] - 1) ** 0.5)
+        image_resolution = vision_patch_size * grid_size
+    else:
+        counts: list = [len(set(k.split(".")[2] for k in state_dict if k.startswith(f"visual.layer{b}"))) for b in [1, 2, 3, 4]]
+        vision_layers = tuple(counts)
+        vision_width = state_dict["visual.layer1.0.conv1.weight"].shape[0]
+        output_width = round((state_dict["visual.attnpool.positional_embedding"].shape[0] - 1) ** 0.5)
+        vision_patch_size = None
+        assert output_width ** 2 + 1 == state_dict["visual.attnpool.positional_embedding"].shape[0]
+        image_resolution = output_width * 32
+
+    embed_dim = state_dict["text_projection"].shape[1]
+    context_length = state_dict["positional_embedding"].shape[0]
+    vocab_size = state_dict["token_embedding.weight"].shape[0]
+    transformer_width = state_dict["ln_final.weight"].shape[0]
+    transformer_heads = transformer_width // 64
+    transformer_layers = len(set(k.split(".")[2] for k in state_dict if k.startswith("transformer.resblocks")))
+
+    model = CLIP(
+        embed_dim,
+        image_resolution, vision_layers, vision_width, vision_patch_size,
+        context_length, vocab_size, transformer_width, transformer_heads, transformer_layers
+    )
+
+    for key in ["input_resolution", "context_length", "vocab_size"]:
+        if key in state_dict:
+            del state_dict[key]
+
+    convert_weights(model)
+    model.load_state_dict(state_dict, strict=False)
+    return model.eval()

models/clip/_clip/prepare.py

@@ -0,0 +1,95 @@
+# Prepare the models to speed up loading them later
+import torch
+from torch import nn, Tensor
+import os
+from tqdm import tqdm
+import json
+
+from .utils import load
+
+
+model_name_map = {
+    "RN50": "resnet50",
+    "RN101": "resnet101",
+    "RN50x4": "resnet50x4",
+    "RN50x16": "resnet50x16",
+    "RN50x64": "resnet50x64",
+    "ViT-B/32": "vit_b_32",
+    "ViT-B/16": "vit_b_16",
+    "ViT-L/14": "vit_l_14",
+    "ViT-L/14@336px": "vit_l_14_336px",
+}
+
+
+class CLIPTextEncoderTemp(nn.Module):
+    def __init__(
+        self,
+        clip: nn.Module,
+    ) -> None:
+        super().__init__()
+        self.context_length = clip.context_length
+        self.vocab_size = clip.vocab_size
+        self.dtype = clip.dtype
+        self.token_embedding = clip.token_embedding
+        self.positional_embedding = clip.positional_embedding
+        self.transformer = clip.transformer
+        self.ln_final = clip.ln_final
+        self.text_projection = clip.text_projection
+
+    def forward(self, text: Tensor) -> None:
+        pass
+
+
+def prepare() -> None:
+    print("Preparing CLIP models...")
+    curr_dir = os.path.dirname(os.path.abspath(__file__))
+    weight_dir = os.path.join(curr_dir, "weights")
+    config_dir = os.path.join(curr_dir, "configs")
+    os.makedirs(weight_dir, exist_ok=True)
+    os.makedirs(config_dir, exist_ok=True)
+    device = torch.device("cpu")
+
+    for model_name in tqdm(["RN50", "RN101", "RN50x4", "RN50x16", "RN50x64", "ViT-B/32", "ViT-B/16", "ViT-L/14", "ViT-L/14@336px"]):
+        model = load(model_name, device=device).to(device)
+        image_encoder = model.visual.to(device)
+        text_encoder = CLIPTextEncoderTemp(model).to(device)
+        torch.save(model.state_dict(), os.path.join(weight_dir, f"clip_{model_name_map[model_name]}.pth"))
+        torch.save(image_encoder.state_dict(), os.path.join(weight_dir, f"clip_image_encoder_{model_name_map[model_name]}.pth"))
+        torch.save(text_encoder.state_dict(), os.path.join(weight_dir, f"clip_text_encoder_{model_name_map[model_name]}.pth"))
+        model_config = {
+            "embed_dim": model.embed_dim,
+            # vision
+            "image_resolution": model.image_resolution,
+            "vision_layers": model.vision_layers,
+            "vision_width": model.vision_width,
+            "vision_patch_size": model.vision_patch_size,
+            # text
+            "context_length": model.context_length,
+            "vocab_size": model.vocab_size,
+            "transformer_width": model.transformer_width,
+            "transformer_heads": model.transformer_heads,
+            "transformer_layers": model.transformer_layers,
+        }
+        image_encoder_config = {
+            "embed_dim": model.embed_dim,
+            "image_resolution": model.image_resolution,
+            "vision_layers": model.vision_layers,
+            "vision_width": model.vision_width,
+            "vision_patch_size": model.vision_patch_size,
+            "vision_heads": model.vision_heads,
+        }
+        text_encoder_config = {
+            "embed_dim": model.embed_dim,
+            "context_length": model.context_length,
+            "vocab_size": model.vocab_size,
+            "transformer_width": model.transformer_width,
+            "transformer_heads": model.transformer_heads,
+            "transformer_layers": model.transformer_layers,
+        }
+        with open(os.path.join(config_dir, f"clip_{model_name_map[model_name]}.json"), "w") as f:
+            json.dump(model_config, f, indent=4)
+        with open(os.path.join(config_dir, f"clip_image_encoder_{model_name_map[model_name]}.json"), "w") as f:
+            json.dump(image_encoder_config, f, indent=4)
+        with open(os.path.join(config_dir, f"clip_text_encoder_{model_name_map[model_name]}.json"), "w") as f:
+            json.dump(text_encoder_config, f, indent=4)
+    print("Done!")

models/clip/_clip/simple_tokenizer.py

@@ -0,0 +1,132 @@
+import gzip
+import html
+import os
+from functools import lru_cache
+
+import ftfy
+import regex as re
+
+
+@lru_cache()
+def default_bpe():
+    return os.path.join(os.path.dirname(os.path.abspath(__file__)), "bpe_simple_vocab_16e6.txt.gz")
+
+
+@lru_cache()
+def bytes_to_unicode():
+    """
+    Returns list of utf-8 byte and a corresponding list of unicode strings.
+    The reversible bpe codes work on unicode strings.
+    This means you need a large # of unicode characters in your vocab if you want to avoid UNKs.
+    When you're at something like a 10B token dataset you end up needing around 5K for decent coverage.
+    This is a significant percentage of your normal, say, 32K bpe vocab.
+    To avoid that, we want lookup tables between utf-8 bytes and unicode strings.
+    And avoids mapping to whitespace/control characters the bpe code barfs on.
+    """
+    bs = list(range(ord("!"), ord("~")+1))+list(range(ord("¡"), ord("¬")+1))+list(range(ord("®"), ord("ÿ")+1))
+    cs = bs[:]
+    n = 0
+    for b in range(2**8):
+        if b not in bs:
+            bs.append(b)
+            cs.append(2**8+n)
+            n += 1
+    cs = [chr(n) for n in cs]
+    return dict(zip(bs, cs))
+
+
+def get_pairs(word):
+    """Return set of symbol pairs in a word.
+    Word is represented as tuple of symbols (symbols being variable-length strings).
+    """
+    pairs = set()
+    prev_char = word[0]
+    for char in word[1:]:
+        pairs.add((prev_char, char))
+        prev_char = char
+    return pairs
+
+
+def basic_clean(text):
+    text = ftfy.fix_text(text)
+    text = html.unescape(html.unescape(text))
+    return text.strip()
+
+
+def whitespace_clean(text):
+    text = re.sub(r'\s+', ' ', text)
+    text = text.strip()
+    return text
+
+
+class SimpleTokenizer(object):
+    def __init__(self, bpe_path: str = default_bpe()):
+        self.byte_encoder = bytes_to_unicode()
+        self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
+        merges = gzip.open(bpe_path).read().decode("utf-8").split('\n')
+        merges = merges[1:49152-256-2+1]
+        merges = [tuple(merge.split()) for merge in merges]
+        vocab = list(bytes_to_unicode().values())
+        vocab = vocab + [v+'</w>' for v in vocab]
+        for merge in merges:
+            vocab.append(''.join(merge))
+        vocab.extend(['<|startoftext|>', '<|endoftext|>'])
+        self.encoder = dict(zip(vocab, range(len(vocab))))
+        self.decoder = {v: k for k, v in self.encoder.items()}
+        self.bpe_ranks = dict(zip(merges, range(len(merges))))
+        self.cache = {'<|startoftext|>': '<|startoftext|>', '<|endoftext|>': '<|endoftext|>'}
+        self.pat = re.compile(r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""", re.IGNORECASE)
+
+    def bpe(self, token):
+        if token in self.cache:
+            return self.cache[token]
+        word = tuple(token[:-1]) + ( token[-1] + '</w>',)
+        pairs = get_pairs(word)
+
+        if not pairs:
+            return token+'</w>'
+
+        while True:
+            bigram = min(pairs, key = lambda pair: self.bpe_ranks.get(pair, float('inf')))
+            if bigram not in self.bpe_ranks:
+                break
+            first, second = bigram
+            new_word = []
+            i = 0
+            while i < len(word):
+                try:
+                    j = word.index(first, i)
+                    new_word.extend(word[i:j])
+                    i = j
+                except:
+                    new_word.extend(word[i:])
+                    break
+
+                if word[i] == first and i < len(word)-1 and word[i+1] == second:
+                    new_word.append(first+second)
+                    i += 2
+                else:
+                    new_word.append(word[i])
+                    i += 1
+            new_word = tuple(new_word)
+            word = new_word
+            if len(word) == 1:
+                break
+            else:
+                pairs = get_pairs(word)
+        word = ' '.join(word)
+        self.cache[token] = word
+        return word
+
+    def encode(self, text):
+        bpe_tokens = []
+        text = whitespace_clean(basic_clean(text)).lower()
+        for token in re.findall(self.pat, text):
+            token = ''.join(self.byte_encoder[b] for b in token.encode('utf-8'))
+            bpe_tokens.extend(self.encoder[bpe_token] for bpe_token in self.bpe(token).split(' '))
+        return bpe_tokens
+
+    def decode(self, tokens):
+        text = ''.join([self.decoder[token] for token in tokens])
+        text = bytearray([self.byte_decoder[c] for c in text]).decode('utf-8', errors="replace").replace('</w>', ' ')
+        return text

models/clip/_clip/text_encoder.py

@@ -0,0 +1,53 @@
+import torch
+from torch import nn, Tensor
+
+from .blocks import LayerNorm, Transformer
+
+
+class CLIPTextEncoder(nn.Module):
+    def __init__(
+        self,
+        embed_dim: int,
+        context_length: int,
+        vocab_size: int,
+        transformer_width: int,
+        transformer_heads: int,
+        transformer_layers: int,
+    ) -> None:
+        super().__init__()
+        self.context_length = context_length
+        self.token_embedding = nn.Embedding(vocab_size, transformer_width)
+        self.transformer = Transformer(
+            width=transformer_width,
+            layers=transformer_layers,
+            heads=transformer_heads,
+            attn_mask=self.build_attention_mask(),
+        )
+        self.vocab_size = vocab_size
+        self.token_embedding = nn.Embedding(vocab_size, transformer_width)
+        self.positional_embedding = nn.Parameter(torch.empty(self.context_length, transformer_width))
+        self.ln_final = LayerNorm(transformer_width)
+
+        self.text_projection = nn.Parameter(torch.empty(transformer_width, embed_dim))
+
+    def build_attention_mask(self):
+        # lazily create causal attention mask, with full attention between the vision tokens
+        # pytorch uses additive attention mask; fill with -inf
+        mask = torch.empty(self.context_length, self.context_length)
+        mask.fill_(float("-inf"))
+        mask.triu_(1)  # zero out the lower diagonal
+        return mask
+    
+    @property
+    def dtype(self):
+        return self.transformer.resblocks[0].attn.in_proj_weight.dtype
+
+    def forward(self, text: Tensor):
+        x = self.token_embedding(text).type(self.dtype)
+        x = x + self.positional_embedding.type(self.dtype)
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.transformer(x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        x = self.ln_final(x).type(self.dtype)
+        x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
+        return x

models/clip/_clip/utils.py

@@ -0,0 +1,249 @@
+import hashlib
+import os
+import urllib
+import warnings
+from typing import Union, List
+from pkg_resources import packaging
+
+from PIL import Image
+from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
+import torch
+
+from typing import List, Union
+from tqdm import tqdm
+
+from .model import build_model
+from .simple_tokenizer import SimpleTokenizer as _Tokenizer
+
+try:
+    from torchvision.transforms import InterpolationMode
+    BICUBIC = InterpolationMode.BICUBIC
+except ImportError:
+    BICUBIC = Image.BICUBIC
+
+
+if packaging.version.parse(torch.__version__) < packaging.version.parse("1.7.1"):
+    warnings.warn("PyTorch version 1.7.1 or higher is recommended")
+
+
+__all__ = ["available_models", "load", "tokenize"]
+_tokenizer = _Tokenizer()
+
+
+
+_MODELS = {
+    "RN50": "https://openaipublic.azureedge.net/clip/models/afeb0e10f9e5a86da6080e35cf09123aca3b358a0c3e3b6c78a7b63bc04b6762/RN50.pt",
+    "RN101": "https://openaipublic.azureedge.net/clip/models/8fa8567bab74a42d41c5915025a8e4538c3bdbe8804a470a72f30b0d94fab599/RN101.pt",
+    "RN50x4": "https://openaipublic.azureedge.net/clip/models/7e526bd135e493cef0776de27d5f42653e6b4c8bf9e0f653bb11773263205fdd/RN50x4.pt",
+    "RN50x16": "https://openaipublic.azureedge.net/clip/models/52378b407f34354e150460fe41077663dd5b39c54cd0bfd2b27167a4a06ec9aa/RN50x16.pt",
+    "RN50x64": "https://openaipublic.azureedge.net/clip/models/be1cfb55d75a9666199fb2206c106743da0f6468c9d327f3e0d0a543a9919d9c/RN50x64.pt",
+    "ViT-B/32": "https://openaipublic.azureedge.net/clip/models/40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt",
+    "ViT-B/16": "https://openaipublic.azureedge.net/clip/models/5806e77cd80f8b59890b7e101eabd078d9fb84e6937f9e85e4ecb61988df416f/ViT-B-16.pt",
+    "ViT-L/14": "https://openaipublic.azureedge.net/clip/models/b8cca3fd41ae0c99ba7e8951adf17d267cdb84cd88be6f7c2e0eca1737a03836/ViT-L-14.pt",
+    "ViT-L/14@336px": "https://openaipublic.azureedge.net/clip/models/3035c92b350959924f9f00213499208652fc7ea050643e8b385c2dac08641f02/ViT-L-14-336px.pt",
+}
+
+
+def _download(url: str, root: str):
+    os.makedirs(root, exist_ok=True)
+    filename = os.path.basename(url)
+
+    expected_sha256 = url.split("/")[-2]
+    download_target = os.path.join(root, filename)
+
+    if os.path.exists(download_target) and not os.path.isfile(download_target):
+        raise RuntimeError(f"{download_target} exists and is not a regular file")
+
+    if os.path.isfile(download_target):
+        if hashlib.sha256(open(download_target, "rb").read()).hexdigest() == expected_sha256:
+            return download_target
+        else:
+            warnings.warn(f"{download_target} exists, but the SHA256 checksum does not match; re-downloading the file")
+
+    with urllib.request.urlopen(url) as source, open(download_target, "wb") as output:
+        with tqdm(total=int(source.info().get("Content-Length")), ncols=80, unit='iB', unit_scale=True, unit_divisor=1024) as loop:
+            while True:
+                buffer = source.read(8192)
+                if not buffer:
+                    break
+
+                output.write(buffer)
+                loop.update(len(buffer))
+
+    if hashlib.sha256(open(download_target, "rb").read()).hexdigest() != expected_sha256:
+        raise RuntimeError("Model has been downloaded but the SHA256 checksum does not not match")
+
+    return download_target
+
+
+def _convert_image_to_rgb(image):
+    return image.convert("RGB")
+
+
+def transform(n_px):
+    return Compose([
+        Resize(n_px, interpolation=BICUBIC),
+        CenterCrop(n_px),
+        _convert_image_to_rgb,
+        ToTensor(),
+        Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
+    ])
+
+
+def available_models() -> List[str]:
+    """Returns the names of available CLIP models"""
+    return list(_MODELS.keys())
+
+
+def load(name: str, device: Union[str, torch.device] = "cuda" if torch.cuda.is_available() else "cpu", jit: bool = False, download_root: str = None):
+    """Load a CLIP model
+
+    Parameters
+    ----------
+    name : str
+        A model name listed by `clip.available_models()`, or the path to a model checkpoint containing the state_dict
+
+    device : Union[str, torch.device]
+        The device to put the loaded model
+
+    jit : bool
+        Whether to load the optimized JIT model or more hackable non-JIT model (default).
+
+    download_root: str
+        path to download the model files; by default, it uses "~/.cache/clip"
+
+    Returns
+    -------
+    model : torch.nn.Module
+        The CLIP model
+
+    preprocess : Callable[[PIL.Image], torch.Tensor]
+        A torchvision transform that converts a PIL image into a tensor that the returned model can take as its input
+    """
+    if name in _MODELS:
+        model_path = _download(_MODELS[name], download_root or os.path.expanduser("~/.cache/clip"))
+    elif os.path.isfile(name):
+        model_path = name
+    else:
+        raise RuntimeError(f"Model {name} not found; available models = {available_models()}")
+
+    with open(model_path, 'rb') as opened_file:
+        try:
+            # loading JIT archive
+            model = torch.jit.load(opened_file, map_location=device if jit else "cpu").eval()
+            state_dict = None
+        except RuntimeError:
+            # loading saved state dict
+            if jit:
+                warnings.warn(f"File {model_path} is not a JIT archive. Loading as a state dict instead")
+                jit = False
+            state_dict = torch.load(opened_file, map_location="cpu")
+
+    if not jit:
+        model = build_model(state_dict or model.state_dict()).to(device)
+        if str(device) == "cpu":
+            model.float()
+        return model
+
+    # patch the device names
+    device_holder = torch.jit.trace(lambda: torch.ones([]).to(torch.device(device)), example_inputs=[])
+    device_node = [n for n in device_holder.graph.findAllNodes("prim::Constant") if "Device" in repr(n)][-1]
+
+    def _node_get(node: torch._C.Node, key: str):
+        """Gets attributes of a node which is polymorphic over return type.
+        
+        From https://github.com/pytorch/pytorch/pull/82628
+        """
+        sel = node.kindOf(key)
+        return getattr(node, sel)(key)
+
+    def patch_device(module):
+        try:
+            graphs = [module.graph] if hasattr(module, "graph") else []
+        except RuntimeError:
+            graphs = []
+
+        if hasattr(module, "forward1"):
+            graphs.append(module.forward1.graph)
+
+        for graph in graphs:
+            for node in graph.findAllNodes("prim::Constant"):
+                if "value" in node.attributeNames() and str(_node_get(node, "value")).startswith("cuda"):
+                    node.copyAttributes(device_node)
+
+    model.apply(patch_device)
+    patch_device(model.encode_image)
+    patch_device(model.encode_text)
+
+    # patch dtype to float32 on CPU
+    if str(device) == "cpu":
+        float_holder = torch.jit.trace(lambda: torch.ones([]).float(), example_inputs=[])
+        float_input = list(float_holder.graph.findNode("aten::to").inputs())[1]
+        float_node = float_input.node()
+
+        def patch_float(module):
+            try:
+                graphs = [module.graph] if hasattr(module, "graph") else []
+            except RuntimeError:
+                graphs = []
+
+            if hasattr(module, "forward1"):
+                graphs.append(module.forward1.graph)
+
+            for graph in graphs:
+                for node in graph.findAllNodes("aten::to"):
+                    inputs = list(node.inputs())
+                    for i in [1, 2]:  # dtype can be the second or third argument to aten::to()
+                        if _node_get(inputs[i].node(), "value") == 5:
+                            inputs[i].node().copyAttributes(float_node)
+
+        model.apply(patch_float)
+        patch_float(model.encode_image)
+        patch_float(model.encode_text)
+
+        model.float()
+
+    return model
+
+
+def tokenize(texts: Union[str, List[str]], context_length: int = 77, truncate: bool = False) -> Union[torch.IntTensor, torch.LongTensor]:
+    """
+    Returns the tokenized representation of given input string(s)
+
+    Parameters
+    ----------
+    texts : Union[str, List[str]]
+        An input string or a list of input strings to tokenize
+
+    context_length : int
+        The context length to use; all CLIP models use 77 as the context length
+
+    truncate: bool
+        Whether to truncate the text in case its encoding is longer than the context length
+
+    Returns
+    -------
+    A two-dimensional tensor containing the resulting tokens, shape = [number of input strings, context_length].
+    We return LongTensor when torch version is <1.8.0, since older index_select requires indices to be long.
+    """
+    if isinstance(texts, str):
+        texts = [texts]
+
+    sot_token = _tokenizer.encoder["<|startoftext|>"]
+    eot_token = _tokenizer.encoder["<|endoftext|>"]
+    all_tokens = [[sot_token] + _tokenizer.encode(text) + [eot_token] for text in texts]
+    if packaging.version.parse(torch.__version__) < packaging.version.parse("1.8.0"):
+        result = torch.zeros(len(all_tokens), context_length, dtype=torch.long)
+    else:
+        result = torch.zeros(len(all_tokens), context_length, dtype=torch.int)
+
+    for i, tokens in enumerate(all_tokens):
+        if len(tokens) > context_length:
+            if truncate:
+                tokens = tokens[:context_length]
+                tokens[-1] = eot_token
+            else:
+                raise RuntimeError(f"Input {texts[i]} is too long for context length {context_length}")
+        result[i, :len(tokens)] = torch.tensor(tokens)
+
+    return result

models/clip/model.py

@@ -0,0 +1,331 @@
+import torch
+from torch import nn, Tensor
+import torch.nn.functional as F
+import numpy as np
+import os
+import math
+from typing import List, Tuple, Union, Optional
+
+from . import _clip
+from ..utils import _init_weights, make_resnet_layers, Bottleneck, BasicBlock
+from .utils import format_count
+
+curr_dir = os.path.abspath(os.path.dirname(__file__))
+
+
+# resnet50: reduction, channels, embed_dim = 32, 2048, 1024
+# resnet101: reduction, channels, embed_dim = 32, 2048, 512
+# resnet50x4: reduction, channels, embed_dim = 32, 2560, 640
+# resnet50x16: reduction, channels, embed_dim = 32, 3072, 768
+# resnet50x64: reduction, channels, embed_dim = 32, 4096, 1024
+# vit_b_32: reduction, channels, embed_dim = 32, 768, 512
+# vit_b_16: reduction, channels, embed_dim = 16, 768, 512
+# vit_l_14: reduction, channels, embed_dim = 14, 1024, 768
+# vit_l_14_336px: reduction, channels, embed_dim = 14, 1024, 768
+
+resnet_backbones = ["resnet50", "resnet101", "resnet50x4", "resnet50x16", "resnet50x64"]
+vit_backbones = ["vit_b_16", "vit_b_32", "vit_l_14", "vit_l_14_336px"]
+
+
+class CLIP_EBC(nn.Module):
+    def __init__(
+        self,
+        backbone: str,
+        bins: List[Tuple[float, float]],
+        anchor_points: List[float],
+        reduction: Optional[int] = None,
+        freeze_text_encoder: bool = True,
+        prompt_type: str = "number",
+        input_size: Optional[int] = None,
+        num_vpt: Optional[int] = None,
+        deep_vpt: Optional[bool] = None,
+        vpt_drop: Optional[float] = None,
+        decoder_block: Optional[nn.Module] = None,
+        decoder_cfg: Optional[List[Union[str, int]]] = None,
+    ) -> None:
+        super().__init__()
+        assert backbone in resnet_backbones + vit_backbones, f"Backbone should be in {resnet_backbones + vit_backbones}, got {backbone}"
+        self.backbone = backbone
+
+        # Image encoder
+        if backbone in resnet_backbones:
+            self.image_encoder = getattr(_clip, f"{backbone}_img")(features_only=True, out_indices=(-1,), reduction=reduction)
+
+        else:
+            assert input_size is not None, "Expected input_size to be an integer, got None."
+            assert num_vpt is not None, "Expected num_vpt to be an integer, got None."
+            assert deep_vpt is not None, "Expected deep_vpt to be a boolean, got None."
+            assert vpt_drop is not None, "Expected vpt_drop to be a float, got None."
+
+            self.image_encoder = getattr(_clip, f"{backbone}_img")(features_only=True, input_size=input_size)
+            self.image_encoder_depth = len(self.image_encoder.transformer.resblocks)
+
+            # Use VPT. Freeze the image encoder.
+            for param in self.image_encoder.parameters():
+                param.requires_grad = False
+
+            self.num_vpt = num_vpt
+            self.deep_vpt = deep_vpt
+
+            patch_size = self.image_encoder.patch_size[0]
+            val = math.sqrt(6. / float(3 * patch_size + self.image_encoder.channels))
+
+            for idx in range(self.image_encoder_depth if self.deep_vpt else 1):
+                setattr(self, f"vpt_{idx}", nn.Parameter(torch.empty(self.num_vpt, self.image_encoder.channels)))
+                nn.init.uniform_(getattr(self, f"vpt_{idx}"), -val, val)
+                setattr(self, f"vpt_drop_{idx}", nn.Dropout(vpt_drop) if vpt_drop > 0 else nn.Identity())
+
+        self.encoder_reduction = self.image_encoder.reduction
+        self.reduction = self.encoder_reduction if reduction is None else reduction
+        self.channels = self.image_encoder.channels
+        self.clip_embed_dim = self.image_encoder.clip_embed_dim
+
+        if decoder_cfg is not None:
+            assert decoder_block is not None, "Expected decoder_block to be a nn.Module, got None."
+            self.image_decoder = make_resnet_layers(decoder_block, decoder_cfg, in_channels=self.channels, expansion=1, dilation=1)
+            self.image_decoder.apply(_init_weights)
+            self.channels = decoder_cfg[-1]
+        else:
+            self.image_decoder = nn.Identity()
+
+        if self.channels != self.clip_embed_dim:
+            self.projection = nn.Conv2d(in_channels=self.channels, out_channels=self.clip_embed_dim, kernel_size=1)
+            self.projection.apply(_init_weights)
+        else:
+            self.projection = nn.Identity()
+
+        # Text encoder
+        assert prompt_type in ["number", "word"], f"Expected prompt_type to be 'number' or 'word', got {prompt_type}"
+        self.prompt_type = prompt_type
+        self.text_encoder = getattr(_clip, f"{backbone}_txt")()
+        self.freeze_text_encoder = freeze_text_encoder
+        if self.freeze_text_encoder:
+            for param in self.text_encoder.parameters():
+                param.requires_grad = False
+
+        self.bins = bins
+        self.anchor_points = torch.tensor(anchor_points, dtype=torch.float32, requires_grad=False).view(1, -1, 1, 1)
+
+        self._get_text_prompts()
+        self._tokenize_text_prompts()
+
+        if self.freeze_text_encoder:
+            self._extract_text_features()
+        else:
+            self.text_features = None
+
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07), requires_grad=True)
+
+    def _get_text_prompts(self) -> None:
+        bins = [b[0] if b[0] == b[1] else b for b in self.bins]
+        self.text_prompts = [format_count(b, self.prompt_type) for b in bins]
+        print(f"Initialized model with text prompts: {self.text_prompts}")
+
+    def _tokenize_text_prompts(self) -> None:
+        self.text_prompts = _clip.tokenize(self.text_prompts)
+
+    def _extract_text_features(self) -> None:
+        with torch.no_grad():
+            self.text_features = self.text_encoder(self.text_prompts)
+
+    def _prepare_vpt(self, layer: int, batch_size: int, device: torch.device) -> Tensor:
+        if not self.deep_vpt:
+            assert layer == 0, f"Expected layer to be 0 when using Shallow Visual Prompt Tuning, got {layer}"
+
+        vpt = getattr(self, f"vpt_{layer}").to(device)
+        vpt = vpt.unsqueeze(0).expand(batch_size, -1, -1)
+        vpt = getattr(self, f"vpt_drop_{layer}")(vpt)
+        vpt = vpt.permute(1, 0, 2)  # (num_vpt, batch_size, hidden_dim)
+        assert vpt.shape[1] == batch_size, f"Expected the VPT to have the shape [L_vis B C], got {vpt.shape}."
+        return vpt
+
+    def _forward_vpt(self, x: Tensor) -> Tuple[Tensor]:
+        device = x.device
+        batch_size, _, height, width = x.shape
+        num_h_patches, num_w_patches = height // self.image_encoder.patch_size[0], width // self.image_encoder.patch_size[1]
+
+        image_features = self.image_encoder.conv1(x)
+        image_features = image_features.reshape(batch_size, image_features.shape[1], -1)
+        image_features = image_features.permute(0, 2, 1)  # (B, num_patches, C)
+        image_features = torch.cat([
+            self.image_encoder.class_embedding + torch.zeros(batch_size, 1, image_features.shape[-1], dtype=image_features.dtype, device=device),
+            image_features,
+        ], dim=1)  # (B, num_patches + 1, C)
+
+        pos_embedding = self.image_encoder._interpolate_pos_embed(num_h_patches, num_w_patches)
+        image_features = image_features + pos_embedding
+        image_features = self.image_encoder.ln_pre(image_features)
+        image_features = image_features.permute(1, 0, 2)  # (num_patches + 1, B, C)
+        assert image_features.shape[0] == num_h_patches * num_w_patches + 1 and image_features.shape[1] == batch_size, f"Expected image_features to have shape [num_patches + 1, B, C], got {image_features.shape}."
+
+        vpt = self._prepare_vpt(0, batch_size, device)
+        for idx in range(self.image_encoder_depth):
+            # assemble
+            image_features = torch.cat([
+                image_features[:1, :, :],  # CLS token
+                vpt,
+                image_features[1:, :, :],
+            ], dim=0)
+
+            # transformer
+            image_features = self.image_encoder.transformer.resblocks[idx](image_features)
+
+            # disassemble
+            if idx < self.image_encoder_depth - 1:
+                if self.deep_vpt:
+                    vpt = self._prepare_vpt(idx + 1, batch_size, device)
+                else:
+                    vpt = image_features[1: (self.num_vpt + 1), :, :]
+
+            image_features = torch.cat([
+                image_features[:1, :, :],  # CLS token
+                image_features[(self.num_vpt + 1):, :, :],
+            ], dim=0)
+            
+        image_features = image_features.permute(1, 0, 2)  # (B, num_patches + 1, C)
+        image_features = self.image_encoder.ln_post(image_features)
+        image_features = image_features[:, 1:, :].permute(0, 2, 1)  # (B, C, num_patches)
+        image_features = image_features.reshape(batch_size, -1, num_h_patches, num_w_patches)
+        return image_features
+
+    def _forward(self, x: Tensor) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+        device = x.device
+
+        x = self.image_encoder(x) if self.backbone in resnet_backbones else self._forward_vpt(x)
+        if self.reduction != self.encoder_reduction:
+            x = F.interpolate(x, scale_factor=self.encoder_reduction / self.reduction, mode="bilinear")
+        x = self.image_decoder(x)
+        x = self.projection(x)
+
+        image_features = x.permute(0, 2, 3, 1)  # shape (B, H, W, C)
+        text_features = self.text_encoder(self.text_prompts.to(device)) if self.text_features is None else self.text_features.to(device)  # shape (N, C)
+
+        image_features = F.normalize(image_features, p=2, dim=-1)
+        text_features = F.normalize(text_features, p=2, dim=-1)
+
+        # cosine similarity as logits
+        logit_scale = self.logit_scale.exp()
+        logits = logit_scale * image_features @ text_features.t()  # (B, H, W, N), logits per image
+        logits = logits.permute(0, 3, 1, 2)  # (B, N, H, W)
+
+        probs = logits.softmax(dim=1)
+        exp = (probs * self.anchor_points.to(x.device)).sum(dim=1, keepdim=True)  # (B, 1, H, W)
+
+        if self.training:
+            return logits, exp
+        else:
+            return exp
+    
+    def forward(self, x: Tensor) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+        assert len(x.shape) == 4, f"Expected input to have shape (B C H W), got {x.shape}."
+        if "vit" in self.backbone:
+            image_height, image_width = x.shape[2], x.shape[3]
+            window_height, window_width = self.image_encoder.input_resolution
+
+            if self.training:
+                assert (image_height, image_width) == (window_height, window_width), f"Expected input to have shape ({window_height} {window_width}), got ({image_height} {image_width})."
+                return self._forward(x)
+            
+            elif (image_height, image_width) == (window_height, window_width):  # evaluation, input size = training size
+                return self._forward(x)
+            
+            else:  # evaluation, input_size != training size, use sliding window prediction
+                stride_height, stride_width = window_height, window_width
+                reduction = self.reduction
+
+                num_rows = int(np.ceil((image_height - window_height) / stride_height) + 1)
+                num_cols = int(np.ceil((image_width - window_width) / stride_width) + 1)
+
+                windows = []
+                for i in range(num_rows):
+                    for j in range(num_cols):
+                        x_start, y_start = i * stride_height, j * stride_width
+                        x_end, y_end = x_start + window_height, y_start + window_width
+                        if x_end > image_height:
+                            x_start, x_end = image_height - window_height, image_height
+                        if y_end > image_width:
+                            y_start, y_end = image_width - window_width, image_width
+
+                        window = x[:, :, x_start:x_end, y_start:y_end]
+                        windows.append(window)
+
+                windows = torch.cat(windows, dim=0).to(x.device)  # batched windows, shape: (num_windows, c, h, w)
+
+                preds = self._forward(windows)
+                preds = preds.cpu().detach().numpy()
+
+                # assemble the density map
+                pred_map = np.zeros((preds.shape[1], image_height // reduction, image_width // reduction), dtype=np.float32)
+                count_map = np.zeros((preds.shape[1], image_height // reduction, image_width // reduction), dtype=np.float32)
+                idx = 0
+                for i in range(num_rows):
+                    for j in range(num_cols):
+                        x_start, y_start = i * stride_height, j * stride_width
+                        x_end, y_end = x_start + window_height, y_start + window_width
+                        if x_end > image_height:
+                            x_start, x_end = image_height - window_height, image_height
+                        if y_end > image_width:
+                            y_start, y_end = image_width - window_width, image_width
+
+                        pred_map[:, (x_start // reduction): (x_end // reduction), (y_start // reduction): (y_end // reduction)] += preds[idx, :, :, :]
+                        count_map[:, (x_start // reduction): (x_end // reduction), (y_start // reduction): (y_end // reduction)] += 1.
+                        idx += 1
+
+                pred_map /= count_map  # average the overlapping regions
+                return torch.tensor(pred_map).unsqueeze(0)  # shape: (1, 1, h // reduction, w // reduction)
+        
+        else:
+            return self._forward(x)
+
+
+def _clip_ebc(
+    backbone: str,
+    bins: List[Tuple[float, float]],
+    anchor_points: List[float],
+    reduction: Optional[int] = None,
+    freeze_text_encoder: bool = True,
+    prompt_type: str = "number",
+    input_size: Optional[int] = None,
+    num_vpt: Optional[int] = None,
+    deep_vpt: Optional[bool] = None,
+    vpt_drop: Optional[float] = None,
+    decoder_block: Optional[nn.Module] = None,
+    decoder_cfg: Optional[List[Union[str, int]]] = None
+) -> CLIP_EBC:
+    if backbone in resnet_backbones:
+        decoder_block = Bottleneck
+        if decoder_cfg is None:
+            if backbone == "resnet50":
+                decoder_cfg = [2048]
+            elif backbone == "resnet50x4":
+                decoder_cfg = [1280]
+            elif backbone == "resnet50x16":
+                decoder_cfg = [1536]
+            elif backbone == "resnet50x64":
+                decoder_cfg = [2048]
+            else:  # backbone == "resnet101"
+                decoder_cfg = [2048, 1024]
+    else:
+        decoder_block = BasicBlock
+        if decoder_cfg is None:
+            if backbone == "vit_b_16":
+                decoder_cfg = [768]
+            elif backbone == "vit_b_32":
+                decoder_cfg = [768]
+            else:  # backbone == "vit_l_14"
+                decoder_cfg = [1024]
+
+    return CLIP_EBC(
+        backbone=backbone,
+        bins=bins,
+        anchor_points=anchor_points,
+        reduction=reduction,
+        freeze_text_encoder=freeze_text_encoder,
+        prompt_type=prompt_type,
+        input_size=input_size,
+        num_vpt=num_vpt,
+        deep_vpt=deep_vpt,
+        vpt_drop=vpt_drop,
+        decoder_block=decoder_block,
+        decoder_cfg=decoder_cfg,
+    )

models/clip/utils.py

@@ -0,0 +1,40 @@
+from typing import Union, Tuple
+
+
+num_to_word = {
+    "0": "zero", "1": "one", "2": "two", "3": "three", "4": "four", "5": "five", "6": "six", "7": "seven", "8": "eight", "9": "nine", 
+    "10": "ten", "11": "eleven", "12": "twelve", "13": "thirteen", "14": "fourteen", "15": "fifteen", "16": "sixteen", "17": "seventeen", "18": "eighteen", "19": "nineteen", 
+    "20": "twenty", "21": "twenty-one", "22": "twenty-two", "23": "twenty-three", "24": "twenty-four", "25": "twenty-five", "26": "twenty-six", "27": "twenty-seven", "28": "twenty-eight", "29": "twenty-nine",
+    "30": "thirty", "31": "thirty-one", "32": "thirty-two", "33": "thirty-three", "34": "thirty-four", "35": "thirty-five", "36": "thirty-six", "37": "thirty-seven", "38": "thirty-eight", "39": "thirty-nine",
+    "40": "forty", "41": "forty-one", "42": "forty-two", "43": "forty-three", "44": "forty-four", "45": "forty-five", "46": "forty-six", "47": "forty-seven", "48": "forty-eight", "49": "forty-nine",
+    "50": "fifty", "51": "fifty-one", "52": "fifty-two", "53": "fifty-three", "54": "fifty-four", "55": "fifty-five", "56": "fifty-six", "57": "fifty-seven", "58": "fifty-eight", "59": "fifty-nine",
+    "60": "sixty", "61": "sixty-one", "62": "sixty-two", "63": "sixty-three", "64": "sixty-four", "65": "sixty-five", "66": "sixty-six", "67": "sixty-seven", "68": "sixty-eight", "69": "sixty-nine",
+    "70": "seventy", "71": "seventy-one", "72": "seventy-two", "73": "seventy-three", "74": "seventy-four", "75": "seventy-five", "76": "seventy-six", "77": "seventy-seven", "78": "seventy-eight", "79": "seventy-nine",
+    "80": "eighty", "81": "eighty-one", "82": "eighty-two", "83": "eighty-three", "84": "eighty-four", "85": "eighty-five", "86": "eighty-six", "87": "eighty-seven", "88": "eighty-eight", "89": "eighty-nine",
+    "90": "ninety", "91": "ninety-one", "92": "ninety-two", "93": "ninety-three", "94": "ninety-four", "95": "ninety-five", "96": "ninety-six", "97": "ninety-seven", "98": "ninety-eight", "99": "ninety-nine",
+    "100": "one hundred", "200": "two hundred", "300": "three hundred", "400": "four hundred", "500": "five hundred", "600": "six hundred", "700": "seven hundred", "800": "eight hundred", "900": "nine hundred",
+    "1000": "one thousand"
+}
+
+
+def num2word(num: Union[int, str]) -> str:
+    """
+    Convert the input number to the corresponding English word. For example, 1 -> "one", 2 -> "two", etc.
+    """
+    num = str(int(num))
+    return num_to_word.get(num, num)
+
+
+def format_count(count: Union[float, Tuple[float, float]], prompt_type: str = "word") -> str:
+    if count == 0:
+        return "There is no person." if prompt_type == "word" else "There is 0 person."
+    elif count == 1:
+        return "There is one person." if prompt_type == "word" else "There is 1 person."
+    elif isinstance(count, (int, float)):
+        return f"There are {num2word(int(count))} people." if prompt_type == "word" else f"There are {int(count)} people."
+    elif count[1] == float("inf"):
+        return f"There are more than {num2word(int(count[0]))} people." if prompt_type == "word" else f"There are more than {int(count[0])} people."
+    else:  # count is a tuple of finite numbers
+        left, right = int(count[0]), int(count[1])
+        left, right = num2word(left), num2word(right) if prompt_type == "word" else left, right
+        return f"There are between {left} and {right} people."

models/encoder/__init__.py

@@ -0,0 +1,10 @@
+from .vgg import vgg11, vgg11_bn, vgg13, vgg13_bn, vgg16, vgg16_bn, vgg19, vgg19_bn
+from .vit import vit_b_16, vit_b_32, vit_l_16, vit_l_32, vit_h_14
+from .timm_models import _timm_encoder
+
+
+__all__ = [
+    "vgg11", "vgg11_bn", "vgg13", "vgg13_bn", "vgg16", "vgg16_bn", "vgg19", "vgg19_bn",
+    "vit_b_16", "vit_b_32", "vit_l_16", "vit_l_32", "vit_h_14",
+    "_timm_encoder",
+]

models/encoder/timm_models.py

@@ -0,0 +1,54 @@
+from timm import create_model, list_models
+from torch import nn, Tensor
+import torch.nn.functional as F
+from typing import Optional
+
+from warnings import warn
+
+
+class TIMMEncoder(nn.Module):
+    def __init__(
+        self,
+        backbone: str,
+        reduction: Optional[int] = None,
+    ) -> None:
+        super().__init__()
+        assert backbone in list_models(), f"Backbone {backbone} not available in timm"
+        encoder = create_model(backbone, pretrained=True, features_only=True, out_indices=[-1])
+        encoder_reduction = encoder.feature_info.reduction()[-1]
+
+        if reduction <= 16:
+            if "resnet" in backbone:
+                if "resnet18" in backbone or "resnet34" in backbone:
+                    encoder.layer4[0].conv1.stride = (1, 1)
+                    encoder.layer4[0].downsample[0].stride = (1, 1)
+                else:
+                    encoder.layer4[0].conv2.stride = (1, 1)
+                    encoder.layer4[0].downsample[0].stride = (1, 1)
+                encoder_reduction = encoder_reduction // 2
+
+            elif "mobilenetv2" in backbone:
+                encoder.blocks[5][0].conv_dw.stride = (1, 1)
+                encoder_reduction = encoder_reduction // 2
+
+            elif "densenet" in backbone:
+                encoder.features_transition3.pool = nn.Identity()
+                encoder_reduction = encoder_reduction // 2
+
+            else:
+                warn(f"Reduction for {backbone} not handled. Using default reduction of {encoder_reduction}")
+
+        self.encoder = encoder
+        self.encoder_reduction = encoder_reduction
+        self.reduction = self.encoder_reduction if reduction is None else reduction
+        self.channels = self.encoder.feature_info.channels()[-1]
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.encoder(x)[-1]
+        if self.encoder_reduction != self.reduction:
+            x = F.interpolate(x, scale_factor=self.encoder_reduction / self.reduction, mode="bilinear")
+        return x
+
+
+def _timm_encoder(backbone: str, reduction: Optional[int] = None) -> TIMMEncoder:
+    return TIMMEncoder(backbone, reduction)

models/encoder/vgg.py

@@ -0,0 +1,69 @@
+from torch import nn, Tensor
+import torch.nn.functional as F
+from torch.hub import load_state_dict_from_url
+from typing import Optional
+
+from ..utils import make_vgg_layers, vgg_cfgs, vgg_urls
+
+
+class VGG(nn.Module):
+    def __init__(
+        self,
+        features: nn.Module,
+        reduction: Optional[int] = None,
+    ) -> None:
+        super().__init__()
+        self.features = features
+        self.encoder_reduction = 16
+        self.reduction = self.encoder_reduction if reduction is None else reduction
+        self.channels = 512
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.features(x)
+        if self.encoder_reduction != self.reduction:
+            x = F.interpolate(x, scale_factor=self.encoder_reduction / self.reduction, mode="bilinear")
+        return x
+
+
+def _load_weights(model: VGG, url: str) -> VGG:
+    state_dict = load_state_dict_from_url(url)
+    missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
+    print("Loading pre-trained weights")
+    if len(missing_keys) > 0:
+        print(f"Missing keys: {missing_keys}")
+    if len(unexpected_keys) > 0:
+        print(f"Unexpected keys: {unexpected_keys}")
+    return model
+
+
+def vgg11(reduction: int = 8) -> VGG:
+    model = VGG(make_vgg_layers(vgg_cfgs["A"]), reduction=reduction)
+    return _load_weights(model, vgg_urls["vgg11"])
+
+def vgg11_bn(reduction: int = 8) -> VGG:
+    model = VGG(make_vgg_layers(vgg_cfgs["A"], batch_norm=True), reduction=reduction)
+    return _load_weights(model, vgg_urls["vgg11_bn"])
+
+def vgg13(reduction: int = 8) -> VGG:
+    model = VGG(make_vgg_layers(vgg_cfgs["B"]), reduction=reduction)
+    return _load_weights(model, vgg_urls["vgg13"])
+
+def vgg13_bn(reduction: int = 8) -> VGG:
+    model = VGG(make_vgg_layers(vgg_cfgs["B"], batch_norm=True), reduction=reduction)
+    return _load_weights(model, vgg_urls["vgg13_bn"])
+
+def vgg16(reduction: int = 8) -> VGG:
+    model = VGG(make_vgg_layers(vgg_cfgs["D"]), reduction=reduction)
+    return _load_weights(model, vgg_urls["vgg16"])
+
+def vgg16_bn(reduction: int = 8) -> VGG:
+    model = VGG(make_vgg_layers(vgg_cfgs["D"], batch_norm=True), reduction=reduction)
+    return _load_weights(model, vgg_urls["vgg16_bn"])
+
+def vgg19(reduction: int = 8) -> VGG:
+    model = VGG(make_vgg_layers(vgg_cfgs["E"]), reduction=reduction)
+    return _load_weights(model, vgg_urls["vgg19"])
+
+def vgg19_bn(reduction: int = 8) -> VGG:
+    model = VGG(make_vgg_layers(vgg_cfgs["E"], batch_norm=True), reduction=reduction)
+    return _load_weights(model, vgg_urls["vgg19_bn"])

models/encoder/vit.py

@@ -0,0 +1,526 @@
+import math
+from collections import OrderedDict
+from functools import partial
+from typing import Any, Callable, List, NamedTuple, Optional, Tuple
+
+import torch
+from torch import nn, Tensor
+import torch.nn.functional as F
+from torch.hub import load_state_dict_from_url
+from einops import rearrange
+
+from ..utils import Conv2dNormActivation, MLP
+from ..utils import _log_api_usage_once
+
+
+weights = {
+    "vit_b_16": "https://download.pytorch.org/models/vit_b_16-c867db91.pth",
+    "vit_b_32": "https://download.pytorch.org/models/vit_b_32-d86f8d99.pth",
+    "vit_l_16": "https://download.pytorch.org/models/vit_l_16-852ce7e3.pth",
+    "vit_l_32": "https://download.pytorch.org/models/vit_l_32-c7638314.pth",
+    "vit_h_14": "https://download.pytorch.org/models/vit_h_14-6kbcf7eb.pth",
+}
+
+
+class ConvStemConfig(NamedTuple):
+    out_channels: int
+    kernel_size: int
+    stride: int
+    norm_layer: Callable[..., nn.Module] = nn.BatchNorm2d
+    activation_layer: Callable[..., nn.Module] = nn.ReLU
+
+
+class MLPBlock(MLP):
+    """Transformer MLP block."""
+
+    _version = 2
+
+    def __init__(self, in_dim: int, mlp_dim: int, dropout: float):
+        super().__init__(in_dim, [mlp_dim, in_dim], activation_layer=nn.GELU, inplace=None, dropout=dropout)
+
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.xavier_uniform_(m.weight)
+                if m.bias is not None:
+                    nn.init.normal_(m.bias, std=1e-6)
+
+    def _load_from_state_dict(
+        self,
+        state_dict,
+        prefix,
+        local_metadata,
+        strict,
+        missing_keys,
+        unexpected_keys,
+        error_msgs,
+    ):
+        version = local_metadata.get("version", None)
+
+        if version is None or version < 2:
+            # Replacing legacy MLPBlock with MLP. See https://github.com/pytorch/vision/pull/6053
+            for i in range(2):
+                for type in ["weight", "bias"]:
+                    old_key = f"{prefix}linear_{i+1}.{type}"
+                    new_key = f"{prefix}{3*i}.{type}"
+                    if old_key in state_dict:
+                        state_dict[new_key] = state_dict.pop(old_key)
+
+        super()._load_from_state_dict(
+            state_dict,
+            prefix,
+            local_metadata,
+            strict,
+            missing_keys,
+            unexpected_keys,
+            error_msgs,
+        )
+
+
+class EncoderBlock(nn.Module):
+    """Transformer encoder block."""
+
+    def __init__(
+        self,
+        num_heads: int,
+        hidden_dim: int,
+        mlp_dim: int,
+        dropout: float,
+        attention_dropout: float,
+        norm_layer: Callable[..., torch.nn.Module] = partial(nn.LayerNorm, eps=1e-6),
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+
+        # Attention block
+        self.ln_1 = norm_layer(hidden_dim)
+        self.self_attention = nn.MultiheadAttention(hidden_dim, num_heads, dropout=attention_dropout, batch_first=True)
+        self.dropout = nn.Dropout(dropout)
+
+        # MLP block
+        self.ln_2 = norm_layer(hidden_dim)
+        self.mlp = MLPBlock(hidden_dim, mlp_dim, dropout)
+
+    def forward(self, input: Tensor):
+        torch._assert(input.dim() == 3, f"Expected (batch_size, seq_length, hidden_dim) got {input.shape}")
+        x = self.ln_1(input)
+        x, _ = self.self_attention(x, x, x, need_weights=False)
+        x = self.dropout(x)
+        x = x + input
+
+        y = self.ln_2(x)
+        y = self.mlp(y)
+        return x + y
+
+
+class Encoder(nn.Module):
+    """Transformer Model Encoder for sequence to sequence translation."""
+    def __init__(
+        self,
+        num_h_patches: int,
+        num_w_patches: int,
+        num_layers: int,
+        num_heads: int,
+        hidden_dim: int,
+        mlp_dim: int,
+        dropout: float,
+        attention_dropout: float,
+        norm_layer: Callable[..., torch.nn.Module] = partial(nn.LayerNorm, eps=1e-6),
+    ):
+        super().__init__()
+        self.num_h_patches = num_h_patches
+        self.num_w_patches = num_w_patches
+
+        # Note that batch_size is on the first dim because
+        # we have batch_first=True in nn.MultiAttention() by default
+        seq_length = num_h_patches * num_w_patches + 1  # +1 for the class token
+        self.pos_embedding = nn.Parameter(torch.empty(1, seq_length, hidden_dim).normal_(std=0.02))  # from BERT
+        self.dropout = nn.Dropout(dropout)
+        layers: OrderedDict[str, nn.Module] = OrderedDict()
+        for i in range(num_layers):
+            layers[f"encoder_layer_{i}"] = EncoderBlock(
+                num_heads,
+                hidden_dim,
+                mlp_dim,
+                dropout,
+                attention_dropout,
+                norm_layer,
+            )
+        self.layers = nn.Sequential(layers)
+        self.ln = norm_layer(hidden_dim)
+
+    def _get_pos_embedding(self, n_h: int, n_w: int) -> Tensor:
+        if n_h == self.num_h_patches and n_w == self.num_w_patches:
+            return self.pos_embedding
+        else:
+            pos_embedding = self.pos_embedding[:, 1:, :]
+            pos_embedding = rearrange(pos_embedding, "1 (h w) d -> 1 d h w", h=self.num_h_patches, w=self.num_w_patches)
+            pos_embedding = F.interpolate(pos_embedding, size=(n_h, n_w), mode="bicubic")
+            pos_embedding = rearrange(pos_embedding, "1 d h w -> 1 (h w) d")
+            return torch.cat([self.pos_embedding[:, :1, :], pos_embedding], dim=1)
+
+    def forward(self, input: Tensor, n_h: int, n_w: int) -> Tensor:
+        torch._assert(input.dim() == 3, f"Expected (batch_size, seq_length, hidden_dim) got {input.shape}")
+        input = input + self._get_pos_embedding(n_h, n_w)
+        return self.ln(self.layers(self.dropout(input)))
+
+
+class VisionTransformer(nn.Module):
+    """Vision Transformer as a feature extractor."""
+
+    def __init__(
+        self,
+        image_size: int,
+        patch_size: int,
+        num_layers: int,
+        num_heads: int,
+        hidden_dim: int,
+        mlp_dim: int,
+        dropout: float = 0.0,
+        attention_dropout: float = 0.0,
+        # num_classes: int = 1000,  # No need for the classification head as we only need the features
+        reduction: Optional[int] = None,
+        representation_size: Optional[int] = None,
+        norm_layer: Callable[..., torch.nn.Module] = partial(nn.LayerNorm, eps=1e-6),
+        conv_stem_configs: Optional[List[ConvStemConfig]] = None,
+    ):
+        super().__init__()
+        _log_api_usage_once(self)
+        torch._assert(image_size % patch_size == 0, "Input shape indivisible by patch size!")
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.hidden_dim = hidden_dim
+        self.mlp_dim = mlp_dim
+        self.attention_dropout = attention_dropout
+        self.dropout = dropout
+        # self.num_classes = num_classes
+        self.representation_size = representation_size
+        self.norm_layer = norm_layer
+
+        if conv_stem_configs is not None:
+            # As per https://arxiv.org/abs/2106.14881
+            seq_proj = nn.Sequential()
+            prev_channels = 3
+            for i, conv_stem_layer_config in enumerate(conv_stem_configs):
+                seq_proj.add_module(
+                    f"conv_bn_relu_{i}",
+                    Conv2dNormActivation(
+                        in_channels=prev_channels,
+                        out_channels=conv_stem_layer_config.out_channels,
+                        kernel_size=conv_stem_layer_config.kernel_size,
+                        stride=conv_stem_layer_config.stride,
+                        norm_layer=conv_stem_layer_config.norm_layer,
+                        activation_layer=conv_stem_layer_config.activation_layer,
+                    ),
+                )
+                prev_channels = conv_stem_layer_config.out_channels
+            seq_proj.add_module(
+                "conv_last", nn.Conv2d(in_channels=prev_channels, out_channels=hidden_dim, kernel_size=1)
+            )
+            self.conv_proj: nn.Module = seq_proj
+        else:
+            self.conv_proj = nn.Conv2d(
+                in_channels=3, out_channels=hidden_dim, kernel_size=patch_size, stride=patch_size
+            )
+
+        seq_length = (image_size // patch_size) ** 2
+
+        # Add a class token
+        self.class_token = nn.Parameter(torch.zeros(1, 1, hidden_dim))
+        seq_length += 1
+
+        self.encoder = Encoder(
+            image_size // patch_size,
+            image_size // patch_size,
+            num_layers,
+            num_heads,
+            hidden_dim,
+            mlp_dim,
+            dropout,
+            attention_dropout,
+            norm_layer,
+        )
+        self.seq_length = seq_length
+
+        # heads_layers: OrderedDict[str, nn.Module] = OrderedDict()
+        # if representation_size is None:
+        #     heads_layers["head"] = nn.Linear(hidden_dim, num_classes)
+        # else:
+        #     heads_layers["pre_logits"] = nn.Linear(hidden_dim, representation_size)
+        #     heads_layers["act"] = nn.Tanh()
+        #     heads_layers["head"] = nn.Linear(representation_size, num_classes)
+
+        # self.heads = nn.Sequential(heads_layers)
+
+        if isinstance(self.conv_proj, nn.Conv2d):
+            # Init the patchify stem
+            fan_in = self.conv_proj.in_channels * self.conv_proj.kernel_size[0] * self.conv_proj.kernel_size[1]
+            nn.init.trunc_normal_(self.conv_proj.weight, std=math.sqrt(1 / fan_in))
+            if self.conv_proj.bias is not None:
+                nn.init.zeros_(self.conv_proj.bias)
+        elif self.conv_proj.conv_last is not None and isinstance(self.conv_proj.conv_last, nn.Conv2d):
+            # Init the last 1x1 conv of the conv stem
+            nn.init.normal_(
+                self.conv_proj.conv_last.weight, mean=0.0, std=math.sqrt(2.0 / self.conv_proj.conv_last.out_channels)
+            )
+            if self.conv_proj.conv_last.bias is not None:
+                nn.init.zeros_(self.conv_proj.conv_last.bias)
+
+        # if hasattr(self.heads, "pre_logits") and isinstance(self.heads.pre_logits, nn.Linear):
+        #     fan_in = self.heads.pre_logits.in_features
+        #     nn.init.trunc_normal_(self.heads.pre_logits.weight, std=math.sqrt(1 / fan_in))
+        #     nn.init.zeros_(self.heads.pre_logits.bias)
+
+        # if isinstance(self.heads.head, nn.Linear):
+        #     nn.init.zeros_(self.heads.head.weight)
+        #     nn.init.zeros_(self.heads.head.bias)
+
+        self.encoder_reduction = self.patch_size
+        self.reduction = self.encoder_reduction if reduction is None else reduction
+        self.channels = hidden_dim
+
+    def _process_input(self, x: Tensor) -> Tuple[Tensor, int, int, int]:
+        # (n, c, h, w) -> (n, hidden_dim, n_h, n_w)
+        x = self.conv_proj(x)
+        n, _, n_h, n_w = x.shape
+        # (n, hidden_dim, n_h, n_w) -> (n, hidden_dim, (n_h * n_w))
+        x = x.reshape(n, self.hidden_dim, n_h * n_w)
+
+        # (n, hidden_dim, (n_h * n_w)) -> (n, (n_h * n_w), hidden_dim)
+        # The self attention layer expects inputs in the format (N, S, E)
+        # where S is the source sequence length, N is the batch size, E is the
+        # embedding dimension
+        x = x.permute(0, 2, 1)
+
+        return x, n, n_h, n_w
+
+    def forward(self, x: Tensor) -> Tensor:
+        # Reshape and permute the input tensor
+        x, n, n_h, n_w = self._process_input(x)
+
+        # Expand the class token to the full batch
+        batch_class_token = self.class_token.expand(n, -1, -1)
+        x = torch.cat([batch_class_token, x], dim=1)
+
+        x = self.encoder(x, n_h, n_w)  # Allows input image to be of any size.
+
+        # Classifier "token" as used by standard language architectures
+        # x = x[:, 0]
+
+        # x = self.heads(x)
+
+        x = x[:, 1:, :]
+        x = rearrange(x, "n (h w) d -> n d h w", h=n_h, w=n_w)
+        if self.encoder_reduction != self.reduction:
+            x = F.interpolate(x, scale_factor=self.encoder_reduction / self.reduction, mode="bilinear")
+        return x  # To be consistent with timm models
+
+
+def _vision_transformer(
+    patch_size: int,
+    num_layers: int,
+    num_heads: int,
+    hidden_dim: int,
+    mlp_dim: int,
+    weights: str,
+    **kwargs: Any,
+) -> VisionTransformer:
+    image_size = kwargs.pop("image_size", 224)
+
+    model = VisionTransformer(
+        image_size=image_size,
+        patch_size=patch_size,
+        num_layers=num_layers,
+        num_heads=num_heads,
+        hidden_dim=hidden_dim,
+        mlp_dim=mlp_dim,
+        **kwargs,
+    )
+
+    if weights is not None:
+        weights = load_state_dict_from_url(weights, progress=kwargs.get("progress", True))
+        missing_keys, unexpected_keys = model.load_state_dict(weights, strict=False)
+        if len(missing_keys) > 0:
+            print(f"Missing keys: {missing_keys}")
+        if len(unexpected_keys) > 0:
+            print(f"Unexpected keys: {unexpected_keys}")
+
+    return model
+
+
+def interpolate_embeddings(
+    image_size: int,
+    patch_size: int,
+    pos_embedding: Tensor,
+    interpolation_mode: str = "bicubic",
+) -> Tensor:
+    """This function helps interpolate positional embeddings during checkpoint loading,
+    especially when you want to apply a pre-trained model on images with different resolution.
+
+    Args:
+        image_size (int): Image size of the new model.
+        patch_size (int): Patch size of the new model.
+        model_state (OrderedDict[str, Tensor]): State dict of the pre-trained model.
+        interpolation_mode (str): The algorithm used for upsampling. Default: bicubic.
+        reset_heads (bool): If true, not copying the state of heads. Default: False.
+
+    Returns:
+        Tensor: The interpolated positional embedding.
+    """
+    # Shape of pos_embedding is (1, seq_length, hidden_dim)
+    n, seq_length, hidden_dim = pos_embedding.shape
+    if n != 1:
+        raise ValueError(f"Unexpected position embedding shape: {pos_embedding.shape}")
+
+    new_seq_length = (image_size // patch_size) ** 2 + 1
+
+    # Need to interpolate the weights for the position embedding.
+    # We do this by reshaping the positions embeddings to a 2d grid, performing
+    # an interpolation in the (h, w) space and then reshaping back to a 1d grid.
+    if new_seq_length != seq_length:
+        # The class token embedding shouldn't be interpolated, so we split it up.
+        seq_length -= 1
+        new_seq_length -= 1
+        pos_embedding_token = pos_embedding[:, :1, :]
+        pos_embedding_img = pos_embedding[:, 1:, :]
+
+        # (1, seq_length, hidden_dim) -> (1, hidden_dim, seq_length)
+        pos_embedding_img = pos_embedding_img.permute(0, 2, 1)
+        seq_length_1d = int(math.sqrt(seq_length))
+        if seq_length_1d * seq_length_1d != seq_length:
+            raise ValueError(
+                f"seq_length is not a perfect square! Instead got seq_length_1d * seq_length_1d = {seq_length_1d * seq_length_1d } and seq_length = {seq_length}"
+            )
+
+        # (1, hidden_dim, seq_length) -> (1, hidden_dim, seq_l_1d, seq_l_1d)
+        pos_embedding_img = pos_embedding_img.reshape(1, hidden_dim, seq_length_1d, seq_length_1d)
+        new_seq_length_1d = image_size // patch_size
+
+        # Perform interpolation.
+        # (1, hidden_dim, seq_l_1d, seq_l_1d) -> (1, hidden_dim, new_seq_l_1d, new_seq_l_1d)
+        new_pos_embedding_img = nn.functional.interpolate(
+            pos_embedding_img,
+            size=new_seq_length_1d,
+            mode=interpolation_mode,
+        )
+
+        # (1, hidden_dim, new_seq_l_1d, new_seq_l_1d) -> (1, hidden_dim, new_seq_length)
+        new_pos_embedding_img = new_pos_embedding_img.reshape(1, hidden_dim, new_seq_length)
+
+        # (1, hidden_dim, new_seq_length) -> (1, new_seq_length, hidden_dim)
+        new_pos_embedding_img = new_pos_embedding_img.permute(0, 2, 1)
+        new_pos_embedding = torch.cat([pos_embedding_token, new_pos_embedding_img], dim=1)
+
+        return new_pos_embedding
+
+    return pos_embedding
+
+
+def vit_b_16(
+    image_size: int = 224,
+    reduction: int = 16,
+    **kwargs: Any,
+) -> VisionTransformer:
+    vit = _vision_transformer(
+        patch_size=16,
+        num_layers=12,
+        num_heads=12,
+        hidden_dim=768,
+        mlp_dim=3072,
+        weights=weights["vit_b_16"],
+        reduction=reduction,
+        **kwargs,
+    )
+    if image_size != 224:
+        vit.image_size = image_size
+        new_pos_embedding = interpolate_embeddings(image_size, 16, vit.state_dict()["encoder.pos_embedding"], "bicubic")
+        vit.encoder.pos_embedding = nn.Parameter(new_pos_embedding, requires_grad=True)
+    return vit
+
+
+def vit_b_32(
+    image_size: int = 224,
+    reduction: int = 32,
+    **kwargs: Any,
+) -> VisionTransformer:
+    vit = _vision_transformer(
+        patch_size=32,
+        num_layers=12,
+        num_heads=12,
+        hidden_dim=768,
+        mlp_dim=3072,
+        weights=weights["vit_b_32"],
+        reduction=reduction,
+        **kwargs,
+    )
+    if image_size != 224:
+        vit.image_size = image_size
+        new_pos_embedding = interpolate_embeddings(image_size, 32, vit.state_dict()["encoder.pos_embedding"], "bicubic")
+        vit.encoder.pos_embedding = nn.Parameter(new_pos_embedding, requires_grad=True)
+    return vit
+
+
+def vit_l_16(
+    image_size: int = 224,
+    reduction: int = 16,
+    **kwargs: Any,
+) -> VisionTransformer:
+    vit = _vision_transformer(
+        patch_size=16,
+        num_layers=24,
+        num_heads=16,
+        hidden_dim=1024,
+        mlp_dim=4096,
+        weights=weights["vit_l_16"],
+        reduction=reduction,
+        **kwargs,
+    )
+    if image_size != 224:
+        vit.image_size = image_size
+        new_pos_embedding = interpolate_embeddings(image_size, 16, vit.state_dict()["encoder.pos_embedding"], "bicubic")
+        vit.encoder.pos_embedding = nn.Parameter(new_pos_embedding, requires_grad=True)
+    return vit
+
+
+def vit_l_32(
+    image_size: int = 224,
+    reduction: int = 32,
+    **kwargs: Any,
+) -> VisionTransformer:
+    vit = _vision_transformer(
+        patch_size=32,
+        num_layers=24,
+        num_heads=16,
+        hidden_dim=1024,
+        mlp_dim=4096,
+        weights=weights["vit_l_32"],
+        reduction=reduction,
+        **kwargs,
+    )
+    if image_size != 224:
+        vit.image_size = image_size
+        new_pos_embedding = interpolate_embeddings(image_size, 32, vit.state_dict()["encoder.pos_embedding"], "bicubic")
+        vit.encoder.pos_embedding = nn.Parameter(new_pos_embedding, requires_grad=True)
+    return vit
+
+
+def vit_h_14(
+    image_size: int = 224,
+    reduction: int = 14,
+    **kwargs: Any,
+) -> VisionTransformer:
+    vit = _vision_transformer(
+        patch_size=14,
+        num_layers=32,
+        num_heads=16,
+        hidden_dim=1280,
+        mlp_dim=5120,
+        weights=weights["vit_h_14"],
+        reduction=reduction,
+        **kwargs,
+    )
+    if image_size != 224:
+        vit.image_size = image_size
+        new_pos_embedding = interpolate_embeddings(image_size, 14, vit.state_dict()["encoder.pos_embedding"], "bicubic")
+        vit.encoder.pos_embedding = nn.Parameter(new_pos_embedding, requires_grad=True)
+    return vit
+

models/encoder_decoder/__init__.py

@@ -0,0 +1,17 @@
+from .vgg import vgg11 as vgg11_ae, vgg11_bn as vgg11_bn_ae
+from .vgg import vgg13 as vgg13_ae, vgg13_bn as vgg13_bn_ae
+from .vgg import vgg16 as vgg16_ae, vgg16_bn as vgg16_bn_ae
+from .vgg import vgg19 as vgg19_ae, vgg19_bn as vgg19_bn_ae
+from .resnet import resnet18 as resnet18_ae, resnet34 as resnet34_ae
+from .resnet import resnet50 as resnet50_ae, resnet101 as resnet101_ae, resnet152 as resnet152_ae
+
+from .cannet import cannet, cannet_bn
+from .csrnet import csrnet, csrnet_bn
+
+
+__all__ = [
+    "vgg11_ae", "vgg11_bn_ae", "vgg13_ae", "vgg13_bn_ae", "vgg16_ae", "vgg16_bn_ae", "vgg19_ae", "vgg19_bn_ae",
+    "resnet18_ae", "resnet34_ae", "resnet50_ae", "resnet101_ae", "resnet152_ae",
+    "cannet", "cannet_bn",
+    "csrnet", "csrnet_bn",
+]

models/encoder_decoder/cannet.py

@@ -0,0 +1,85 @@
+import torch
+from torch import nn, Tensor
+import torch.nn.functional as F
+
+from typing import List, Optional
+
+from ..utils import _init_weights
+from .csrnet import CSRNet, csrnet, csrnet_bn
+
+EPS = 1e-6
+
+
+class ContextualModule(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int = 512,
+        sizes: List[int] = [1, 2, 3, 6],
+    ) -> None:
+        super().__init__()
+        self.scales = nn.ModuleList([self.__make_scale__(in_channels, size) for size in sizes])
+        self.bottleneck = nn.Conv2d(in_channels * 2, out_channels, kernel_size=1)
+        self.relu = nn.ReLU(inplace=True)
+        self.weight_net = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+
+    def __make_weight__(self, feature: Tensor, scale_feature: Tensor) -> Tensor:
+        weight_feature = feature - scale_feature
+        weight_feature = self.weight_net(weight_feature)
+        return F.sigmoid(weight_feature)
+    
+    def __make_scale__(self, channels: int, size: int) -> nn.Module:
+        return nn.Sequential(
+            nn.AdaptiveAvgPool2d(output_size=(size, size)),
+            nn.Conv2d(channels, channels, kernel_size=1, bias=False),
+        )
+
+    def forward(self, feature: Tensor) -> Tensor:
+        h, w = feature.shape[-2:]
+        multi_scales = [F.interpolate(input=scale(feature), size=(h, w), mode="bilinear") for scale in self.scales]
+        weights = [self.__make_weight__(feature, scale_feature) for scale_feature in multi_scales]
+        multi_scales = sum([multi_scales[i] * weights[i] for i in range(len(weights))]) / (sum(weights) + EPS)
+        overall_features = torch.cat([multi_scales, feature], dim=1)
+        overall_features = self.bottleneck(overall_features)
+        overall_features = self.relu(overall_features)
+        return overall_features
+
+
+class CANNet(nn.Module):
+    def __init__(
+        self,
+        csrnet: CSRNet,
+        sizes: List[int] = [1, 2, 3, 6],
+        reduction: Optional[int] = 8,
+    ) -> None:
+        super().__init__()
+        assert isinstance(csrnet, CSRNet), f"csrnet should be an instance of CSRNet, got {type(csrnet)}."
+        assert isinstance(sizes, (tuple, list)), f"sizes should be a list or tuple, got {type(sizes)}."
+        assert len(sizes) > 0, f"Expected at least one size, got {len(sizes)}."
+        assert all([isinstance(size, int) for size in sizes]), f"Expected all size to be int, got {sizes}."
+        self.sizes = sizes
+        self.encoder_reduction = csrnet.encoder_reduction
+        self.reduction = self.encoder_reduction if reduction is None else reduction
+
+        self.features = csrnet.features
+        self.decoder = csrnet.decoder
+        self.decoder.apply(_init_weights)
+        self.context = ContextualModule(512, 512, self.sizes)
+        self.context.apply(_init_weights)
+
+        self.channels = csrnet.channels
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.features(x)
+        x = self.context(x)
+        if self.encoder_reduction != self.reduction:
+            x = F.interpolate(x, scale_factor=self.encoder_reduction / self.reduction, mode="bilinear")
+        x = self.decoder(x)
+        return x
+
+
+def cannet(sizes=[1, 2, 3, 6], reduction: int = 8) -> CANNet:
+    return CANNet(csrnet(), sizes=sizes, reduction=reduction)
+
+def cannet_bn(sizes=[1, 2, 3, 6], reduction: int = 8) -> CANNet:
+    return CANNet(csrnet_bn(), sizes=sizes, reduction=reduction)

models/encoder_decoder/csrnet.py

@@ -0,0 +1,54 @@
+from torch import nn, Tensor
+import torch.nn.functional as F
+from typing import Optional
+
+from ..utils import _init_weights, make_vgg_layers, vgg_urls
+from .vgg import _load_weights
+
+EPS = 1e-6
+
+
+encoder_cfg = [64, 64, "M", 128, 128, "M", 256, 256, 256, "M", 512, 512, 512]
+decoder_cfg = [512, 512, 512, 256, 128, 64]
+
+
+class CSRNet(nn.Module):
+    def __init__(
+        self,
+        features: nn.Module,
+        decoder: nn.Module,
+        reduction: Optional[int] = None,
+    ) -> None:
+        super().__init__()
+        self.features = features
+        self.features.apply(_init_weights)
+        self.decoder = decoder
+        self.decoder.apply(_init_weights)
+
+        self.encoder_reduction = 8
+        self.reduction = self.encoder_reduction if reduction is None else reduction
+        self.channels = 64
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.features(x)
+        if self.encoder_reduction != self.reduction:
+            x = F.interpolate(x, scale_factor=self.encoder_reduction / self.reduction, mode="bilinear")
+        x = self.decoder(x)
+        return x
+
+
+def csrnet(reduction: int = 8) -> CSRNet:
+    model = CSRNet(
+        make_vgg_layers(encoder_cfg, in_channels=3, batch_norm=False, dilation=1),
+        make_vgg_layers(decoder_cfg, in_channels=512, batch_norm=False, dilation=2),
+        reduction=reduction
+    )
+    return _load_weights(model, vgg_urls["vgg16"])
+
+def csrnet_bn(reduction: int = 8) -> CSRNet:
+    model = CSRNet(
+        make_vgg_layers(encoder_cfg, in_channels=3, batch_norm=True, dilation=1),
+        make_vgg_layers(decoder_cfg, in_channels=512, batch_norm=True, dilation=2),
+        reduction=reduction
+    )
+    return _load_weights(model, vgg_urls["vgg16"])

models/encoder_decoder/resnet.py

@@ -0,0 +1,95 @@
+from torch import nn, Tensor
+import torch.nn.functional as F
+import timm
+from typing import Union, Optional
+
+from ..utils import BasicBlock, Bottleneck, make_resnet_layers
+from ..utils import _init_weights
+
+
+model_configs = {
+    "resnet18.tv_in1k": {
+        "decoder_channels": [512, 256, 128],
+    },
+    "resnet34.tv_in1k": {
+        "decoder_channels": [512, 256, 128],
+    },
+    "resnet50.tv_in1k": {
+        "decoder_channels": [512, 256, 256, 128],
+    },
+    "resnet101.tv_in1k": {
+        "decoder_channels": [512, 512, 256, 256, 128],
+    },
+    "resnet152.tv_in1k": {
+        "decoder_channels": [512, 512, 512, 256, 256, 128],
+    },
+}
+
+
+class ResNet(nn.Module):
+    def __init__(
+        self,
+        decoder_block: Union[BasicBlock, Bottleneck],
+        backbone: str = "resnet34.tv_in1k",
+        reduction: Optional[int] = None,
+    ) -> None:
+        super().__init__()
+        assert backbone in model_configs.keys(), f"Backbone should be in {model_configs.keys()}"
+        config = model_configs[backbone]
+        encoder = timm.create_model(backbone, pretrained=True, features_only=True, out_indices=(-1,))
+        encoder_reduction = encoder.feature_info.reduction()[-1]
+
+        if reduction <= 16:
+            if "resnet18" in backbone or "resnet34" in backbone:
+                encoder.layer4[0].conv1.stride = (1, 1)
+                encoder.layer4[0].downsample[0].stride = (1, 1)
+            else:
+                encoder.layer4[0].conv2.stride = (1, 1)
+                encoder.layer4[0].downsample[0].stride = (1, 1)
+            encoder_reduction = encoder_reduction // 2
+
+        self.encoder = encoder
+        self.encoder_reduction = encoder_reduction
+
+        encoder_out_channels = self.encoder.feature_info.channels()[-1]
+
+        decoder_channels = config["decoder_channels"]
+        self.decoder = make_resnet_layers(
+            block=decoder_block,
+            cfg=decoder_channels,
+            in_channels=encoder_out_channels,
+            dilation=1,
+            expansion=1,
+        )
+        self.decoder.apply(_init_weights)
+
+        self.reduction = self.encoder_reduction if reduction is None else reduction
+        self.channels = decoder_channels[-1]
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.encoder(x)[-1]
+        if self.encoder_reduction != self.reduction:
+            x = F.interpolate(x, scale_factor=self.encoder_reduction / self.reduction, mode="bilinear")
+        x = self.decoder(x)
+
+        return x
+
+
+def resnet18(reduction: int = 32) -> ResNet:
+    return ResNet(decoder_block=BasicBlock, backbone="resnet18.tv_in1k", reduction=reduction)
+
+
+def resnet34(reduction: int = 32) -> ResNet:
+    return ResNet(decoder_block=BasicBlock, backbone="resnet34.tv_in1k", reduction=reduction)
+
+
+def resnet50(reduction: int = 32) -> ResNet:
+    return ResNet(decoder_block=Bottleneck, backbone="resnet50.tv_in1k", reduction=reduction)
+
+
+def resnet101(reduction: int = 32) -> ResNet:
+    return ResNet(decoder_block=Bottleneck, backbone="resnet101.tv_in1k", reduction=reduction)
+
+
+def resnet152(reduction: int = 32) -> ResNet:
+    return ResNet(decoder_block=Bottleneck, backbone="resnet152.tv_in1k", reduction=reduction)

models/encoder_decoder/vgg.py

@@ -0,0 +1,85 @@
+# The model used in the paper Distribution Matching for Crowd Counting.
+# Code adapted from https://github.com/cvlab-stonybrook/DM-Count/blob/master/models.py
+from torch import nn, Tensor
+import torch.nn.functional as F
+from torch.hub import load_state_dict_from_url
+from typing import Optional
+
+from ..utils import make_vgg_layers, vgg_cfgs, vgg_urls
+from ..utils import _init_weights
+
+
+
+class VGG(nn.Module):
+    def __init__(
+        self,
+        features: nn.Module,
+        reduction: Optional[int] = None,
+    ) -> None:
+        super().__init__()
+        self.features = features
+        self.reg_layer = nn.Sequential(
+            nn.Conv2d(512, 256, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(256, 128, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+        )
+
+        self.reg_layer.apply(_init_weights)
+        # Remove the density layer, as the output from this model is not final and will be further processed.
+        # self.density_layer = nn.Sequential(nn.Conv2d(128, 1, 1), nn.ReLU())
+        self.encoder_reduction = 16
+        self.reduction = self.encoder_reduction if reduction is None else reduction
+        self.channels = 128
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.features(x)
+        if self.encoder_reduction != self.reduction:
+            x = F.interpolate(x, scale_factor=self.encoder_reduction / self.reduction, mode="bilinear")
+        x = self.reg_layer(x)
+        # x = self.density_layer(x)
+        return x
+
+
+def _load_weights(model: VGG, url: str) -> VGG:
+    state_dict = load_state_dict_from_url(url)
+    missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
+    print("Loading pre-trained weights")
+    if len(missing_keys) > 0:
+        print(f"Missing keys: {missing_keys}")
+    if len(unexpected_keys) > 0:
+        print(f"Unexpected keys: {unexpected_keys}")
+    return model
+
+
+def vgg11(reduction: int = 8) -> VGG:
+    model = VGG(make_vgg_layers(vgg_cfgs["A"]), reduction=reduction)
+    return _load_weights(model, vgg_urls["vgg11"])
+
+def vgg11_bn(reduction: int = 8) -> VGG:
+    model = VGG(make_vgg_layers(vgg_cfgs["A"], batch_norm=True), reduction=reduction)
+    return _load_weights(model, vgg_urls["vgg11_bn"])
+
+def vgg13(reduction: int = 8) -> VGG:
+    model = VGG(make_vgg_layers(vgg_cfgs["B"]), reduction=reduction)
+    return _load_weights(model, vgg_urls["vgg13"])
+
+def vgg13_bn(reduction: int = 8) -> VGG:
+    model = VGG(make_vgg_layers(vgg_cfgs["B"], batch_norm=True), reduction=reduction)
+    return _load_weights(model, vgg_urls["vgg13_bn"])
+
+def vgg16(reduction: int = 8) -> VGG:
+    model = VGG(make_vgg_layers(vgg_cfgs["D"]), reduction=reduction)
+    return _load_weights(model, vgg_urls["vgg16"])
+
+def vgg16_bn(reduction: int = 8) -> VGG:
+    model = VGG(make_vgg_layers(vgg_cfgs["D"], batch_norm=True), reduction=reduction)
+    return _load_weights(model, vgg_urls["vgg16_bn"])
+
+def vgg19(reduction: int = 8) -> VGG:
+    model = VGG(make_vgg_layers(vgg_cfgs["E"]), reduction=reduction)
+    return _load_weights(model, vgg_urls["vgg19"])
+
+def vgg19_bn(reduction: int = 8) -> VGG:
+    model = VGG(make_vgg_layers(vgg_cfgs["E"], batch_norm=True), reduction=reduction)
+    return _load_weights(model, vgg_urls["vgg19_bn"])

models/model.py

@@ -0,0 +1,112 @@
+import torch
+from torch import nn, Tensor
+import os
+from typing import List, Tuple, Union, Callable
+from functools import partial
+
+from .utils import _init_weights
+
+from . import encoder
+from . import encoder_decoder
+from .encoder import _timm_encoder
+
+
+curr_dir = os.path.abspath(os.path.dirname(__file__))
+
+
+class Regressor(nn.Module):
+    def __init__(self, backbone: nn.Module) -> None:
+        super().__init__()
+        self.backbone = backbone
+        self.reduction = backbone.reduction
+
+        self.regressor = nn.Sequential(
+            nn.Conv2d(backbone.channels, 1, kernel_size=1),
+            nn.ReLU(inplace=True),
+        )
+        self.regressor.apply(_init_weights)
+        self.bins = None
+        self.anchor_points = None
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.backbone(x)
+        x = self.regressor(x)
+        return x
+
+
+class Classifier(nn.Module):
+    def __init__(
+        self,
+        backbone: nn.Module,
+        bins: List[Tuple[float, float]],
+        anchor_points: List[float],
+    ) -> None:
+        super().__init__()
+        self.backbone = backbone
+        self.reduction = backbone.reduction
+
+        assert len(bins) == len(anchor_points), f"Expected bins and anchor_points to have the same length, got {len(bins)} and {len(anchor_points)}"
+        assert all(len(b) == 2 for b in bins), f"Expected bins to be a list of tuples of length 2, got {bins}"
+        assert all(bin[0] <= p <= bin[1] for bin, p in zip(bins, anchor_points)), f"Expected anchor_points to be within the range of the corresponding bin, got {bins} and {anchor_points}"
+
+        self.bins = bins
+        self.anchor_points = torch.tensor(anchor_points, dtype=torch.float32, requires_grad=False).view(1, -1, 1, 1)
+
+        if backbone.channels > 512:
+            self.classifier = nn.Sequential(
+                nn.Conv2d(backbone.channels, 512, kernel_size=1),  # serves as a linear layer for feature vectors at each pixel
+                nn.ReLU(inplace=True),
+                nn.Conv2d(512, len(self.bins), kernel_size=1),
+            )
+        else:
+            self.classifier = nn.Conv2d(backbone.channels, len(self.bins), kernel_size=1)
+
+        self.classifier.apply(_init_weights)
+
+    def forward(self, x: Tensor) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+        x = self.backbone(x)
+        x = self.classifier(x)  # shape (B, C, H, W), where C = len(bins), x is the logits
+
+        probs = x.softmax(dim=1)  # shape (B, C, H, W)
+        exp = (probs * self.anchor_points.to(x.device)).sum(dim=1, keepdim=True)  # shape (B, 1, H, W)
+        if self.training:
+            return x, exp
+        else:
+            return exp
+
+
+def _get_backbone(backbone: str, input_size: int, reduction: int) -> Callable:
+    assert "clip" not in backbone, f"This function does not support CLIP model, got {backbone}"
+
+    if backbone in ["vit_b_16", "vit_b_32", "vit_l_16", "vit_l_32", "vit_h_14"]:
+        return partial(getattr(encoder, backbone), image_size=input_size, reduction=reduction)
+    elif backbone in ["vgg11", "vgg11_bn", "vgg13", "vgg13_bn", "vgg16", "vgg16_bn", "vgg19", "vgg19_bn"]:
+        return partial(getattr(encoder, backbone), reduction=reduction)
+    elif backbone in ["vgg11_ae", "vgg11_bn_ae", "vgg13_ae", "vgg13_bn_ae", "vgg16_ae", "vgg16_bn_ae", "vgg19_ae", "vgg19_bn_ae"]:
+        return partial(getattr(encoder_decoder, backbone), reduction=reduction)
+    elif backbone in ["resnet18_ae", "resnet34_ae", "resnet50_ae", "resnet101_ae", "resnet152_ae"]:
+        return partial(getattr(encoder_decoder, backbone), reduction=reduction)
+    elif backbone in ["cannet", "cannet_bn", "csrnet", "csrnet_bn"]:
+        return partial(getattr(encoder_decoder, backbone), reduction=reduction)
+    else:
+        return partial(_timm_encoder, backbone=backbone, reduction=reduction)
+
+
+def _regressor(
+    backbone: str,
+    input_size: int,
+    reduction: int,
+) -> Regressor:
+    backbone = _get_backbone(backbone.lower(), input_size, reduction)
+    return Regressor(backbone())
+
+
+def _classifier(
+    backbone: nn.Module,
+    input_size: int,
+    reduction: int,
+    bins: List[Tuple[float, float]],
+    anchor_points: List[float],
+) -> Classifier:
+    backbone = _get_backbone(backbone.lower(), input_size, reduction)
+    return Classifier(backbone(), bins, anchor_points)

models/utils.py

@@ -0,0 +1,444 @@
+import torch
+from torch import nn, Tensor
+import torch.nn.functional as F
+from functools import partial
+from typing import Callable, Optional, Sequence, Tuple, Union, Any, List, TypeVar, List
+from types import FunctionType
+from itertools import repeat
+import warnings
+import os
+from collections.abc import Iterable
+
+V = TypeVar("V")
+curr_dir = os.path.dirname(os.path.abspath(__file__))
+
+
+vgg_urls = {
+    "vgg11": "https://download.pytorch.org/models/vgg11-8a719046.pth",
+    "vgg11_bn": "https://download.pytorch.org/models/vgg11_bn-6002323d.pth",
+    "vgg13": "https://download.pytorch.org/models/vgg13-19584684.pth",
+    "vgg13_bn": "https://download.pytorch.org/models/vgg13_bn-abd245e5.pth",
+    "vgg16": "https://download.pytorch.org/models/vgg16-397923af.pth",
+    "vgg16_bn": "https://download.pytorch.org/models/vgg16_bn-6c64b313.pth",
+    "vgg19": "https://download.pytorch.org/models/vgg19-dcbb9e9d.pth",
+    "vgg19_bn": "https://download.pytorch.org/models/vgg19_bn-c79401a0.pth",
+}
+
+vgg_cfgs = {
+    "A": [64, "M", 128, "M", 256, 256, "M", 512, 512, "M", 512, 512],
+    "B": [64, 64, "M", 128, 128, "M", 256, 256, "M", 512, 512, "M", 512, 512],
+    "D": [64, 64, "M", 128, 128, "M", 256, 256, 256, "M", 512, 512, 512, "M", 512, 512, 512],
+    "E": [64, 64, "M", 128, 128, "M", 256, 256, 256, 256, "M", 512, 512, 512, 512, "M", 512, 512, 512, 512]
+}
+
+
+def _log_api_usage_once(obj: Any) -> None:
+
+    """
+    Logs API usage(module and name) within an organization.
+    In a large ecosystem, it's often useful to track the PyTorch and
+    TorchVision APIs usage. This API provides the similar functionality to the
+    logging module in the Python stdlib. It can be used for debugging purpose
+    to log which methods are used and by default it is inactive, unless the user
+    manually subscribes a logger via the `SetAPIUsageLogger method <https://github.com/pytorch/pytorch/blob/eb3b9fe719b21fae13c7a7cf3253f970290a573e/c10/util/Logging.cpp#L114>`_.
+    Please note it is triggered only once for the same API call within a process.
+    It does not collect any data from open-source users since it is no-op by default.
+    For more information, please refer to
+    * PyTorch note: https://pytorch.org/docs/stable/notes/large_scale_deployments.html#api-usage-logging;
+    * Logging policy: https://github.com/pytorch/vision/issues/5052;
+
+    Args:
+        obj (class instance or method): an object to extract info from.
+    """
+    module = obj.__module__
+    if not module.startswith("torchvision"):
+        module = f"torchvision.internal.{module}"
+    name = obj.__class__.__name__
+    if isinstance(obj, FunctionType):
+        name = obj.__name__
+    torch._C._log_api_usage_once(f"{module}.{name}")
+
+
+def _make_ntuple(x: Any, n: int) -> Tuple[Any, ...]:
+    """
+    Make n-tuple from input x. If x is an iterable, then we just convert it to tuple.
+    Otherwise, we will make a tuple of length n, all with value of x.
+    reference: https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/utils.py#L8
+
+    Args:
+        x (Any): input value
+        n (int): length of the resulting tuple
+    """
+    if isinstance(x, Iterable):
+        return tuple(x)
+    return tuple(repeat(x, n))
+
+
+class ConvNormActivation(torch.nn.Sequential):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: Union[int, Tuple[int, ...]] = 3,
+        stride: Union[int, Tuple[int, ...]] = 1,
+        padding: Optional[Union[int, Tuple[int, ...], str]] = None,
+        groups: int = 1,
+        norm_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.BatchNorm2d,
+        activation_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.ReLU,
+        dilation: Union[int, Tuple[int, ...]] = 1,
+        inplace: Optional[bool] = True,
+        bias: Optional[bool] = None,
+        conv_layer: Callable[..., torch.nn.Module] = torch.nn.Conv2d,
+    ) -> None:
+
+        if padding is None:
+            if isinstance(kernel_size, int) and isinstance(dilation, int):
+                padding = (kernel_size - 1) // 2 * dilation
+            else:
+                _conv_dim = len(kernel_size) if isinstance(kernel_size, Sequence) else len(dilation)
+                kernel_size = _make_ntuple(kernel_size, _conv_dim)
+                dilation = _make_ntuple(dilation, _conv_dim)
+                padding = tuple((kernel_size[i] - 1) // 2 * dilation[i] for i in range(_conv_dim))
+        if bias is None:
+            bias = norm_layer is None
+
+        layers = [
+            conv_layer(
+                in_channels,
+                out_channels,
+                kernel_size,
+                stride,
+                padding,
+                dilation=dilation,
+                groups=groups,
+                bias=bias,
+            )
+        ]
+
+        if norm_layer is not None:
+            layers.append(norm_layer(out_channels))
+
+        if activation_layer is not None:
+            params = {} if inplace is None else {"inplace": inplace}
+            layers.append(activation_layer(**params))
+        super().__init__(*layers)
+        _log_api_usage_once(self)
+        self.out_channels = out_channels
+
+        if self.__class__ == ConvNormActivation:
+            warnings.warn(
+                "Don't use ConvNormActivation directly, please use Conv2dNormActivation and Conv3dNormActivation instead."
+            )
+
+
+class Conv2dNormActivation(ConvNormActivation):
+    """
+    Configurable block used for Convolution2d-Normalization-Activation blocks.
+
+    Args:
+        in_channels (int): Number of channels in the input image
+        out_channels (int): Number of channels produced by the Convolution-Normalization-Activation block
+        kernel_size: (int, optional): Size of the convolving kernel. Default: 3
+        stride (int, optional): Stride of the convolution. Default: 1
+        padding (int, tuple or str, optional): Padding added to all four sides of the input. Default: None, in which case it will be calculated as ``padding = (kernel_size - 1) // 2 * dilation``
+        groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1
+        norm_layer (Callable[..., torch.nn.Module], optional): Norm layer that will be stacked on top of the convolution layer. If ``None`` this layer won't be used. Default: ``torch.nn.BatchNorm2d``
+        activation_layer (Callable[..., torch.nn.Module], optional): Activation function which will be stacked on top of the normalization layer (if not None), otherwise on top of the conv layer. If ``None`` this layer won't be used. Default: ``torch.nn.ReLU``
+        dilation (int): Spacing between kernel elements. Default: 1
+        inplace (bool): Parameter for the activation layer, which can optionally do the operation in-place. Default ``True``
+        bias (bool, optional): Whether to use bias in the convolution layer. By default, biases are included if ``norm_layer is None``.
+
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: Union[int, Tuple[int, int]] = 3,
+        stride: Union[int, Tuple[int, int]] = 1,
+        padding: Optional[Union[int, Tuple[int, int], str]] = None,
+        groups: int = 1,
+        norm_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.BatchNorm2d,
+        activation_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.ReLU,
+        dilation: Union[int, Tuple[int, int]] = 1,
+        inplace: Optional[bool] = True,
+        bias: Optional[bool] = None,
+    ) -> None:
+
+        super().__init__(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride,
+            padding,
+            groups,
+            norm_layer,
+            activation_layer,
+            dilation,
+            inplace,
+            bias,
+            torch.nn.Conv2d,
+        )
+
+
+class MLP(torch.nn.Sequential):
+    """This block implements the multi-layer perceptron (MLP) module.
+
+    Args:
+        in_channels (int): Number of channels of the input
+        hidden_channels (List[int]): List of the hidden channel dimensions
+        norm_layer (Callable[..., torch.nn.Module], optional): Norm layer that will be stacked on top of the linear layer. If ``None`` this layer won't be used. Default: ``None``
+        activation_layer (Callable[..., torch.nn.Module], optional): Activation function which will be stacked on top of the normalization layer (if not None), otherwise on top of the linear layer. If ``None`` this layer won't be used. Default: ``torch.nn.ReLU``
+        inplace (bool, optional): Parameter for the activation layer, which can optionally do the operation in-place.
+            Default is ``None``, which uses the respective default values of the ``activation_layer`` and Dropout layer.
+        bias (bool): Whether to use bias in the linear layer. Default ``True``
+        dropout (float): The probability for the dropout layer. Default: 0.0
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        hidden_channels: List[int],
+        norm_layer: Optional[Callable[..., torch.nn.Module]] = None,
+        activation_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.ReLU,
+        inplace: Optional[bool] = None,
+        bias: bool = True,
+        dropout: float = 0.0,
+    ):
+        # The addition of `norm_layer` is inspired from the implementation of TorchMultimodal:
+        # https://github.com/facebookresearch/multimodal/blob/5dec8a/torchmultimodal/modules/layers/mlp.py
+        params = {} if inplace is None else {"inplace": inplace}
+
+        layers = []
+        in_dim = in_channels
+        for hidden_dim in hidden_channels[:-1]:
+            layers.append(torch.nn.Linear(in_dim, hidden_dim, bias=bias))
+            if norm_layer is not None:
+                layers.append(norm_layer(hidden_dim))
+            layers.append(activation_layer(**params))
+            layers.append(torch.nn.Dropout(dropout, **params))
+            in_dim = hidden_dim
+
+        layers.append(torch.nn.Linear(in_dim, hidden_channels[-1], bias=bias))
+        layers.append(torch.nn.Dropout(dropout, **params))
+
+        super().__init__(*layers)
+        _log_api_usage_once(self)
+
+
+def conv3x3(
+    in_channels: int,
+    out_channels: int,
+    stride: int = 1,
+    groups: int = 1,
+    dilation: int = 1,
+) -> nn.Conv2d:
+    """3x3 convolution with padding"""
+    return nn.Conv2d(
+        in_channels,
+        out_channels,
+        kernel_size=3,
+        stride=stride,
+        padding=dilation,
+        groups=groups,
+        bias=False,
+        dilation=dilation,
+    )
+
+
+def conv1x1(in_channels: int, out_channels: int, stride: int = 1) -> nn.Conv2d:
+    """1x1 convolution"""
+    return nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False)
+
+
+class BasicBlock(nn.Module):
+    expansion: int = 1
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        stride: int = 1,
+        groups: int = 1,
+        base_width: int = 64,
+        dilation: int = 1,
+        norm_layer: Optional[Callable[..., nn.Module]] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        if groups != 1 or base_width != 64:
+            raise ValueError("BasicBlock only supports groups=1 and base_width=64")
+        if dilation > 1:
+            raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
+        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv3x3(in_channels, out_channels, stride)
+        self.bn1 = norm_layer(out_channels)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(out_channels, out_channels)
+        self.bn2 = norm_layer(out_channels)
+        self.stride = stride
+        if in_channels != out_channels:
+            self.downsample = nn.Sequential(
+                conv1x1(in_channels, out_channels),
+                nn.BatchNorm2d(out_channels),
+            )
+        else:
+            self.downsample = nn.Identity()
+
+    def forward(self, x: Tensor) -> Tensor:
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        out += self.downsample(identity)
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    # Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2)
+    # while original implementation places the stride at the first 1x1 convolution(self.conv1)
+    # according to "Deep residual learning for image recognition" https://arxiv.org/abs/1512.03385.
+    # This variant is also known as ResNet V1.5 and improves accuracy according to
+    # https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch.
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        stride: int = 1,
+        groups: int = 1,
+        base_width: int = 64,
+        dilation: int = 1,
+        expansion: int = 4,
+        norm_layer: Optional[Callable[..., nn.Module]] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        width = int(out_channels * (base_width / 64.0)) * groups
+        self.expansion = expansion
+        # Both self.conv2 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv1x1(in_channels, width)
+        self.bn1 = norm_layer(width)
+        self.conv2 = conv3x3(width, width, stride, groups, dilation)
+        self.bn2 = norm_layer(width)
+        self.conv3 = conv1x1(width, out_channels * self.expansion)
+        self.bn3 = norm_layer(out_channels * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.stride = stride
+        if in_channels != out_channels:
+            self.downsample = nn.Sequential(
+                conv1x1(in_channels, out_channels),
+                nn.BatchNorm2d(out_channels),
+            )
+        else:
+            self.downsample = nn.Identity()
+
+    def forward(self, x: Tensor) -> Tensor:
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        out += self.downsample(identity)
+        out = self.relu(out)
+
+        return out
+    
+
+def _init_weights(model: nn.Module) -> None:
+    for m in model.modules():
+        if isinstance(m, nn.Conv2d):
+            nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0.)
+        elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+            nn.init.constant_(m.weight, 1.)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0.)
+        elif isinstance(m, nn.Linear):
+            nn.init.normal_(m.weight, std=0.01)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0.)
+
+
+class Upsample(nn.Module):
+    def __init__(
+        self,
+        size: Union[int, Tuple[int, int]] = None,
+        scale_factor: Union[float, Tuple[float, float]] = None,
+        mode: str = "nearest",
+        align_corners: bool = False,
+        antialias: bool = False,
+    ) -> None:
+        super().__init__()
+        self.interpolate = partial(
+            F.interpolate,
+            size=size,
+            scale_factor=scale_factor,
+            mode=mode,
+            align_corners=align_corners,
+            antialias=antialias,
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.interpolate(x)
+
+
+def make_vgg_layers(cfg: List[Union[str, int]], in_channels: int = 3, batch_norm: bool = False, dilation: int = 1) -> nn.Sequential:
+    layers = []
+    for v in cfg:
+        if v == "M":
+            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
+        elif v == "U":
+            layers += [Upsample(scale_factor=2, mode="bilinear")]
+        else:
+            conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=dilation, dilation=dilation)
+            if batch_norm:
+                layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
+            else:
+                layers += [conv2d, nn.ReLU(inplace=True)]
+            in_channels = v
+    return nn.Sequential(*layers)
+
+
+def make_resnet_layers(
+    block: Union[BasicBlock, Bottleneck],
+    cfg: List[Union[int, str]],
+    in_channels: int,
+    dilation: int = 1,
+    expansion: int = 1,
+) -> nn.Sequential:
+    layers = []
+    for v in cfg:
+        if v == "U":
+            layers.append(Upsample(scale_factor=2, mode="bilinear"))
+        else:
+            layers.append(block(
+                in_channels=in_channels,
+                out_channels=v,
+                dilation=dilation,
+                expansion=expansion,
+            ))
+            in_channels = v
+
+    layers = nn.Sequential(*layers)
+    layers.apply(_init_weights)
+    return layers