infer_model.py

import abc
import math

import torch
import torch.nn.functional as F
from sentence_transformers import SentenceTransformer
from timm.models.vision_transformer import (
    VisionTransformer,
    build_model_with_cfg,
    checkpoint_filter_fn,
    checkpoint_seq,
    resolve_pretrained_cfg,
)
from torch import Tensor, nn


class BlankLayer(nn.Module):
    pass


class CustomViT(VisionTransformer):
    def __init__(
            self,
            *args,
            image_pooling="gmp",
            **kwargs,
    ):
        super(CustomViT, self).__init__(
            *args, **kwargs
        )
        self.image_pooling = image_pooling

    def forward_head(self, x, pre_logits: bool = False):
        if self.image_pooling:
            if self.image_pooling == "gap":
                x = x[:, self.num_prefix_tokens:].mean(dim=1)
            elif self.image_pooling == "gmp":
                x = x[:, self.num_prefix_tokens:].max(dim=-2)[0]
            elif self.image_pooling == "all":
                x = x[:, self.num_prefix_tokens:]
            else:  # cls by default
                x = x[:, 0]
        x = self.fc_norm(x)
        return x if pre_logits else self.head(x)

    def forward(self, x, get_pos_tokens=False):
        x = self.forward_features(x, get_pos_tokens=get_pos_tokens)
        if get_pos_tokens:
            return self.fc_norm(x[:, self.num_prefix_tokens:])
        x = self.forward_head(x)
        return x

    def forward_features(self, x, get_pos_tokens=False):
        _, nc, h, w = x.shape
        x = self.patch_embed(x)
        x = self._pos_embed(x, w, h)
        if self.grad_checkpointing and not torch.jit.is_scripting():
            x = checkpoint_seq(self.blocks, x)
        else:
            x = self.blocks(x)
        x = self.norm(x)
        return x

    def _pos_embed(self, x, w, h):
        if self.no_embed_class:
            # deit-3, updated JAX (big vision)
            # position embedding does not overlap with class token, add then concat
            x = x + self.pos_embed
            if self.cls_token is not None:
                x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1)
        else:
            # original timm, JAX, and deit vit impl
            # pos_embed has entry for class token, concat then add
            if self.cls_token is not None:
                x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1)
            x = x + self._interpolate_pos_encoding(x, w, h)
        return self.pos_drop(x)

    def _interpolate_pos_encoding(self, x, w, h):
        npatch = x.shape[1] - 1
        N = self.pos_embed.shape[1] - 1
        if npatch == N and w == h:
            return self.pos_embed
        class_pos_embed = self.pos_embed[:, 0]
        patch_pos_embed = self.pos_embed[:, 1:]
        dim = x.shape[-1]
        w0 = w // self.patch_embed.patch_size[0]
        h0 = h // self.patch_embed.patch_size[1]
        # we add a small number to avoid floating point error in the interpolation
        # see discussion at https://github.com/facebookresearch/dino/issues/8
        w0, h0 = w0 + 0.1, h0 + 0.1
        patch_pos_embed = nn.functional.interpolate(
            patch_pos_embed.reshape(1, int(math.sqrt(N)), int(math.sqrt(N)), dim).permute(
                0, 3, 1, 2
            ),
            scale_factor=(w0 / math.sqrt(N), h0 / math.sqrt(N)),
            mode="bicubic",
        )
        assert int(w0) == patch_pos_embed.shape[-2] and int(h0) == patch_pos_embed.shape[-1]
        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)


def _create_vision_transformer(variant, pretrained=False, **kwargs):
    if kwargs.get("features_only", None):
        raise RuntimeError("features_only not implemented for Vision Transformer models.")

    pretrained_cfg = resolve_pretrained_cfg(
        variant, pretrained_cfg=kwargs.pop("pretrained_cfg", None)
    )
    model = build_model_with_cfg(
        CustomViT,
        variant,
        pretrained,
        pretrained_cfg=pretrained_cfg,
        pretrained_filter_fn=checkpoint_filter_fn,
        pretrained_custom_load="npz" in pretrained_cfg["url"],
        **kwargs,
    )
    return model


def vit_base_patch16_224(pretrained=False, variant="vit_base_patch16_224_dino", **kwargs):
    """ViT-Base (ViT-B/16) /w DINO pretrained weights (no head) - https://arxiv.org/abs/2104.14294"""
    model_kwargs = dict(patch_size=16, embed_dim=768, depth=12, num_heads=12, **kwargs)
    model = _create_vision_transformer(variant, pretrained=pretrained, **model_kwargs)
    return model


class CLIPpyModel(abc.ABC, torch.nn.Module):
    """ Implements code for running inference with pre-trained CLIPpy model.

    NOTE: weights used are for a model trained with lower batch-size leading to results below those in paper.
    """

    def __init__(
            self,
            image_pooling: str = "cls",
            text_pooling: str = "gap",
    ):
        super().__init__()

        self.visual = BlankLayer()

        self.visual.trunk = vit_base_patch16_224(True, image_pooling=image_pooling)

        self.text = SentenceTransformer("sentence-transformers/sentence-t5-base")
        self.logit_scale = nn.Parameter(torch.ones([]) * math.log(1 / 0.07))
        self.set_text_pooling(text_pooling)

        self._divisor_eps = 1e-4
        self._image_pooling = image_pooling
        self._text_pooling = text_pooling

    def forward(
            self,
            images: Tensor,
            input_ids: Tensor,
            input_id_masks: Tensor,
            get_pos_tokens: bool = False,
            **kwargs,
    ):

        image_encodings = self.encode_image(images, get_pos_tokens=get_pos_tokens)

        if get_pos_tokens:
            return {
                image_encodings: image_encodings,
            }

        text_encodings = self.encode_text(input_ids, input_id_masks)

        return {
            image_encodings: image_encodings,
            text_encodings: text_encodings,
        }

    def encode_text(self, input_ids: Tensor, input_id_masks: Tensor = None, **kwargs):
        output = self.text({"input_ids": input_ids, "attention_mask": input_id_masks})[
            "sentence_embedding"
        ]
        return self.text_head(output)

    def text_head(self, hidden_states: Tensor, input_id_masks: Tensor = None, **kwargs):
        return F.normalize(hidden_states, dim=-1, eps=self._divisor_eps).float()

    def encode_image(self, images: Tensor, get_pos_tokens: bool = False, **kwargs):
        output = self.visual.trunk(images, get_pos_tokens)
        return self.image_head(output, get_pos_tokens=get_pos_tokens)

    def image_head(self, hidden_states: Tensor, get_pos_tokens: bool = False, **kwargs):
        return F.normalize(hidden_states, dim=-1, eps=self._divisor_eps).float()

    def set_text_pooling(self, pooling):
        """ Converts pooling in the Hugging Face model to be max or average pooling"""
        if pooling == "gmp":
            self.text[1].pooling_mode_mean_tokens = False
            self.text[1].pooling_mode_max_tokens = True
        elif pooling == "gap":
            pass
        else:
            raise NotImplementedError(f"{pooling} not implemented")