An implementation of "CLIP4STR: A Simple Baseline for Scene Text Recognition with Pre-trained Vision-Language Model".
#!/usr/bin/env python3
# Scene Text Recognition Model Hub
# Copyright 2022 Darwin Bautista
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
import argparse

import torch

from PIL import Image

from strhub.data.module import SceneTextDataModule
from strhub.models.utils import load_from_checkpoint, parse_model_args


@torch.inference_mode()
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('checkpoint', help="Model checkpoint (or 'pretrained=<model_id>')")
    parser.add_argument('--images_path', type=str, help='Images to read')
    parser.add_argument('--device', default='cuda')
    args, unknown = parser.parse_known_args()
    kwargs = parse_model_args(unknown)
    print(f'Additional keyword arguments: {kwargs}')

    model = load_from_checkpoint(args.checkpoint, **kwargs).eval().to(args.device)
    img_transform = SceneTextDataModule.get_transform(model.hparams.img_size)

    files = sorted([x for x in os.listdir(args.images_path) if x.endswith('png') or x.endswith('jpeg') or x.endswith('jpg')])

    for fname in files:
        # Load image and prepare for input
        filename = os.path.join(args.images_path, fname)
        image = Image.open(filename).convert('RGB')
        image = img_transform(image).unsqueeze(0).to(args.device)

        p = model(image).softmax(-1)
        pred, p = model.tokenizer.decode(p)
        print(f'{fname}: {pred[0]}')


if __name__ == '__main__':
    main()

read.py
VamosC/CLIP4STR
import hashlib
import os
import urllib
import warnings
from typing import Any, Union, List
from pkg_resources import packaging

import torch
from PIL import Image
from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
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"))
        raise NotImplementedError("Please pass a checkpoint file")
    elif os.path.isfile(name):
        model_path = name
    else:
        raise RuntimeError(f"Model {name} not found; available models = {available_models()}")
    print("loading checkpoint from {}".format(model_path))

    try:
        with open(model_path, 'rb') as opened_file:
            # 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")
        state_dict = torch.load(model_path, map_location="cpu")['model']

    if not jit:
        model = build_model(state_dict or model.state_dict()).to(device)
        if str(device) == "cpu":
            model.float()
        return model, _transform(model.visual.input_resolution)

    # 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 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["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 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, _transform(model.input_resolution.item())


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

strhub/clip/clip.py
VamosC/CLIP4STR
from collections import OrderedDict
from typing import Tuple, Union

import numpy as np
import torch
import torch.nn.functional as F
from torch import nn


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: torch.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 ** 2 + 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)


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, output_dim, heads, input_resolution=224, width=64):
        super().__init__()
        self.output_dim = output_dim
        self.input_resolution = input_resolution

        # 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=2)

        embed_dim = width * 32  # the ResNet feature dimension
        self.attnpool = AttentionPool2d(input_resolution // 32, embed_dim, heads, 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 forward(self, x):
        def stem(x):
            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

        x = x.type(self.conv1.weight.dtype)
        x = stem(x)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.attnpool(x)

        return x


class LayerNorm(nn.LayerNorm):
    """Subclass torch's LayerNorm to handle fp16."""

    def forward(self, x: torch.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: torch.Tensor):
        return x * torch.sigmoid(1.702 * x)


class ResidualAttentionBlock(nn.Module):
    def __init__(self, d_model: int, n_head: int, attn_mask: torch.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: torch.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: torch.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: torch.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: torch.Tensor):
        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 = nn.Parameter(scale * torch.randn(width))
        self.positional_embedding = nn.Parameter(scale * torch.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 = nn.Parameter(scale * torch.randn(width, output_dim))

    def forward(self, x: torch.Tensor, cls=False, projection=True):
        """
        Args:
            cls: if False, return all features of all tokens
            projection: 
        """
        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 = 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)  # 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

        if cls:
            x = self.ln_post(x[:, 0, :])
            if self.proj is not None and projection:
                x = x @ self.proj
        else:
            # NEW ADDED CODE
            # here we need to output all tokens 
            x = self.ln_post(x)
            if self.proj is not None and projection:
                x = torch.matmul(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

        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
            )
        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
            )

        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))
        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, cls=False, projection=True):
        return self.visual(image.type(self.dtype), cls=cls, projection=projection)

    def encode_text(self, text, eot=False, projection=True):
        x = self.token_embedding(text).type(self.dtype)  # [batch_size, n_ctx, d_model]

        # x = x + self.positional_embedding.type(self.dtype)
        pos_emd = self.positional_embedding[:x.size(1), :]
        x = x + pos_emd
        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]
        if eot:
            # take features from the eot embedding (eot_token is the highest number in each sequence)
            if projection:
                x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
            else:
                x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)]
        else:
            if projection:
                x = torch.matmul(x, self.text_projection)

        return x

    def forward(self, image, text):
        # NOTE: original CLIP return index 0 token, but we return all tokens
        image_features = self.encode_image(image, cls=True)
        text_features = self.encode_text(text, eot=True)

        # 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)
    return model.eval()

strhub/clip/model.py
VamosC/CLIP4STR
# Scene Text Recognition Model Hub
# Copyright 2022 Darwin Bautista
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import io
import glob
import lmdb
import unicodedata
from PIL import Image
from pathlib import Path, PurePath
from typing import Callable, Optional, Union
from torch.utils.data import Dataset, ConcatDataset
from pytorch_lightning.utilities import rank_zero_info

from strhub.data.utils import CharsetAdapter


def build_tree_dataset(root: Union[PurePath, str], *args, **kwargs):
    try:
        kwargs.pop('root')  # prevent 'root' from being passed via kwargs
    except KeyError:
        pass
    root = Path(root).absolute()
    rank_zero_info(f'dataset root:\t{root}')
    datasets = []
    for mdb in glob.glob(str(root / '**/data.mdb'), recursive=True):
        mdb = Path(mdb)
        ds_name = str(mdb.parent.relative_to(root))
        ds_root = str(mdb.parent.absolute())
        dataset = LmdbDataset(ds_root, *args, **kwargs)
        rank_zero_info(f'\tlmdb:\t{ds_name}\tnum samples: {len(dataset)}')
        datasets.append(dataset)
    return ConcatDataset(datasets)


class LmdbDataset(Dataset):
    """Dataset interface to an LMDB database.

    It supports both labelled and unlabelled datasets. For unlabelled datasets, the image index itself is returned
    as the label. Unicode characters are normalized by default. Case-sensitivity is inferred from the charset.
    Labels are transformed according to the charset.
    """

    def __init__(self, root: str, charset: str, max_label_len: int, min_image_dim: int = 0,
                 remove_whitespace: bool = True, normalize_unicode: bool = True,
                 unlabelled: bool = False, transform: Optional[Callable] = None):
        self._env = None
        self.root = root
        self.unlabelled = unlabelled
        self.transform = transform
        self.labels = []
        self.filtered_index_list = []
        self.num_samples = self._preprocess_labels(charset, remove_whitespace, normalize_unicode,
                                                   max_label_len, min_image_dim)
        # log.info('The number of samples is {}'.format(self.num_samples))

    def __del__(self):
        if self._env is not None:
            self._env.close()
            self._env = None

    def _create_env(self):
        return lmdb.open(self.root, max_readers=1, readonly=True, create=False,
                         readahead=False, meminit=False, lock=False)

    @property
    def env(self):
        if self._env is None:
            self._env = self._create_env()
        return self._env

    def _preprocess_labels(self, charset, remove_whitespace, normalize_unicode, max_label_len, min_image_dim):
        charset_adapter = CharsetAdapter(charset)
        with self._create_env() as env, env.begin() as txn:
            num_samples = int(txn.get('num-samples'.encode()))
            if self.unlabelled:
                return num_samples
            for index in range(num_samples):
                index += 1  # lmdb starts with 1
                label_key = f'label-{index:09d}'.encode()
                label = txn.get(label_key).decode()
                # Normally, whitespace is removed from the labels.
                if remove_whitespace:
                    label = ''.join(label.split())
                # Normalize unicode composites (if any) and convert to compatible ASCII characters
                if normalize_unicode:
                    label = unicodedata.normalize('NFKD', label).encode('ascii', 'ignore').decode()
                # Filter by length before removing unsupported characters. The original label might be too long.
                if len(label) > max_label_len:
                    continue
                label = charset_adapter(label)
                # We filter out samples which don't contain any supported characters
                if not label:
                    continue
                # Filter images that are too small.
                if min_image_dim > 0:
                    img_key = f'image-{index:09d}'.encode()
                    buf = io.BytesIO(txn.get(img_key))
                    w, h = Image.open(buf).size
                    if w < self.min_image_dim or h < self.min_image_dim:
                        continue
                self.labels.append(label)
                self.filtered_index_list.append(index)
        return len(self.labels)

    def __len__(self):
        return self.num_samples

    def __getitem__(self, index):
        if self.unlabelled:
            label = index
        else:
            label = self.labels[index]
            index = self.filtered_index_list[index]

        img_key = f'image-{index:09d}'.encode()
        with self.env.begin() as txn:
            imgbuf = txn.get(img_key)
        buf = io.BytesIO(imgbuf)
        img = Image.open(buf).convert('RGB')

        if self.transform is not None:
            img = self.transform(img)

        return img, label


if __name__ == "__main__":
    import os

    output_path = "."
    root = "str_dataset/test/CUTE80/data.mdb"
    charset = "0123456789abcdefghijklmnopqrstuvwxyz"
    max_label_len = 25
    dataset = LmdbDataset(root, charset, max_label_len)
    # 512 * 3 = 1536
    id1, id2 = 1539, 1568
    id1, id2 = 1576, 1706
    # image1, label1 = dataset[1539]
    # image2, label2 = dataset[1568]
    image1, label1 = dataset[id1]
    image2, label2 = dataset[id2]

    print(label1, label2)
    image1.save(os.path.join(output_path, "imiage_{}.jpg".format(id1)))
    image2.save(os.path.join(output_path, "imiage_{}.jpg".format(id2)))

strhub/data/dataset.py
VamosC/CLIP4STR
# Scene Text Recognition Model Hub
# Copyright 2022 Darwin Bautista
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import pytorch_lightning as pl
from pathlib import PurePath
from torch.utils.data import DataLoader
from torchvision import transforms as T
from typing import Optional, Callable, Sequence, Tuple
from pytorch_lightning.utilities import rank_zero_info

from .dataset import build_tree_dataset, LmdbDataset


class SceneTextDataModule(pl.LightningDataModule):
    # TEST_BENCHMARK_SUB = ('IIIT5k', 'SVT', 'IC13_857', 'IC15_1811', 'SVTP', 'CUTE80')
    TEST_BENCHMARK_SUB = ('IIIT5k', 'SVT', 'IC13_1015', 'IC15_1811', 'IC15_2077', 'SVTP', 'CUTE80', 'HOST', 'WOST')
    TEST_BENCHMARK = ('IIIT5k', 'SVT', 'IC13_1015', 'IC15_1811', 'IC15_2077', 'SVTP', 'CUTE80', 'HOST', 'WOST')
    # TEST_BENCHMARK_SUB = ('HOST',)
    # TEST_BENCHMARK = ('HOST',)
    TEST_NEW = ('ArT', 'COCOv1.4', 'Uber')
    TEST_ALL = tuple(set(TEST_BENCHMARK_SUB + TEST_BENCHMARK + TEST_NEW))

    def __init__(self, root_dir: str, train_dir: str, img_size: Sequence[int], max_label_length: int,
                 charset_train: str, charset_test: str, batch_size: int, num_workers: int, augment: bool,
                 remove_whitespace: bool = True, normalize_unicode: bool = True,
                 min_image_dim: int = 0, rotation: int = 0, collate_fn: Optional[Callable] = None,
                 output_url: str = None, openai_meanstd: bool = True,):
        super().__init__()
        self.root_dir = root_dir
        self.train_dir = train_dir
        self.img_size = tuple(img_size)
        self.max_label_length = max_label_length
        self.charset_train = charset_train
        self.charset_test = charset_test
        self.batch_size = batch_size
        self.num_workers = num_workers
        self.augment = augment
        self.remove_whitespace = remove_whitespace
        self.normalize_unicode = normalize_unicode
        self.min_image_dim = min_image_dim
        self.rotation = rotation
        self.collate_fn = collate_fn
        self._train_dataset = None
        self._val_dataset = None
        # https://github.com/mlfoundations/open_clip/blob/b4cf9269b0b11c0eea47cb16039369a46bd67449/src/open_clip/constants.py
        # OPENAI_DATASET_MEAN = (0.48145466, 0.4578275, 0.40821073)
        # OPENAI_DATASET_STD = (0.26862954, 0.26130258, 0.27577711)

        self.mean = (0.48145466, 0.4578275, 0.40821073) if openai_meanstd else 0.5
        self.std = (0.26862954, 0.26130258, 0.27577711) if openai_meanstd else 0.5
        rank_zero_info("[dataset] mean {}, std {}".format(self.mean, self.std))

    @staticmethod
    def get_transform(img_size: Tuple[int], augment: bool = False, rotation: int = 0, mean=0.5, std=0.5):
        transforms = []
        if augment:
            from .augment import rand_augment_transform
            transforms.append(rand_augment_transform())
        if rotation:
            transforms.append(lambda img: img.rotate(rotation, expand=True))
        transforms.extend([
            T.Resize(img_size, T.InterpolationMode.BICUBIC),
            T.ToTensor(),
            T.Normalize(mean, std)
        ])
        return T.Compose(transforms)

    @property
    def train_dataset(self):
        if self._train_dataset is None:
            transform = self.get_transform(self.img_size, self.augment, mean=self.mean, std=self.std)
            root = PurePath(self.root_dir, 'train', self.train_dir)
            self._train_dataset = build_tree_dataset(root, self.charset_train, self.max_label_length,
                                                     self.min_image_dim, self.remove_whitespace, self.normalize_unicode,
                                                     transform=transform)
            rank_zero_info('\tlmdb: The number of training samples is {}'.format(len(self._train_dataset)))
        return self._train_dataset

    @property
    def val_dataset(self):
        if self._val_dataset is None:
            transform = self.get_transform(self.img_size, mean=self.mean, std=self.std)
            root = PurePath(self.root_dir, 'val')
            self._val_dataset = build_tree_dataset(root, self.charset_test, self.max_label_length,
                                                   self.min_image_dim, self.remove_whitespace, self.normalize_unicode,
                                                   transform=transform)
            rank_zero_info('\tlmdb: The number of validation samples is {}'.format(len(self._val_dataset)))
        return self._val_dataset

    def train_dataloader(self):
        return DataLoader(self.train_dataset, batch_size=self.batch_size, shuffle=True,
                          num_workers=self.num_workers, persistent_workers=self.num_workers > 0,
                          pin_memory=True, collate_fn=self.collate_fn)

    def val_dataloader(self):
        return DataLoader(self.val_dataset, batch_size=self.batch_size,
                          num_workers=self.num_workers, persistent_workers=self.num_workers > 0,
                          pin_memory=True, collate_fn=self.collate_fn)

    def test_dataloaders(self, subset):
        transform = self.get_transform(self.img_size, rotation=self.rotation)
        root = PurePath(self.root_dir, 'test')
        datasets = {s: LmdbDataset(str(root / s), self.charset_test, self.max_label_length,
                                   self.min_image_dim, self.remove_whitespace, self.normalize_unicode,
                                   transform=transform) for s in subset}
        return {k: DataLoader(v, batch_size=self.batch_size, num_workers=self.num_workers,
                              pin_memory=True, collate_fn=self.collate_fn)
                for k, v in datasets.items()}

strhub/data/module.py
VamosC/CLIP4STR
# Scene Text Recognition Model Hub
# Copyright 2022 Darwin Bautista
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import math
import re
from abc import ABC, abstractmethod
from itertools import groupby
from typing import List, Optional, Tuple

import torch
from torch import Tensor
from torch.nn.utils.rnn import pad_sequence


class CharsetAdapter:
    """Transforms labels according to the target charset."""

    def __init__(self, target_charset) -> None:
        super().__init__()
        self.lowercase_only = target_charset == target_charset.lower()
        self.uppercase_only = target_charset == target_charset.upper()
        self.unsupported = f'[^{re.escape(target_charset)}]'

    def __call__(self, label):
        if self.lowercase_only:
            label = label.lower()
        elif self.uppercase_only:
            label = label.upper()
        # Remove unsupported characters
        label = re.sub(self.unsupported, '', label)
        return label


class BaseTokenizer(ABC):

    def __init__(self, charset: str, specials_first: tuple = (), specials_last: tuple = ()) -> None:
        self._itos = specials_first + tuple(charset) + specials_last
        self._stoi = {s: i for i, s in enumerate(self._itos)}

    def __len__(self):
        return len(self._itos)

    def _tok2ids(self, tokens: str) -> List[int]:
        return [self._stoi[s] for s in tokens]

    def _ids2tok(self, token_ids: List[int], join: bool = True) -> str:
        tokens = [self._itos[i] for i in token_ids]
        return ''.join(tokens) if join else tokens

    @abstractmethod
    def encode(self, labels: List[str], device: Optional[torch.device] = None) -> Tensor:
        """Encode a batch of labels to a representation suitable for the model.

        Args:
            labels: List of labels. Each can be of arbitrary length.
            device: Create tensor on this device.

        Returns:
            Batched tensor representation padded to the max label length. Shape: N, L
        """
        raise NotImplementedError

    @abstractmethod
    def _filter(self, probs: Tensor, ids: Tensor) -> Tuple[Tensor, List[int]]:
        """Internal method which performs the necessary filtering prior to decoding."""
        raise NotImplementedError

    def decode(self, token_dists: Tensor, raw: bool = False) -> Tuple[List[str], List[Tensor]]:
        """Decode a batch of token distributions.

        Args:
            token_dists: softmax probabilities over the token distribution. Shape: N, L, C
            raw: return unprocessed labels (will return list of list of strings)

        Returns:
            list of string labels (arbitrary length) and
            their corresponding sequence probabilities as a list of Tensors
        """
        batch_tokens = []
        batch_probs = []
        for dist in token_dists:
            probs, ids = dist.max(-1)  # greedy selection
            # probs, ids = beam_search_decoder(dist, 5)
            if not raw:
                probs, ids = self._filter(probs, ids)
            tokens = self._ids2tok(ids, not raw)
            batch_tokens.append(tokens)
            batch_probs.append(probs)
        return batch_tokens, batch_probs

    def decode_for_refine(self, token_dists: Tensor, raw: bool = False) -> Tuple[List[str], List[Tensor]]:
        raise NotImplementedError


class Tokenizer(BaseTokenizer):
    BOS = '[B]'
    EOS = '[E]'
    PAD = '[P]'

    def __init__(self, charset: str) -> None:
        specials_first = (self.EOS,)
        specials_last = (self.BOS, self.PAD)
        super().__init__(charset, specials_first, specials_last)
        self.eos_id, self.bos_id, self.pad_id = [self._stoi[s] for s in specials_first + specials_last]

    def encode(self, labels: List[str], device: Optional[torch.device] = None) -> Tensor:
        batch = [torch.as_tensor([self.bos_id] + self._tok2ids(y) + [self.eos_id], dtype=torch.long, device=device)
                 for y in labels]
        return pad_sequence(batch, batch_first=True, padding_value=self.pad_id)

    def _filter(self, probs: Tensor, ids: Tensor) -> Tuple[Tensor, List[int]]:
        ids = ids.tolist()
        try:
            eos_idx = ids.index(self.eos_id)
        except ValueError:
            eos_idx = len(ids)  # Nothing to truncate.
        # Truncate after EOS
        ids = ids[:eos_idx]
        probs = probs[:eos_idx + 1]  # but include prob. for EOS (if it exists)
        return probs, ids

    def decode_fast(self, token_dists: Tensor, raw: bool = False, charset_adapter = None) -> Tuple[List[str], List[Tensor]]:
        """fast version of decode"""
        bs_tokens = []
        bs_probs = []
        batch_probs, batch_ids = token_dists.max(-1)
        batch_ids = batch_ids.tolist()

        for i, ids in enumerate(batch_ids):
            tokens = self._ids2tok(ids, not raw)
            # Truncate after EOS
            tokens = tokens.partition(self.EOS)[0]
            # but include prob. for EOS (if it exists)
            bs_probs.append(batch_probs[i, :len(tokens) + 1])

            if charset_adapter is not None:
                tokens = charset_adapter(tokens)
            bs_tokens.append(tokens)

        return bs_tokens, bs_probs

    def decode_for_refine(self, token_dists: Tensor, raw: bool = False) -> Tuple[List[str], List[Tensor]]:
        """Decode a batch of token distributions.

        Args:
            token_dists: softmax probabilities over the token distribution. Shape: N, L, C
            raw: return unprocessed labels (will return list of list of strings)

        Returns:
            # list of string labels (arbitrary length) and
            # their corresponding sequence probabilities as a list of Tensors
            lisft of sequence probabilities
        """
        batch_dist = []
        for dist in token_dists:
            # greedy selection
            max_probs, ids = dist.max(-1)
            if not raw:
                ids = ids.tolist()
                try:
                    eos_idx = ids.index(self.eos_id)
                except ValueError:
                    eos_idx = len(ids)
                # trunucate after EOS
                ids = ids[:eos_idx]
                # do not include EOS
                trunc_dist = dist[:eos_idx, ...]
                # set probability of [EOS] to zero
                trunc_dist[:, self.eos_id] = torch.zeros_like(trunc_dist[:, self.eos_id])
                batch_dist.append(trunc_dist)
            else:
                batch_dist.append(dist)

        return batch_dist


class CTCTokenizer(BaseTokenizer):
    BLANK = '[B]'

    def __init__(self, charset: str) -> None:
        # BLANK uses index == 0 by default
        super().__init__(charset, specials_first=(self.BLANK,))
        self.blank_id = self._stoi[self.BLANK]

    def encode(self, labels: List[str], device: Optional[torch.device] = None) -> Tensor:
        # We use a padded representation since we don't want to use CUDNN's CTC implementation
        batch = [torch.as_tensor(self._tok2ids(y), dtype=torch.long, device=device) for y in labels]
        return pad_sequence(batch, batch_first=True, padding_value=self.blank_id)

    def _filter(self, probs: Tensor, ids: Tensor) -> Tuple[Tensor, List[int]]:
        # Best path decoding:
        ids = list(zip(*groupby(ids.tolist())))[0]  # Remove duplicate tokens
        ids = [x for x in ids if x != self.blank_id]  # Remove BLANKs
        # `probs` is just pass-through since all positions are considered part of the path
        return probs, ids


def beam_search_decoder(data, k):
    """a simple beam search algorithm"""
    device = data.device
    data_numpy = data.cpu().numpy()

    sequences = [[list(), 0.0]]
    # walk over each step in sequence
    for row in data_numpy:
        all_candidates = list()
        # expand each current candidate
        for i in range(len(sequences)):
            seq, score = sequences[i]
            for j in range(len(row)):
                candidate = [seq + [j], score - math.log(row[j])]
                all_candidates.append(candidate)
        # order all candidates by score
        ordered = sorted(all_candidates, key=lambda tup:tup[1])
        # select k best
        sequences = ordered[:k]

    ids = torch.tensor(sequences[0][0], dtype=torch.long, device=data.device)
    probs = data[torch.arange(len(ids), dtype=torch.long, device=device), ids]
    return probs, ids


if __name__ == "__main__":
    _char_94_full = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~"
    t = Tokenizer(_char_94_full)
    for x in _char_94_full:
        print(t._tok2ids(x))
    print(t._itos)
    print(t._stoi)
    print(t.bos_id)
    print(t.eos_id)
    print(t.pad_id)

