ADD file via upload

jclip/model.py

@@ -0,0 +1,286 @@
+from typing import Tuple, Union
+
+import numpy as np
+import jittor as jt
+from jittor import nn, init
+from .mha import MultiheadAttention
+
+
+def normal_(module, mean=0, std=1, bias=0):
+    if hasattr(module, 'weight') and module.weight is not None:
+        init.gauss_(module.weight, mean, std)
+    if hasattr(module, 'bias') and isinstance(
+            module.bias, jt.Var) and module.bias is not None:
+        init.constant_(module.bias, bias)
+
+
+class LayerNorm(nn.LayerNorm):
+
+    def execute(self, x):
+        ret = super().execute(x)
+        return ret
+
+
+class QuickGELU(nn.Module):
+
+    def execute(self, x):
+        return x * jt.sigmoid(1.702 * x)
+
+
+class MLP(nn.Module):
+
+    def __init__(self, d_model):
+        super().__init__()
+        self.c_fc = nn.Linear(d_model, d_model * 4)
+        self.gelu = QuickGELU()
+        self.c_proj = nn.Linear(d_model * 4, d_model)
+
+    def execute(self, x):
+        return self.c_proj(self.gelu(self.c_fc(x)))
+
+
+class ResidualAttentionBlock(nn.Module):
+
+    def __init__(self, d_model, n_head, attn_mask):
+        super().__init__()
+
+        self.attn = MultiheadAttention(d_model, n_head)
+        self.ln_1 = LayerNorm(d_model)
+        self.mlp = MLP(d_model)
+        self.ln_2 = LayerNorm(d_model)
+        self.attn_mask = attn_mask
+
+    def attention(self, x):
+        self.attn_mask = self.attn_mask.to(
+            dtype=x.dtype) 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 execute(self, x):
+        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, layers, heads, attn_mask=None):
+        super().__init__()
+        self.width = width
+        self.layers = layers
+        self.resblocks = nn.Sequential(*[
+            ResidualAttentionBlock(width, heads, attn_mask)
+            for _ in range(layers)
+        ])
+
+    def execute(self, x):
+        return self.resblocks(x)
+
+
+class VisionTransformer(nn.Module):
+
+    def __init__(self, input_resolution: int, patch_size: int, width: int,
+                 layers: int, heads: int, output_dim: int):
+        super().__init__()
+        self.input_resolution = input_resolution
+        self.output_dim = output_dim
+        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 = scale * jt.randn((width))
+        self.positional_embedding = scale * jt.randn(
+            ((input_resolution // patch_size)**2 + 1, width))
+        self.ln_pre = LayerNorm(width)
+
+        self.transformer = Transformer(width, layers, heads)
+
+        self.ln_post = LayerNorm(width)
+        self.proj = scale * jt.randn((width, output_dim))
+
+    def execute(self, x):
+        x = self.conv1(x)  # shape = [*, width, grid, grid]
+        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 = jt.concat([
+            self.class_embedding.to(x.dtype) + jt.zeros(
+                (x.shape[0], 1, x.shape[-1]), dtype=x.dtype), x
+        ],
+                      dim=1)  # shape = [*, grid ** 2 + 1, width]
+        x = x + self.positional_embedding.to(x.dtype)
+        x = self.ln_pre(x)
+
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.transformer(x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+
+        x = self.ln_post(x[:, 0, :])
+
+        if self.proj is not None:
+            x = x @ self.proj
+
+        return x
+
+
+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):
+        super().__init__()
+
+        self.context_length = context_length
+
+        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)
+
+        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 = jt.empty(
+            (self.context_length, transformer_width))
+        self.ln_final = LayerNorm(transformer_width)
+
+        self.text_projection = jt.empty((transformer_width, embed_dim))
+        self.logit_scale = jt.ones([]) * np.log(1 / 0.07)
+
+        self.initialize_parameters()
+
+    def initialize_parameters(self):
+        normal_(self.token_embedding.weight, std=0.02)
+        normal_(self.positional_embedding, std=0.01)
+
+        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:
+            normal_(block.attn.in_proj_weight, std=attn_std)
+            normal_(block.attn.out_proj.weight, std=proj_std)
+            normal_(block.mlp.c_fc.weight, std=fc_std)
+            normal_(block.mlp.c_proj.weight, std=proj_std)
+
+        if self.text_projection is not None:
+            normal_(self.text_projection, std=self.transformer.width**-0.5)
+
+    def build_attention_mask(self):
+        mask = jt.empty((self.context_length, self.context_length))
+        mask.fill_(float("-inf"))
+        mask = jt.triu(mask, 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)
+
+    def encode_text(self, text):
+        x = self.token_embedding(text)
+
+        x = x + self.positional_embedding
+        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)
+
+        # 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[jt.arange(x.shape[0]),
+              text.argmax(dim=-1)[0]] @ self.text_projection
+        return x
+
+    def execute(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 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]
+
+    model.load_parameters(state_dict)
+    return model.eval()