# {'[E]': 0, '0': 1, '1': 2, '2': 3, '3': 4, '4': 5, '5': 6, '6': 7, '7': 8, '8': 9, '9': 10, 
# 'a': 11, 'b': 12, 'c': 13, 'd': 14, 'e': 15, 'f': 16, 'g': 17, 'h': 18, 'i': 19, 'j': 20, 
# 'k': 21, 'l': 22, 'm': 23, 'n': 24, 'o': 25, 'p': 26, 'q': 27, 'r': 28, 's': 29, 't': 30, 
# 'u': 31, 'v': 32, 'w': 33, 'x': 34, 'y': 35, 'z': 36, 
# 'A': 37, 'B': 38, 'C': 39, 'D': 40, 'E': 41, 'F': 42, 'G': 43, 'H': 44, 'I': 45, 'J': 46, 
# 'K': 47, 'L': 48, 'M': 49, 'N': 50, 'O': 51, 'P': 52, 'Q': 53, 'R': 54, 'S': 55, 'T': 56, 
# 'U': 57, 'V': 58, 'W': 59, 'X': 60, 'Y': 61, 'Z': 62, 
# '!': 63, '"': 64, '#': 65, '$': 66, '%': 67, '&': 68, "'": 69, '(': 70, ')': 71, '*': 72, 
# '+': 73, ',': 74, '-': 75, '.': 76, '/': 77, ':': 78, ';': 79, '<': 80, '=': 81, '>': 82, 
# '?': 83, '@': 84, '[': 85, '\\': 86, ']': 87, '^': 88, '_': 89, '`': 90, '{': 91, '|': 92, 
# '}': 93, '~': 94, '[B]': 95, '[P]': 96}
# bos_id 95
# eos_id 0
# pad_id 96
strhub/data/utils.py
VamosC/CLIP4STR
# coding=utf-8
import os
import logging
import torch.distributed as dist
from pytorch_lightning.utilities import rank_zero_only


def is_dist_avail_and_initialized():
    if not dist.is_available():
        return False
    if not dist.is_initialized():
        return False
    return True


def get_world_size():
    if not is_dist_avail_and_initialized():
        return 1
    return dist.get_world_size()


def get_rank():
    if not is_dist_avail_and_initialized():
        return 0
    return dist.get_rank()


def is_main_process():
    return get_rank() == 0


def setup_logger(log_file=None, include_host=False):
    """
    Ref: https://github.com/mlfoundations/open_clip/blob/db338b0bb36c15ae12fcd37e86120414903df1ef/src/training/logger.
    """
    if include_host:
        import socket
        hostname = socket.gethostname()
        formatter = logging.Formatter(
            f'%(asctime)s |  {hostname} | %(levelname)s | %(message)s', datefmt='%Y-%m-%d,%H:%M:%S')
    else:
        formatter = logging.Formatter('%(asctime)s | %(levelname)s | %(message)s', datefmt='%Y-%m-%d,%H:%M:%S')
    level = logging.INFO if is_main_process() else logging.WARN

    logging.root.setLevel(level)
    loggers = [logging.getLogger(name) for name in logging.root.manager.loggerDict]
    for logger in loggers:
        logger.setLevel(level)

    stream_handler = logging.StreamHandler()
    stream_handler.setFormatter(formatter)
    logging.root.addHandler(stream_handler)

    if log_file:
        file_handler = logging.FileHandler(filename=log_file)
        file_handler.setFormatter(formatter)
        logging.root.addHandler(file_handler)


@rank_zero_only
def copy_remote(local_dir, remote_dir=None):
    if remote_dir is not None:
        print('=' * 50)
        os.chdir(local_dir)
        os.chdir('..')
        base_name = os.path.basename(local_dir)
        tar_filename = '{}.tar'.format(base_name)
        if os.path.exists(os.path.join(os.getcwd(), tar_filename)):
            print('remove existing tarfile and create a new tarfile')
            os.system("rm {}".format(tar_filename))

        os.system("tar -zcvf {} {}".format(tar_filename, base_name))
        os.system("rsync -rvP {} {}".format(tar_filename, remote_dir))
        os.system("ls -lah {}".format(remote_dir))
        print("Copy success!")

strhub/dist_utils.py
VamosC/CLIP4STR
r"""
Fang, Shancheng, Hongtao, Xie, Yuxin, Wang, Zhendong, Mao, and Yongdong, Zhang.
"Read Like Humans: Autonomous, Bidirectional and Iterative Language Modeling for Scene Text Recognition." .
In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) (pp. 7098-7107).2021.

https://arxiv.org/abs/2103.06495

All source files, except `system.py`, are based on the implementation listed below,
and hence are released under the license of the original.

Source: https://github.com/FangShancheng/ABINet
License: 2-clause BSD License (see included LICENSE file)
"""

strhub/models/abinet/__init__.py
VamosC/CLIP4STR
import torch
import torch.nn as nn

from .transformer import PositionalEncoding


class Attention(nn.Module):
    def __init__(self, in_channels=512, max_length=25, n_feature=256):
        super().__init__()
        self.max_length = max_length

        self.f0_embedding = nn.Embedding(max_length, in_channels)
        self.w0 = nn.Linear(max_length, n_feature)
        self.wv = nn.Linear(in_channels, in_channels)
        self.we = nn.Linear(in_channels, max_length)

        self.active = nn.Tanh()
        self.softmax = nn.Softmax(dim=2)

    def forward(self, enc_output):
        enc_output = enc_output.permute(0, 2, 3, 1).flatten(1, 2)
        reading_order = torch.arange(self.max_length, dtype=torch.long, device=enc_output.device)
        reading_order = reading_order.unsqueeze(0).expand(enc_output.size(0), -1)  # (S,) -> (B, S)
        reading_order_embed = self.f0_embedding(reading_order)  # b,25,512

        t = self.w0(reading_order_embed.permute(0, 2, 1))  # b,512,256
        t = self.active(t.permute(0, 2, 1) + self.wv(enc_output))  # b,256,512

        attn = self.we(t)  # b,256,25
        attn = self.softmax(attn.permute(0, 2, 1))  # b,25,256
        g_output = torch.bmm(attn, enc_output)  # b,25,512
        return g_output, attn.view(*attn.shape[:2], 8, 32)


def encoder_layer(in_c, out_c, k=3, s=2, p=1):
    return nn.Sequential(nn.Conv2d(in_c, out_c, k, s, p),
                         nn.BatchNorm2d(out_c),
                         nn.ReLU(True))


def decoder_layer(in_c, out_c, k=3, s=1, p=1, mode='nearest', scale_factor=None, size=None):
    align_corners = None if mode == 'nearest' else True
    return nn.Sequential(nn.Upsample(size=size, scale_factor=scale_factor,
                                     mode=mode, align_corners=align_corners),
                         nn.Conv2d(in_c, out_c, k, s, p),
                         nn.BatchNorm2d(out_c),
                         nn.ReLU(True))


class PositionAttention(nn.Module):
    def __init__(self, max_length, in_channels=512, num_channels=64,
                 h=8, w=32, mode='nearest', **kwargs):
        super().__init__()
        self.max_length = max_length
        self.k_encoder = nn.Sequential(
            encoder_layer(in_channels, num_channels, s=(1, 2)),
            encoder_layer(num_channels, num_channels, s=(2, 2)),
            encoder_layer(num_channels, num_channels, s=(2, 2)),
            encoder_layer(num_channels, num_channels, s=(2, 2))
        )
        self.k_decoder = nn.Sequential(
            decoder_layer(num_channels, num_channels, scale_factor=2, mode=mode),
            decoder_layer(num_channels, num_channels, scale_factor=2, mode=mode),
            decoder_layer(num_channels, num_channels, scale_factor=2, mode=mode),
            decoder_layer(num_channels, in_channels, size=(h, w), mode=mode)
        )

        self.pos_encoder = PositionalEncoding(in_channels, dropout=0., max_len=max_length)
        self.project = nn.Linear(in_channels, in_channels)

    def forward(self, x):
        N, E, H, W = x.size()
        k, v = x, x  # (N, E, H, W)

        # calculate key vector
        features = []
        for i in range(0, len(self.k_encoder)):
            k = self.k_encoder[i](k)
            features.append(k)
        for i in range(0, len(self.k_decoder) - 1):
            k = self.k_decoder[i](k)
            k = k + features[len(self.k_decoder) - 2 - i]
        k = self.k_decoder[-1](k)

        # calculate query vector
        # TODO q=f(q,k)
        zeros = x.new_zeros((self.max_length, N, E))  # (T, N, E)
        q = self.pos_encoder(zeros)     # (T, N, E)
        q = q.permute(1, 0, 2)          # (N, T, E)
        q = self.project(q)             # (N, T, E)

        # calculate attention
        attn_scores = torch.bmm(q, k.flatten(2, 3))  # (N, T, (H*W))
        attn_scores = attn_scores / (E ** 0.5)
        attn_scores = torch.softmax(attn_scores, dim=-1)

        v = v.permute(0, 2, 3, 1).view(N, -1, E)  # (N, (H*W), E)
        attn_vecs = torch.bmm(attn_scores, v)  # (N, T, E)

        return attn_vecs, attn_scores.view(N, -1, H, W)


class PositionAttentionForSequence(nn.Module):
    def __init__(self, max_length, in_channels=512, num_channels=64,
                 h=14, w=14, mode='nearest', **kwargs):
        """this module deals with sequence input"""
        super().__init__()
        self.max_length = max_length
        self.k_encoder = nn.Sequential(
            encoder_layer(in_channels, num_channels, s=(1, 1)),
            encoder_layer(num_channels, num_channels, s=(1, 1)),
            encoder_layer(num_channels, num_channels, s=(1, 1)),
            encoder_layer(num_channels, num_channels, s=(1, 1))
        )
        self.k_decoder = nn.Sequential(
            decoder_layer(num_channels, num_channels, scale_factor=1, mode=mode),
            decoder_layer(num_channels, num_channels, scale_factor=1, mode=mode),
            decoder_layer(num_channels, num_channels, scale_factor=1, mode=mode),
            decoder_layer(num_channels, in_channels, size=(h, w), mode=mode)
        )

        self.pos_encoder = PositionalEncoding(in_channels, dropout=0., max_len=max_length)
        self.project = nn.Linear(in_channels, in_channels)

    def forward(self, x):
        """
        Args:
            x [N, L, E]
        """
        N, L, E = x.size()
        v = x                                           # (N, L, E)

        H, W = int((L - 1) ** 0.5), int((L - 1) ** 0.5)
        cls_token = x[:, 0:1, :]
        # [N, C, H, W]
        k = x[:, 1:, :].reshape(N, H, W, -1).permute(0, 3, 1, 2).contiguous()

        # calculate key vector
        features = []
        for i in range(0, len(self.k_encoder)):
            k = self.k_encoder[i](k)
            features.append(k)
        for i in range(0, len(self.k_decoder) - 1):
            k = self.k_decoder[i](k)
            k = k + features[len(self.k_decoder) - 2 - i]
        k = self.k_decoder[-1](k)
        k = k.flatten(2, 3).transpose(-1, -2)           # (N, (H*W), E)
        k = torch.cat([cls_token, k], dim=1)            # (N, L, E)

        # calculate query vector
        zeros = x.new_zeros((self.max_length, N, E))    # (T, N, E)
        q = self.pos_encoder(zeros)                     # (T, N, E)
        q = q.permute(1, 0, 2)                          # (N, T, E)
        q = self.project(q)                             # (N, T, E)

        # calculate attention
        attn_scores = torch.bmm(q, k.transpose(-1, -2)) # (N, T, L)
        attn_scores = attn_scores / (E ** 0.5)
        attn_scores = torch.softmax(attn_scores, dim=-1)

        attn_vecs = torch.bmm(attn_scores, v)           # (N, T, E)

        return attn_vecs, attn_scores

strhub/models/abinet/attention.py
VamosC/CLIP4STR
import torch.nn as nn
from torch.nn import TransformerEncoderLayer, TransformerEncoder

from .resnet import resnet45
from .transformer import PositionalEncoding


class ResTranformer(nn.Module):
    def __init__(self, d_model=512, nhead=8, d_inner=2048, dropout=0.1, activation='relu', backbone_ln=2):
        super().__init__()
        self.resnet = resnet45()
        self.pos_encoder = PositionalEncoding(d_model, max_len=8 * 32)
        encoder_layer = TransformerEncoderLayer(d_model=d_model, nhead=nhead,
                                                dim_feedforward=d_inner, dropout=dropout, activation=activation)
        self.transformer = TransformerEncoder(encoder_layer, backbone_ln)

    def forward(self, images):
        feature = self.resnet(images)
        n, c, h, w = feature.shape
        feature = feature.view(n, c, -1).permute(2, 0, 1)
        feature = self.pos_encoder(feature)
        feature = self.transformer(feature)
        feature = feature.permute(1, 2, 0).view(n, c, h, w)
        return feature

strhub/models/abinet/backbone.py
VamosC/CLIP4STR
import torch
import torch.nn as nn


class Model(nn.Module):

    def __init__(self, dataset_max_length: int, null_label: int):
        super().__init__()
        self.max_length = dataset_max_length + 1  # additional stop token
        self.null_label = null_label

    def _get_length(self, logit, dim=-1):
        """ Greed decoder to obtain length from logit"""
        out = (logit.argmax(dim=-1) == self.null_label)
        abn = out.any(dim)
        out = ((out.cumsum(dim) == 1) & out).max(dim)[1]
        out = out + 1  # additional end token
        out = torch.where(abn, out, out.new_tensor(logit.shape[1], device=out.device))
        return out

    @staticmethod
    def _get_padding_mask(length, max_length):
        length = length.unsqueeze(-1)
        grid = torch.arange(0, max_length, device=length.device).unsqueeze(0)
        return grid >= length

    @staticmethod
    def _get_location_mask(sz, device=None):
        mask = torch.eye(sz, device=device)
        mask = mask.float().masked_fill(mask == 1, float('-inf'))
        return mask

strhub/models/abinet/model.py
VamosC/CLIP4STR
import torch
from torch import nn

from .model_alignment import BaseAlignment
from .model_language import BCNLanguage
from .model_vision import BaseVision


class ABINetIterModel(nn.Module):
    def __init__(self, dataset_max_length, null_label, num_classes, iter_size=1,
                 d_model=512, nhead=8, d_inner=2048, dropout=0.1, activation='relu',
                 v_loss_weight=1., v_attention='position', v_attention_mode='nearest',
                 v_backbone='transformer', v_num_layers=2,
                 l_loss_weight=1., l_num_layers=4, l_detach=True, l_use_self_attn=False,
                 a_loss_weight=1.):
        super().__init__()
        self.iter_size = iter_size
        self.vision = BaseVision(dataset_max_length, null_label, num_classes, v_attention, v_attention_mode,
                                 v_loss_weight, d_model, nhead, d_inner, dropout, activation, v_backbone, v_num_layers)
        self.language = BCNLanguage(dataset_max_length, null_label, num_classes, d_model, nhead, d_inner, dropout,
                                    activation, l_num_layers, l_detach, l_use_self_attn, l_loss_weight)
        self.alignment = BaseAlignment(dataset_max_length, null_label, num_classes, d_model, a_loss_weight)

    def forward(self, images):
        v_res = self.vision(images)
        a_res = v_res
        all_l_res, all_a_res = [], []
        for _ in range(self.iter_size):
            tokens = torch.softmax(a_res['logits'], dim=-1)
            lengths = a_res['pt_lengths']
            lengths.clamp_(2, self.language.max_length)  # TODO:move to langauge model
            l_res = self.language(tokens, lengths)
            all_l_res.append(l_res)
            a_res = self.alignment(l_res['feature'], v_res['feature'])
            all_a_res.append(a_res)
        if self.training:
            return all_a_res, all_l_res, v_res
        else:
            return a_res, all_l_res[-1], v_res

strhub/models/abinet/model_abinet_iter.py
VamosC/CLIP4STR
import torch
import torch.nn as nn

from .model import Model


class BaseAlignment(Model):
    def __init__(self, dataset_max_length, null_label, num_classes, d_model=512, loss_weight=1.0):
        super().__init__(dataset_max_length, null_label)
        self.loss_weight = loss_weight
        self.w_att = nn.Linear(2 * d_model, d_model)
        self.cls = nn.Linear(d_model, num_classes)

    def forward(self, l_feature, v_feature):
        """
        Args:
            l_feature: (N, T, E) where T is length, N is batch size and d is dim of model
            v_feature: (N, T, E) shape the same as l_feature 
        """
        f = torch.cat((l_feature, v_feature), dim=2)
        f_att = torch.sigmoid(self.w_att(f))
        output = f_att * v_feature + (1 - f_att) * l_feature

        logits = self.cls(output)  # (N, T, C)
        pt_lengths = self._get_length(logits)

        return {'logits': logits, 'pt_lengths': pt_lengths, 'loss_weight': self.loss_weight,
                'name': 'alignment'}

strhub/models/abinet/model_alignment.py
VamosC/CLIP4STR
import torch.nn as nn
from torch.nn import TransformerDecoder

from .model import Model
from .transformer import PositionalEncoding, TransformerDecoderLayer


class BCNLanguage(Model):
    def __init__(self, dataset_max_length, null_label, num_classes, d_model=512, nhead=8, d_inner=2048, dropout=0.1,
                 activation='relu', num_layers=4, detach=True, use_self_attn=False, loss_weight=1.0,
                 global_debug=False):
        super().__init__(dataset_max_length, null_label)
        self.detach = detach
        self.loss_weight = loss_weight
        self.proj = nn.Linear(num_classes, d_model, False)
        self.token_encoder = PositionalEncoding(d_model, max_len=self.max_length)
        self.pos_encoder = PositionalEncoding(d_model, dropout=0, max_len=self.max_length)
        decoder_layer = TransformerDecoderLayer(d_model, nhead, d_inner, dropout,
                                                activation, self_attn=use_self_attn, debug=global_debug)
        self.model = TransformerDecoder(decoder_layer, num_layers)
        self.cls = nn.Linear(d_model, num_classes)

    def forward(self, tokens, lengths):
        """
        Args:
            tokens: (N, T, C) where T is length, N is batch size and C is classes number
            lengths: (N,)
        """
        if self.detach:
            tokens = tokens.detach()
        embed = self.proj(tokens)  # (N, T, E)
        embed = embed.permute(1, 0, 2)  # (T, N, E)
        embed = self.token_encoder(embed)  # (T, N, E)
        padding_mask = self._get_padding_mask(lengths, self.max_length)

        zeros = embed.new_zeros(*embed.shape)
        qeury = self.pos_encoder(zeros)
        location_mask = self._get_location_mask(self.max_length, tokens.device)
        output = self.model(qeury, embed,
                            tgt_key_padding_mask=padding_mask,
                            memory_mask=location_mask,
                            memory_key_padding_mask=padding_mask)  # (T, N, E)
        output = output.permute(1, 0, 2)  # (N, T, E)

        logits = self.cls(output)  # (N, T, C)
        pt_lengths = self._get_length(logits)

        res = {'feature': output, 'logits': logits, 'pt_lengths': pt_lengths,
               'loss_weight': self.loss_weight, 'name': 'language'}
        return res

strhub/models/abinet/model_language.py
VamosC/CLIP4STR
from torch import nn

from .attention import PositionAttention, Attention
from .backbone import ResTranformer
from .model import Model
from .resnet import resnet45


class BaseVision(Model):
    def __init__(self, dataset_max_length, null_label, num_classes,
                 attention='position', attention_mode='nearest', loss_weight=1.0,
                 d_model=512, nhead=8, d_inner=2048, dropout=0.1, activation='relu',
                 backbone='transformer', backbone_ln=2):
        super().__init__(dataset_max_length, null_label)
        self.loss_weight = loss_weight
        self.out_channels = d_model

        if backbone == 'transformer':
            self.backbone = ResTranformer(d_model, nhead, d_inner, dropout, activation, backbone_ln)
        else:
            self.backbone = resnet45()

        if attention == 'position':
            self.attention = PositionAttention(
                max_length=self.max_length,
                mode=attention_mode
            )
        elif attention == 'attention':
            self.attention = Attention(
                max_length=self.max_length,
                n_feature=8 * 32,
            )
        else:
            raise ValueError(f'invalid attention: {attention}')

        self.cls = nn.Linear(self.out_channels, num_classes)

    def forward(self, images):
        features = self.backbone(images)  # (N, E, H, W)
        attn_vecs, attn_scores = self.attention(features)  # (N, T, E), (N, T, H, W)
        logits = self.cls(attn_vecs)  # (N, T, C)
        pt_lengths = self._get_length(logits)

        return {'feature': attn_vecs, 'logits': logits, 'pt_lengths': pt_lengths,
                'attn_scores': attn_scores, 'loss_weight': self.loss_weight, 'name': 'vision'}

strhub/models/abinet/model_vision.py
VamosC/CLIP4STR
import math
from typing import Optional, Callable

import torch.nn as nn
from torchvision.models import resnet


class BasicBlock(resnet.BasicBlock):

    def __init__(self, inplanes: int, planes: int, stride: int = 1, downsample: Optional[nn.Module] = None,
                 groups: int = 1, base_width: int = 64, dilation: int = 1,
                 norm_layer: Optional[Callable[..., nn.Module]] = None) -> None:
        super().__init__(inplanes, planes, stride, downsample, groups, base_width, dilation, norm_layer)
        self.conv1 = resnet.conv1x1(inplanes, planes)
        self.conv2 = resnet.conv3x3(planes, planes, stride)


class ResNet(nn.Module):

    def __init__(self, block, layers):
        super().__init__()
        self.inplanes = 32
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1,
                               bias=False)
        self.bn1 = nn.BatchNorm2d(32)
        self.relu = nn.ReLU(inplace=True)

        self.layer1 = self._make_layer(block, 32, layers[0], stride=2)
        self.layer2 = self._make_layer(block, 64, layers[1], stride=1)
        self.layer3 = self._make_layer(block, 128, layers[2], stride=2)
        self.layer4 = self._make_layer(block, 256, layers[3], stride=1)
        self.layer5 = self._make_layer(block, 512, layers[4], stride=1)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()

    def _make_layer(self, block, planes, blocks, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes, planes * block.expansion,
                          kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(planes * block.expansion),
            )

        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample))
        self.inplanes = planes * block.expansion
        for i in range(1, blocks):
            layers.append(block(self.inplanes, planes))

        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.layer5(x)
        return x


def resnet45():
    return ResNet(BasicBlock, [3, 4, 6, 6, 3])

strhub/models/abinet/resnet.py
VamosC/CLIP4STR
# Scene Text Recognition Model Hub
# Copyright 2022 Darwin Bautista
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import logging
import math
from typing import Any, Tuple, List, Optional

import torch
import torch.nn.functional as F
from torch import Tensor, nn
from torch.optim import AdamW
from torch.optim.lr_scheduler import OneCycleLR

from pytorch_lightning.utilities.types import STEP_OUTPUT
from timm.optim.optim_factory import param_groups_weight_decay

from strhub.models.base import CrossEntropySystem
from strhub.models.utils import init_weights
from .model_abinet_iter import ABINetIterModel as Model

log = logging.getLogger(__name__)


class ABINet(CrossEntropySystem):

    def __init__(self, charset_train: str, charset_test: str, max_label_length: int,
                 batch_size: int, lr: float, warmup_pct: float, weight_decay: float,
                 iter_size: int, d_model: int, nhead: int, d_inner: int, dropout: float, activation: str,
                 v_loss_weight: float, v_attention: str, v_attention_mode: str, v_backbone: str, v_num_layers: int,
                 l_loss_weight: float, l_num_layers: int, l_detach: bool, l_use_self_attn: bool,
                 l_lr: float, a_loss_weight: float, lm_only: bool = False, **kwargs) -> None:
        super().__init__(charset_train, charset_test, batch_size, lr, warmup_pct, weight_decay)
        self.scheduler = None
        self.save_hyperparameters()
        self.max_label_length = max_label_length
        self.num_classes = len(self.tokenizer) - 2  # We don't predict <bos> nor <pad>
        self.model = Model(max_label_length, self.eos_id, self.num_classes, iter_size, d_model, nhead, d_inner,
                           dropout, activation, v_loss_weight, v_attention, v_attention_mode, v_backbone, v_num_layers,
                           l_loss_weight, l_num_layers, l_detach, l_use_self_attn, a_loss_weight)
        self.model.apply(init_weights)
        # FIXME: doesn't support resumption from checkpoint yet
        self._reset_alignment = True
        self._reset_optimizers = True
        self.l_lr = l_lr
        self.lm_only = lm_only
        # Train LM only. Freeze other submodels.
        if lm_only:
            self.l_lr = lr  # for tuning
            self.model.vision.requires_grad_(False)
            self.model.alignment.requires_grad_(False)

    @property
    def _pretraining(self):
        # In the original work, VM was pretrained for 8 epochs while full model was trained for an additional 10 epochs.
        total_steps = self.trainer.estimated_stepping_batches * self.trainer.accumulate_grad_batches
        return self.global_step < (8 / (8 + 10)) * total_steps

    @torch.jit.ignore
    def no_weight_decay(self):
        return {'model.language.proj.weight'}

    def _add_weight_decay(self, model: nn.Module, skip_list=()):
        if self.weight_decay:
            return param_groups_weight_decay(model, self.weight_decay, skip_list)
        else:
            return [{'params': model.parameters()}]

    def configure_optimizers(self):
        agb = self.trainer.accumulate_grad_batches
        # Linear scaling so that the effective learning rate is constant regardless of the number of GPUs used with DDP.
        lr_scale = agb * math.sqrt(self.trainer.num_devices) * self.batch_size / 256.
        lr = lr_scale * self.lr
        l_lr = lr_scale * self.l_lr
        params = []
        params.extend(self._add_weight_decay(self.model.vision))
        params.extend(self._add_weight_decay(self.model.alignment))
        # We use a different learning rate for the LM.
        for p in self._add_weight_decay(self.model.language, ('proj.weight',)):
            p['lr'] = l_lr
            params.append(p)
        max_lr = [p.get('lr', lr) for p in params]
        optim = AdamW(params, lr)
        self.scheduler = OneCycleLR(optim, max_lr, self.trainer.estimated_stepping_batches,
                                    pct_start=self.warmup_pct, cycle_momentum=False)
        return {'optimizer': optim, 'lr_scheduler': {'scheduler': self.scheduler, 'interval': 'step'}}

    def forward(self, images: Tensor, max_length: Optional[int] = None) -> Tensor:
        max_length = self.max_label_length if max_length is None else min(max_length, self.max_label_length)
        logits = self.model.forward(images)[0]['logits']
        return logits[:, :max_length + 1]  # truncate

    def calc_loss(self, targets, *res_lists) -> Tensor:
        total_loss = 0
        for res_list in res_lists:
            loss = 0
            if isinstance(res_list, dict):
                res_list = [res_list]
            for res in res_list:
                logits = res['logits'].flatten(end_dim=1)
                loss += F.cross_entropy(logits, targets.flatten(), ignore_index=self.pad_id)
            loss /= len(res_list)
            self.log('loss_' + res_list[0]['name'], loss)
            total_loss += res_list[0]['loss_weight'] * loss
        return total_loss

    def on_train_batch_start(self, batch: Any, batch_idx: int) -> None:
        if not self._pretraining and self._reset_optimizers:
            log.info('Pretraining ends. Updating base LRs.')
            self._reset_optimizers = False
            # Make base_lr the same for all groups
            base_lr = self.scheduler.base_lrs[0]  # base_lr of group 0 - VM
            self.scheduler.base_lrs = [base_lr] * len(self.scheduler.base_lrs)

    def _prepare_inputs_and_targets(self, labels):
        # Use dummy label to ensure sequence length is constant.
        dummy = ['0' * self.max_label_length]
        targets = self.tokenizer.encode(dummy + list(labels), self.device)[1:]
        targets = targets[:, 1:]  # remove <bos>. Unused here.
        # Inputs are padded with eos_id
        inputs = torch.where(targets == self.pad_id, self.eos_id, targets)
        inputs = F.one_hot(inputs, self.num_classes).float()
        lengths = torch.as_tensor(list(map(len, labels)), device=self.device) + 1  # +1 for eos
        return inputs, lengths, targets

    def training_step(self, batch, batch_idx) -> STEP_OUTPUT:
        images, labels = batch
        inputs, lengths, targets = self._prepare_inputs_and_targets(labels)
        if self.lm_only:
            l_res = self.model.language(inputs, lengths)
            loss = self.calc_loss(targets, l_res)
        # Pretrain submodels independently first
        elif self._pretraining:
            # Vision
            v_res = self.model.vision(images)
            # Language
            l_res = self.model.language(inputs, lengths)
            # We also train the alignment model to 'satisfy' DDP requirements (all parameters should be used).
            # We'll reset its parameters prior to joint training.
            a_res = self.model.alignment(l_res['feature'].detach(), v_res['feature'].detach())
            loss = self.calc_loss(targets, v_res, l_res, a_res)
        else:
            # Reset alignment model's parameters once prior to full model training.
            if self._reset_alignment:
                log.info('Pretraining ends. Resetting alignment model.')
                self._reset_alignment = False
                self.model.alignment.apply(init_weights)
            all_a_res, all_l_res, v_res = self.model.forward(images)
            loss = self.calc_loss(targets, v_res, all_l_res, all_a_res)
        self.log('loss', loss)
        return loss

    def forward_logits_loss(self, images: Tensor, labels: List[str]) -> Tuple[Tensor, Tensor, int]:
        if self.lm_only:
            inputs, lengths, targets = self._prepare_inputs_and_targets(labels)
            l_res = self.model.language(inputs, lengths)
            loss = self.calc_loss(targets, l_res)
            loss_numel = (targets != self.pad_id).sum()
            return l_res['logits'], loss, loss_numel
        else:
            return super().forward_logits_loss(images, labels)

strhub/models/abinet/system.py
VamosC/CLIP4STR
import math

import torch
import torch.nn.functional as F
from torch import nn
from torch.nn.modules.transformer import _get_activation_fn


class TransformerDecoderLayer(nn.Module):
    r"""TransformerDecoderLayer is made up of self-attn, multi-head-attn and feedforward network.
    This standard decoder layer is based on the paper "Attention Is All You Need".
    Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez,
    Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in
    Neural Information Processing Systems, pages 6000-6010. Users may modify or implement
    in a different way during application.

    Args:
        d_model: the number of expected features in the input (required).
        nhead: the number of heads in the multiheadattention models (required).
        dim_feedforward: the dimension of the feedforward network model (default=2048).
        dropout: the dropout value (default=0.1).
        activation: the activation function of intermediate layer, relu or gelu (default=relu).

    Examples::
        >>> decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8)
        >>> memory = torch.rand(10, 32, 512)
        >>> tgt = torch.rand(20, 32, 512)
        >>> out = decoder_layer(tgt, memory)
    """

    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1,
                 activation="relu", self_attn=True, siamese=False, debug=False):
        super().__init__()
        self.has_self_attn, self.siamese = self_attn, siamese
        self.debug = debug
        if self.has_self_attn:
            self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
            self.norm1 = nn.LayerNorm(d_model)
            self.dropout1 = nn.Dropout(dropout)
        self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
        # Implementation of Feedforward model
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)

        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout2 = nn.Dropout(dropout)
        self.dropout3 = nn.Dropout(dropout)
        if self.siamese:
            self.multihead_attn2 = nn.MultiheadAttention(d_model, nhead, dropout=dropout)

        self.activation = _get_activation_fn(activation)

    def __setstate__(self, state):
        if 'activation' not in state:
            state['activation'] = F.relu
        super().__setstate__(state)

    def forward(self, tgt, memory, tgt_mask=None, memory_mask=None,
                tgt_key_padding_mask=None, memory_key_padding_mask=None,
                memory2=None, memory_mask2=None, memory_key_padding_mask2=None):
        # type: (Tensor, Tensor, Optional[Tensor], Optional[Tensor], Optional[Tensor], Optional[Tensor]) -> Tensor
        r"""Pass the inputs (and mask) through the decoder layer.

        Args:
            tgt: the sequence to the decoder layer (required).
            memory: the sequence from the last layer of the encoder (required).
            tgt_mask: the mask for the tgt sequence (optional).
            memory_mask: the mask for the memory sequence (optional).
            tgt_key_padding_mask: the mask for the tgt keys per batch (optional).
            memory_key_padding_mask: the mask for the memory keys per batch (optional).

        Shape:
            see the docs in Transformer class.
        """
        if self.has_self_attn:
            tgt2, attn = self.self_attn(tgt, tgt, tgt, attn_mask=tgt_mask,
                                        key_padding_mask=tgt_key_padding_mask)
            tgt = tgt + self.dropout1(tgt2)
            tgt = self.norm1(tgt)
            if self.debug: self.attn = attn
        tgt2, attn2 = self.multihead_attn(tgt, memory, memory, attn_mask=memory_mask,
                                          key_padding_mask=memory_key_padding_mask)
        if self.debug: self.attn2 = attn2

        if self.siamese:
            tgt3, attn3 = self.multihead_attn2(tgt, memory2, memory2, attn_mask=memory_mask2,
                                               key_padding_mask=memory_key_padding_mask2)
            tgt = tgt + self.dropout2(tgt3)
            if self.debug: self.attn3 = attn3

        tgt = tgt + self.dropout2(tgt2)
        tgt = self.norm2(tgt)
        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
        tgt = tgt + self.dropout3(tgt2)
        tgt = self.norm3(tgt)

        return tgt


class PositionalEncoding(nn.Module):
    r"""Inject some information about the relative or absolute position of the tokens
        in the sequence. The positional encodings have the same dimension as
        the embeddings, so that the two can be summed. Here, we use sine and cosine
        functions of different frequencies.
    .. math::
        \text{PosEncoder}(pos, 2i) = sin(pos/10000^(2i/d_model))
        \text{PosEncoder}(pos, 2i+1) = cos(pos/10000^(2i/d_model))
        \text{where pos is the word position and i is the embed idx)
    Args:
        d_model: the embed dim (required).
        dropout: the dropout value (default=0.1).
        max_len: the max. length of the incoming sequence (default=5000).
    Examples:
        >>> pos_encoder = PositionalEncoding(d_model)
    """

    def __init__(self, d_model, dropout=0.1, max_len=5000):
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)

        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0).transpose(0, 1)
        # self.register_buffer('pe', pe)
        # https://github.com/Lightning-AI/lightning/discussions/14377
        # https://github.com/pytorch/pytorch/issues/68407
        self.register_parameter('pe', nn.Parameter(pe, requires_grad=False))

    def forward(self, x):
        r"""Inputs of forward function
        Args:
            x: the sequence fed to the positional encoder model (required).
        Shape:
            x: [sequence length, batch size, embed dim]
            output: [sequence length, batch size, embed dim]
        Examples:
            >>> output = pos_encoder(x)
        """

        x = x + self.pe[:x.size(0), :]
        return self.dropout(x)

strhub/models/abinet/transformer.py
VamosC/CLIP4STR
# Scene Text Recognition Model Hub
# Copyright 2022 Darwin Bautista
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import math
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Optional, Tuple, List

import pytorch_lightning as pl
import torch
import torch.nn.functional as F
from nltk import edit_distance
from pytorch_lightning.utilities.types import EPOCH_OUTPUT, STEP_OUTPUT
from timm.optim import create_optimizer_v2
from torch import Tensor
from torch.optim import Optimizer
from torch.optim.lr_scheduler import OneCycleLR

from strhub.data.utils import CharsetAdapter, CTCTokenizer, Tokenizer, BaseTokenizer
from strhub.models.post_process import clip_post_process
from strhub.clip import clip


@dataclass
class BatchResult:
    num_samples: int
    correct: int
    ned: float
    confidence: float
    label_length: int
    loss: Tensor
    loss_numel: int


class BaseSystem(pl.LightningModule, ABC):

    def __init__(self, tokenizer: BaseTokenizer, charset_test: str,
                 batch_size: int, lr: float, warmup_pct: float, weight_decay: float) -> None:
        super().__init__()
        self.tokenizer = tokenizer
        self.charset_adapter = CharsetAdapter(charset_test)
        self.batch_size = batch_size
        self.lr = lr
        self.warmup_pct = warmup_pct
        self.weight_decay = weight_decay

    @abstractmethod
    def forward(self, images: Tensor, max_length: Optional[int] = None) -> Tensor:
        """Inference

        Args:
            images: Batch of images. Shape: N, Ch, H, W
            max_length: Max sequence length of the output. If None, will use default.

        Returns:
            logits: N, L, C (L = sequence length, C = number of classes, typically len(charset_train) + num specials)
        """
        raise NotImplementedError

    @abstractmethod
    def forward_logits_loss(self, images: Tensor, labels: List[str]) -> Tuple[Tensor, Tensor, int]:
        """Like forward(), but also computes the loss (calls forward() internally).

        Args:
            images: Batch of images. Shape: N, Ch, H, W
            labels: Text labels of the images

        Returns:
            logits: N, L, C (L = sequence length, C = number of classes, typically len(charset_train) + num specials)
            loss: mean loss for the batch
            loss_numel: number of elements the loss was calculated from
        """
        raise NotImplementedError

    def configure_optimizers(self):
        agb = self.trainer.accumulate_grad_batches
        # Linear scaling so that the effective learning rate is constant regardless of the number of GPUs used with DDP.
        lr_scale = agb * math.sqrt(self.trainer.num_devices) * self.batch_size / 256.
        lr = lr_scale * self.lr
        optim = create_optimizer_v2(self, 'adamw', lr, self.weight_decay)
        sched = OneCycleLR(optim, lr, self.trainer.estimated_stepping_batches, pct_start=self.warmup_pct,
                           cycle_momentum=False)
        return {'optimizer': optim, 'lr_scheduler': {'scheduler': sched, 'interval': 'step'}}

    def optimizer_zero_grad(self, epoch: int, batch_idx: int, optimizer: Optimizer, optimizer_idx: int):
        optimizer.zero_grad(set_to_none=True)

    def _eval_step(self, batch, validation: bool, clip_refine: bool = False, **kwargs) -> Optional[STEP_OUTPUT]:
        images, labels = batch

        correct = 0
        total = 0
        ned = 0
        confidence = 0
        label_length = 0
        if validation:
            logits, loss, loss_numel = self.forward_logits_loss(images, labels)
        else:
            # At test-time, we shouldn't specify a max_label_length because the test-time charset used
            # might be different from the train-time charset. max_label_length in eval_logits_loss() is computed
            # based on the transformed label, which could be wrong if the actual gt label contains characters existing
            # in the train-time charset but not in the test-time charset. For example, "aishahaleyes.blogspot.com"
            # is exactly 25 characters, but if processed by CharsetAdapter for the 36-char set, it becomes 23 characters
            # long only, which sets max_label_length = 23. This will cause the model prediction to be truncated.
            logits = self.forward(images)
            loss = loss_numel = None  # Only used for validation; not needed at test-time.

        probs = logits.softmax(-1)

        if clip_refine:
            # use clip to refine the final results
            batch_probs = self.tokenizer.decode_for_refine(probs)
            preds, probs = [], []
            for i, prob in enumerate(batch_probs):
                post_prob, post_ids, post_pred = \
                    clip_post_process(self.clip_model,
                                images[i, ...].unsqueeze(0),
                                prob, self.charset_adapter, self.tokenizer,
                                K=5 if not "sample_K" in kwargs else kwargs["sample_K"],
                                K2=0 if not "sample_K2" in kwargs else kwargs["sample_K2"],
                                num_samples=50 if not "sample_total" in kwargs else kwargs["sample_total"],
                                prompt=None if not "sample_prompt" in kwargs else kwargs["sample_prompt"],
                                alpha=0.1 if not "alpha" in kwargs else kwargs["alpha"]
                                )
                preds.append(post_pred)
                probs.append(post_prob)
        else:
            preds, probs = self.tokenizer.decode(probs)

        for pred, prob, gt in zip(preds, probs, labels):
            confidence += prob.prod().item()
            # adapt for the test charset
            pred = self.charset_adapter(pred) if not clip_refine else pred
            # Follow ICDAR 2019 definition of N.E.D.
            ned += edit_distance(pred, gt) / max(len(pred), len(gt))
            if pred == gt:
                correct += 1
            # else:
            #     # check for the wrong case
            #     print(total, pred, gt)
            total += 1
            label_length += len(pred)

        return dict(output=BatchResult(total, correct, ned, confidence, label_length, loss, loss_numel))

    @staticmethod
    def _aggregate_results(outputs: EPOCH_OUTPUT) -> Tuple[float, float, float]:
        if not outputs:
            return 0., 0., 0.
        total_loss = 0
        total_loss_numel = 0
        total_n_correct = 0
        total_norm_ED = 0
        total_size = 0
        for result in outputs:
            result = result['output']
            total_loss += result.loss_numel * result.loss
            total_loss_numel += result.loss_numel
            total_n_correct += result.correct
            total_norm_ED += result.ned
            total_size += result.num_samples
        acc = total_n_correct / total_size
        ned = (1 - total_norm_ED / total_size)
        loss = total_loss / total_loss_numel
        return acc, ned, loss

    def validation_step(self, batch, batch_idx) -> Optional[STEP_OUTPUT]:
        return self._eval_step(batch, True, clip_refine=False)

    def validation_epoch_end(self, outputs: EPOCH_OUTPUT) -> None:
        acc, ned, loss = self._aggregate_results(outputs)
        self.log('val_accuracy', 100 * acc, sync_dist=True)
        self.log('val_NED', 100 * ned, sync_dist=True)
        self.log('val_loss', loss, sync_dist=True)
        self.log('hp_metric', acc, sync_dist=True)

    def test_step(self, batch, batch_idx, clip_refine: bool = False, **kwargs) -> Optional[STEP_OUTPUT]:
        if clip_refine and self.clip_model is None:
            print('[CLIP] creating CLIP model for post refinement')
            model, clip_transform = clip.load(kwargs["clip_model_path"], device=batch[0].device)
            self.clip_model = model.float()

        return self._eval_step(batch, False, clip_refine=clip_refine, **kwargs)


class CrossEntropySystem(BaseSystem):

    def __init__(self, charset_train: str, charset_test: str,
                 batch_size: int, lr: float, warmup_pct: float, weight_decay: float) -> None:
        tokenizer = Tokenizer(charset_train)
        super().__init__(tokenizer, charset_test, batch_size, lr, warmup_pct, weight_decay)
        self.bos_id = tokenizer.bos_id
        self.eos_id = tokenizer.eos_id
        self.pad_id = tokenizer.pad_id
        self.clip_model = None

    def forward_logits_loss(self, images: Tensor, labels: List[str]) -> Tuple[Tensor, Tensor, int]:
        targets = self.tokenizer.encode(labels, self.device)
        targets = targets[:, 1:]            # Discard <bos>
        max_len = targets.shape[1] - 1      # exclude <eos> from count
        logits = self.forward(images, max_len)
        loss = F.cross_entropy(logits.flatten(end_dim=1), targets.flatten(), ignore_index=self.pad_id)
        loss_numel = (targets != self.pad_id).sum()

        return logits, loss, loss_numel


class CTCSystem(BaseSystem):

    def __init__(self, charset_train: str, charset_test: str,
                 batch_size: int, lr: float, warmup_pct: float, weight_decay: float) -> None:
        tokenizer = CTCTokenizer(charset_train)
        super().__init__(tokenizer, charset_test, batch_size, lr, warmup_pct, weight_decay)
        self.blank_id = tokenizer.blank_id

    def forward_logits_loss(self, images: Tensor, labels: List[str]) -> Tuple[Tensor, Tensor, int]:
        targets = self.tokenizer.encode(labels, self.device)
        logits = self.forward(images)
        log_probs = logits.log_softmax(-1).transpose(0, 1)  # swap batch and seq. dims
        T, N, _ = log_probs.shape
        input_lengths = torch.full(size=(N,), fill_value=T, dtype=torch.long, device=self.device)
        target_lengths = torch.as_tensor(list(map(len, labels)), dtype=torch.long, device=self.device)
        loss = F.ctc_loss(log_probs, targets, input_lengths, target_lengths, blank=self.blank_id, zero_infinity=True)
        return logits, loss, N

strhub/models/base.py
VamosC/CLIP4STR
# Scene Text Recognition Model Hub
# Copyright 2022 Darwin Bautista
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import math
from functools import partial
from itertools import permutations
from typing import Sequence, Any, Optional

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor

from pytorch_lightning.utilities.types import STEP_OUTPUT
from timm.models.helpers import named_apply

from strhub.models.base import CrossEntropySystem
from strhub.models.utils import init_weights
from strhub.models.vitstr.model import load_pretrained_weight
from .modules import DecoderLayer, Decoder, Encoder, TokenEmbedding


class PARSeq(CrossEntropySystem):

    def __init__(self, charset_train: str, charset_test: str, max_label_length: int,
                 batch_size: int, lr: float, warmup_pct: float, weight_decay: float,
                 img_size: Sequence[int], patch_size: Sequence[int], embed_dim: int,
                 enc_num_heads: int, enc_mlp_ratio: int, enc_depth: int,
                 dec_num_heads: int, dec_mlp_ratio: int, dec_depth: int,
                 perm_num: int, perm_forward: bool, perm_mirrored: bool,
                 decode_ar: bool, refine_iters: int, dropout: float, **kwargs: Any) -> None:
        super().__init__(charset_train, charset_test, batch_size, lr, warmup_pct, weight_decay)
        self.save_hyperparameters()

        self.max_label_length = max_label_length
        self.decode_ar = decode_ar
        self.refine_iters = refine_iters

        self.encoder = Encoder(img_size, patch_size, embed_dim=embed_dim, depth=enc_depth, num_heads=enc_num_heads,
                               mlp_ratio=enc_mlp_ratio)
        # load pretrained weights for encoder
        if embed_dim == 768 and "imagenet21k_pretrained" in kwargs.keys():
            load_pretrained_weight(self.encoder, kwargs["imagenet21k_pretrained"])

        if embed_dim > 512:
            self.proj = nn.Linear(embed_dim, 512)
            embed_dim = 512
        else:
            self.proj = nn.Identity()

        decoder_layer = DecoderLayer(embed_dim, dec_num_heads, embed_dim * dec_mlp_ratio, dropout)
        self.decoder = Decoder(decoder_layer, num_layers=dec_depth, norm=nn.LayerNorm(embed_dim))

        # Perm/attn mask stuff
        self.rng = np.random.default_rng()
        self.max_gen_perms = perm_num // 2 if perm_mirrored else perm_num
        self.perm_forward = perm_forward
        self.perm_mirrored = perm_mirrored

        # We don't predict <bos> nor <pad>
        self.head = nn.Linear(embed_dim, len(self.tokenizer) - 2)
        self.text_embed = TokenEmbedding(len(self.tokenizer), embed_dim)

        # +1 for <eos>
        self.pos_queries = nn.Parameter(torch.Tensor(1, max_label_length + 1, embed_dim))
        self.dropout = nn.Dropout(p=dropout)
        # Encoder has its own init.
        named_apply(partial(init_weights, exclude=['encoder']), self)
        nn.init.trunc_normal_(self.pos_queries, std=.02)

    @torch.jit.ignore
    def no_weight_decay(self):
        param_names = {'text_embed.embedding.weight', 'pos_queries'}
        enc_param_names = {'encoder.' + n for n in self.encoder.no_weight_decay()}
        return param_names.union(enc_param_names)

    def encode(self, img: torch.Tensor):
        return self.proj(self.encoder(img))

    def decode(self, tgt: torch.Tensor, memory: torch.Tensor, tgt_mask: Optional[Tensor] = None,
               tgt_padding_mask: Optional[Tensor] = None, tgt_query: Optional[Tensor] = None,
               tgt_query_mask: Optional[Tensor] = None):
        N, L = tgt.shape
        # <bos> stands for the null context. We only supply position information for characters after <bos>.
        null_ctx = self.text_embed(tgt[:, :1])
        tgt_emb = self.pos_queries[:, :L - 1] + self.text_embed(tgt[:, 1:])
        tgt_emb = self.dropout(torch.cat([null_ctx, tgt_emb], dim=1))
        if tgt_query is None:
            tgt_query = self.pos_queries[:, :L].expand(N, -1, -1)
        tgt_query = self.dropout(tgt_query)
        return self.decoder(tgt_query, tgt_emb, memory, tgt_query_mask, tgt_mask, tgt_padding_mask)

    def forward(self, images: Tensor, max_length: Optional[int] = None) -> Tensor:
        testing = max_length is None
        max_length = self.max_label_length if max_length is None else min(max_length, self.max_label_length)
        bs = images.shape[0]
        # +1 for <eos> at end of sequence.
        num_steps = max_length + 1
        memory = self.encode(images)

        # Query positions up to `num_steps`
        pos_queries = self.pos_queries[:, :num_steps].expand(bs, -1, -1)

        # Special case for the forward permutation. Faster than using `generate_attn_masks()`
        tgt_mask = query_mask = torch.triu(torch.full((num_steps, num_steps), float('-inf'), device=self._device), 1)

        if self.decode_ar:
            tgt_in = torch.full((bs, num_steps), self.pad_id, dtype=torch.long, device=self._device)
            tgt_in[:, 0] = self.bos_id

            logits = []
            for i in range(num_steps):
                j = i + 1  # next token index
                # Efficient decoding:
                # Input the context up to the ith token. We use only one query (at position = i) at a time.
                # This works because of the lookahead masking effect of the canonical (forward) AR context.
                # Past tokens have no access to future tokens, hence are fixed once computed.
                tgt_out = self.decode(tgt_in[:, :j], memory, tgt_mask[:j, :j], tgt_query=pos_queries[:, i:j],
                                      tgt_query_mask=query_mask[i:j, :j])
                # the next token probability is in the output's ith token position
                p_i = self.head(tgt_out)
                logits.append(p_i)
                if j < num_steps:
                    # greedy decode. add the next token index to the target input
                    tgt_in[:, j] = p_i.squeeze().argmax(-1)
                    # Efficient batch decoding: If all output words have at least one EOS token, end decoding.
                    if testing and (tgt_in == self.eos_id).any(dim=-1).all():
                        break

            logits = torch.cat(logits, dim=1)
        else:
            # No prior context, so input is just <bos>. We query all positions.
            tgt_in = torch.full((bs, 1), self.bos_id, dtype=torch.long, device=self._device)
            tgt_out = self.decode(tgt_in, memory, tgt_query=pos_queries)
            logits = self.head(tgt_out)

        if self.refine_iters:
            # For iterative refinement, we always use a 'cloze' mask.
            # We can derive it from the AR forward mask by unmasking the token context to the right.
            query_mask[torch.triu(torch.ones(num_steps, num_steps, dtype=torch.bool, device=self._device), 2)] = 0
            bos = torch.full((bs, 1), self.bos_id, dtype=torch.long, device=self._device)
            for i in range(self.refine_iters):
                # Prior context is the previous output.
                tgt_in = torch.cat([bos, logits[:, :-1].argmax(-1)], dim=1)
                tgt_padding_mask = ((tgt_in == self.eos_id).cumsum(-1) > 0)  # mask tokens beyond the first EOS token.
                tgt_out = self.decode(tgt_in, memory, tgt_mask, tgt_padding_mask,
                                      tgt_query=pos_queries, tgt_query_mask=query_mask[:, :tgt_in.shape[1]])
                logits = self.head(tgt_out)

        return logits

    def gen_tgt_perms(self, tgt):
        """Generate shared permutations for the whole batch.
           This works because the same attention mask can be used for the shorter sequences
           because of the padding mask.
        """
        # We don't permute the position of BOS, we permute EOS separately
        max_num_chars = tgt.shape[1] - 2
        # Special handling for 1-character sequences
        if max_num_chars == 1:
            return torch.arange(3, device=self._device).unsqueeze(0)
        perms = [torch.arange(max_num_chars, device=self._device)] if self.perm_forward else []
        # Additional permutations if needed
        max_perms = math.factorial(max_num_chars)
        if self.perm_mirrored:
            max_perms //= 2
        num_gen_perms = min(self.max_gen_perms, max_perms)
        # For 4-char sequences and shorter, we generate all permutations and sample from the pool to avoid collisions
        # Note that this code path might NEVER get executed since the labels in a mini-batch typically exceed 4 chars.
        if max_num_chars < 5:
            # Pool of permutations to sample from. We only need the first half (if complementary option is selected)
            # Special handling for max_num_chars == 4 which correctly divides the pool into the flipped halves
            if max_num_chars == 4 and self.perm_mirrored:
                selector = [0, 3, 4, 6, 9, 10, 12, 16, 17, 18, 19, 21]
            else:
                selector = list(range(max_perms))
            perm_pool = torch.as_tensor(list(permutations(range(max_num_chars), max_num_chars)), device=self._device)[selector]
            # If the forward permutation is always selected, no need to add it to the pool for sampling
            if self.perm_forward:
                perm_pool = perm_pool[1:]
            perms = torch.stack(perms)
            if len(perm_pool):
                i = self.rng.choice(len(perm_pool), size=num_gen_perms - len(perms), replace=False)
                perms = torch.cat([perms, perm_pool[i]])
        else:
            perms.extend([torch.randperm(max_num_chars, device=self._device) for _ in range(num_gen_perms - len(perms))])
            perms = torch.stack(perms)
        if self.perm_mirrored:
            # Add complementary pairs
            comp = perms.flip(-1)
            # Stack in such a way that the pairs are next to each other.
            perms = torch.stack([perms, comp]).transpose(0, 1).reshape(-1, max_num_chars)
        # NOTE:
        # The only meaningful way of permuting the EOS position is by moving it one character position at a time.
        # However, since the number of permutations = T! and number of EOS positions = T + 1, the number of possible EOS
        # positions will always be much less than the number of permutations (unless a low perm_num is set).
        # Thus, it would be simpler to just train EOS using the full and null contexts rather than trying to evenly
        # distribute it across the chosen number of permutations.
        # Add position indices of BOS and EOS
        bos_idx = perms.new_zeros((len(perms), 1))
        eos_idx = perms.new_full((len(perms), 1), max_num_chars + 1)
        perms = torch.cat([bos_idx, perms + 1, eos_idx], dim=1)
        # Special handling for the reverse direction. This does two things:
        # 1. Reverse context for the characters
        # 2. Null context for [EOS] (required for learning to predict [EOS] in NAR mode)
        if len(perms) > 1:
            perms[1, 1:] = max_num_chars + 1 - torch.arange(max_num_chars + 1, device=self._device)
        return perms

    def generate_attn_masks(self, perm):
        """Generate attention masks given a sequence permutation (includes pos. for bos and eos tokens)
        :param perm: the permutation sequence. i = 0 is always the BOS
        :return: lookahead attention masks
        """
        sz = perm.shape[0]
        mask = torch.zeros((sz, sz), device=self._device)
        for i in range(sz):
            query_idx = perm[i]
            masked_keys = perm[i + 1:]
            mask[query_idx, masked_keys] = float('-inf')
        content_mask = mask[:-1, :-1].clone()
        mask[torch.eye(sz, dtype=torch.bool, device=self._device)] = float('-inf')  # mask "self"
        query_mask = mask[1:, :-1]
        return content_mask, query_mask

    def training_step(self, batch, batch_idx) -> STEP_OUTPUT:
        images, labels = batch
        tgt = self.tokenizer.encode(labels, self._device)

        # Encode the source sequence (i.e. the image codes)
        memory = self.encode(images)

        # Prepare the target sequences (input and output)
        tgt_perms = self.gen_tgt_perms(tgt)
        tgt_in = tgt[:, :-1]
        tgt_out = tgt[:, 1:]
        # The [EOS] token is not depended upon by any other token in any permutation ordering
        tgt_padding_mask = (tgt_in == self.pad_id) | (tgt_in == self.eos_id)

        loss = 0
        loss_numel = 0
        n = (tgt_out != self.pad_id).sum().item()
        for i, perm in enumerate(tgt_perms):
            tgt_mask, query_mask = self.generate_attn_masks(perm)
            out = self.decode(tgt_in, memory, tgt_mask, tgt_padding_mask, tgt_query_mask=query_mask)
            logits = self.head(out).flatten(end_dim=1)
            loss += n * F.cross_entropy(logits, tgt_out.flatten(), ignore_index=self.pad_id)
            loss_numel += n
            # After the second iteration (i.e. done with canonical and reverse orderings),
            # remove the [EOS] tokens for the succeeding perms
            if i == 1:
                tgt_out = torch.where(tgt_out == self.eos_id, self.pad_id, tgt_out)
                n = (tgt_out != self.pad_id).sum().item()
        loss /= loss_numel

        self.log('loss', loss)
        return loss

strhub/models/parseq/system.py
VamosC/CLIP4STR
# coding = utf-8
import torch
from strhub.clip import clip


@torch.no_grad()
def clip_post_process(clip_model, image, probs, charset_adapter, char_tokenizer,
                        K=10, K2=5, num_samples=50, prompt=None, alpha=0.3):
    """using CLIP to do post-refinement
    Args:
        clip_model: the clip model
        image: input image, by default, after transforms
        probs: [L, C], probability distribution
        K: the sampled captions
        K2: if K2 > 0, before sampling, choose the top-K2 probabilities
        charset_adapter: charset adapter
    """
    max_p_y, max_ids = probs.max(-1)
    max_p_y_pord = max_p_y.prod(dim=0)
    p_y_prod, p_y, p_y_ind = top_k_candidate_v2(probs, K=K, num_samples=num_samples, K2=K2)

    # if the prediction with max probability not sampled, mannully add it
    if p_y_prod.max() < max_p_y_pord:
        try:
            p_y_prod = torch.cat([p_y_prod, max_p_y_pord.unsqueeze(0)], dim=0)
            p_y = torch.cat([p_y, max_p_y.unsqueeze(0)], dim=0)
            p_y_ind = torch.cat([p_y_ind, max_ids.unsqueeze(0)], dim=0)
        except:
            print(p_y.shape, p_y_ind.shape)

    text, preds = [], []
    for ind in p_y_ind:
        ids = ind.tolist()
        # maybe a single id
        ids = ids if isinstance(ids, list) else [ids]
        tokens = char_tokenizer._ids2tok(ids, join=True)

        # charset_adapter is necessary
        string = charset_adapter(tokens)
        # add prompt or not
        if prompt is not None:
            text.append(prompt + " " + string)
        else:
            text.append(string)
        preds.append(string)

    if len(text) > 1:
        with torch.no_grad():
            text_input = clip.tokenize(text).to(image.device)
            logits_from_image, logits_from_text = clip_model(image, text_input)
            clip_prob = logits_from_image.softmax(dim=-1)
    else:
        clip_prob = p_y_prod.softmax(-1)

    # combine the probabilities
    final_prob = p_y_prod.softmax(-1) + alpha * clip_prob
    p, ind = final_prob.max(dim=-1)

    final_p_y, final_ind, final_pred = p_y[ind, ...], p_y_ind[ind, ...], preds[ind.item()]

    return final_p_y, final_ind, final_pred


def top_k_candidate(prob, K=10, num_samples=30):
    """
    sample the top-k candidate from the probability distribution
    prob: shape [length of the string, num_classes], probability distribution
    K: the number of candidates
    Return:
        p_y_prod: [K], the joint probability
        p_y: [K, L], all probabilities
        p_y_ind: [K, L], correspoding index of p_y
    """
    p_y, p_y_ind, p_y_prod = [], [], []
    for i in range(num_samples):
        # sample as multinomial distribution
        ind = torch.multinomial(prob, 1)
        sample = torch.gather(prob, 1, ind).squeeze(dim=-1)
        prod = sample.prod(dim=0).item()
        if prod in p_y_prod:
            continue
        else:
            p_y.append(sample)
            p_y_ind.append(ind.squeeze(dim=-1))
            p_y_prod.append(prod)

    p_y = torch.stack(p_y, dim=0)
    p_y_ind = torch.stack(p_y_ind, dim=0)
    p_y_prod = p_y.prod(dim=1)

    k_prod, k_ind = torch.topk(p_y_prod, min(K, len(p_y_prod)), dim=0)

    return p_y_prod[k_ind], p_y[k_ind, ...], p_y_ind[k_ind, ...]


def top_k_candidate_v2(prob, K=10, num_samples=30, K2=5):
    """
    sample the top-k candidate from the probability distribution.
    In this version, we first choose the top-K samples of each prediciton of char, then
    sample from the choosed top-K predictios. 
    prob: shape [length of the string, num_classes], probability distribution
    K: the number of candidates
    K2: the number of choosed top-K candidates
    """
    if K2 <= 0 or K2 >= prob.shape[-1]:
        return top_k_candidate(prob, K, num_samples)
    else:
        # [L, C] -> [L, K2]
        k_prob, k_ind = torch.topk(prob, k=K2, dim=-1)
        # k_p_y_ind: [K, L]
        k_p_y_prod, k_p_y, k_p_y_ind = top_k_candidate(k_prob, K, num_samples)
        # map k_p_y_ind -> k_ind
        p_y_ind = torch.gather(k_ind, 1, k_p_y_ind.transpose(0, 1)).transpose(0, 1)

        return k_p_y_prod, k_p_y, p_y_ind


if __name__ == "__main__":
    # sum_x = 0
    # for i in range(1000):
    #     a = torch.tensor([  [0.1, 0.2, 0.7],
    #                         [0.2, 0.7, 0.1] ], dtype=torch.float)
    #     x = torch.multinomial(a, 1)
    #     # print(x)
    #     sum_x += x

    # print(sum_x / 1000)
    probs = torch.tensor([  [0.1, 0.2, 0.7],
                            [0.2, 0.7, 0.1],
                            [0.7, 0.1, 0.2] ], dtype=torch.float)
    # probs = torch.tensor([  [0.1, 0.2, 0.7] ], dtype=torch.float)
    p, p_y, p_y_ind = top_k_candidate_v2(probs, K=4, num_samples=30, K2=2)
    print("input probability:\n", probs)
    print("joint probability:\n", p)
    print("probabilities of discrete random variable:\n", p_y)
    print("index:\n", p_y_ind)

strhub/models/post_process.py
VamosC/CLIP4STR
# coding=utf-8
import os
import copy
import math
from functools import partial
from collections import OrderedDict
from typing import Optional, Sequence, Callable, Tuple
from timm.models.helpers import named_apply

import torch
from torch import nn as nn, Tensor
from torch.nn import functional as F
from torch.nn.functional import _in_projection_packed
from torch.nn.modules import transformer
from strhub.clip.model import ResidualAttentionBlock, LayerNorm, QuickGELU
from strhub.models.utils import init_weights


class DecoderLayer(nn.Module):
    """A Transformer decoder layer supporting two-stream attention (XLNet)
       This implements a pre-LN decoder, as opposed to the post-LN default in PyTorch."""

    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, activation='gelu',
                 layer_norm_eps=1e-5):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=True)
        self.cross_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=True)
        # Implementation of Feedforward model
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)

        self.norm1 = nn.LayerNorm(d_model, eps=layer_norm_eps)
        self.norm2 = nn.LayerNorm(d_model, eps=layer_norm_eps)
        self.norm_q = nn.LayerNorm(d_model, eps=layer_norm_eps)
        self.norm_c = nn.LayerNorm(d_model, eps=layer_norm_eps)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.dropout3 = nn.Dropout(dropout)

        self.activation = transformer._get_activation_fn(activation)

    def __setstate__(self, state):
        if 'activation' not in state:
            state['activation'] = F.gelu
        super().__setstate__(state)

    def forward_stream(self, tgt: Tensor, tgt_norm: Tensor, tgt_kv: Tensor, memory: Tensor, tgt_mask: Optional[Tensor],
                       tgt_key_padding_mask: Optional[Tensor]):
        """Forward pass for a single stream (i.e. content or query)
        tgt_norm is just a LayerNorm'd tgt. Added as a separate parameter for efficiency.
        Both tgt_kv and memory are expected to be LayerNorm'd too.
        memory is LayerNorm'd by ViT.
        """
        tgt2, sa_weights = self.self_attn(tgt_norm, tgt_kv, tgt_kv, attn_mask=tgt_mask,
                                          key_padding_mask=tgt_key_padding_mask)
        tgt = tgt + self.dropout1(tgt2)

        tgt2, ca_weights = self.cross_attn(self.norm1(tgt), memory, memory)
        tgt = tgt + self.dropout2(tgt2)

        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(self.norm2(tgt)))))
        tgt = tgt + self.dropout3(tgt2)
        return tgt, sa_weights, ca_weights

    def forward(self, query, content, memory, query_mask: Optional[Tensor] = None, content_mask: Optional[Tensor] = None,
                content_key_padding_mask: Optional[Tensor] = None, update_content: bool = True):
        query_norm = self.norm_q(query)
        content_norm = self.norm_c(content)
        query = self.forward_stream(query, query_norm, content_norm, memory, query_mask, content_key_padding_mask)[0]
        if update_content:
            content = self.forward_stream(content, content_norm, content_norm, memory, content_mask,
                                          content_key_padding_mask)[0]
        return query, content


class Decoder(nn.Module):
    __constants__ = ['norm']

    def __init__(self, decoder_layer, num_layers, norm):
        super().__init__()
        self.layers = transformer._get_clones(decoder_layer, num_layers)
        self.num_layers = num_layers
        self.norm = norm

    def forward(self, query, content, memory, query_mask: Optional[Tensor] = None, content_mask: Optional[Tensor] = None,
                content_key_padding_mask: Optional[Tensor] = None):
        for i, mod in enumerate(self.layers):
            last = i == len(self.layers) - 1
            query, content = mod(query, content, memory, query_mask, content_mask, content_key_padding_mask,
                                 update_content=not last)
        query = self.norm(query)
        return query


class TokenEmbedding(nn.Module):

    def __init__(self, charset_size: int, embed_dim: int):
        super().__init__()
        self.embedding = nn.Embedding(charset_size, embed_dim)
        self.embed_dim = embed_dim

    def forward(self, tokens: torch.Tensor):
        return math.sqrt(self.embed_dim) * self.embedding(tokens)


class Hook():
    # A simple hook class that returns the input and output of a layer during forward/backward pass
    def __init__(self, module, backward=False):
        if backward == False:
            self.hook = module.register_forward_hook(self.hook_fn)
        else:
            self.hook = module.register_backward_hook(self.hook_fn)
    
    def hook_fn(self, module, input, output):
        self.input = input
        self.output = output

    def close(self):
        self.hook.remove()


class IndentityAdapter(nn.Module):
    def __init__(self, *args, **kwargs):
        super(IndentityAdapter, self).__init__()
        self.identity = nn.Identity()

    def forward(self, x):
        return self.identity(x)


class LinearAdapter(nn.Module):
    """
    CLIP adapter from
    https://github.com/gaopengcuhk/CLIP-Adapter/blob/main/clip_adapter.py
    """
    def __init__(self, c_in, reduction=2, ratio=0.2):
        super(LinearAdapter, self).__init__()
        self.fc = nn.Sequential(
            nn.Linear(c_in, c_in // reduction, bias=False),
            nn.ReLU(inplace=True),
            nn.Linear(c_in // reduction, c_in, bias=False),
            nn.ReLU(inplace=True)
        )
        self.ratio = ratio

        # init
        for m in self.modules():
            init_weights(m)

    def forward(self, x):
        x = self.ratio * self.fc(x) + (1.0 - self.ratio) * x

        return x


class AttentionAdapter(nn.Module):
    def __init__(self, c_in, n_head, ratio=0.5):
        super(AttentionAdapter, self).__init__()
        self.attn = ResidualAttentionBlock(d_model=c_in, n_head=n_head)
        self.ratio = ratio
        # init
        named_apply(partial(init_weights, exclude=['none']), self)

    def forward(self, x):
        # x [N, L, D]
        x = x.permute(1, 0, 2)      # NLD -> LND
        x = self.ratio * self.attn(x) + (1.0 - self.ratio) * x
        x = x.permute(1, 0, 2)      # LND -> NLD

        return x


class LadderSideAdapter(nn.Module):
    def __init__(self, clip_model, block_ids=[0, 4, 8, 11], T=0.1) -> None:
        """
        A Ladder Side Adapter
        Args:
            transformer: a transformer object defined by CLIP
            block_ids: ids of blocks in the transformer
            T: temperature
        """
        super(LadderSideAdapter, self).__init__()
        self.T = T
        self.block_ids = block_ids

        self.alpha = nn.Parameter(torch.zeros(len(block_ids)), requires_grad=True)
        resblocks = []
        self.hooks = []
        for i in block_ids:
            self.hooks.append(Hook(clip_model.visual.transformer.resblocks[i], backward=False))
            resblocks.append(copy.deepcopy(clip_model.visual.transformer.resblocks[i]))
        self.resblocks = nn.ModuleList(resblocks)
        self.ln_post = copy.deepcopy(clip_model.visual.ln_post)
        self.proj = copy.deepcopy(clip_model.visual.proj)

    def forward(self, memory=None):
        # x: LND in hooks
        
        # ladder side forward
        W = torch.sigmoid(self.alpha / self.T)
        for i, j in enumerate(self.block_ids):
            if i == 0:
                x = self.resblocks[i](self.hooks[i].input[0])
            else:
                x = W[i] * self.hooks[i].input[0] + (1.0 - W[i]) * x
                x = self.resblocks[i](x)
        x = x.permute(1, 0, 2)                  # LND -> NLD

        # output transformation
        x = self.ln_post(x)
        if self.proj is not None:
            x = torch.matmul(x, self.proj)

        if memory is not None:
            x = W[0] * memory + (1.0 - W[0]) * x

        return x


class LinearLadderSideAdapter(nn.Module):
    def __init__(self, clip_model, block_ids=[1, 3, 5, 7, 9, 11], T=0.1) -> None:
        """
        A Ladder Side Adapter
        Args:
            transformer: a transformer object defined by CLIP
            block_ids: ids of blocks in the transformer
            T: temperature
        """
        super(LinearLadderSideAdapter, self).__init__()
        self.T = T
        self.block_ids = block_ids

        self.alpha = nn.Parameter(torch.zeros(len(block_ids)), requires_grad=True)
        width = clip_model.visual.class_embedding.shape[0]
        resblocks = []
        self.hooks = []
        for i in block_ids:
            self.hooks.append(Hook(clip_model.visual.transformer.resblocks[i], backward=False))
            resblocks.append(LinearAdapter(width, reduction=4, ratio=0.5))
        self.resblocks = nn.ModuleList(resblocks)
        self.ln_post = copy.deepcopy(clip_model.visual.ln_post)
        self.proj = copy.deepcopy(clip_model.visual.proj)
        # init
        named_apply(partial(init_weights, exclude=['ln_post', 'proj', 'alpha']), self)

    def forward(self, memory=None):
        # x: LND in hooks
        
        # ladder side forward
        W = torch.sigmoid(self.alpha / self.T)
        for i, j in enumerate(self.block_ids):
            if i == 0:
                x = self.resblocks[i](self.hooks[i].input[0])
            else:
                x = W[i] * self.hooks[i].input[0] + (1.0 - W[i]) * x
                x = self.resblocks[i](x)
        x = x.permute(1, 0, 2)                  # LND -> NLD

        # output transformation
        x = self.ln_post(x)
        if self.proj is not None:
            x = torch.matmul(x, self.proj)

        if memory is not None:
            x = W[0] * memory + (1.0 - W[0]) * x

        return x


class LadderSideAdapterPruning(nn.Module):
    def __init__(self, clip_model, block_ids=[1, 3, 5, 7, 9, 11], T=0.1, reduction=4.0) -> None:
        """
        A Ladder Side Adapter using the pruned model as the side network
        Args:
            clip_model: clip model
            block_ids: ids of blocks in the transformer
            T: temperature
            reduction: prune the row or columns of weight matrix to (1 / reduction) of the original
        """
        super(LadderSideAdapterPruning, self).__init__()
        self.T = T
        self.block_ids = block_ids
        self.alpha = nn.Parameter(torch.zeros(len(block_ids)), requires_grad=True)
        
        d_model = clip_model.visual.transformer.resblocks[0].attn.embed_dim
        n_head = clip_model.visual.transformer.resblocks[0].attn.num_heads
        resblocks = []
        downsamples = []
        self.hooks = []
        for i in block_ids:
            self.hooks.append(Hook(clip_model.visual.transformer.resblocks[i], backward=False))
    
            # prune and init
            # block = ResidualAttentionBlock(int(d_model / reduction), n_head)
            block = ResidualAttentionBlockCustom(d_model, n_head, reduction=reduction)
            state_dict = clip_model.visual.transformer.resblocks[i].state_dict()
            # new_state_dict = self.prune(state_dict, reduction=reduction)
            new_state_dict = self.prune_v2(state_dict, reduction=reduction)
            block.load_state_dict(new_state_dict)
            resblocks.append(block)
            # downsamples.append(nn.Linear(d_model, int(d_model / reduction), bias=False))
            downsamples.append(nn.Identity())

        self.resblocks = nn.ModuleList(resblocks)
        self.downsamples = nn.ModuleList(downsamples)
        # self.upsample = nn.Linear(int(d_model / reduction), d_model, bias=False)
        self.upsample = nn.Identity()
        self.ln_post = copy.deepcopy(clip_model.visual.ln_post)
        self.proj = copy.deepcopy(clip_model.visual.proj)
        # init
        named_apply(partial(init_weights, exclude=['resblocks', 'ln_post', 'proj', 'alpha']), self)

    def forward(self, memory=None):
        # x: LND in hooks
        
        # ladder side forward
        W = torch.sigmoid(self.alpha / self.T)
        for i, j in enumerate(self.block_ids):
            res_x = self.downsamples[i](self.hooks[i].input[0])
            if i == 0:
                x = self.resblocks[i](res_x)
            else:
                x = W[i] * res_x + (1.0 - W[i]) * x
                x = self.resblocks[i](x)
        x = self.upsample(x)
        x = x.permute(1, 0, 2)                  # LND -> NLD

        # output transformation
        x = self.ln_post(x)
        if self.proj is not None:
            x = torch.matmul(x, self.proj)

        if memory is not None:
            x = W[0] * memory + (1.0 - W[0]) * x

        return x

    def prune(self, state_dcit, reduction):
        new_sd = {}
        for k, v in state_dcit.items():
            if "in_proj" in k:
                v_q, v_k, v_v = torch.chunk(v, 3, dim=0)
                new_v_q = ln_pruning(v_q, reduction=reduction, prune_col=True)
                new_v_k = ln_pruning(v_k, reduction=reduction, prune_col=True)
                new_v_v = ln_pruning(v_v, reduction=reduction, prune_col=True)
                new_v = torch.cat([new_v_q, new_v_k, new_v_v], dim=0).contiguous()
            else:
                new_v = ln_pruning(v, reduction=reduction, prune_col=True)
            new_sd[k] = new_v

        return new_sd

    def prune_v2(self, state_dcit, reduction):
        new_sd = {}
        for k, v in state_dcit.items():
            if "in_proj" in k:
                v_q, v_k, v_v = torch.chunk(v, 3, dim=0)
                new_v_q = ln_pruning(v_q, reduction=reduction, prune_col=False)
                new_v_k = ln_pruning(v_k, reduction=reduction, prune_col=False)
                new_v_v = ln_pruning(v_v, reduction=reduction, prune_col=False)
                new_v = torch.cat([new_v_q, new_v_k, new_v_v], dim=0).contiguous()
            
            elif "out_proj.weight" in k or "mlp.c_proj.weight" in k:
                # only prune in_feature dimension
                new_v = ln_pruning(v.transpose(0, 1), reduction=reduction, prune_col=False)
                new_v = new_v.transpose(0, 1)
            
            elif "mlp.c_fc" in k:
                new_v = ln_pruning(v, reduction=reduction, prune_col=False)

            else:
                new_v = v

            new_sd[k] = new_v

        return new_sd


def ln_pruning(W, reduction=4.0, p=1, prune_col=False):
    if W.ndim == 1:
        n = W.numel()
        ln_norm = torch.abs(W)
        n_to_prune = int(n / reduction)
        _, indices = torch.topk(ln_norm, k=n_to_prune, sorted=False)
        s_indices, _ = torch.sort(indices, descending=False)
        pruned_W = W[s_indices]

    elif W.ndim == 2:
        row, col = W.shape
        # first prune the row
        ln_norm = torch.norm(W, p=p, dim=1)
        n_to_prune = int(row / reduction)
        _, indices = torch.topk(ln_norm, k=n_to_prune, sorted=False)
        s_indices, _ = torch.sort(indices, descending=False)
        pruned_W = W[s_indices, ...]

        if prune_col:
            # then prune the column
            W_ = pruned_W.transpose(0, 1)
            ln_norm = torch.norm(W_, p=p, dim=1)
            n_to_prune = int(col / reduction)
            _, indices = torch.topk(ln_norm, k=n_to_prune, sorted=False)
            s_indices, _ = torch.sort(indices, descending=False)
            pruned_W = W_[s_indices, ...]
            pruned_W = pruned_W.transpose(0, 1)

    else:
        raise NotImplementedError

    return pruned_W


class ResidualAttentionBlockCustom(nn.Module):
    def __init__(self, d_model: int, n_head: int, attn_mask: torch.Tensor = None, reduction = 4):
        super().__init__()

        self.attn = MultiheadAttentionCustom(d_model, n_head, reduction=reduction)
        self.ln_1 = LayerNorm(d_model)
        self.mlp = nn.Sequential(OrderedDict([
            ("c_fc", nn.Linear(d_model, d_model * 4 // reduction)),
            ("gelu", QuickGELU()),
            ("c_proj", nn.Linear(d_model * 4 // reduction, d_model))
        ]))
        self.ln_2 = LayerNorm(d_model)
        self.attn_mask = attn_mask

    def attention(self, x: torch.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: torch.Tensor):
        x = x + self.attention(self.ln_1(x))
        x = x + self.mlp(self.ln_2(x))
        return x


class MultiheadAttentionCustom(nn.Module):
    def __init__(self, d_model: int, n_head: int, bias = True, attn_mask: torch.Tensor = None, reduction = 4):
        super().__init__()
        self.embed_dim = d_model
        self.num_heads = n_head
        self.in_proj_weight = nn.Parameter(torch.empty((3 * d_model // reduction, d_model)))
        self.in_proj_bias = nn.Parameter(torch.empty(3 * d_model // reduction))
        self.out_proj = nn.Linear(d_model // reduction, d_model, bias=bias)
        self._reset_parameters()

    def _reset_parameters(self):
        nn.init.xavier_uniform_(self.in_proj_weight)
        if self.in_proj_bias is not None:
            nn.init.constant_(self.in_proj_bias, 0.)
            nn.init.constant_(self.out_proj.bias, 0.)

    def forward(self, query: Tensor, key: Tensor, value: Tensor, key_padding_mask: Optional[Tensor] = None,
                need_weights: bool = True, attn_mask: Optional[Tensor] = None,
                average_attn_weights: bool = True) -> Tuple[Tensor, Optional[Tensor]]:
        # by default, length frsit in input, e.g., [L, N D]
        num_heads = self.num_heads

        # set up shape vars
        tgt_len, bsz, embed_dim = query.shape
        src_len, _, _ = key.shape

        # we assume q,k,v are the same
        q, k, v = _in_projection_packed(query, key, value, self.in_proj_weight, self.in_proj_bias)

        head_dim = q.shape[-1] // num_heads

        #
        # reshape q, k, v for multihead attention and make em batch first
        #
        q = q.contiguous().view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)
        k = k.contiguous().view(k.shape[0], bsz * num_heads, head_dim).transpose(0, 1)
        v = v.contiguous().view(v.shape[0], bsz * num_heads, head_dim).transpose(0, 1)
        
        # update source sequence length after adjustments
        src_len = k.size(1)

        dropout_p = 0.0
        #
        # (deep breath) calculate attention and out projection
        #
        attn_output, attn_output_weights = _scaled_dot_product_attention(q, k, v, attn_mask, dropout_p)
        attn_output = attn_output.transpose(0, 1).contiguous().view(tgt_len, bsz, head_dim * num_heads)
        attn_output = F.linear(attn_output, self.out_proj.weight, self.out_proj.bias)

        if need_weights:
            # average attention weights over heads
            attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)
            return attn_output, attn_output_weights.sum(dim=1) / num_heads
        else:
            return attn_output, None

# https://github.com/pytorch/pytorch/blob/71f889c7d265b9636b93ede9d651c0a9c4bee191/torch/nn/functional.py#L4809
def _scaled_dot_product_attention(
    q: Tensor,
    k: Tensor,
    v: Tensor,
    attn_mask: Optional[Tensor] = None,
    dropout_p: float = 0.0,
) -> Tuple[Tensor, Tensor]:
    r"""
    Computes scaled dot product attention on query, key and value tensors, using
    an optional attention mask if passed, and applying dropout if a probability
    greater than 0.0 is specified.
    Returns a tensor pair containing attended values and attention weights.
    Args:
        q, k, v: query, key and value tensors. See Shape section for shape details.
        attn_mask: optional tensor containing mask values to be added to calculated
            attention. May be 2D or 3D; see Shape section for details.
        dropout_p: dropout probability. If greater than 0.0, dropout is applied.
    Shape:
        - q: :math:`(B, Nt, E)` where B is batch size, Nt is the target sequence length,
            and E is embedding dimension.
        - key: :math:`(B, Ns, E)` where B is batch size, Ns is the source sequence length,
            and E is embedding dimension.
        - value: :math:`(B, Ns, E)` where B is batch size, Ns is the source sequence length,
            and E is embedding dimension.
        - attn_mask: either a 3D tensor of shape :math:`(B, Nt, Ns)` or a 2D tensor of
            shape :math:`(Nt, Ns)`.
        - Output: attention values have shape :math:`(B, Nt, E)`; attention weights
            have shape :math:`(B, Nt, Ns)`
    """
    B, Nt, E = q.shape
    q = q / math.sqrt(E)
    # (B, Nt, E) x (B, E, Ns) -> (B, Nt, Ns)
    attn = torch.bmm(q, k.transpose(-2, -1))
    if attn_mask is not None:
        attn += attn_mask
    attn = F.softmax(attn, dim=-1)
    if dropout_p > 0.0:
        attn = F.dropout(attn, p=dropout_p)
    # (B, Nt, Ns) x (B, Ns, E) -> (B, Nt, E)
    output = torch.bmm(attn, v)

    return output, attn


if __name__ == "__main__":
    # W = torch.rand(4, 4)
    # W = torch.rand(4)
    # b = ln_pruning(W, reduction=2.0, p=1)
    # print(W)
    # print(b)

    m = MultiheadAttentionCustom(768, 12, reduction=4)
    m = ResidualAttentionBlockCustom(768, 12, reduction=4)
    num = 0
    for n, p in m.named_parameters():
        num += p.numel()
    print(num)

    x = torch.rand(16, 4, 768)
    # y = m(x, x, x)
    y = m(x)
    print(y[0].shape)

strhub/models/str_adapter/modules.py
VamosC/CLIP4STR
# CLIP Adapters for STR
import os
import math
import warnings
from functools import partial
from itertools import permutations
from typing import Sequence, Any, Optional

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from torch.optim.lr_scheduler import OneCycleLR
from pytorch_lightning.utilities import rank_zero_info
from pytorch_lightning.utilities.types import STEP_OUTPUT
from timm.models.helpers import named_apply

from strhub.clip import clip
from strhub.models.utils import init_weights
from strhub.models.base import CrossEntropySystem
from .modules import DecoderLayer, Decoder, TokenEmbedding
from .modules import IndentityAdapter, LinearAdapter, AttentionAdapter
from .modules import LadderSideAdapter, LinearLadderSideAdapter, LadderSideAdapterPruning


# an alternative choice when the input argument is not valid 
CLIP_PATH = '/PUT/YOUR/PATH/HERE/pretrained/clip'


class STRAdapter(CrossEntropySystem):

    def __init__(self, charset_train: str, charset_test: str, max_label_length: int,
                 batch_size: int, lr: float, warmup_pct: float, weight_decay: float,
                 img_size: Sequence[int], patch_size: Sequence[int],
                 dec_num_heads: int, dec_mlp_ratio: int, dec_depth: int,
                 perm_num: int, perm_forward: bool, perm_mirrored: bool,
                 decode_ar: bool, refine_iters: int, dropout: float, **kwargs: Any) -> None:
        super().__init__(charset_train, charset_test, batch_size, lr, warmup_pct, weight_decay)
        self.save_hyperparameters()

        # some args
        self.max_label_length = max_label_length
        self.decode_ar = decode_ar
        self.refine_iters = refine_iters
        self.coef_lr = kwargs["coef_lr"] if "coef_lr" in kwargs.keys() else 1.0
        self.coef_wd = kwargs["coef_wd"] if "coef_wd" in kwargs.keys() else 1.0
        self.freeze_backbone = kwargs["freeze_backbone"] if "freeze_backbone" in kwargs.keys() else False
        self.adapter_type = kwargs["adapter_type"] if "adapter_type" in kwargs.keys() else "linear"
        self.block_ids = kwargs["block_ids"] if "block_ids" in kwargs.keys() else [0, 4, 8, 11]
        self.prune_reduction = kwargs["prune_reduction"] if "prune_reduction" in kwargs.keys() else 4

        rank_zero_info("[STRAdapter] freeze_backbone {}, adapter_type {}, block_ids {}".format(
                            self.freeze_backbone, self.adapter_type, self.block_ids))

        # load CLIP model
        assert "clip_pretrained" in kwargs.keys()
        clip_model, _ = clip.load(name=kwargs["clip_pretrained"], device='cpu')
        self.clip_model = clip_model.float()

        assert "clip_pretrained" in kwargs.keys()
        if not os.path.exists(kwargs["clip_pretrained"]):
            kwargs["clip_pretrained"] = os.path.join(CLIP_PATH, os.path.basename(kwargs["clip_pretrained"]))
            assert os.path.exists(kwargs["clip_pretrained"])

        embed_dim = self.clip_model.text_projection.shape[-1]
        if self.adapter_type == "linear":
            self.adapter = LinearAdapter(embed_dim, reduction=2, ratio=0.2)
        elif self.adapter_type == "attention":
            self.adapter = AttentionAdapter(embed_dim, dec_num_heads, ratio=0.5)
        elif self.adapter_type == "ladder":
            self.adapter = LadderSideAdapter(self.clip_model, block_ids=self.block_ids)
        elif self.adapter_type == "linear_ladder":
            self.adapter = LinearLadderSideAdapter(self.clip_model, block_ids=self.block_ids)
        elif self.adapter_type == "ladder_pruning":
            self.adapter = LadderSideAdapterPruning(self.clip_model, block_ids=self.block_ids, reduction=self.prune_reduction)
        else:
            warnings.warn("Warning: you are using IndentityAdapter")
            self.adapter = IndentityAdapter()

        # decoder module
        decoder_layer = DecoderLayer(embed_dim, dec_num_heads, embed_dim * dec_mlp_ratio, dropout)
        self.decoder = Decoder(decoder_layer, num_layers=dec_depth, norm=nn.LayerNorm(embed_dim))

        # Perm/attn mask stuff
        self.rng = np.random.default_rng()
        self.max_gen_perms = perm_num // 2 if perm_mirrored else perm_num
        self.perm_forward = perm_forward
        self.perm_mirrored = perm_mirrored

        # We don't predict <bos> nor <pad>
        self.head = nn.Linear(embed_dim, len(self.tokenizer) - 2, bias=True)
        self.text_embed = TokenEmbedding(len(self.tokenizer), embed_dim)

        # +1 for <eos>
        self.pos_queries = nn.Parameter(torch.Tensor(1, max_label_length + 1, embed_dim))
        self.dropout = nn.Dropout(p=dropout)

        # init decoder
        named_apply(partial(init_weights, exclude=['clip_model', 'adapter']), self)
        nn.init.trunc_normal_(self.pos_queries, std=.02)

    def encode(self, img: torch.Tensor):
        if self.freeze_backbone:
            self.clip_model.eval()
            with torch.no_grad():
                memory = self.clip_model.encode_image(img)
        else:
            memory = self.clip_model.encode_image(img)

        memory = self.adapter(memory)

        return memory

    def decode(self, tgt: torch.Tensor, memory: torch.Tensor, tgt_mask: Optional[Tensor] = None,
               tgt_padding_mask: Optional[Tensor] = None, tgt_query: Optional[Tensor] = None,
               tgt_query_mask: Optional[Tensor] = None):
        N, L = tgt.shape
        # <bos> stands for the null context. We only supply position information for characters after <bos>.
        null_ctx = self.text_embed(tgt[:, :1])
        tgt_emb = self.pos_queries[:, :L - 1] + self.text_embed(tgt[:, 1:])
        tgt_emb = self.dropout(torch.cat([null_ctx, tgt_emb], dim=1))
        if tgt_query is None:
            tgt_query = self.pos_queries[:, :L].expand(N, -1, -1)
        tgt_query = self.dropout(tgt_query)
        return self.decoder(tgt_query, tgt_emb, memory, tgt_query_mask, tgt_mask, tgt_padding_mask)

    def forward(self, images: Tensor, max_length: Optional[int] = None) -> Tensor:
        testing = max_length is None
        max_length = self.max_label_length if max_length is None else min(max_length, self.max_label_length)
        bs = images.shape[0]
        # +1 for <eos> at end of sequence.
        num_steps = max_length + 1
        memory = self.encode(images)

        # Query positions up to `num_steps`
        pos_queries = self.pos_queries[:, :num_steps].expand(bs, -1, -1)

        # Special case for the forward permutation. Faster than using `generate_attn_masks()`
        tgt_mask = query_mask = torch.triu(torch.full((num_steps, num_steps), float('-inf'), device=self._device), 1)

        if self.decode_ar:
            tgt_in = torch.full((bs, num_steps), self.pad_id, dtype=torch.long, device=self._device)
            tgt_in[:, 0] = self.bos_id

            logits = []
            for i in range(num_steps):
                j = i + 1  # next token index
                # Efficient decoding:
                # Input the context up to the ith token. We use only one query (at position = i) at a time.
                # This works because of the lookahead masking effect of the canonical (forward) AR context.
                # Past tokens have no access to future tokens, hence are fixed once computed.
                tgt_out = self.decode(tgt_in[:, :j], memory, tgt_mask[:j, :j], tgt_query=pos_queries[:, i:j],
                                      tgt_query_mask=query_mask[i:j, :j])
                # the next token probability is in the output's ith token position
                p_i = self.head(tgt_out)
                logits.append(p_i)
                if j < num_steps:
                    # greedy decode. add the next token index to the target input
                    tgt_in[:, j] = p_i.squeeze().argmax(-1)
                    # Efficient batch decoding: If all output words have at least one EOS token, end decoding.
                    if testing and (tgt_in == self.eos_id).any(dim=-1).all():
                        break

            logits = torch.cat(logits, dim=1)
        else:
            # No prior context, so input is just <bos>. We query all positions.
            tgt_in = torch.full((bs, 1), self.bos_id, dtype=torch.long, device=self._device)
            tgt_out = self.decode(tgt_in, memory, tgt_query=pos_queries)
            logits = self.head(tgt_out)

        if self.refine_iters:
            # For iterative refinement, we always use a 'cloze' mask.
            # We can derive it from the AR forward mask by unmasking the token context to the right.
            query_mask[torch.triu(torch.ones(num_steps, num_steps, dtype=torch.bool, device=self._device), 2)] = 0
            bos = torch.full((bs, 1), self.bos_id, dtype=torch.long, device=self._device)
            for i in range(self.refine_iters):
                # Prior context is the previous output.
                tgt_in = torch.cat([bos, logits[:, :-1].argmax(-1)], dim=1)
                tgt_padding_mask = ((tgt_in == self.eos_id).cumsum(-1) > 0)  # mask tokens beyond the first EOS token.
                tgt_out = self.decode(tgt_in, memory, tgt_mask, tgt_padding_mask,
                                      tgt_query=pos_queries, tgt_query_mask=query_mask[:, :tgt_in.shape[1]])
                logits = self.head(tgt_out)

        return logits

    def gen_tgt_perms(self, tgt):
        """Generate shared permutations for the whole batch.
           This works because the same attention mask can be used for the shorter sequences
           because of the padding mask.
        """
        # We don't permute the position of BOS, we permute EOS separately
        max_num_chars = tgt.shape[1] - 2
        # Special handling for 1-character sequences
        if max_num_chars == 1:
            return torch.arange(3, device=self._device).unsqueeze(0)
        perms = [torch.arange(max_num_chars, device=self._device)] if self.perm_forward else []
        # Additional permutations if needed
        max_perms = math.factorial(max_num_chars)
        if self.perm_mirrored:
            max_perms //= 2
        num_gen_perms = min(self.max_gen_perms, max_perms)
        # For 4-char sequences and shorter, we generate all permutations and sample from the pool to avoid collisions
        # Note that this code path might NEVER get executed since the labels in a mini-batch typically exceed 4 chars.
        if max_num_chars < 5:
            # Pool of permutations to sample from. We only need the first half (if complementary option is selected)
            # Special handling for max_num_chars == 4 which correctly divides the pool into the flipped halves
            if max_num_chars == 4 and self.perm_mirrored:
                selector = [0, 3, 4, 6, 9, 10, 12, 16, 17, 18, 19, 21]
            else:
                selector = list(range(max_perms))
            perm_pool = torch.as_tensor(list(permutations(range(max_num_chars), max_num_chars)), device=self._device)[selector]
            # If the forward permutation is always selected, no need to add it to the pool for sampling
            if self.perm_forward:
                perm_pool = perm_pool[1:]
            perms = torch.stack(perms)
            if len(perm_pool):
                i = self.rng.choice(len(perm_pool), size=num_gen_perms - len(perms), replace=False)
                perms = torch.cat([perms, perm_pool[i]])
        else:
            perms.extend([torch.randperm(max_num_chars, device=self._device) for _ in range(num_gen_perms - len(perms))])
            perms = torch.stack(perms)
        if self.perm_mirrored:
            # Add complementary pairs
            comp = perms.flip(-1)
            # Stack in such a way that the pairs are next to each other.
            perms = torch.stack([perms, comp]).transpose(0, 1).reshape(-1, max_num_chars)
        # NOTE:
        # The only meaningful way of permuting the EOS position is by moving it one character position at a time.
        # However, since the number of permutations = T! and number of EOS positions = T + 1, the number of possible EOS
        # positions will always be much less than the number of permutations (unless a low perm_num is set).
        # Thus, it would be simpler to just train EOS using the full and null contexts rather than trying to evenly
        # distribute it across the chosen number of permutations.
        # Add position indices of BOS and EOS
        bos_idx = perms.new_zeros((len(perms), 1))
        eos_idx = perms.new_full((len(perms), 1), max_num_chars + 1)
        perms = torch.cat([bos_idx, perms + 1, eos_idx], dim=1)
        # Special handling for the reverse direction. This does two things:
        # 1. Reverse context for the characters
        # 2. Null context for [EOS] (required for learning to predict [EOS] in NAR mode)
        if len(perms) > 1:
            perms[1, 1:] = max_num_chars + 1 - torch.arange(max_num_chars + 1, device=self._device)
        return perms

    def generate_attn_masks(self, perm):
        """Generate attention masks given a sequence permutation (includes pos. for bos and eos tokens)
        :param perm: the permutation sequence. i = 0 is always the BOS
        :return: lookahead attention masks
        """
        sz = perm.shape[0]
        mask = torch.zeros((sz, sz), device=self._device)
        for i in range(sz):
            query_idx = perm[i]
            masked_keys = perm[i + 1:]
            mask[query_idx, masked_keys] = float('-inf')
        content_mask = mask[:-1, :-1].clone()
        mask[torch.eye(sz, dtype=torch.bool, device=self._device)] = float('-inf')  # mask "self"
        query_mask = mask[1:, :-1]
        return content_mask, query_mask

    def training_step(self, batch, batch_idx) -> STEP_OUTPUT:
        images, labels = batch
        tgt = self.tokenizer.encode(labels, self._device)

        # Encode the source sequence (i.e. the image codes)
        memory = self.encode(images)

        # Prepare the target sequences (input and output)
        tgt_perms = self.gen_tgt_perms(tgt)
        tgt_in = tgt[:, :-1]
        tgt_out = tgt[:, 1:]
        # The [EOS] token is not depended upon by any other token in any permutation ordering
        tgt_padding_mask = (tgt_in == self.pad_id) | (tgt_in == self.eos_id)

        loss = 0
        loss_numel = 0
        n = (tgt_out != self.pad_id).sum().item()
        for i, perm in enumerate(tgt_perms):
            tgt_mask, query_mask = self.generate_attn_masks(perm)
            out = self.decode(tgt_in, memory, tgt_mask, tgt_padding_mask, tgt_query_mask=query_mask)
            logits = self.head(out).flatten(end_dim=1)
            loss += n * F.cross_entropy(logits, tgt_out.flatten(), ignore_index=self.pad_id)
            loss_numel += n
            # After the second iteration (i.e. done with canonical and reverse orderings),
            # remove the [EOS] tokens for the succeeding perms
            if i == 1:
                tgt_out = torch.where(tgt_out == self.eos_id, self.pad_id, tgt_out)
                n = (tgt_out != self.pad_id).sum().item()
        loss /= loss_numel

        self.log('loss', loss)
        return loss

    def configure_optimizers(self):
        agb = self.trainer.accumulate_grad_batches
        # Linear scaling so that the effective learning rate is constant regardless of the number of GPUs used with DDP.
        lr_scale = agb * math.sqrt(self.trainer.num_devices) * self.batch_size / 256.
        lr = lr_scale * self.lr

        # optim = create_optimizer_v2(self, 'adamw', lr, self.weight_decay)
        # https://github.com/mlfoundations/open_clip/blob/b4cf9269b0b11c0eea47cb16039369a46bd67449/src/training/main.py#L171
        exclude = lambda n, p: p.ndim < 2 or "bn" in n or "ln" in n or "bias" in n or 'logit_scale' in n
        include = lambda n, p: not exclude(n, p)

        # encoder parameters
        if self.adapter_type in ["ladder", "ladder_pruning"]:
            encoder_params = list(self.adapter.named_parameters())
            enc_gain_or_bias_params = [p for n, p in encoder_params if exclude(n, p) and p.requires_grad]
            enc_rest_params = [p for n, p in encoder_params if include(n, p) and p.requires_grad]
        else:
            enc_gain_or_bias_params = []
            enc_rest_params = []

        # decoder and adapter parameters
        decoder_params = list(self.decoder.named_parameters())
        if self.adapter_type not in ["ladder", "ladder_pruning"]:
            decoder_params += list(self.adapter.named_parameters())
        dec_gain_or_bias_params = [p for n, p in decoder_params if exclude(n, p) and p.requires_grad]
        dec_rest_params = [p for n, p in decoder_params if include(n, p) and p.requires_grad]

        # following PARSEQ pos_queries, no weight decay
        dec_gain_or_bias_params.append(self.pos_queries)
        dec_gain_or_bias_params.append(self.head.bias)
        dec_gain_or_bias_params.append(self.text_embed.embedding.weight)
        dec_rest_params.append(self.head.weight)

        rank_zero_info("[STRAdapter] The length of encoder params with and without weight decay is {} and {}, respectively.".format(
            len(enc_rest_params), len(enc_gain_or_bias_params)
        ))
        rank_zero_info("[STRAdapter] The length of decoder params with and without weight decay is {} and {}, respectively.".format(
            len(dec_rest_params), len(dec_gain_or_bias_params)
        ))

        optimizer = torch.optim.AdamW(
            [
                {"params": enc_gain_or_bias_params, "weight_decay": 0., 'lr': lr},
                {"params": enc_rest_params, "weight_decay": self.weight_decay, 'lr': lr},
                {"params": dec_gain_or_bias_params, "weight_decay": 0., 'lr': lr * self.coef_lr},
                {"params": dec_rest_params, "weight_decay": self.weight_decay * self.coef_wd, 'lr': lr * self.coef_lr},
            ],
            lr=lr, betas=(0.9, 0.98), eps=1.0e-6,
            )
        sched = OneCycleLR(optimizer, [lr, lr, lr * self.coef_lr, lr * self.coef_lr],
                            self.trainer.estimated_stepping_batches, pct_start=self.warmup_pct,
                            cycle_momentum=False)

        return {'optimizer': optimizer, 'lr_scheduler': {'scheduler': sched, 'interval': 'step'}}

strhub/models/str_adapter/system.py
VamosC/CLIP4STR
from pathlib import PurePath
from typing import Sequence

import torch
from torch import nn

import yaml
from torchvision.ops.misc import FrozenBatchNorm2d


class InvalidModelError(RuntimeError):
    """Exception raised for any model-related error (creation, loading)"""


_WEIGHTS_URL = {
    'parseq-tiny': 'https://github.com/baudm/parseq/releases/download/v1.0.0/parseq_tiny-e7a21b54.pt',
    'parseq': 'https://github.com/baudm/parseq/releases/download/v1.0.0/parseq-bb5792a6.pt',
    'abinet': 'https://github.com/baudm/parseq/releases/download/v1.0.0/abinet-1d1e373e.pt',
    'trba': 'https://github.com/baudm/parseq/releases/download/v1.0.0/trba-cfaed284.pt',
    'vitstr': 'https://github.com/baudm/parseq/releases/download/v1.0.0/vitstr-26d0fcf4.pt',
    'crnn': 'https://github.com/baudm/parseq/releases/download/v1.0.0/crnn-679d0e31.pt',
}


def _get_config(experiment: str, **kwargs):
    """Emulates hydra config resolution"""
    root = PurePath(__file__).parents[2]
    with open(root / 'configs/main.yaml', 'r') as f:
        config = yaml.load(f, yaml.Loader)['model']
    with open(root / f'configs/charset/94_full.yaml', 'r') as f:
        config.update(yaml.load(f, yaml.Loader)['model'])
    with open(root / f'configs/experiment/{experiment}.yaml', 'r') as f:
        exp = yaml.load(f, yaml.Loader)
    # Apply base model config
    model = exp['defaults'][0]['override /model']
    with open(root / f'configs/model/{model}.yaml', 'r') as f:
        config.update(yaml.load(f, yaml.Loader))
    # Apply experiment config
    if 'model' in exp:
        config.update(exp['model'])
    config.update(kwargs)
    return config


def _get_model_class(key):
    exp = 'parseq'
    if 'abinet' in key:
        from .abinet.system import ABINet as ModelClass
        exp = 'abinet'
    elif 'crnn' in key:
        from .crnn.system import CRNN as ModelClass
        exp = 'crnn'
    elif 'parseq' in key:
        from .parseq.system import PARSeq as ModelClass
        exp = 'parseq'
    elif 'trba' in key:
        from .trba.system import TRBA as ModelClass
    elif 'trbc' in key:
        from .trba.system import TRBC as ModelClass
    elif 'vitstr' in key:
        from .vitstr.system import ViTSTR as ModelClass
        exp = 'vitstr'
    elif 'vl4str' in key or 'clip4str' in key:
        from .vl_str.system import VL4STR as ModelClass
        exp = 'vl4str'
        if 'large' in key:
            exp = 'vl4str-large'
        if 'base32x32' in key:
            exp = 'vl4str-base32'

    elif 'str_adapter' in key:
        from .str_adapter.system import STRAdapter as ModelClass
        exp = 'str_adapter'
    else:
        raise InvalidModelError("Unable to find model class for '{}'".format(key))

    return ModelClass, exp


def create_model(experiment: str, pretrained: bool = False, **kwargs):
    try:
        config = _get_config(experiment, **kwargs)
    except FileNotFoundError:
        raise InvalidModelError("No configuration found for '{}'".format(experiment)) from None
    ModelClass = _get_model_class(experiment)
    model = ModelClass(**config)
    if pretrained:
        try:
            url = _WEIGHTS_URL[experiment]
        except KeyError:
            raise InvalidModelError("No pretrained weights found for '{}'".format(experiment)) from None
        checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location='cpu', check_hash=True)
        model.load_state_dict(checkpoint)
    return model


def load_from_checkpoint(checkpoint_path: str, **kwargs):
    # if checkpoint_path.startswith('pretrained='):
    #     model_id = checkpoint_path.split('=', maxsplit=1)[1]
    #     model = create_model(model_id, True, **kwargs)
    # else:
    #     ModelClass, _ = _get_model_class(checkpoint_path)
    #     model = ModelClass.load_from_checkpoint(checkpoint_path, **kwargs)
    try:
        ModelClass, _ = _get_model_class(checkpoint_path)
        model = ModelClass.load_from_checkpoint(checkpoint_path, **kwargs)
    except:
        ModelClass, experiment = _get_model_class(checkpoint_path)
        try:
            config = _get_config(experiment, **kwargs)
        except FileNotFoundError:
            raise InvalidModelError("No configuration found for '{}'".format(experiment)) from None
        model = ModelClass(**config)
        checkpoint = torch.load(checkpoint_path, map_location='cpu')
        model.load_state_dict(checkpoint)

    return model


def parse_model_args(args):
    kwargs = {}
    arg_types = {t.__name__: t for t in [int, float, str]}
    arg_types['bool'] = lambda v: v.lower() == 'true'  # special handling for bool
    for arg in args:
        name, value = arg.split('=', maxsplit=1)
        name, arg_type = name.split(':', maxsplit=1)
        kwargs[name] = arg_types[arg_type](value)

    return kwargs


def init_weights(module: nn.Module, name: str = '', exclude: Sequence[str] = ()):
    """Initialize the weights using the typical initialization schemes used in SOTA models."""
    if any(map(name.startswith, exclude)):
        return
    if isinstance(module, nn.Linear):
        nn.init.trunc_normal_(module.weight, std=.02)
        if module.bias is not None:
            nn.init.zeros_(module.bias)
    elif isinstance(module, nn.Embedding):
        nn.init.trunc_normal_(module.weight, std=.02)
        if module.padding_idx is not None:
            module.weight.data[module.padding_idx].zero_()
    elif isinstance(module, nn.Conv2d):
        nn.init.kaiming_normal_(module.weight, mode='fan_out', nonlinearity='relu')
        if module.bias is not None:
            nn.init.zeros_(module.bias)
    elif isinstance(module, (nn.LayerNorm, nn.BatchNorm2d, nn.GroupNorm)):
        nn.init.ones_(module.weight)
        nn.init.zeros_(module.bias)


def freeze_batch_norm_2d(module, module_match={}, name=""):
    """
    Converts all `BatchNorm2d` and `SyncBatchNorm` layers of provided module into `FrozenBatchNorm2d`. If `module` is
    itself an instance of either `BatchNorm2d` or `SyncBatchNorm`, it is converted into `FrozenBatchNorm2d` and
    returned. Otherwise, the module is walked recursively and submodules are converted in place.
    Args:
        module (torch.nn.Module): Any PyTorch module.
        module_match (dict): Dictionary of full module names to freeze (all if empty)
        name (str): Full module name (prefix)
    Returns:
        torch.nn.Module: Resulting module
    Inspired by https://github.com/pytorch/pytorch/blob/a5895f85be0f10212791145bfedc0261d364f103/torch/nn/modules/batchnorm.py#L762
    """
    res = module
    is_match = True
    if module_match:
        is_match = name in module_match
    if is_match and isinstance(
        module, (nn.modules.batchnorm.BatchNorm2d, nn.modules.batchnorm.SyncBatchNorm)
    ):
        res = FrozenBatchNorm2d(module.num_features)
        res.num_features = module.num_features
        res.affine = module.affine
        if module.affine:
            res.weight.data = module.weight.data.clone().detach()
            res.bias.data = module.bias.data.clone().detach()
        res.running_mean.data = module.running_mean.data
        res.running_var.data = module.running_var.data
        res.eps = module.eps
    else:
        for child_name, child in module.named_children():
            full_child_name = ".".join([name, child_name]) if name else child_name
            new_child = freeze_batch_norm_2d(child, module_match, full_child_name)
            if new_child is not child:
                res.add_module(child_name, new_child)
    return res


if __name__ == "__main__":
    import numpy as np

    # define a sequence of 10 words over a vocab of 5 words
    data = [[0.1, 0.2, 0.3, 0.4, 0.5],
    [0.5, 0.4, 0.3, 0.2, 0.1],
    [0.1, 0.2, 0.3, 0.4, 0.5],
    [0.5, 0.4, 0.3, 0.2, 0.1],
    [0.1, 0.2, 0.3, 0.4, 0.5],
    [0.5, 0.4, 0.3, 0.2, 0.1],
    [0.1, 0.2, 0.3, 0.4, 0.5],
    [0.5, 0.4, 0.3, 0.2, 0.1],
    [0.1, 0.2, 0.3, 0.4, 0.5],
    [0.5, 0.4, 0.3, 0.2, 0.1]]
    data = np.array(data)
    # # decode sequence
    # result = beam_search_decoder(data, 3)
    # # print result
    # for seq in result:
    #     print(seq)

strhub/models/utils.py
VamosC/CLIP4STR
"""
Implementation of ViTSTR based on timm VisionTransformer.

TODO:
1) distilled deit backbone
2) base deit backbone

Copyright 2021 Rowel Atienza
"""
import os
import torch
import torch.nn as nn
from pytorch_lightning.utilities import rank_zero_info
from timm.models.vision_transformer import VisionTransformer


class ViTSTR(VisionTransformer):
    """
    ViTSTR is basically a ViT that uses DeiT weights.
    Modified head to support a sequence of characters prediction for STR.
    """

    def forward(self, x, seqlen: int = 25):
        x = self.forward_features(x)
        x = x[:, :seqlen]

        # batch, seqlen, embsize
        b, s, e = x.size()
        x = x.reshape(b * s, e)
        x = self.head(x).view(b, s, self.num_classes)
        return x


def load_pretrained_weight(model, pretrained):
    if not os.path.exists(pretrained):
        rank_zero_info("[ViT] pretrained weight ({}) does not exist".format(pretrained))
    rank_zero_info("[ViT] Load pretrained weights from ImageNet-21K pretrained model ({})".format(pretrained))
    state_dict = torch.load(pretrained, map_location='cpu')
    if 'model' in state_dict:
        state_dict = state_dict['model']
    if 'state_dict' in state_dict:
        state_dict = state_dict['state_dict']

    # diff shapes
    # pos_embed torch.Size([1, 129, 768])
    # pos_embed torch.Size([1, 197, 768])
    # patch_embed.proj.weight torch.Size([768, 3, 4, 8])
    # patch_embed.proj.weight torch.Size([768, 3, 16, 16])
    # head.weight torch.Size([25, 768])
    # head.weight torch.Size([21843, 768])
    # head.bias torch.Size([25])
    # head.bias torch.Size([21843])
    # convert to fp32
    new_state_dict = {}
    for keys, value in state_dict.items():
        new_state_dict[keys] = value.float()

    # position embedding
    old_pe = new_state_dict['pos_embed']
    new_pe_len = model.pos_embed.shape[1]
    new_state_dict['pos_embed'] = old_pe[:, :new_pe_len, ...]

    # first conv
    old_conv1_w = new_state_dict['patch_embed.proj.weight']
    new_conv1_w_shape = model.patch_embed.proj.weight.shape
    if old_conv1_w.shape != new_conv1_w_shape:
        rank_zero_info("[ViT] averageing conv1 weight (patch_embed.proj) of ViT")
        # average across [H, W]
        kernel_size = (old_conv1_w.shape[-2] // new_conv1_w_shape[-2], old_conv1_w.shape[-1] // new_conv1_w_shape[-1])
        new_state_dict['patch_embed.proj.weight'] = nn.AvgPool2d(kernel_size, stride=kernel_size)(old_conv1_w)

    # head
    old_proj_w = new_state_dict['head.weight']
    old_proj_bias = new_state_dict['head.bias']
    if not isinstance(model.head, nn.Identity):
        new_output_dim = model.head.weight.shape[0]
        if new_output_dim != old_proj_w.shape[1]:
            rank_zero_info("[ViT] downsmaple the output head of ViT")
            new_state_dict['head.weight'] = nn.AdaptiveAvgPool1d(new_output_dim)(old_proj_w.transpose(0, 1)).transpose(0,1)
            if hasattr(model.head, 'bias'):
                new_state_dict['head.bias'] = nn.AdaptiveAvgPool1d(new_output_dim)(old_proj_bias.unsqueeze(0)).squeeze()
    else:
        del new_state_dict['head.weight']
        del new_state_dict['head.bias']

    # cls token
    if model.cls_token is None:
        del new_state_dict['cls_token']

    # load
    missing_keys, unexpected_keys = model.load_state_dict(new_state_dict, strict=False)
    print("missing_keys: ", missing_keys)
    print("unexpected_keys: ", unexpected_keys)

# model1 = VisionTransformer(img_size=224, patch_size=16, depth=12, mlp_ratio=4, qkv_bias=True,
#                            embed_dim=768, num_heads=12, num_classes=21843)
# model1.load_pretrained("B_16-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.0-sd_0.0.npz")
# sd1 = model1.state_dict()
# torch.save(sd1, "B_16-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.0-sd_0.0.pth")

strhub/models/vitstr/model.py
VamosC/CLIP4STR
# Scene Text Recognition Model Hub
# Copyright 2022 Darwin Bautista
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Sequence, Any, Optional

import torch
from pytorch_lightning.utilities.types import STEP_OUTPUT
from pytorch_lightning.utilities import rank_zero_info
from torch import Tensor

from strhub.models.base import CrossEntropySystem
from strhub.models.utils import init_weights
from .model import ViTSTR as Model, load_pretrained_weight


class ViTSTR(CrossEntropySystem):

    def __init__(self, charset_train: str, charset_test: str, max_label_length: int,
                 batch_size: int, lr: float, warmup_pct: float, weight_decay: float,
                 img_size: Sequence[int], patch_size: Sequence[int], embed_dim: int, num_heads: int,
                 **kwargs: Any) -> None:
        super().__init__(charset_train, charset_test, batch_size, lr, warmup_pct, weight_decay)
        self.save_hyperparameters()
        self.max_label_length = max_label_length
        # We don't predict <bos> nor <pad>
        self.model = Model(img_size=img_size, patch_size=patch_size, depth=12, mlp_ratio=4, qkv_bias=True,
                           embed_dim=embed_dim, num_heads=num_heads, num_classes=len(self.tokenizer) - 2)
        # Non-zero weight init for the head
        self.model.head.apply(init_weights)
        # load pretrained weights
        load_pretrained_weight(self.model, kwargs["imagenet21k_pretrained"])

    @torch.jit.ignore
    def no_weight_decay(self):
        return {'model.' + n for n in self.model.no_weight_decay()}

    def forward(self, images: Tensor, max_length: Optional[int] = None) -> Tensor:
        max_length = self.max_label_length if max_length is None else min(max_length, self.max_label_length)
        logits = self.model.forward(images, max_length + 2)  # +2 tokens for [GO] and [s]
        # Truncate to conform to other models. [GO] in ViTSTR is actually used as the padding (therefore, ignored).
        # First position corresponds to the class token, which is unused and ignored in the original work.
        logits = logits[:, 1:]
        return logits

    def training_step(self, batch, batch_idx) -> STEP_OUTPUT:
        images, labels = batch
        loss = self.forward_logits_loss(images, labels)[1]
        self.log('loss', loss)
        return loss

strhub/models/vitstr/system.py
VamosC/CLIP4STR
"""
 Copyright (c) 2022, salesforce.com, inc.
 All rights reserved.
 SPDX-License-Identifier: BSD-3-Clause
 For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
"""
import torch
import torch.distributed.nn
from torch import distributed as dist, nn as nn
from torch.nn import functional as F


def gather_features(
    image_features,
    text_features,
    local_loss=False,
    gather_with_grad=False,
    rank=0,
    world_size=1,
    use_horovod=False,
):
    if use_horovod:
        raise NotImplementedError
    else:
        # We gather tensors from all gpus
        if gather_with_grad:
            all_image_features = torch.cat(
                torch.distributed.nn.all_gather(image_features), dim=0
            )
            all_text_features = torch.cat(
                torch.distributed.nn.all_gather(text_features), dim=0
            )
        else:
            gathered_image_features = [
                torch.zeros_like(image_features) for _ in range(world_size)
            ]
            gathered_text_features = [
                torch.zeros_like(text_features) for _ in range(world_size)
            ]
            dist.all_gather(gathered_image_features, image_features)
            dist.all_gather(gathered_text_features, text_features)
            if not local_loss:
                # ensure grads for local rank when all_* features don't have a gradient
                gathered_image_features[rank] = image_features
                gathered_text_features[rank] = text_features
            all_image_features = torch.cat(gathered_image_features, dim=0)
            all_text_features = torch.cat(gathered_text_features, dim=0)

    return all_image_features, all_text_features


class ClipLoss(nn.Module):
    def __init__(
        self,
        local_loss=False,
        gather_with_grad=False,
        cache_labels=False,
        rank=0,
        world_size=1,
        use_horovod=False,
    ):
        super().__init__()
        self.local_loss = local_loss
        self.gather_with_grad = gather_with_grad
        self.cache_labels = cache_labels
        self.rank = rank
        self.world_size = world_size
        self.use_horovod = use_horovod

        # cache state
        self.prev_num_logits = 0
        self.labels = {}

    def forward(self, image_features, text_features, logit_scale):
        device = image_features.device
        if self.world_size > 1:
            all_image_features, all_text_features = gather_features(
                image_features,
                text_features,
                self.local_loss,
                self.gather_with_grad,
                self.rank,
                self.world_size,
                self.use_horovod,
            )

            if self.local_loss:
                logits_per_image = logit_scale * image_features @ all_text_features.T
                logits_per_text = logit_scale * text_features @ all_image_features.T
            else:
                logits_per_image = (
                    logit_scale * all_image_features @ all_text_features.T
                )
                logits_per_text = logits_per_image.T
        else:
            logits_per_image = logit_scale * image_features @ text_features.T
            logits_per_text = logit_scale * text_features @ image_features.T

        # calculated ground-truth and cache if enabled
        num_logits = logits_per_image.shape[0]
        if self.prev_num_logits != num_logits or device not in self.labels:
            labels = torch.arange(num_logits, device=device, dtype=torch.long)
            if self.world_size > 1 and self.local_loss:
                labels = labels + num_logits * self.rank
            if self.cache_labels:
                self.labels[device] = labels
                self.prev_num_logits = num_logits
        else:
            labels = self.labels[device]

        total_loss = (
            F.cross_entropy(logits_per_image, labels)
            + F.cross_entropy(logits_per_text, labels)
        ) / 2
        return total_loss
strhub/models/vl_str/loss.py
VamosC/CLIP4STR
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
import math
from functools import partial
from typing import Optional, Sequence, Callable

import torch
from torch import nn as nn, Tensor
from torch.nn import functional as F
from torch.nn.modules import transformer
from timm.models.helpers import named_apply
from strhub.models.utils import init_weights


class DecoderLayer(nn.Module):
    """A Transformer decoder layer supporting two-stream attention (XLNet)
       This implements a pre-LN decoder, as opposed to the post-LN default in PyTorch."""

    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, activation='gelu',
                 layer_norm_eps=1e-5):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=True)
        self.cross_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=True)
        # Implementation of Feedforward model
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)

        self.norm1 = nn.LayerNorm(d_model, eps=layer_norm_eps)
        self.norm2 = nn.LayerNorm(d_model, eps=layer_norm_eps)
        self.norm_q = nn.LayerNorm(d_model, eps=layer_norm_eps)
        self.norm_c = nn.LayerNorm(d_model, eps=layer_norm_eps)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.dropout3 = nn.Dropout(dropout)

        self.activation = transformer._get_activation_fn(activation)

    def __setstate__(self, state):
        if 'activation' not in state:
            state['activation'] = F.gelu
        super().__setstate__(state)

    def forward_stream(self, tgt: Tensor, tgt_norm: Tensor, tgt_kv: Tensor, memory: Tensor, attn_mask: Optional[Tensor],
                       key_padding_mask: Optional[Tensor]):
        """Forward pass for a single stream (i.e. content or query)
        tgt_norm is just a LayerNorm'd tgt. Added as a separate parameter for efficiency.
        Both tgt_kv and memory are expected to be LayerNorm'd too.
        memory is LayerNorm'd by ViT.
        """
        tgt2, sa_weights = self.self_attn(tgt_norm, tgt_kv, tgt_kv, attn_mask=attn_mask, key_padding_mask=key_padding_mask)
        tgt = tgt + self.dropout1(tgt2)

        tgt2, ca_weights = self.cross_attn(self.norm1(tgt), memory, memory)
        tgt = tgt + self.dropout2(tgt2)

        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(self.norm2(tgt)))))
        tgt = tgt + self.dropout3(tgt2)
        return tgt, sa_weights, ca_weights

    def forward(self, query, content, memory, tgt_query_mask: Optional[Tensor] = None, content_mask: Optional[Tensor] = None,
                content_key_padding_mask: Optional[Tensor] = None, update_content: bool = True):
        query_norm = self.norm_q(query)
        content_norm = self.norm_c(content)
        query = self.forward_stream(query, query_norm, content_norm, memory, tgt_query_mask, content_key_padding_mask)[0]
        if update_content:
            content = self.forward_stream(content, content_norm, content_norm, memory, content_mask, content_key_padding_mask)[0]
        return query, content


class Decoder(nn.Module):
    __constants__ = ['norm']

    def __init__(self, decoder_layer, num_layers, norm, embed_dim=512, dropout=0.0, num_classes=94,
                        charset_size=94, max_label_length=25):
        """a self-contained decoder for character extraction"""
        super().__init__()
        self.layers = transformer._get_clones(decoder_layer, num_layers)
        self.text_embed = TokenEmbedding(charset_size, embed_dim)
        self.pos_queries = nn.Parameter(torch.Tensor(1, max_label_length, embed_dim))
        self.num_layers = num_layers
        self.norm = norm
        self.dropout = nn.Dropout(p=dropout)
        self.head = nn.Linear(embed_dim, num_classes, bias=True)

        named_apply(partial(init_weights, exclude=['none']), self)
        nn.init.trunc_normal_(self.pos_queries, std=.02)

    def forward(self, tgt, memory,
                tgt_query: Optional[Tensor] = None,
                tgt_query_mask: Optional[Tensor] = None,
                content_mask: Optional[Tensor] = None,
                content_key_padding_mask: Optional[Tensor] = None):
        N, L = tgt.shape
        # <bos> stands for the null context. We only supply position information for characters after <bos>.
        null_ctx = self.text_embed(tgt[:, :1])
        tgt_emb = self.pos_queries[:, :L - 1] + self.text_embed(tgt[:, 1:])
        content = self.dropout(torch.cat([null_ctx, tgt_emb], dim=1))
        if tgt_query is None:
            tgt_query = self.pos_queries[:, :L].expand(N, -1, -1)
        query = self.dropout(tgt_query)

        # forward layers
        for i, mod in enumerate(self.layers):
            last = i == len(self.layers) - 1
            query, content = mod(query, content, memory, tgt_query_mask, content_mask, content_key_padding_mask,
                                 update_content=not last)
        query = self.norm(query)

        # prediction
        logits = self.head(query)

        # return prediction and feature
        return logits, query


class TokenEmbedding(nn.Module):

    def __init__(self, charset_size: int, embed_dim: int):
        super().__init__()
        self.embedding = nn.Embedding(charset_size, embed_dim)
        self.embed_dim = embed_dim

    def forward(self, tokens: torch.Tensor):
        return math.sqrt(self.embed_dim) * self.embedding(tokens)


class Hook():
    # A simple hook class that returns the input and output of a layer during forward/backward pass
    def __init__(self, module, backward=False):
        if backward == False:
            self.hook = module.register_forward_hook(self.hook_fn)
        else:
            self.hook = module.register_backward_hook(self.hook_fn)
    
    def hook_fn(self, module, input, output):
        self.input = input
        self.output = output

    def close(self):
        self.hook.remove()


def modify_attn_mask(m, context_length=10):
    if hasattr(m, "attn_mask"):
        if m.attn_mask is not None:
            m.attn_mask = m.attn_mask[:context_length, :context_length]


class FusionFC(nn.Module):
    def __init__(self, d_model, num_classes):
        super().__init__()
        self.w_att = nn.Linear(2 * d_model, d_model)
        self.cls = nn.Linear(d_model, num_classes)
        # init
        for m in self.modules():
            init_weights(m)

    def forward(self, feature1, feature2, detach=True):
        """
        Args:
            feature1: (N, T, E) where T is length, N is batch size and d is dim of model
            feature2: (N, T, E) shape the same as l_feature 
        """
        if detach:
            feature1 = feature1.detach().clone()
            feature2 = feature2.detach().clone()
        f = torch.cat((feature1, feature2), dim=2)
        f_att = torch.sigmoid(self.w_att(f))
        output = f_att * feature2 + (1.0 - f_att) * feature1
        # (N, T, C)
        logits = self.cls(output)

        return logits

strhub/models/vl_str/modules.py
VamosC/CLIP4STR
# coding=utf-8

import torch

def generate_attn_masks(perm):
    """Generate attention masks given a sequence permutation (includes pos. for bos and eos tokens)
    :param perm: the permutation sequence. i = 0 is always the BOS
    :return: lookahead attention masks
    """
    sz = perm.shape[0]
    mask = torch.zeros((sz, sz))
    for i in range(sz):
        query_idx = perm[i]
        masked_keys = perm[i + 1:]
        mask[query_idx, masked_keys] = float('-inf')
    content_mask = mask[:-1, :-1].clone()
    mask[torch.eye(sz, dtype=torch.bool)] = float('-inf')  # mask "self"
    query_mask = mask[1:, :-1]
    return content_mask, query_mask

# perms = torch.tensor([[0, 1, 2, 3, 4, 5, 6],
#                         [0, 6, 5, 4, 3, 2, 1],
#                         [0, 3, 1, 5, 2, 4, 6],
#                         [0, 1, 4, 2, 3, 5, 6],
#                         [0, 2, 3, 5, 1, 4, 6],
#                         [0, 2, 1, 3, 5, 4, 6]])

# for perm in perms:
#     sz = perm.shape[0]
#     mask = torch.zeros((sz, sz))

#     for i in range(sz):
#         q_idx = perm[i]
#         masked_keys = perm[i + 1:]
#         mask[q_idx, masked_keys] = float('-inf')
#     print(mask)

L= 7
mask = torch.triu(torch.full((L, L), float('-inf')), 1)
content_mask = mask[:-1, :-1].clone()
strhub/models/vl_str/perm_test.py
VamosC/CLIP4STR
# coding=utf-8
import os
import math
import numpy as np
from itertools import permutations
from typing import Sequence, Any, Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from torch.optim.lr_scheduler import OneCycleLR
from pytorch_lightning.utilities import rank_zero_info
from pytorch_lightning.utilities.types import STEP_OUTPUT

from strhub.clip import clip
from strhub.models.base import CrossEntropySystem
from .modules import DecoderLayer, Decoder, modify_attn_mask


# an alternative choice when the input argument is not valid 
CLIP_PATH = '/PUT/YOUR/PATH/HERE/pretrained/clip'


class VL4STR(CrossEntropySystem):

    def __init__(self, charset_train: str, charset_test: str, max_label_length: int,
                 batch_size: int, lr: float, warmup_pct: float, weight_decay: float,
                 img_size: Sequence[int], patch_size: Sequence[int], embed_dim: int,
                 enc_num_heads: int, enc_mlp_ratio: int, enc_depth: int,
                 dec_num_heads: int, dec_mlp_ratio: int, dec_depth: int,
                 perm_num: int, perm_forward: bool, perm_mirrored: bool,
                 decode_ar: bool, refine_iters: int, dropout: float, **kwargs: Any) -> None:
        super().__init__(charset_train, charset_test, batch_size, lr, warmup_pct, weight_decay)
        self.save_hyperparameters()

        self.max_label_length = max_label_length
        self.decode_ar = decode_ar
        self.refine_iters = refine_iters
        self.coef_lr = kwargs["coef_lr"] if "coef_lr" in kwargs.keys() else 1.0
        self.coef_wd = kwargs["coef_wd"] if "coef_wd" in kwargs.keys() else 1.0
        self.image_freeze_nlayer = kwargs["image_freeze_nlayer"] if "image_freeze_nlayer" in kwargs.keys() else -1
        self.text_freeze_nlayer = kwargs["text_freeze_nlayer"] if "text_freeze_nlayer" in kwargs.keys() else -1
        self.freeze_language_backbone = self.text_freeze_nlayer >= 12
        self.freeze_image_backbone = self.image_freeze_nlayer >= 12
        self.use_language_model = kwargs["use_language_model"] if "use_language_model" in kwargs.keys() else False
        self.context_length = kwargs["context_length"] if "context_length" in kwargs.keys() else 20
        self.cross_loss_w = kwargs["cross_loss_w"] if "cross_loss_w" in kwargs.keys() else 1.0
        self.use_share_dim = kwargs["use_share_dim"] if "use_share_dim" in kwargs.keys() else True
        self.cross_gt_context = kwargs["cross_gt_context"] if "cross_gt_context" in kwargs.keys() else True
        self.cross_cloze_mask = kwargs["cross_cloze_mask"] if "cross_cloze_mask" in kwargs.keys() else False
        self.cross_correct_once = kwargs["cross_correct_once"] if "cross_correct_once" in kwargs.keys() else False
        self.image_detach = kwargs["image_detach"] if "image_detach" in kwargs.keys() else True
        self.cross_token_embeding = kwargs["cross_token_embeding"] if "cross_token_embeding" in kwargs.keys() else False
        self.cross_fast_decode = False

        rank_zero_info("\n config of VL4STR: \n"
                "\t image_freeze_nlayer: {}, text_freeze_nlayer: {}, freeze_language_backbone: {}, freeze_image_backbone: {} \n"
                "\t use_language_model: {}, context_length: {}, cross_token_embeding: {}, cross_loss_weight: {} \n"
                "\t use_share_dim: {}, image_detach: {} \n"
                "\t cross_gt_context: {}, cross_cloze_mask: {}, cross_fast_decode: {} \n".format(
                self.image_freeze_nlayer, self.text_freeze_nlayer, self.freeze_language_backbone, self.freeze_image_backbone,
                self.use_language_model, self.context_length,  self.cross_token_embeding, self.cross_loss_w,
                self.use_share_dim, self.image_detach, self.cross_gt_context, self.cross_cloze_mask, self.cross_fast_decode)
                )

        assert "clip_pretrained" in kwargs.keys()
        if not os.path.exists(kwargs["clip_pretrained"]):
            kwargs["clip_pretrained"] = os.path.join(CLIP_PATH, os.path.basename(kwargs["clip_pretrained"]))
            assert os.path.exists(kwargs["clip_pretrained"])
        # load CLIP model
        clip_model, _ = clip.load(name=kwargs["clip_pretrained"], device='cpu')
        self.clip_model = clip_model.float()

        # modify the attention mask according to context length
        self.clip_model.transformer.apply(lambda m: modify_attn_mask(m, context_length=self.context_length))
        self.freeze_cip_layers(self.image_freeze_nlayer, self.text_freeze_nlayer)

        # visual deoder
        vis_embed_dim = self.clip_model.text_projection.shape[-1] if self.use_share_dim else self.clip_model.visual.proj.shape[0]
        rank_zero_info("The dimension of the visual decoder is {}.".format(vis_embed_dim))
        decoder_layer = DecoderLayer(vis_embed_dim, dec_num_heads, vis_embed_dim * dec_mlp_ratio, dropout)
        # We don't predict <bos> nor <pad>, so num_classes=len(self.tokenizer) - 2, label length + 1 for <eos>
        self.visual_decoder = Decoder(decoder_layer, dec_depth, norm=nn.LayerNorm(vis_embed_dim),
                                        embed_dim=vis_embed_dim, 
                                        dropout=dropout,
                                        num_classes=len(self.tokenizer) - 2,
                                        charset_size=len(self.tokenizer),
                                        max_label_length=max_label_length + 1)
        # cross-modal decoder
        if self.use_language_model:
            cross_embed_dim = self.clip_model.text_projection.shape[-1]
            decoder_layer = DecoderLayer(cross_embed_dim, dec_num_heads, cross_embed_dim * dec_mlp_ratio, dropout)
            self.cross_decoder = Decoder(decoder_layer, dec_depth, norm=nn.LayerNorm(cross_embed_dim),
                                        embed_dim=cross_embed_dim,
                                        dropout=dropout,
                                        num_classes=len(self.tokenizer) - 2,
                                        charset_size=len(self.tokenizer),
                                        max_label_length=max_label_length + 1)
        # Perm/attn mask stuff
        self.rng = np.random.default_rng()
        self.max_gen_perms = perm_num // 2 if perm_mirrored else perm_num
        self.perm_forward = perm_forward
        self.perm_mirrored = perm_mirrored

    def encode(self, img: torch.Tensor):
        """extract CLIP image features"""
        if self.freeze_image_backbone:
            with torch.no_grad():
                memory = self.clip_model.encode_image(img, projection=self.use_share_dim)
        else:
            memory = self.clip_model.encode_image(img, projection=self.use_share_dim)

        return memory

    def visual_decode(self, tgt: torch.Tensor, memory: torch.Tensor,
                        tgt_query: Optional[Tensor] = None, tgt_query_mask: Optional[Tensor] = None,
                        content_mask: Optional[Tensor] = None, tgt_padding_mask: Optional[Tensor] = None, ):
        return self.visual_decoder(tgt, memory, tgt_query, tgt_query_mask, content_mask, tgt_padding_mask)

    def encoder_cross_modal_feature(self, prev_logits, image_feat):
        prev_logits = prev_logits.detach().clone()
        image_features = image_feat.detach().clone() if self.image_detach else image_feat
        if not self.use_share_dim:
            image_features = torch.matmul(image_features, self.clip_model.visual.proj)

        # get previous predictions
        probs = prev_logits.softmax(-1)
        # adapt for the test charset, CLIP is not sensitive to uppercase or symbols
        captions, _ = self.tokenizer.decode_fast(probs, charset_adapter=self.charset_adapter)
        text = clip.tokenize(captions, context_length=self.context_length, truncate=True).to(image_feat.device)

        # return all text features
        if self.freeze_language_backbone:
            with torch.no_grad():
                text_features = self.clip_model.token_embedding(text) if self.cross_token_embeding else \
                                    self.clip_model.encode_text(text)
        else:
            text_features = self.clip_model.token_embedding(text) if self.cross_token_embeding else \
                                    self.clip_model.encode_text(text)        

        return torch.cat([image_features, text_features], dim=1)

    def cross_decode(self, prev_logits, tgt: torch.Tensor, memory: torch.Tensor,
                        tgt_query: Optional[Tensor] = None, tgt_query_mask: Optional[Tensor] = None, 
                        content_mask: Optional[Tensor] = None, tgt_padding_mask: Optional[Tensor] = None,
                        cross_memory = None):
        if cross_memory is None:
            cross_memory = self.encoder_cross_modal_feature(prev_logits, memory)

        return self.cross_decoder(tgt, cross_memory, tgt_query, tgt_query_mask, content_mask, tgt_padding_mask)

    def forward(self, images: Tensor, max_length: Optional[int] = None) -> Tensor:
        testing = max_length is None
        max_length = self.max_label_length if max_length is None else min(max_length, self.max_label_length)
        bs = images.shape[0]
        # +1 for <eos> at end of sequence.
        num_steps = max_length + 1
        memory = self.encode(images)

        # Query positions up to `num_steps`
        vis_pos_queries = self.visual_decoder.pos_queries[:, :num_steps].expand(bs, -1, -1)
        crs_pos_queries = self.cross_decoder.pos_queries[:, :num_steps].expand(bs, -1, -1) if self.use_language_model else None

        # a left-to-right auto-regressive mask, special case for the forward permutation
        content_mask = query_mask = torch.triu(torch.full((num_steps, num_steps), float('-inf'), device=self._device), 1)
        bos = torch.full((bs, 1), self.bos_id, dtype=torch.long, device=self._device)

        if self.decode_ar:
            tgt_in = torch.full((bs, num_steps), self.pad_id, dtype=torch.long, device=self._device)
            tgt_in[:, 0] = self.bos_id

            logits = []
            all_visual_vec = []
            for i in range(num_steps):
                j = i + 1  # next token index
                # Efficient decoding: Input the context up to the ith token. We use only one query (at position = i) at a time.
                # This works because of the lookahead masking effect of the canonical (forward) AR context.
                # Past tokens have no access to future tokens, hence are fixed once computed.
                p_i, visual_vec = self.visual_decode(tgt_in[:, :j], memory, tgt_query=vis_pos_queries[:, i:j],
                                            tgt_query_mask=query_mask[i:j, :j], content_mask=content_mask[:j, :j],)
                # the next token probability is in the output's ith token position
                logits.append(p_i)
                all_visual_vec.append(visual_vec.clone())
                if j < num_steps:
                    # greedy decode. add the next token index to the target input
                    tgt_in[:, j] = p_i.squeeze().argmax(-1)
                    if testing and (tgt_in == self.eos_id).any(dim=-1).all():
                        break
            logits = torch.cat(logits, dim=1)
            visual_vec = torch.cat(all_visual_vec, dim=1)

        else:
            # No prior context, so input is just <bos>. We query all positions.
            # tgt_in = torch.full((bs, 1), self.bos_id, dtype=torch.long, device=self._device)
            logits, visual_vec = self.visual_decode(bos, memory, tgt_query=vis_pos_queries)

        if self.use_language_model:
            crs_num_steps = logits.shape[1]
            if self.cross_fast_decode:
                # just use visual output as input context
                cross_logits, cross_vec = self.cross_decode(logits, tgt_in[:, :crs_num_steps], memory, tgt_query=crs_pos_queries[:, :crs_num_steps],
                                            tgt_query_mask=query_mask[:crs_num_steps, :crs_num_steps],
                                            content_mask=content_mask[:crs_num_steps, :crs_num_steps],)
            else:
                # prediction of cross-modal branch as input context
                cross_memory = self.encoder_cross_modal_feature(logits, memory)
                cross_logits = []
                all_cross_vec = []
                for i in range(crs_num_steps):
                    j = i + 1  # next token index
                    p_i, cross_vec = self.cross_decode(logits, tgt_in[:, :j], memory, tgt_query=crs_pos_queries[:, i:j],
                                                tgt_query_mask=query_mask[i:j, :j], content_mask=content_mask[:j, :j], cross_memory=cross_memory)                
                    cross_logits.append(p_i)
                    all_cross_vec.append(cross_vec.clone())
                    if j < crs_num_steps:
                        tgt_in[:, j] = p_i.squeeze().argmax(-1)
                cross_logits = torch.cat(cross_logits, dim=1)
                cross_vec = torch.cat(all_cross_vec, dim=1)

        if self.refine_iters:
            # For iterative refinement, we always use a 'cloze' mask.
            # We can derive it from the AR forward mask by unmasking the token context to the right.
            query_mask[torch.triu(torch.ones(num_steps, num_steps, dtype=torch.bool, device=self._device), 2)] = 0
            bos = torch.full((bs, 1), self.bos_id, dtype=torch.long, device=self._device)
            for i in range(self.refine_iters):
                # Prior context is the previous output.
                tgt_in = torch.cat([bos, logits[:, :-1].argmax(-1)], dim=1)
                tgt_padding_mask = ((tgt_in == self.eos_id).cumsum(-1) > 0)  # mask tokens beyond the first EOS token.
                logits, visual_vec = self.visual_decode(tgt_in, memory,
                                                    tgt_query=vis_pos_queries, tgt_query_mask=query_mask[:, :tgt_in.shape[1]],
                                                    content_mask=content_mask, tgt_padding_mask=tgt_padding_mask,)
                if self.use_language_model:
                    tgt_in = torch.cat([bos, cross_logits[:, :-1].argmax(-1)], dim=1)
                    tgt_padding_mask = ((tgt_in == self.eos_id).cumsum(-1) > 0)
                    cross_logits, cross_vec = self.cross_decode(logits, tgt_in, memory,
                                                    tgt_query=crs_pos_queries, tgt_query_mask=query_mask[:, :tgt_in.shape[1]],
                                                    content_mask=content_mask, tgt_padding_mask=tgt_padding_mask,)

        # TODO: how to fuse the final predictions
        logits = cross_logits
        return logits

    def training_step(self, batch, batch_idx) -> STEP_OUTPUT:
        images, labels = batch
        tgt = self.tokenizer.encode(labels, self._device)

        # Encode the source sequence (i.e. the image codes)
        memory = self.encode(images)

        # Prepare the target sequences (input and output)
        tgt_perms = self.gen_tgt_perms(tgt)
        tgt_in = tgt[:, :-1]        # remove [EOS] token
        tgt_out = tgt[:, 1:]        # remove [BOS] token
        # cross_target = tgt[:, 1:]
        # max_len = min(cross_target.shape[1] - 1, self.max_label_length)
        # The [EOS] token is not depended upon by any other token in any permutation ordering
        tgt_padding_mask = (tgt_in == self.pad_id) | (tgt_in == self.eos_id)
        bs = images.shape[0]
        bos = torch.full((bs, 1), self.bos_id, dtype=torch.long, device=self._device)
        # try to use a cloze_mask for cross decoding
        L = tgt_in.shape[1]
        cloze_content_mask = cloze_query_mask = torch.triu(torch.full((L, L), float('-inf'), device=self._device), 1)
        cloze_query_mask[torch.triu(torch.ones(L, L, dtype=torch.bool, device=self._device), 2)] = 0

        loss = 0
        loss_numel = 0
        all_inter_vars = []
        n = (tgt_out != self.pad_id).sum().item()
        for i, perm in enumerate(tgt_perms):
            content_mask, query_mask = self.generate_attn_masks(perm)
            visual_logits, visual_vec = self.visual_decode(tgt_in, memory, tgt_query_mask=query_mask,
                                                content_mask=content_mask, tgt_padding_mask=tgt_padding_mask, )
            loss += n * F.cross_entropy(visual_logits.flatten(end_dim=1), tgt_out.flatten(), ignore_index=self.pad_id)
            all_inter_vars.append((visual_logits.detach().clone(), query_mask, content_mask))

            # cross forward
            if self.use_language_model:
                # 1. use the prediction of visual branch as context, keep the shape of cross_tgt_in the same as `tgt_in`
                # 2. use the GT as the context
                cross_tgt_in = tgt_in if self.cross_gt_context else torch.cat([bos, visual_logits[:, :-1].argmax(-1)], dim=1)
                cross_query_mask = cloze_query_mask if self.cross_cloze_mask else query_mask
                cross_content_mask = cloze_content_mask if self.cross_cloze_mask else content_mask
                cross_logits, cross_vec = self.cross_decode(visual_logits, cross_tgt_in, memory, tgt_query_mask=cross_query_mask,
                                                content_mask=cross_content_mask, tgt_padding_mask=tgt_padding_mask, )
                loss += self.cross_loss_w * n * F.cross_entropy(cross_logits.flatten(end_dim=1), tgt_out.flatten(), ignore_index=self.pad_id)

            loss_numel += n
            # After the second iteration (i.e. done with canonical and reverse orderings),
            # remove the [EOS] tokens for the succeeding perms
            if i == 1:
                tgt_out = torch.where(tgt_out == self.eos_id, self.pad_id, tgt_out)
                n = (tgt_out != self.pad_id).sum().item()

        loss /= loss_numel

        self.log('loss', loss)
        return loss

    def gen_tgt_perms(self, tgt):
        """Generate shared permutations for the whole batch.
           This works because the same attention mask can be used for the shorter sequences
           because of the padding mask.
           An example fo perms with string length 5
           >>> tensor([[0, 1, 2, 3, 4, 5, 6],   # canonical order
                    [0, 6, 5, 4, 3, 2, 1],      # reverse order
                    [0, 3, 1, 5, 2, 4, 6],
                    [0, 1, 4, 2, 3, 5, 6],
                    [0, 2, 3, 5, 1, 4, 6],
                    [0, 2, 1, 3, 5, 4, 6]])
        """
        # We don't permute the position of BOS, we permute EOS separately
        max_num_chars = tgt.shape[1] - 2
        # Special handling for 1-character sequences
        if max_num_chars == 1:
            return torch.arange(3, device=self._device).unsqueeze(0)
        perms = [torch.arange(max_num_chars, device=self._device)] if self.perm_forward else []
        # Additional permutations if needed
        max_perms = math.factorial(max_num_chars)
        if self.perm_mirrored:
            max_perms //= 2
        num_gen_perms = min(self.max_gen_perms, max_perms)
        # For 4-char sequences and shorter, we generate all permutations and sample from the pool to avoid collisions
        # Note that this code path might NEVER get executed since the labels in a mini-batch typically exceed 4 chars.
        if max_num_chars < 5:
            # Pool of permutations to sample from. We only need the first half (if complementary option is selected)
            # Special handling for max_num_chars == 4 which correctly divides the pool into the flipped halves
            if max_num_chars == 4 and self.perm_mirrored:
                selector = [0, 3, 4, 6, 9, 10, 12, 16, 17, 18, 19, 21]
            else:
                selector = list(range(max_perms))
            perm_pool = torch.as_tensor(list(permutations(range(max_num_chars), max_num_chars)), device=self._device)[selector]
            # If the forward permutation is always selected, no need to add it to the pool for sampling
            if self.perm_forward:
                perm_pool = perm_pool[1:]
            perms = torch.stack(perms)
            if len(perm_pool):
                i = self.rng.choice(len(perm_pool), size=num_gen_perms - len(perms), replace=False)
                perms = torch.cat([perms, perm_pool[i]])
        else:
            perms.extend([torch.randperm(max_num_chars, device=self._device) for _ in range(num_gen_perms - len(perms))])
            perms = torch.stack(perms)
        if self.perm_mirrored:
            # Add complementary pairs
            comp = perms.flip(-1)
            # Stack in such a way that the pairs are next to each other.
            perms = torch.stack([perms, comp]).transpose(0, 1).reshape(-1, max_num_chars)
        # NOTE:
        # The only meaningful way of permuting the EOS position is by moving it one character position at a time.
        # However, since the number of permutations = T! and number of EOS positions = T + 1, the number of possible EOS
        # positions will always be much less than the number of permutations (unless a low perm_num is set).
        # Thus, it would be simpler to just train EOS using the full and null contexts rather than trying to evenly
        # distribute it across the chosen number of permutations.
        # Add position indices of BOS and EOS
        bos_idx = perms.new_zeros((len(perms), 1))
        eos_idx = perms.new_full((len(perms), 1), max_num_chars + 1)
        perms = torch.cat([bos_idx, perms + 1, eos_idx], dim=1)
        # Special handling for the reverse direction. This does two things:
        # 1. Reverse context for the characters
        # 2. Null context for [EOS] (required for learning to predict [EOS] in NAR mode)
        if len(perms) > 1:
            perms[1, 1:] = max_num_chars + 1 - torch.arange(max_num_chars + 1, device=self._device)
        return perms

    def generate_attn_masks(self, perm):
        """Generate attention masks given a sequence permutation (includes pos. for bos and eos tokens)
        :param perm: the permutation sequence. i = 0 is always the BOS
        :return: lookahead attention masks
        """
        sz = perm.shape[0]
        mask = torch.zeros((sz, sz), device=self._device)
        for i in range(sz):
            query_idx = perm[i]
            masked_keys = perm[i + 1:]
            mask[query_idx, masked_keys] = float('-inf')
        content_mask = mask[:-1, :-1].clone()
        mask[torch.eye(sz, dtype=torch.bool, device=self._device)] = float('-inf')  # mask "self"
        query_mask = mask[1:, :-1]
        return content_mask, query_mask

    def configure_optimizers(self):
        agb = self.trainer.accumulate_grad_batches
        # Linear scaling so that the effective learning rate is constant regardless of the number of GPUs used with DDP.
        # lr_scale = agb * math.sqrt(self.trainer.num_devices) * self.batch_size / 256.
        lr_scale = agb * (self.trainer.num_devices * self.batch_size) / 512
        lr = lr_scale * self.lr

        # https://github.com/mlfoundations/open_clip/blob/b4cf9269b0b11c0eea47cb16039369a46bd67449/src/training/main.py#L171
        exclude = lambda n, p: p.ndim < 2 or "bn" in n or "ln" in n or "bias" in n or 'logit_scale' in n \
                            or "pos_queries" in n or "text_embed" in n
        include = lambda n, p: not exclude(n, p)

        # encoder parameters
        encoder_params = list(self.clip_model.named_parameters())
        enc_gain_or_bias_params = [p for n, p in encoder_params if exclude(n, p) and p.requires_grad]
        enc_rest_params = [p for n, p in encoder_params if include(n, p) and p.requires_grad]

        # decoder parameters
        decoder_params = [(n, p) for n, p in list(self.named_parameters()) if "clip_model" not in n]
        dec_gain_or_bias_params = [p for n, p in decoder_params if exclude(n, p) and p.requires_grad]
        dec_rest_params = [p for n, p in decoder_params if include(n, p) and p.requires_grad]

        rank_zero_info("[VL4STR] The length of encoder params with and without weight decay is {} and {}, respectively.".format(
            len(enc_rest_params), len(enc_gain_or_bias_params)
        ))
        rank_zero_info("[VL4STR] The length of decoder params with and without weight decay is {} and {}, respectively.".format(
            len(dec_rest_params), len(dec_gain_or_bias_params)
        ))

        optimizer = torch.optim.AdamW(
            [
                {"params": enc_gain_or_bias_params, "weight_decay": 0., 'lr': lr},
                {"params": enc_rest_params, "weight_decay": self.weight_decay, 'lr': lr},
                {"params": dec_gain_or_bias_params, "weight_decay": 0., 'lr': lr * self.coef_lr},
                {"params": dec_rest_params, "weight_decay": self.weight_decay * self.coef_wd, 'lr': lr * self.coef_lr},
            ],
            lr=lr, betas=(0.9, 0.98), eps=1.0e-6,
            )
        sched = OneCycleLR(optimizer, [lr, lr, lr * self.coef_lr, lr * self.coef_lr],
                            self.trainer.estimated_stepping_batches, pct_start=self.warmup_pct,
                            cycle_momentum=False)

        return {'optimizer': optimizer, 'lr_scheduler': {'scheduler': sched, 'interval': 'step'}}

    def _freeze_backbones(self):
        """set frozen backbones to eval mode"""
        for name, mod in self.clip_model.named_modules():
            if name.startswith("visual.transformer.resblocks."):
                layer_num = int(name.split(".resblocks.")[1].split(".")[0])
                if layer_num < self.image_freeze_nlayer:
                    mod.eval()

                # if self.image_freeze_layer_divisor > 0 and (layer_num + 1) % self.image_freeze_layer_divisor == 0:
                #     mod.eval()
            elif name.startswith("transformer.resblocks."):
                layer_num = int(name.split(".resblocks.")[1].split(".")[0])
                if layer_num < self.text_freeze_nlayer:
                    mod.eval()

        if self.freeze_language_backbone:
            self.clip_model.transformer.eval()
            self.clip_model.ln_final.eval()

    def freeze_cip_layers(self, image_freeze_nlayer, text_freeze_nlayer, image_freeze_layer_divisor=-1, image_only_fc=False):
        """
        freeze the parameters of layers with No.layer < image_freeze_nlayer or text_freeze_nlayer,
        """
        assert image_freeze_nlayer <= 12 and text_freeze_nlayer <=12 and image_freeze_layer_divisor <= 12
        if hasattr(self, "clip_model"):
            if image_freeze_nlayer > -1:
                for name, param in self.clip_model.visual.named_parameters():
                    # top layers always need to train
                    if name.startswith("ln_post.") or name.startswith("proj") or name.startswith("conv1") or name.startswith("ln_pre"):
                        continue
                    elif name.startswith("transformer.resblocks."):
                        layer_num = int(name.split(".resblocks.")[1].split(".")[0])
                        if layer_num >= image_freeze_nlayer:
                            continue
                    param.requires_grad = False

            #### freeze the layers which the index can be divided by image_freeze_layer_divisor
            # if image_freeze_layer_divisor > 0:
            #     for name, param in self.clip_model.visual.named_parameters():
            #         if name.startswith("transformer.resblocks."):
            #             layer_num = int(name.split(".resblocks.")[1].split(".")[0])
            #             if (layer_num + 1) % image_freeze_layer_divisor == 0:
            #                 param.requires_grad = False

            #### only train the top fc layer
            # if image_only_fc:
            #     for name, param in self.clip_model.visual.named_parameters():
            #         if "out_proj" in name or "conv1" in name or name.startswith("ln_post.") or name.startswith("proj"):
            #             continue
            #         param.requires_grad = False

            if text_freeze_nlayer > -1:
                for name, param in self.clip_model.named_parameters():
                    # top layers always need to train
                    if name.startswith("ln_final.") or name.startswith("text_projection") or name.startswith("visual"):
                        continue
                    elif name.startswith("transformer.resblocks."):
                        layer_num = int(name.split(".resblocks.")[1].split(".")[0])
                        if layer_num >= text_freeze_nlayer:
                            continue
                    param.requires_grad = False

            # freeze the whole backbones and related parameters
            if text_freeze_nlayer >= 12:
                for n, p in self.clip_model.named_parameters():
                    # exclude visual parameters
                    if "visual" not in n:                
                        if "transformer" in n or "token_embedding" in n or "ln_final" in n or "text_projection" in n:
                            p.requires_grad = False
            
            if image_freeze_nlayer >= 12:
                for n, p in self.clip_model.visual.named_parameters():
                    p.requires_grad = False               

    def train(self, mode=True):
        """
        Override the default train() to freeze the BN parameters
        """
        super().train(mode)
        if mode:
            self._freeze_backbones()

        return self

strhub/models/vl_str/system.py
VamosC/CLIP4STR
#!/usr/bin/env python3
# Scene Text Recognition Model Hub
# Copyright 2022 Darwin Bautista
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import sys
import time
import string
import argparse
from dataclasses import dataclass
from typing import List

import torch

from tqdm import tqdm

from strhub.data.module import SceneTextDataModule
from strhub.models.utils import load_from_checkpoint, parse_model_args


@dataclass
class Result:
    dataset: str
    num_samples: int
    accuracy: float
    ned: float
    confidence: float
    label_length: float


def print_results_table(results: List[Result], file=None):
    w = max(map(len, map(getattr, results, ['dataset'] * len(results))))
    w = max(w, len('Dataset'), len('Combined'))
    print('| {:<{w}} | # samples | Accuracy | 1 - NED | Confidence | Label Length |'.format('Dataset', w=w), file=file)
    print('|:{:-<{w}}:|----------:|---------:|--------:|-----------:|-------------:|'.format('----', w=w), file=file)
    c = Result('Combined', 0, 0, 0, 0, 0)
    for res in results:
        c.num_samples += res.num_samples
        c.accuracy += res.num_samples * res.accuracy
        c.ned += res.num_samples * res.ned
        c.confidence += res.num_samples * res.confidence
        c.label_length += res.num_samples * res.label_length
        print(f'| {res.dataset:<{w}} | {res.num_samples:>9} | {res.accuracy:>8.2f} | {res.ned:>7.2f} '
              f'| {res.confidence:>10.2f} | {res.label_length:>12.2f} |', file=file)
    c.accuracy /= c.num_samples
    c.ned /= c.num_samples
    c.confidence /= c.num_samples
    c.label_length /= c.num_samples
    print('|-{:-<{w}}-|-----------|----------|---------|------------|--------------|'.format('----', w=w), file=file)
    print(f'| {c.dataset:<{w}} | {c.num_samples:>9} | {c.accuracy:>8.2f} | {c.ned:>7.2f} '
          f'| {c.confidence:>10.2f} | {c.label_length:>12.2f} |', file=file)


@torch.inference_mode()
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('checkpoint', help="Model checkpoint (or 'pretrained=<model_id>')")
    parser.add_argument('--data_root', default='data')
    parser.add_argument('--batch_size', type=int, default=512)
    parser.add_argument('--num_workers', type=int, default=4)
    parser.add_argument('--cased', action='store_true', default=False, help='Cased comparison')
    parser.add_argument('--punctuation', action='store_true', default=False, help='Check punctuation')
    parser.add_argument('--new', action='store_true', default=False, help='Evaluate on new benchmark datasets')
    parser.add_argument('--rotation', type=int, default=0, help='Angle of rotation (counter clockwise) in degrees.')
    parser.add_argument('--device', default='cuda')

    parser.add_argument('--clip_model_path', type=str,
                            default="path/to/clip/model",
                            help='path to the clip model')
    parser.add_argument('--clip_refine', action='store_true', default=False,
                            help='use clip to refine the predicted results')
    parser.add_argument('--sample_K', type=int, default=5,
                            help='K of top-K when performing CLIP post-refinement')
    parser.add_argument('--sample_K2', type=int, default=5,
                            help='K of top-K choosed predictions when performing CLIP post-refinement')
    parser.add_argument('--sample_total', type=int, default=50,
                            help='the number of samples when sample from the predicted probability distribution')
    parser.add_argument('--sample_prompt', type=str, default=None,
                            help='prompt for CLIP')
    parser.add_argument('--alpha', type=float, default=0.1,
                            help='prompt for CLIP')                                

    args, unknown = parser.parse_known_args()
    kwargs = parse_model_args(unknown)
    print(args)

    if os.path.isdir(args.checkpoint):
        ckpts = [x for x in os.listdir(args.checkpoint) if 'val' in x]
        assert len(ckpts) >= 1
        val_acc = [float(x.split('-')[-2].split('=')[-1]) for x in ckpts]
        best_ckpt = os.path.join(args.checkpoint, ckpts[val_acc.index(max(val_acc))])
        best_epoch = int(best_ckpt.split('/')[-1].split('-')[0].split('=')[-1])
        print('The val accuracy is best {}-{}e'.format(max(val_acc), best_epoch))
        args.checkpoint = best_ckpt

    charset_test = string.digits + string.ascii_lowercase
    if args.cased:
        charset_test += string.ascii_uppercase
    if args.punctuation:
        charset_test += string.punctuation
    kwargs.update({'charset_test': charset_test})
    print(f'Additional keyword arguments: {kwargs}')

    print('load weights from checkpoint {}'.format(args.checkpoint))
    model = load_from_checkpoint(args.checkpoint, **kwargs).eval().to(args.device)
    hp = model.hparams
    datamodule = SceneTextDataModule(args.data_root, '_unused_', hp.img_size, hp.max_label_length, hp.charset_train,
                                     hp.charset_test, args.batch_size, args.num_workers, False, rotation=args.rotation)

    test_set = SceneTextDataModule.TEST_BENCHMARK_SUB + SceneTextDataModule.TEST_BENCHMARK
    if args.new:
        test_set += SceneTextDataModule.TEST_NEW
    test_set = sorted(set(test_set))

    torch.cuda.synchronize()
    start_time = time.time()
    results = {}
    max_width = max(map(len, test_set))
    all_total = 0
    for name, dataloader in datamodule.test_dataloaders(test_set).items():
        total = 0
        correct = 0
        ned = 0
        confidence = 0
        label_length = 0
        for imgs, labels in tqdm(iter(dataloader), desc=f'{name:>{max_width}}'):
            res = model.test_step((imgs.to(model.device), labels), -1,
                                    clip_model_path=args.clip_model_path,
                                    clip_refine=args.clip_refine,
                                    sample_K=args.sample_K,
                                    sample_K2=args.sample_K2,
                                    sample_total=args.sample_total,
                                    sample_prompt=args.sample_prompt,
                                    alpha=args.alpha
                                    )['output']

            total += res.num_samples
            correct += res.correct
            ned += res.ned
            confidence += res.confidence
            label_length += res.label_length
        all_total += total
        accuracy = 100 * correct / total
        mean_ned = 100 * (1 - ned / total)
        mean_conf = 100 * confidence / total
        mean_label_length = label_length / total
        results[name] = Result(name, total, accuracy, mean_ned, mean_conf, mean_label_length)

    result_groups = {
        'Benchmark (Subset)': SceneTextDataModule.TEST_BENCHMARK_SUB,
        'Benchmark': SceneTextDataModule.TEST_BENCHMARK
    }
    if args.new:
        result_groups.update({'New': SceneTextDataModule.TEST_NEW})

    torch.cuda.synchronize()
    total_time = time.time() - start_time

    if args.clip_refine:
        log_filename = args.checkpoint + '.log_K{}-{}-{}_prompt{}_alpha{}_new{}.txt'.format(
            args.sample_K2, args.sample_K, args.sample_total, args.sample_prompt, args.alpha, str(args.new))
    else:
        log_filename = args.checkpoint + '.log_new{}.txt'.format(str(args.new))

    with open(log_filename, 'w') as f:
        for out in [f, sys.stdout]:
            for group, subset in result_groups.items():
                print(f'{group} set:', file=out)
                print_results_table([results[s] for s in subset], out)
                print('\n', file=out)
            print("Time: Total {}s, Average {}ms. Total samples {}.".format(total_time, total_time * 1000.0 / all_total, all_total), file=out)


if __name__ == '__main__':
    main()

test.py
VamosC/CLIP4STR
#!/usr/bin/env python3
# Scene Text Recognition Model Hub
# Copyright 2022 Darwin Bautista
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
import math
import hydra
import torch
from pathlib import Path
from hydra.core.hydra_config import HydraConfig
from omegaconf import DictConfig, open_dict, OmegaConf
from pytorch_lightning import Trainer
from pytorch_lightning.loggers import TensorBoardLogger
from pytorch_lightning.strategies import DDPStrategy
from pytorch_lightning.utilities import rank_zero_info, rank_zero_only
from pytorch_lightning.utilities.model_summary import summarize
from pytorch_lightning.callbacks import ModelCheckpoint, StochasticWeightAveraging, LearningRateMonitor

from strhub.data.module import SceneTextDataModule
from strhub.models.base import BaseSystem
from strhub.models.utils import create_model
from strhub.dist_utils import copy_remote, is_main_process


@hydra.main(config_path='configs', config_name='main', version_base='1.2')
def main(config: DictConfig):
    trainer_strategy = None
    with open_dict(config):
        # Resolve absolute path to data.root_dir
        config.data.root_dir = hydra.utils.to_absolute_path(config.data.root_dir)
        # Special handling for GPU-affected config
        gpus = config.trainer.get('gpus', 0)
        if gpus:
            # Use mixed-precision training
            config.trainer.precision = 16
        if gpus > 1:
            # Use DDP
            config.trainer.strategy = 'ddp'
            # DDP optimizations
            trainer_strategy = DDPStrategy(find_unused_parameters=getattr(config.model, "find_unused_parameters", False),
                                            gradient_as_bucket_view=True)
            # Scale steps-based config
            config.trainer.val_check_interval //= gpus
            if config.trainer.get('max_steps', -1) > 0:
                config.trainer.max_steps //= gpus

    # Special handling for PARseq
    if config.model.get('perm_mirrored', False):
        assert config.model.perm_num % 2 == 0, 'perm_num should be even if perm_mirrored = True'

    # print config
    rank_zero_info(OmegaConf.to_yaml(config))

    # If specified, use pretrained weights to initialize the model
    if config.pretrained is not None:
        model: BaseSystem = create_model(config.pretrained, True)
    else:
        model: BaseSystem = hydra.utils.instantiate(config.model)
    
    rank_zero_info(summarize(model, max_depth=1 if model.hparams.name.startswith('parseq') else 2))

    datamodule: SceneTextDataModule = hydra.utils.instantiate(config.data)
    checkpoint = ModelCheckpoint(monitor='val_accuracy', mode='max', save_top_k=1, save_last=True,
                                    filename='{epoch}-{step}-{val_accuracy:.4f}-{val_NED:.4f}',
                                    every_n_epochs=1)
    lr_monitor = LearningRateMonitor(logging_interval='step')
    callbacks = [lr_monitor, checkpoint]
    if getattr(config, 'swa', False):
        # set swa lrs
        swa_epoch_start = 0.8
        lr_scale = math.sqrt(torch.cuda.device_count()) * config.data.batch_size / 256.
        lr = lr_scale * config.model.lr
        if "clip" in config.model.name:
            swa_lrs = [lr, lr, config.model.coef_lr * lr, config.model.coef_lr * lr]
        else:
            swa_lrs = [lr,]

        swa_lrs = [x * (1 - swa_epoch_start) for x in swa_lrs]
        swa = StochasticWeightAveraging(swa_lrs=swa_lrs, swa_epoch_start=swa_epoch_start)
        callbacks.append(swa)

    cwd = HydraConfig.get().runtime.output_dir if config.ckpt_path is None else \
            str(Path(config.ckpt_path).parents[1].absolute())
    trainer: Trainer = hydra.utils.instantiate(config.trainer, logger=TensorBoardLogger(cwd, '', '.'),
                                               strategy=trainer_strategy, enable_model_summary=False,
                                               accumulate_grad_batches=config.trainer.accumulate_grad_batches,
                                               callbacks=callbacks)

    trainer.fit(model, datamodule=datamodule, ckpt_path=config.ckpt_path)

    # copy data and perform test
    torch.distributed.barrier()
    if is_main_process():
        copy_remote(cwd, config.data.output_url)
        test_call(cwd, config.data.root_dir, config.model.code_path)

    torch.distributed.barrier()
    torch.distributed.destroy_process_group()


@rank_zero_only
def test_call(cwd, data_dir, code_path=None):
    file = os.path.join(code_path, 'test.py')
    assert os.path.exists(file)
    print("The execute file is {}".format(file))
    ckpts = [x for x in os.listdir(os.path.join(cwd, 'checkpoints')) if 'val' in x]
    val_acc = [float(x.split('-')[-2].split('=')[-1]) for x in ckpts]

    best_ckpt = os.path.join(os.path.join(cwd, 'checkpoints'), ckpts[val_acc.index(max(val_acc))])
    print("The best ckpt is {}".format(best_ckpt))
    best_epoch = int(best_ckpt.split('/')[-1].split('-')[0].split('=')[-1])
    print('The val accuracy is best {}-{}e'.format(max(val_acc), best_epoch))

    # test best
    # print("\n Test results with the best checkpoint")
    # os.system("python {} {} --data_root {} --new".format(file, best_ckpt, data_dir))
    # test last

    print("\n Test results with the last checkpoint")
    last_ckpt = os.path.join(os.path.join(cwd, 'checkpoints'), "last.ckpt")
    os.system("python {} {} --data_root {} --new".format(file, last_ckpt, data_dir))

    # test last with refinement
    # print("\n Test results with the last checkpoint")
    # last_ckpt = os.path.join(os.path.join(cwd, 'checkpoints'), "last.ckpt")
    # os.system("python {} {} --data_root {} --new --clip_refine".format(file, last_ckpt, data_dir))


if __name__ == '__main__':
    main()

train.py
VamosC/CLIP4STR