Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+res/effect.png filter=lfs diff=lfs merge=lfs -text

.mv

@@ -0,0 +1 @@
+Revision:master,CreatedAt:1742457812

README.md

@@ -0,0 +1,94 @@
+---
+tasks:
+- multi-modal-embedding
+- image-text-retrieval
+domain:
+- multi-modal
+frameworks:
+- pytorch
+backbone:
+- transformers
+metrics:
+- R@1
+license: apache-2.0
+tags:
+- Ant Group
+- multi-modal-embedding
+widgets:
+  - inputs:
+      - validator:
+          max_words: 52
+        type: text
+        title: 查询文本
+    output:
+      maximize: false
+    examples:
+      - name: 1
+        inputs:
+          - data: 戴眼镜的猫
+      - name: 2
+        inputs:
+          - data: 一个在逛公园的女孩
+    task: multi-modal-embedding
+---
+
+## 模型描述
+M2-Encoder是强大的中英双语多模态模型，它在我们构建的包含60亿图文对(30亿中文+30亿英文)的BM-6B上训练得到，支持zero-shot 图文跨模态检索(文搜图、图搜文) 以及 zero-shot图片分类 任务。
+
+模型效果如下：
+
+![M2-Encoder](./res/effect.png)
+
+## 期望模型使用方式以及适用范围
+本模型主要用于：
+1. 图片检索文本，或文本检索图片: 以文本检索图片为例，使用M2-Encoder提前对所有图片底库进行特征抽取，给定文本query，使用M2-Encoder对query文本进行特征抽取， 然后和图片底库保存的特征进行相似度计算。
+2. 图片zero-shot开集分类: 给定图像以及对应的标签列表，根据图像和标签相似度，输出与图像最匹配的标签。
+
+
+## 如何使用
+
+### 代码范例
+```
+# 新建环境（Python版本3.8）
+conda create -n m2-encoder python=3.8
+source activate m2-encoder
+
+# clone项目地址
+cd /YourPath/
+git clone https://github.com/alipay/Ant-Multi-Modal-Framework
+
+# 安装包依赖
+cd ./Ant-Multi-Modal-Framework/prj/M2_Encoder/
+pip install -r requirements.txt
+
+# 运行demo，会自动通过model_scope下载对应模型权重
+python run.py
+```
+
+### 模型局限性以及可能的偏差
+模型在数据集上训练，有可能产生一些偏差，请用户自行评测后决定如何使用。
+
+## 训练数据介绍
+BM-6B数据集: 包含60亿清洗后的高质量中英双语图文对数据，其中文和英文数据比例基本保持一致，均为30亿。数据集搜集、构建过程详见[技术报告](https://arxiv.org/abs/2401.15896)。
+
+## 模型训练流程
+暂时不支持通过ModelScope接口进行训练，敬请期待。
+
+
+### 训练
+暂不支持。
+## 数据评估及结果
+zero-shot图文跨模态检索和zero-shot分类任务均达到SOTA.
+
+
+
+### 相关论文以及引用信息
+如果你觉得这个该模型对有所帮助，请考虑引用下面的相关的论文：
+```
+@misc{guo2024m2encoder,
+      title={M2-Encoder: Advancing Bilingual Image-Text Understanding by Large-scale Efficient Pretraining}, 
+      author={Qingpei Guo and Furong Xu and Hanxiao Zhang and Wang Ren and Ziping Ma and Lin Ju and Jian Wang and Jingdong Chen and Ming Yang},
+      year={2024},
+      url={https://arxiv.org/abs/2401.15896},
+}
+```

configuration.json

@@ -0,0 +1 @@
+{"framework":"Pytorch","task":"multi-modal-embeddings","pipeline":{"type":"multi-modal-embedding-pipeline"},"allow_remote": true}

m2_encoder_1B.ckpt

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

ms_wrapper.py

@@ -0,0 +1,219 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+
+import torch
+import os
+
+from modelscope.models.base import TorchModel
+from modelscope.preprocessors.base import Preprocessor
+from modelscope.pipelines.base import Model, Pipeline
+from modelscope.utils.config import Config
+from modelscope.pipelines.builder import PIPELINES
+from modelscope.preprocessors.builder import PREPROCESSORS
+from modelscope.models.builder import MODELS
+from modelscope.preprocessors.image import load_image
+
+
+from vlmo.utils.beit_utils import load_from_config
+
+
+@PIPELINES.register_module(
+    "multi-modal-embeddings", module_name="multi-modal-embedding-pipeline"
+)
+class MyCustomPipeline(Pipeline):
+    """Give simple introduction to this pipeline.
+
+    Examples:
+
+    >>> from modelscope.pipelines import pipeline
+    >>> input = "Hello, ModelScope!"
+    >>> my_pipeline = pipeline('my-task', 'my-model-id')
+    >>> result = my_pipeline(input)
+
+    """
+
+    def __init__(self, model, preprocessor=None, **kwargs):
+        """
+        use `model` and `preprocessor` to create a custom pipeline for prediction
+        Args:
+            model: model id on modelscope hub.
+            preprocessor: the class of method be init_preprocessor
+        """
+        super().__init__(model=model, auto_collate=False)
+        self.model_dir = model
+        self._device = "cuda" if torch.cuda.is_available() else "cpu"
+        # model_config = {
+        #     "loss_names": {"itc": 1},
+        #     "encoder_layers": 9,
+        #     "beit3_vl_layers": 3,
+        #     "tokenizer_type": "GLMChineseTokenizer",
+        #     "tokenizer": os.path.join(self.model_dir, "./vlmo/tokenizer"),
+        #     "vocab_size": 115244,
+        #     "whole_word_masking": True,
+        #     "precision": 32,
+        #     "test_only": True,
+        #     "flash_attn": True,
+        #     "model_path": os.path.join(self.model_dir, "m2_encoder_1B.ckpt"),
+        #     "modelscope": {"model_id": "M2Cognition/M2-Encoder-Large"},
+        #     "model_file": "m2_encoder_1B.ckpt",
+        # }
+        model_config = {
+            "loss_names": {"itc": 1},
+            "beit_version": "large",
+            "encoder_embed_dim": 1024,
+            "out_embed_dim": 1024,
+            "encoder_layers": 21,
+            "beit3_vl_layers": 3,
+            # "image_size": 224,
+            "visual_mask_size": 14,
+            "tokenizer_type": "GLMChineseTokenizer",
+            "tokenizer": os.path.join(self.model_dir, "./vlmo/tokenizer"),
+            "vocab_size": 115244,
+            "whole_word_masking": False,
+            "precision": 32,
+            "test_only": True,
+            "flash_attn": True,
+            "model_path": os.path.join(self.model_dir, "m2_encoder_1B.ckpt"),
+            "modelscope": {
+                "model_id": "M2Cognition/M2_Encoder_Large"
+            },
+            "model_file": "m2_encoder_1B.ckpt"
+        }
+        model, processors = load_from_config(model_config)
+        self.model = model
+        self.model.to(self._device).eval()
+        self._tokenizer, self._img_processor = processors
+
+    def _sanitize_parameters(self, **pipeline_parameters):
+        """
+        this method should sanitize the keyword args to preprocessor params,
+        forward params and postprocess params on '__call__' or '_process_single' method
+        considered to be a normal classmethod with default implementation / output
+
+        Default Returns:
+            Dict[str, str]:  preprocess_params = {}
+            Dict[str, str]:  forward_params = {}
+            Dict[str, str]:  postprocess_params = pipeline_parameters
+        """
+        return {}, pipeline_parameters, {}
+
+    def _check_input(self, inputs):
+        pass
+
+    def _check_output(self, outputs):
+        pass
+
+    def forward(self, forward_params):
+        """Provide default implementation using self.model and user can reimplement it"""
+        # print("forward_params", forward_params)
+        labels = forward_params.get("label_list", "")
+        labels = labels.split(",")
+        if len(labels) > 1 and labels[0] != "":
+            txt_encoding = self._tokenizer(
+                labels,
+                padding="max_length",
+                truncation=True,
+                max_length=self.model.hparams.config["max_text_len"],
+                return_special_tokens_mask=True,
+            )
+            txt_data = {
+                "text_ids": torch.tensor(txt_encoding["input_ids"]).to(self._device),
+                "text_masks": torch.tensor(txt_encoding["attention_mask"]).to(
+                    self._device
+                ),
+                "text_labels": None,
+            }
+            txt_feats = self.model.infer_text(txt_data)["cls_vlffn_feats"]
+            image = forward_params["image"]
+            image = load_image(image)
+            img = self._img_processor(image).unsqueeze(0)
+            img_data = {"image": [img.to(self._device)]}
+            img_feats = self.model.infer_image(img_data)["cls_vlffn_feats"]
+            logits_per_image = self.model.logit_scale.exp() * img_feats @ txt_feats.t()
+            probs = logits_per_image.softmax(dim=-1).detach().cpu()
+            index = probs.max(dim=-1)[1][0]
+            label = labels[index]
+            return {"text": label, "scores": probs.numpy().tolist()[0]}
+        else:
+            rets = {}
+            if "text" in forward_params:
+                text = forward_params.get("text")
+                txt_encoding = self._tokenizer(
+                    text,
+                    padding="max_length",
+                    truncation=True,
+                    max_length=self.model.hparams.config["max_text_len"],
+                    return_special_tokens_mask=True,
+                )
+                txt_data = {
+                    "text_ids": torch.tensor(txt_encoding["input_ids"]).to(
+                        self._device
+                    ),
+                    "text_masks": torch.tensor(txt_encoding["attention_mask"]).to(
+                        self._device
+                    ),
+                    "text_labels": None,
+                }
+                txt_feats = self.model.infer_text(txt_data)["cls_vlffn_feats"]
+                rets.update({"text_embedding": txt_feats.detach()})
+            if "img" in forward_params:
+                input_img = forward_params["img"]
+                img = self._img_processor(input_img).unsqueeze(0)
+                img_data = {"image": [img.to(self._device)]}
+                img_feats = self.model.infer_image(img_data)["cls_vlffn_feats"]
+                rets.update({"img_embedding": img_feats.detach()})
+
+            return rets
+
+    def preprocess(self, inputs):
+        return inputs
+
+    def postprocess(self, inputs):
+        """If current pipeline support model reuse, common postprocess
+            code should be write here.
+
+        Args:
+            inputs:  input data
+
+        Return:
+            dict of results:  a dict containing outputs of model, each
+                output should have the standard output name.
+        """
+        return inputs
+
+
+"""
+# Tips: usr_config_path is the temporary save configuration location， after upload modelscope hub, it is the model_id
+usr_config_path = "/tmp/snapdown/"
+config = Config(
+    {
+        "framework": "pytorch",
+        "task": "multi-modal-embeddings",
+        "model": {"type": "m2-encoder"},
+        "pipeline": {"type": "multi-modal-embedding-pipeline"},
+        "allow_remote": True,
+    }
+)
+config.dump("/tmp/snapdown/" + "configuration.json")
+"""
+
+if __name__ == "__main__":
+    from modelscope.pipelines import pipeline
+    from modelscope.preprocessors.image import load_image
+
+    model = "M2Cognition/M2-Encoder"
+    pipe = pipeline("multi-modal-embeddings", model=model)
+    input = {
+        "image": "https://clip-cn-beijing.oss-cn-beijing.aliyuncs.com/pokemon.jpeg",
+        "label_list": "杰尼龟,妙蛙种子,小火龙,皮卡丘",
+    }
+    demo = pipe(input)
+    print("demo output", demo)
+    inputs = {"text": ["杰尼龟", "妙蛙种子", "小火龙", "皮卡丘"]}
+    output = pipe(inputs)
+    print("text output", output)
+    input_img = load_image(
+        "https://clip-cn-beijing.oss-cn-beijing.aliyuncs.com/pokemon.jpeg"
+    )  # 支持皮卡丘示例图片路径/本地图片 返回PIL.Image
+    inputs = {"img": input_img}
+    img_embedding = pipe(inputs)  # 2D Tensor, [图片数, 特征维度]
+    print("image output", img_embedding)

requirements.txt

@@ -0,0 +1,14 @@
+torch
+pytorch_lightning<=2.0.8
+transformers
+Pillow
+tqdm
+einops
+sacred
+timm
+torchvision
+fairscale
+numpy
+opencv-python
+sentencepiece
+modelscope

vlmo/Encoder_0.4B.json

@@ -0,0 +1,17 @@
+{
+  "loss_names": {"itc": 1},
+  "encoder_layers": 9,
+  "beit3_vl_layers": 3,
+  "tokenizer_type": "GLMChineseTokenizer",
+  "tokenizer": "./vlmo/tokenizer",
+  "vocab_size": 115244,
+  "whole_word_masking": true,
+  "precision": 32,
+  "test_only": true,
+  "flash_attn": true,
+  "model_path": "m2_encoder_0.4B.ckpt",
+  "modelscope": {
+    "model_id": "M2Cognition/M2-Encoder"
+  },
+  "model_file": "m2_encoder_0.4B.ckpt"
+}

vlmo/README.md

@@ -0,0 +1,10 @@
+---
+license: Apache License 2.0
+---
+###### 该模型当前使用的是默认介绍模版，处于“预发布”阶段，页面仅限所有者可见。
+###### 请根据[模型贡献文档说明](https://www.modelscope.cn/docs/%E5%A6%82%E4%BD%95%E6%92%B0%E5%86%99%E5%A5%BD%E7%94%A8%E7%9A%84%E6%A8%A1%E5%9E%8B%E5%8D%A1%E7%89%87)，及时完善模型卡片内容。ModelScope平台将在模型卡片完善后展示。谢谢您的理解。
+
+#### Clone with HTTP
+```bash
+ git clone https://www.modelscope.cn/M2Cognition/M2_Encoder_demo.git
+```

vlmo/config.py

@@ -0,0 +1,165 @@
+from sacred import Experiment
+
+ex = Experiment("VLMo")
+
+
+def _loss_names(d):
+    ret = {
+        "itm": 0,  # image-text matching loss
+        "itc": 0,  # image-text contrastive loss
+        "caption": 0,  # image captioning loss
+        "mvlm": 0,  # masked language modeling loss
+        "textmlm": 0,  # text-only masked language modeling
+        "imagemlm": 0,  # image-only masked language modeling
+        "vqa": 0,
+        "nlvr2": 0,
+        "irtr": 0,  # retrieval task ft
+    }
+    ret.update(d)
+    return ret
+
+
+@ex.config
+def config():
+    exp_name = "vlmo"
+    seed = 1
+    datasets = ["coco", "vg", "sbu", "gcc"]  # dataset name, the definition can refer to: vlmo/datamodules/__init__.py  # noqa
+    loss_names = _loss_names({"itm": 0, "itc": 0, "mvlm": 0})  # training loss
+    batch_size = 1024  # this is a desired batch size; pl trainer will accumulate gradients.
+
+    # BEiT-v3 setting
+    encoder_layers = 12  # the layer number of backbone
+    encoder_embed_dim = 768  # the hidden size of tokenizer
+    out_embed_dim = 768  # the hidden size of output embedding
+    beit_version = "base"  # model size: base(0.4B)|large(1B)|huge(10B)
+    beit3_vl_layers = 3  # the layer number of vl_backbone
+    deepnorm_init = True  # init method
+    share_layer = False  # if share the weight between layer within backbone
+    share_attn = False  # if share the attention weight of different layer
+    one_attn = False  # if share the attention weight of vision and language
+
+    # Image setting
+    train_transform_keys = ["square_transform_randaug"]  # train transform: refer to vlmo/transforms/__init__.py
+    val_transform_keys = ["square_transform"]  # test transform: refer to refer to vlmo/transforms/__init__.py
+    image_size = 224  # image size
+    reclip_image_size = None  # reclip image size
+    patch_size = 16  # patch size
+    draw_false_image = 0  # if get negative image
+    image_only = False  # only input image
+    text_only = False  # # only input text
+
+    # Video setting, video_num_frm is not None means video input
+    video_num_frm = None
+
+    # Visual tokenizer setting based on beit2
+    tokenizer_model = "beit2_visual_tokenizer"
+    codebook_size = 8192
+    codebook_dim = 32
+    visual_mask_size = 14
+    visual_mask_num = 80
+
+    # Text Setting
+    lang = 'cn'  # language for zero-shot imagenet testing: cn|en
+    vqav2_label_size = 3129
+    max_text_len = 40  # the number of characters
+    max_text_len_of_initckpt = 196
+    tokenizer_type = "BertTokenizer"  # Chinese text
+    vocab_size = 21128
+    tokenizer = "./vocab.txt"
+    whole_word_masking = True
+    mlm_prob = 0.15  # language mask ratio
+    draw_false_text = 0
+    mvlm_prob = 0.50  # vision-langurage mlm task
+    mask_ratio = 0  # flip: mask ratio for image
+
+    # cap setting
+    cap_onlytext = False  # default caption image to text
+
+    # imagemlm setting
+    split_data_for_imagemlm = False  # if True, split a batch data to two parts, and the first part for imagemlm.
+
+    # itc setting
+    itc_mask = False  # itc use masked token
+    aggregate_nodes = -1  # aggregate nodes num for compute_itc, default -1 is for all nodes
+
+    # Transformer Setting
+    model_arch = "vlmo_base_patch16"
+    drop_path_rate = 0.1
+
+    # Downstream Setting
+    get_recall_metric = False
+    get_recall_rerank_metric = False
+    get_zeroshot_metric = False
+    get_muge_feat = False
+    get_f30k_feat = False
+    k_test = 32
+
+    # PL Trainer Setting
+    resume_from = None
+    fast_dev_run = False
+    val_check_interval = 1.0
+    test_only = False
+    use_sharded_training = False
+    resume_during_training = False
+    save_top_k = 10
+    every_n_train_steps = 2000  # the step to save checkpoint
+    log_metric_steps = 100  # the step to log metric
+
+    # below params varies with the environment
+    use_pcache = False  # data storage method: pcache or nas
+    pcache_root = ""
+    # main_site: pcache://multimodalproxyi-pool.cz50c.alipay.com:39999/mnt/
+    # public_cloud: pcache://pcache_public_cloud.pcache.local:39999/mnt/abc7c88079a60b45ddfce7afa40720b7/
+    gpu_env = "main_site"  # public_cloud or main_site
+    data_root = ""  # data root for data list
+
+
+    log_dir = "result"
+    per_gpu_batchsize = 4  # you should define this manually with per_gpu_batch_size=#
+    num_gpus = 1
+    num_nodes = 1
+    load_path = ""
+    num_workers = 8
+    precision = 16
+    local_run = True
+    flash_attn = False
+    deepspeed_config = None  # "ds_config.json"
+    coalesce_backbone = False
+    mask_data = "v+l"  # 'v+l':choose input of imagemlm+textmlm task, 'vl': choose input of mvlm task.
+    communication_benchmark = False
+    checkpoint_activations = False
+
+    # dataset setting
+    single_cap = True  # if have only one caption
+    random_one = False  # if choose one caption from caption list
+
+    # ITC setting
+    itc_feats_name = "cls_vlffn_feats"  # feat for itc loss
+    itc_distill = ""
+    itc_distill_dim = 1024
+    itc_teacher_weights = ""
+
+    # mup training setting
+    mup = False
+    base_encoder_embed_dim = 1
+    delta_encoder_embed_dim = 2
+    mup_encoder_attention_heads = 1
+    base_encoder_ffn_embed_dim = 1
+    delta_encoder_ffn_embed_dim = 2
+
+    # atorch
+    atorch_config = None
+    compile_op = False
+    optimizer_state_shard_save = False
+    model_state_shard_save = False
+
+    # itc loss
+    local_loss = False
+    use_dual_softmax = False
+
+    num_frames = 1
+# ----------------------- LMM pretraining config -----------------------
+
+    # norm setting
+    deepnorm = False
+

vlmo/modules/__init__.py

@@ -0,0 +1 @@
+from .vlmo_module import VLMo

vlmo/modules/heads.py

@@ -0,0 +1,24 @@
+import torch.nn as nn
+
+
+class Pooler(nn.Module):
+    def __init__(self, hidden_size):
+        super().__init__()
+        self.dense = nn.Linear(hidden_size, hidden_size)
+        self.activation = nn.Tanh()
+
+    def forward(self, hidden_states):
+        first_token_tensor = hidden_states[:, 0]
+        pooled_output = self.dense(first_token_tensor)
+        pooled_output = self.activation(pooled_output)
+        return pooled_output
+
+
+class ITCHead(nn.Module):
+    def __init__(self, hidden_size, out_size):
+        super().__init__()
+        self.fc = nn.Linear(hidden_size, out_size, bias=False)
+
+    def forward(self, x):
+        x = self.fc(x)
+        return x

vlmo/modules/modeling_utils.py

@@ -0,0 +1,179 @@
+# --------------------------------------------------------
+# Image as a Foreign Language: BEiT Pretraining for Vision and Vision-Language Tasks (https://arxiv.org/abs/2208.10442)
+# Github source: https://github.com/microsoft/unilm/tree/master/beit3
+# Copyright (c) 2023 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------'
+
+import math
+import torch
+import torch.nn as nn
+from timm.models.layers import trunc_normal_ as __call_trunc_normal_
+
+from vlmo.torchscale.model.BEiT3 import BEiT3
+from vlmo.torchscale.architecture.config import EncoderConfig
+
+
+def trunc_normal_(tensor, mean=0.0, std=1.0):
+    __call_trunc_normal_(tensor, mean=mean, std=std, a=-std, b=std)
+
+
+def _get_base_config(
+    img_size=224,
+    patch_size=16,
+    drop_path_rate=0,
+    checkpoint_activations=None,
+    mlp_ratio=4,
+    vocab_size=64010,
+    encoder_layers=12,
+    encoder_embed_dim=768,
+    encoder_attention_heads=12,
+    share_layer=False,
+    share_attn=False,
+    deepnorm=False,
+    mask_ratio=0,
+    max_text_len=52,
+    one_attn=False,
+    **kwargs
+):
+    return EncoderConfig(
+        img_size=img_size,
+        patch_size=patch_size,
+        vocab_size=vocab_size,
+        multiway=True,
+        layernorm_embedding=False,
+        normalize_output=True,
+        no_output_layer=True,
+        drop_path_rate=drop_path_rate,
+        encoder_embed_dim=encoder_embed_dim,
+        encoder_attention_heads=encoder_attention_heads,
+        encoder_layers=encoder_layers,
+        encoder_ffn_embed_dim=int(encoder_embed_dim * mlp_ratio),
+        checkpoint_activations=checkpoint_activations,
+        share_layer=share_layer,
+        share_attn=share_attn,
+        deepnorm=deepnorm,
+        mask_ratio=mask_ratio,
+        max_text_len=max_text_len,
+        one_attn=one_attn,
+    )
+
+
+def _get_large_config(
+    img_size=224,
+    patch_size=16,
+    drop_path_rate=0,
+    checkpoint_activations=None,
+    mlp_ratio=4,
+    vocab_size=64010,
+    encoder_layers=24,
+    encoder_embed_dim=1024,
+    encoder_attention_heads=16,
+    share_layer=False,
+    share_attn=False,
+    deepnorm=False,
+    mask_ratio=0,
+    max_text_len=52,
+    one_attn=False,
+    **kwargs
+):
+    return EncoderConfig(
+        img_size=img_size,
+        patch_size=patch_size,
+        vocab_size=vocab_size,
+        multiway=True,
+        layernorm_embedding=False,
+        normalize_output=True,
+        no_output_layer=True,
+        drop_path_rate=drop_path_rate,
+        encoder_embed_dim=encoder_embed_dim,
+        encoder_attention_heads=encoder_attention_heads,
+        encoder_layers=encoder_layers,
+        encoder_ffn_embed_dim=int(encoder_embed_dim * mlp_ratio),
+        checkpoint_activations=checkpoint_activations,
+        share_layer=share_layer,
+        share_attn=share_attn,
+        deepnorm=deepnorm,
+        mask_ratio=mask_ratio,
+        max_text_len=max_text_len,
+        one_attn=one_attn,
+    )
+
+
+def _get_huge_config(
+    img_size=224,
+    patch_size=16,
+    drop_path_rate=0,
+    checkpoint_activations=None,
+    mlp_ratio=4,
+    vocab_size=30522,
+    encoder_layers=32,
+    encoder_embed_dim=4096,
+    encoder_attention_heads=32,
+    share_layer=False,
+    share_attn=False,
+    deepnorm=False,
+    mask_ratio=0,
+    max_text_len=52,
+    one_attn=False,
+    **kwargs
+):
+    return EncoderConfig(
+        img_size=img_size,
+        patch_size=patch_size,
+        vocab_size=vocab_size,
+        multiway=True,
+        layernorm_embedding=False,
+        normalize_output=True,
+        no_output_layer=True,
+        drop_path_rate=drop_path_rate,
+        encoder_embed_dim=encoder_embed_dim,
+        encoder_attention_heads=encoder_attention_heads,
+        encoder_layers=encoder_layers,
+        encoder_ffn_embed_dim=int(encoder_embed_dim * mlp_ratio),
+        checkpoint_activations=checkpoint_activations,
+        share_layer=share_layer,
+        share_attn=share_attn,
+        deepnorm=deepnorm,
+        mask_ratio=mask_ratio,
+        max_text_len=max_text_len,
+        one_attn=one_attn,
+    )
+
+
+class BEiT3Wrapper(nn.Module):
+    def __init__(self, args, **kwargs):
+        super().__init__()
+        self.args = args
+        self.beit3 = BEiT3(args)
+        self.apply(self._init_weights)
+
+    def fix_init_weight(self):
+        def rescale(param, layer_id):
+            param.div_(math.sqrt(2.0 * layer_id))
+
+        for layer_id, layer in enumerate(self.blocks):
+            rescale(layer.attn.proj.weight.data, layer_id + 1)
+            rescale(layer.mlp.fc2.weight.data, layer_id + 1)
+
+    def get_num_layers(self):
+        return self.beit3.encoder.num_layers
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {
+            "pos_embed",
+            "cls_token",
+            "beit3.encoder.embed_positions.A.weight",
+            "beit3.vision_embed.cls_token",
+            "logit_scale",
+        }
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=0.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)

vlmo/modules/multiway_transformer.py

@@ -0,0 +1,396 @@
+""" Vision Transformer (ViT) in PyTorch
+
+A PyTorch implement of Vision Transformers as described in
+'An Image Is Worth 16 x 16 Words: Transformers for Image Recognition at Scale' - https://arxiv.org/abs/2010.11929
+
+The official jax code is released and available at https://github.com/google-research/vision_transformer
+
+Acknowledgments:
+* The paper authors for releasing code and weights, thanks!
+* I fixed my class token impl based on Phil Wang's https://github.com/lucidrains/vit-pytorch ... check it out
+for some einops/einsum fun
+* Simple transformer style inspired by Andrej Karpathy's https://github.com/karpathy/minGPT
+* Bert reference code checks against Huggingface Transformers and Tensorflow Bert
+
+DeiT model defs and weights from https://github.com/facebookresearch/deit,
+paper `DeiT: Data-efficient Image Transformers` - https://arxiv.org/abs/2012.12877
+
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+from functools import partial
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from timm.models.layers import DropPath, to_2tuple, trunc_normal_
+from timm.models.registry import register_model
+from pytorch_lightning.utilities.rank_zero import rank_zero_info
+
+
+class Mlp(nn.Module):
+    def __init__(
+        self,
+        in_features,
+        hidden_features=None,
+        out_features=None,
+        act_layer=nn.GELU,
+        drop=0.0,
+    ):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        dim,
+        num_heads=8,
+        qkv_bias=False,
+        qk_scale=None,
+        attn_drop=0.0,
+        proj_drop=0.0,
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        # NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights
+        self.scale = qk_scale or head_dim**-0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=False)
+        if qkv_bias:
+            self.q_bias = nn.Parameter(torch.zeros(dim))
+            self.v_bias = nn.Parameter(torch.zeros(dim))
+        else:
+            self.q_bias = None
+            self.v_bias = None
+
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x, mask=None, relative_position_bias=None):
+        B, N, C = x.shape
+
+        qkv_bias = None
+        if self.q_bias is not None:
+            qkv_bias = torch.cat((self.q_bias, torch.zeros_like(self.v_bias, requires_grad=False), self.v_bias))
+        qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias)
+        qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
+
+        q, k, v = (
+            qkv[0],
+            qkv[1],
+            qkv[2],
+        )  # make torchscript happy (cannot use tensor as tuple)
+
+        q = q * self.scale
+        attn = q.float() @ k.float().transpose(-2, -1)
+
+        if relative_position_bias is not None:
+            attn = attn + relative_position_bias.unsqueeze(0)
+
+        if mask is not None:
+            mask = mask.bool()
+            attn = attn.masked_fill(~mask[:, None, None, :], float("-inf"))
+        attn = attn.softmax(dim=-1).type_as(x)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class Block(nn.Module):
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        mlp_ratio=4.0,
+        qkv_bias=False,
+        qk_scale=None,
+        drop=0.0,
+        attn_drop=0.0,
+        drop_path=0.0,
+        act_layer=nn.GELU,
+        norm_layer=nn.LayerNorm,
+        with_vlffn=False,
+        layer_scale_init_values=0.1,
+        max_text_len=40,
+    ):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            qk_scale=qk_scale,
+            attn_drop=attn_drop,
+            proj_drop=drop,
+        )
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+        self.norm2_text = norm_layer(dim)
+        self.norm2_imag = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp_text = Mlp(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            act_layer=act_layer,
+            drop=drop,
+        )
+        self.mlp_imag = Mlp(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            act_layer=act_layer,
+            drop=drop,
+        )
+        self.mlp_vl = None
+        if with_vlffn:
+            self.mlp_vl = Mlp(
+                in_features=dim,
+                hidden_features=mlp_hidden_dim,
+                act_layer=act_layer,
+                drop=drop,
+            )
+            self.norm2_vl = norm_layer(dim)
+
+        self.gamma_1 = (
+            nn.Parameter(layer_scale_init_values * torch.ones((dim)), requires_grad=True)
+            if layer_scale_init_values is not None
+            else 1.0
+        )
+        self.gamma_2 = (
+            nn.Parameter(layer_scale_init_values * torch.ones((dim)), requires_grad=True)
+            if layer_scale_init_values is not None
+            else 1.0
+        )
+
+        self.max_text_len = max_text_len
+
+    def forward(self, x, mask=None, modality_type=None, relative_position_bias=None):
+        x = x + self.drop_path(
+            self.gamma_1 * self.attn(self.norm1(x), mask=mask, relative_position_bias=relative_position_bias)
+        )
+
+        if modality_type == "image":
+            x = x + self.drop_path(self.gamma_2 * self.mlp_imag(self.norm2_imag(x)))
+        elif modality_type == "text":
+            x = x + self.drop_path(self.gamma_2 * self.mlp_text(self.norm2_text(x)))
+        else:
+            if self.mlp_vl is None:
+                x_text = x[:, : self.max_text_len]
+                x_imag = x[:, self.max_text_len :]
+                x_text = x_text + self.drop_path(self.gamma_2 * self.mlp_text(self.norm2_text(x_text)))
+                x_imag = x_imag + self.drop_path(self.gamma_2 * self.mlp_imag(self.norm2_imag(x_imag)))
+                x = torch.cat([x_text, x_imag], dim=1)
+            else:
+                x = x + self.drop_path(self.gamma_2 * self.mlp_vl(self.norm2_vl(x)))
+
+        return x
+
+
+class PatchEmbed(nn.Module):
+    """Image to Patch Embedding"""
+
+    def __init__(
+        self,
+        img_size=224,
+        patch_size=16,
+        in_chans=3,
+        embed_dim=768,
+        no_patch_embed_bias=False,
+    ):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
+        self.patch_shape = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.num_patches = num_patches
+
+        self.proj = nn.Conv2d(
+            in_chans,
+            embed_dim,
+            kernel_size=patch_size,
+            stride=patch_size,
+            bias=False if no_patch_embed_bias else True,
+        )
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        assert (
+            H == self.img_size[0] and W == self.img_size[1]
+        ), f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+        # FIXME look at relaxing size constraints
+        x = self.proj(x)
+        return x
+
+
+class MultiWayTransformer(nn.Module):
+    """Vision Transformer
+
+    A PyTorch impl of : `An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale`  -
+        https://arxiv.org/abs/2010.11929
+    """
+
+    def __init__(
+        self,
+        img_size=224,
+        patch_size=16,
+        in_chans=3,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4.0,
+        qkv_bias=True,
+        qk_scale=None,
+        drop_rate=0.0,
+        attn_drop_rate=0.0,
+        drop_path_rate=0.0,
+        norm_layer=None,
+        need_relative_position_embed=True,
+        use_abs_pos_emb=False,
+        layer_scale_init_values=0.1,
+        vlffn_start_layer_index=10,
+        config=None,
+    ):
+        """
+        Args:
+            img_size (int, tuple): input image size
+            patch_size (int, tuple): patch size
+            in_chans (int): number of input channels
+            num_classes (int): number of classes for classification head
+            embed_dim (int): embedding dimension
+            depth (int): depth of transformer
+            num_heads (int): number of attention heads
+            mlp_ratio (int): ratio of mlp hidden dim to embedding dim
+            qkv_bias (bool): enable bias for qkv if True
+            qk_scale (float): override default qk scale of head_dim ** -0.5 if set
+            drop_rate (float): dropout rate
+            attn_drop_rate (float): attention dropout rate
+            drop_path_rate (float): stochastic depth rate
+            norm_layer: (nn.Module): normalization layer
+            need_relative_position_embed (bool): enable relative position bias on self-attention
+            use_abs_pos_emb (bool): enable abs pos emb
+            layer_scale_init_values (float or None): layer scale init values, set None to disable
+            vlffn_start_layer_index (int): vl-ffn start index
+            config: (dict): other hyper from pytorch-lighting
+        """
+        super().__init__()
+        drop_path_rate = drop_path_rate if config is None else config["drop_path_rate"]
+        rank_zero_info("drop path rate: {}".format(drop_path_rate))
+        self.use_abs_pos_emb = use_abs_pos_emb
+        self.need_relative_position_embed = need_relative_position_embed
+
+        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
+        norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
+
+        self.patch_embed = PatchEmbed(
+            img_size=img_size,
+            patch_size=patch_size,
+            in_chans=in_chans,
+            embed_dim=embed_dim,
+        )
+        num_patches = self.patch_embed.num_patches
+        self.patch_size = patch_size
+        self.num_heads = num_heads
+        self.vlffn_start_layer_index = vlffn_start_layer_index
+        if config["loss_names"]["textmlm"] > 0:
+            self.vlffn_start_layer_index = depth
+            rank_zero_info(
+                "Set vlffn_start_layer_index={} for text-only pretraining".format(self.vlffn_start_layer_index)
+            )
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim)) if self.use_abs_pos_emb else None
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+        self.blocks = nn.ModuleList(
+            [
+                Block(
+                    dim=embed_dim,
+                    num_heads=num_heads,
+                    mlp_ratio=mlp_ratio,
+                    qkv_bias=qkv_bias,
+                    qk_scale=qk_scale,
+                    drop=drop_rate,
+                    attn_drop=attn_drop_rate,
+                    drop_path=dpr[i],
+                    norm_layer=norm_layer,
+                    with_vlffn=(i >= self.vlffn_start_layer_index),
+                    layer_scale_init_values=layer_scale_init_values,
+                    max_text_len=config["max_text_len"],
+                )
+                for i in range(depth)
+            ]
+        )
+        self.norm = norm_layer(embed_dim)
+
+        if self.pos_embed is not None:
+            trunc_normal_(self.pos_embed, std=0.02)
+        trunc_normal_(self.cls_token, std=0.02)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=0.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {"pos_embed", "cls_token"}
+
+    def visual_embed(self, _x):
+        x = self.patch_embed(_x)
+        x = x.flatten(2).transpose(1, 2)
+        B, L, _ = x.shape
+
+        cls_tokens = self.cls_token.expand(B, -1, -1)
+        x = torch.cat((cls_tokens, x), dim=1)
+
+        if self.pos_embed is not None:
+            x = x + self.pos_embed
+        x = self.pos_drop(x)
+
+        x_mask = torch.ones(x.shape[0], x.shape[1])
+
+        return x, x_mask
+
+
+# VLMo base/p16
+@register_model
+def vlmo_base_patch16(pretrained=False, **kwargs):
+    img_size = kwargs.pop("img_size", 224)
+    model = MultiWayTransformer(
+        img_size=img_size,
+        patch_size=16,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4,
+        qkv_bias=True,
+        vlffn_start_layer_index=10,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),
+        **kwargs,
+    )
+    return model

vlmo/modules/objectives.py

@@ -0,0 +1,12 @@
+import torch.nn as nn
+
+
+def init_weights(module):
+    if isinstance(module, (nn.Linear, nn.Embedding)):
+        module.weight.data.normal_(mean=0.0, std=0.02)
+    elif isinstance(module, nn.LayerNorm):
+        module.bias.data.zero_()
+        module.weight.data.fill_(1.0)
+
+    if isinstance(module, nn.Linear) and module.bias is not None:
+        module.bias.data.zero_()

vlmo/modules/vlmo_module.py

@@ -0,0 +1,405 @@
+import math
+import os
+import time
+
+import numpy as np
+import pytorch_lightning as pl
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+from pytorch_lightning.utilities.rank_zero import rank_zero_info
+from timm.models import create_model
+from transformers import AutoTokenizer, BertTokenizer, XLMRobertaTokenizer  # noqa
+from vlmo.modules import heads, objectives, vlmo_utils
+from vlmo.tokenizer.tokenization_glm import GLMChineseTokenizer  # noqa
+from vlmo.torchscale.architecture.encoder import Encoder
+from vlmo.torchscale.model.BEiT3 import BEiT3 as ts_backbone
+from vlmo.transforms.utils import inception_normalize as img_norm
+
+from .modeling_utils import _get_base_config, _get_large_config, _get_huge_config, trunc_normal_  # noqa
+
+
+def convert_pl_ckpt(state_dict, num_visual_token=197):
+    print("start convert_pl_ckpt!!!")
+    new_state_dict = {}
+    for key in state_dict:
+        value = state_dict[key]
+        if "visual_tokenizer" in key:
+            continue
+        elif "backbone.encoder.embed_positions.A.weight" in key:
+            if value.shape[0] < num_visual_token + 2:
+                N = value.shape[0] - 3
+                dim = value.shape[-1]
+                class_pos_embed = value[:3, ]
+                patch_pos_embed = value[3:, ]
+                w0, h0 = int(math.sqrt(num_visual_token - 1)), int(math.sqrt(num_visual_token - 1))
+                patch_pos_embed = patch_pos_embed.float()
+                patch_pos_embed = nn.functional.interpolate(
+                    patch_pos_embed.reshape(1, int(math.sqrt(N)), int(math.sqrt(N)), dim).permute(0, 3, 1, 2),
+                    size=(w0, h0),
+                    mode="area",
+                )
+                patch_pos_embed = patch_pos_embed.to(class_pos_embed.dtype)
+                patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(-1, dim)
+                new_value = torch.cat((class_pos_embed, patch_pos_embed), dim=0)
+                new_state_dict[key] = new_value
+                print("reshape ", key, "raw shape: ", value.shape, "new shape: ", new_value.shape, num_visual_token)
+            elif value.shape[0] > num_visual_token + 2:
+                new_state_dict[key] = value[: num_visual_token + 2, :]
+                print("first ", key, "raw shape: ", value.shape, new_state_dict[key].shape, num_visual_token)
+            else:
+                new_state_dict[key] = value
+                print("raw shape")
+        else:
+            new_state_dict[key] = state_dict[key]
+
+    return new_state_dict
+
+
+def convert_deepspeed_ckpt(state_dict, num_visual_token=197):
+    new_state_dict = {}
+    for key in state_dict:
+        if key.startswith("_forward_module."):
+            new_key = key[len("_forward_module."):]
+            value = state_dict[key]
+            new_state_dict[new_key] = value
+            if "visual_tokenizer.encoder.pos_embed" in new_key or "visual_tokenizer.decoder.pos_embed" in new_key:
+                if value.shape[1] != num_visual_token:
+                    N = value.shape[1] - 1
+                    dim = value.shape[-1]
+                    class_pos_embed = value[:, 0]
+                    patch_pos_embed = value[:, 1:]
+                    w0, h0 = int(math.sqrt(num_visual_token - 1)), int(math.sqrt(num_visual_token - 1))
+                    patch_pos_embed = patch_pos_embed.float()
+                    patch_pos_embed = nn.functional.interpolate(
+                        patch_pos_embed.reshape(1, int(math.sqrt(N)), int(math.sqrt(N)), dim).permute(0, 3, 1, 2),
+                        size=(w0, h0),
+                        mode="area",
+                    )
+                    patch_pos_embed = patch_pos_embed.to(class_pos_embed.dtype)
+                    patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+                    new_value = torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)
+                    new_state_dict[new_key] = new_value
+                    print("reshape ", new_key, "raw shape: ", value.shape, "new_shape: ", new_value.shape)
+            if "backbone.encoder.embed_positions.A.weight" in new_key:
+                if value.shape[1] != num_visual_token + 2:
+                    N = value.shape[0] - 3
+                    dim = value.shape[-1]
+                    class_pos_embed = value[:3, ]
+                    patch_pos_embed = value[3:, ]
+                    w0, h0 = int(math.sqrt(num_visual_token - 1)), int(math.sqrt(num_visual_token - 1))
+                    patch_pos_embed = patch_pos_embed.float()
+                    patch_pos_embed = nn.functional.interpolate(
+                        patch_pos_embed.reshape(1, int(math.sqrt(N)), int(math.sqrt(N)), dim).permute(0, 3, 1, 2),
+                        size=(w0, h0),
+                        mode="area",
+                    )
+                    patch_pos_embed = patch_pos_embed.to(class_pos_embed.dtype)
+                    patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(-1, dim)
+                    new_value = torch.cat((class_pos_embed, patch_pos_embed), dim=0)
+                    new_state_dict[new_key] = new_value
+                    print("reshape ", new_key, "raw shape: ", value.shape, "new_shape: ", new_value.shape)
+
+        else:
+            new_state_dict[key] = state_dict[key]
+
+    return new_state_dict
+
+
+def get_visual_tokenizer(config):
+    tokenizer_name = config["tokenizer_model"]
+    print(f"Creating visual tokenizer: {tokenizer_name}")
+    model = create_model(
+        config["tokenizer_model"],
+        img_size=config["image_size"],
+        n_code=config["codebook_size"],
+        code_dim=config["codebook_dim"],
+    ).eval()
+    return model
+
+
+def get_pretrained_tokenizer(tokenizer_type, from_pretrained):
+    _Tokenizer = eval(f"{tokenizer_type}")
+    if torch.distributed.is_initialized():
+        if torch.distributed.get_rank() == 0:
+            _Tokenizer.from_pretrained(from_pretrained)
+        torch.distributed.barrier()
+    return _Tokenizer.from_pretrained(from_pretrained)
+
+
+class VLMo(pl.LightningModule):
+    def __init__(self, config):
+        super().__init__()
+        self.save_hyperparameters()
+        s_t = time.time()
+
+        # tokenizer & backbone
+        self.img_size = config["image_size"]
+        if not config["test_only"]:
+            self.visual_tokenizer = get_visual_tokenizer(config)
+        kwargs = {}
+        if "encoder_attention_heads" in config:
+            kwargs["encoder_attention_heads"] = config["encoder_attention_heads"]
+        if "atorch_config" in config and config["atorch_config"]:
+            checkpoint_activations = False  # ?
+        else:
+            checkpoint_activations = config["checkpoint_activations"]
+        args = eval(f'_get_{config["beit_version"]}_config')(
+            img_size=config["image_size"],
+            patch_size=config["patch_size"],
+            vocab_size=config["vocab_size"],
+            encoder_layers=config["encoder_layers"],
+            encoder_embed_dim=config["encoder_embed_dim"],
+            checkpoint_activations=checkpoint_activations,
+            share_layer=config["share_layer"],
+            share_attn=config["share_attn"],
+            deepnorm=config["deepnorm"],
+            mask_ratio=config["mask_ratio"],
+            max_text_len=config["max_text_len"],
+            one_attn=config["one_attn"],
+            **kwargs,
+        )
+        self.num_features = args.encoder_embed_dim
+        self.out_features = config["out_embed_dim"]
+        self.cap_onlytext = config["cap_onlytext"]
+        self.lang = config["lang"]
+        self.num_frames = config["num_frames"]
+        self.tokenizer_type = config["tokenizer_type"]
+        self.text_tokenizer = get_pretrained_tokenizer(self.tokenizer_type, from_pretrained=config["tokenizer"])  # noqa
+        print("BEiT args", args.__dict__)
+        self.backbone = ts_backbone(args)
+
+        self.use_vl = config["beit3_vl_layers"] > 0
+        if self.use_vl:
+            args.encoder_layers = config["beit3_vl_layers"]
+            self.backbone_vl = Encoder(args)
+
+        self.norm = nn.LayerNorm(self.num_features, eps=1e-6)
+
+        # task layers
+        self.pooler = heads.Pooler(self.num_features)
+        self.pooler.apply(objectives.init_weights)
+
+        # contrastive loss (or sampling for global hard negative)
+        if config["loss_names"]["itc"] > 0:
+            self.itc_text_proj = heads.ITCHead(self.num_features, self.out_features)
+            self.itc_image_proj = heads.ITCHead(self.num_features, self.out_features)
+            self.itc_text_proj.apply(objectives.init_weights)
+            self.itc_image_proj.apply(objectives.init_weights)
+
+            self.itc_vl_text_proj = heads.ITCHead(self.num_features, self.out_features)
+            self.itc_vl_image_proj = heads.ITCHead(self.num_features, self.out_features)
+            self.itc_vl_text_proj.apply(objectives.init_weights)
+            self.itc_vl_image_proj.apply(objectives.init_weights)
+
+            self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+            self.logit_vl_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+        lp_s_t = time.time()
+
+        self.load_pretrained_weight()
+        load_pretrain_time = time.time() - lp_s_t
+
+        self.current_tasks = list()
+
+        # ===================== load downstream (test_only) ======================
+
+        if self.hparams.config["load_path"] != "" and self.hparams.config["test_only"]:
+            rank_zero_info("Load ckpt from: {}".format(self.hparams.config["load_path"]))
+            ckpt = torch.load(self.hparams.config["load_path"], map_location="cpu")
+
+            state_dict = None
+
+            for state_dict_key in ("state_dict", "module", "model"):
+                if state_dict_key in ckpt:
+                    rank_zero_info("Read state dict from ckpt[%s]. " % state_dict_key)
+                    state_dict = ckpt[state_dict_key]
+                    break
+            if state_dict_key == "module":
+                state_dict = convert_deepspeed_ckpt(state_dict, self.backbone.vision_embed.num_position_embeddings())
+            if state_dict_key == "state_dict":
+                state_dict = convert_pl_ckpt(state_dict, self.backbone.vision_embed.num_position_embeddings())
+            if state_dict is None:
+                if list(ckpt.keys())[0].startswith('_forward_module.'):
+                    rank_zero_info("Read state dict from ckpt with _forward_module prefix. ")
+                    state_dict = convert_deepspeed_ckpt(ckpt, self.backbone.vision_embed.num_position_embeddings())
+                else:
+                    rank_zero_info("Read state dict from ckpt. ")
+                    state_dict = ckpt
+
+            missing_keys, unexpected_keys = self.load_state_dict(state_dict, strict=False)
+            rank_zero_info("missing_keys: {}".format(missing_keys))
+            rank_zero_info("unexpected_keys: {}".format(unexpected_keys))
+
+        construct_time = time.time() - s_t
+        print(
+            f"Process {os.getpid()}. VLMo Constructor time: {construct_time}s;",
+            f"load_pretrain_time: {load_pretrain_time}s",
+            flush=True,
+        )
+        # coalesce backbone calls
+        self._coalesce_backbone = config["coalesce_backbone"]
+        self._mask_data = config["mask_data"]
+        self._backbone_inputs = {}
+        self._backbone_inputs_current_size = 0
+        self._backbone_inputs_keys = {}
+        self._backbone_outputs = None
+        self._default_attn_masks = {}
+        self._itc_group = None
+        self._itc_aggregate_dict = None
+        self._itc_mask = config["itc_mask"]
+        self._local_loss = config["local_loss"]
+        self._aggregate_nodes = config["aggregate_nodes"]
+        self.accumulated_batches_reached = False
+        vlmo_utils.set_task(self)
+        self._only_itc_single_machine = (
+                self._aggregate_nodes > 0 and len(self.current_tasks) == 1 and "itc" in self.current_tasks
+        )
+        self._split_data_for_imagemlm = config["split_data_for_imagemlm"]
+        self.log_metric_steps = config["log_metric_steps"]
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=0.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def fix_init_weight(self):
+        def rescale(param, layer_id):
+            param.div_(math.sqrt(2.0 * layer_id))
+
+        for layer_id, layer in enumerate(self.backbone.encoder.layers):
+            rescale(layer.self_attn.v_proj.A.weight.data, layer_id + 1)
+            rescale(layer.self_attn.v_proj.B.weight.data, layer_id + 1)
+            rescale(layer.self_attn.out_proj.A.weight.data, layer_id + 1)
+            rescale(layer.self_attn.out_proj.B.weight.data, layer_id + 1)
+            rescale(layer.ffn.A.fc2.weight.data, layer_id + 1)
+            rescale(layer.ffn.B.fc2.weight.data, layer_id + 1)
+
+        if self.use_vl:
+            pre_layers = len(self.backbone.encoder.layers) + 1
+            for layer_id, layer in enumerate(self.backbone_vl.layers):
+                rescale(layer.self_attn.v_proj.A.weight.data, layer_id + pre_layers)
+                rescale(layer.self_attn.v_proj.B.weight.data, layer_id + pre_layers)
+                rescale(layer.self_attn.out_proj.A.weight.data, layer_id + pre_layers)
+                rescale(layer.self_attn.out_proj.B.weight.data, layer_id + pre_layers)
+                rescale(layer.ffn.A.fc2.weight.data, layer_id + pre_layers)
+                rescale(layer.ffn.B.fc2.weight.data, layer_id + pre_layers)
+
+    def load_pretrained_weight(self):
+        if self.hparams.config["load_path"] != "" and not self.hparams.config["test_only"]:
+            config = self.hparams.config
+            ckpt = torch.load(self.hparams.config["load_path"], map_location="cpu")
+            rank_zero_info("Load ckpt from: {}".format(self.hparams.config["load_path"]))
+
+            state_dict = None
+
+            for state_dict_key in ("state_dict", "module", "model"):
+                if state_dict_key in ckpt:
+                    rank_zero_info("Read state dict from ckpt[%s]. " % state_dict_key)
+                    state_dict = ckpt[state_dict_key]
+                    break
+            if state_dict_key == "module":
+                state_dict = convert_deepspeed_ckpt(state_dict, self.backbone.vision_embed.num_position_embeddings())
+            if state_dict_key == "state_dict":
+                state_dict = convert_pl_ckpt(state_dict, self.backbone.vision_embed.num_position_embeddings())
+            if state_dict is None:
+                if list(ckpt.keys())[0].startswith('_forward_module.'):
+                    rank_zero_info("Read state dict from ckpt with _forward_module prefix. ")
+                    state_dict = convert_deepspeed_ckpt(ckpt,
+                                                        self.backbone.vision_embed.num_position_embeddings())
+                else:
+                    rank_zero_info("Read state dict from ckpt. ")
+                    state_dict = ckpt
+
+            missing_keys, unexpected_keys = self.load_state_dict(state_dict, strict=False)
+            missing_keys = [k for k in missing_keys if "itc_teacher" not in k]
+            rank_zero_info("missing_keys: {}".format(missing_keys))
+            rank_zero_info("unexpected_keys: {}".format(unexpected_keys))
+
+    def infer_text(
+            self,
+            batch,
+            mask_text=False,
+    ):
+        do_mlm = "_mlm" if mask_text else ""
+        text_ids = batch[f"text_ids{do_mlm}"]
+        text_labels = batch[f"text_labels{do_mlm}"]
+        text_masks = batch[f"text_masks"]
+        text_embed = self.backbone.text_embed(text_ids)
+        text_padding_position = 1 - text_masks
+        lffn_hiddens = self.backbone(
+            textual_tokens=text_ids,
+            text_padding_position=text_padding_position,
+        )["encoder_out"]
+        vlffn_hiddens = self.backbone_vl(
+            src_tokens=None,
+            token_embeddings=lffn_hiddens,
+            encoder_padding_mask=text_padding_position,
+            multiway_split_position=-1,
+        )["encoder_out"]
+
+        cls_feats = self.itc_text_proj(lffn_hiddens[:, 0])
+        cls_feats = cls_feats / cls_feats.norm(dim=-1, keepdim=True)
+
+        cls_vlffn_feats = self.itc_vl_text_proj(vlffn_hiddens[:, 0])
+        cls_vlffn_feats = cls_vlffn_feats / cls_vlffn_feats.norm(dim=-1, keepdim=True)
+
+        ret = {
+            "cls_feats": cls_feats,
+            "cls_vlffn_feats": cls_vlffn_feats,
+            "text_embed": text_embed,
+        }
+
+        return ret
+
+    def infer_image(
+            self,
+            batch,
+            mask_image=False,
+            image_token_type_idx=1,
+            image_embeds=None,
+            image_masks=None,
+    ):
+        if f"image_{image_token_type_idx - 1}" in batch:
+            imgkey = f"image_{image_token_type_idx - 1}"
+        else:
+            imgkey = "image"
+
+        img = batch[imgkey][0]
+        if mask_image:
+            image_masks = batch[f"{imgkey}_masks"][0].flatten(1)
+
+            with torch.no_grad():
+                img = self.visual_tokenizer.pre_process(img)
+                quantize, embed_ind, _ = self.visual_tokenizer.encode(img)
+                image_ids = embed_ind.view(img.shape[0], -1)
+
+                image_labels = torch.full_like(image_ids, -100)
+                bool_masked_pos = image_masks.to(torch.bool)
+                image_labels[bool_masked_pos] = image_ids[bool_masked_pos]
+
+        img_tensor = img_norm(img)
+        vffn_hiddens = self.backbone(visual_tokens=img_tensor)["encoder_out"]
+        vlffn_hiddens = self.backbone_vl(
+            src_tokens=None,
+            token_embeddings=vffn_hiddens,
+            multiway_split_position=-1,
+        )["encoder_out"]
+
+        cls_feats = self.itc_image_proj(vffn_hiddens[:, 0])
+        cls_feats = cls_feats / cls_feats.norm(dim=-1, keepdim=True)
+
+        cls_vlffn_feats = self.itc_vl_image_proj(vlffn_hiddens[:, 0])
+        cls_vlffn_feats = cls_vlffn_feats / cls_vlffn_feats.norm(dim=-1, keepdim=True)
+
+        ret = {
+            "image_feats": vffn_hiddens,
+            "cls_feats": cls_feats,
+            "cls_vlffn_feats": cls_vlffn_feats,
+        }
+
+        return ret

vlmo/modules/vlmo_utils.py

@@ -0,0 +1,12 @@
+def set_task(pl_module):
+    pl_module.current_tasks = [k for k, v in pl_module.hparams.config["loss_names"].items() if v >= 1]
+    return
+
+
+def no_sync_module_apply(module, fn):
+    """FSDP module .apply will use _unshard_params_recurse which will sync params across ranks.
+    using this function when apply fn is unnecessary to sync params across ranks.
+    """
+    for child in module.children():
+        fn(child)
+        no_sync_module_apply(child, fn)

vlmo/tokenizer/__init__.py

@@ -0,0 +1,6 @@
+# coding: utf-8
+# Copyright (c) Antfin, Inc. All rights reserved.
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function

vlmo/tokenizer/sp.model

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

vlmo/tokenizer/tokenization_glm.py

@@ -0,0 +1,307 @@
+import os
+from shutil import copyfile
+from typing import Optional, Tuple, List, Union
+
+import sentencepiece as spm
+import torch
+from transformers import PreTrainedTokenizer
+from transformers.models.auto.tokenization_auto import get_tokenizer_config
+from transformers.tokenization_utils_base import BatchEncoding
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+class GLMBatchEncoding(BatchEncoding):
+    def to(self, device: Union[str, "torch.device"]) -> "BatchEncoding":
+        """
+        Send all values to device by calling `v.to(device)` (PyTorch only).
+
+        Args:
+            device (`str` or `torch.device`): The device to put the tensors on.
+
+        Returns:
+            [`BatchEncoding`]: The same instance after modification.
+        """
+
+        # This check catches things like APEX blindly calling "to" on all inputs to a module
+        # Otherwise it passes the casts down and casts the LongTensor containing the token idxs
+        # into a HalfTensor
+        if isinstance(device, str) or isinstance(device, int):
+        #if isinstance(device, str) or _is_torch_device(device) or isinstance(device, int):
+            self.data = {k: v.to(device=device) if torch.is_tensor(v) else v for k, v in self.data.items()}
+        else:
+            logger.warning(f"Attempting to cast a BatchEncoding to type {str(device)}. This is not supported.")
+        return self
+
+
+class GLMTokenizerMixin:
+    @property
+    def sop_token(self) -> Optional[str]:
+        return "<|startofpiece|>"
+
+    @property
+    def sop_token_id(self) -> Optional[int]:
+        """
+        `Optional[int]`: Id of the start token in the vocabulary, used when training a model with autoregressive blank filling.
+        """
+        return self.convert_tokens_to_ids(self.sop_token)
+
+    @property
+    def eop_token(self) -> Optional[str]:
+        return "<|endofpiece|>"
+
+    @property
+    def eop_token_id(self) -> Optional[int]:
+        """
+        `Optional[int]`: Id of the end token in the vocabulary, used when training a model with autoregressive blank filling.
+        """
+        return self.convert_tokens_to_ids(self.eop_token)
+
+    @property
+    def gmask_token_id(self) -> int:
+        return self.convert_tokens_to_ids("[gMASK]")
+
+    @property
+    def smask_token_id(self) -> int:
+        return self.convert_tokens_to_ids("[sMASK]")
+
+    @property
+    def mask_token_ids(self):
+        return [self.mask_token_id, self.smask_token_id, self.gmask_token_id]
+
+    def _build_input_for_multiple_choice(self, context, choices):
+        context_id = context["input_ids"]
+        if torch.is_tensor(context_id):
+            context_id = context_id.tolist()
+
+        division = len(context_id)
+        mask_position = context_id.index(self.mask_token_id)
+
+        token = torch.tensor(context_id, dtype=torch.long)
+        attention_mask = [context["attention_mask"].expand(division, -1)]
+        position_id = torch.arange(division, dtype=torch.long)
+        block_position_id = torch.zeros(division, dtype=torch.long)
+
+        choice_ids, choice_indices = [], []
+
+        for choice_str in choices:
+            choice = torch.tensor(self(choice_str, add_special_tokens=False, padding=False)['input_ids'],
+                                  dtype=torch.long)
+            choice_ids.append(choice)
+            choice_indices.append(torch.arange(len(token), len(token) + len(choice), dtype=torch.long))
+            attention_mask.append(torch.tril(torch.ones((len(choice), len(choice)), dtype=torch.long)))
+
+            token = torch.cat((token, torch.tensor([self.sop_token_id], dtype=torch.long), choice[:-1]))
+            position_id = torch.cat((position_id, torch.tensor([mask_position] * len(choice), dtype=torch.long)))
+            block_position_id = torch.cat((block_position_id, torch.arange(1, 1 + len(choice), dtype=torch.long)))
+
+        attention_mask = torch.block_diag(*attention_mask)
+        attention_mask[division:, :division] = context["attention_mask"].unsqueeze(0)
+
+        return {
+            "input_ids": token,
+            "position_ids": torch.stack((position_id, block_position_id)),
+            "attention_mask": attention_mask,
+            "choice_ids": choice_ids,
+            "choice_indices": choice_indices
+        }
+
+    def _pad_batch(self, tokens, position_ids, attention_mask, max_seq_length):
+        pad_length = max_seq_length - len(tokens)
+        attention_mask = torch.nn.functional.pad(
+            attention_mask,
+            (0, pad_length, 0, pad_length),
+            mode="constant",
+            value=0,
+        )
+        tokens = torch.cat((tokens, torch.zeros(pad_length, dtype=torch.long)))
+        position_ids = torch.cat((position_ids, position_ids[..., -1:].expand(-1, pad_length)), dim=-1)
+        return tokens, position_ids, attention_mask
+
+    def _collate(self, samples):
+        TILE = 1
+        length_to_pad = (max(map(lambda spl: len(spl["input_ids"]), samples)) + TILE - 1) // TILE * TILE
+
+        token_batch, position_id_batch, attention_mask_batch = [], [], []
+        choices_batch, choice_target_ids_batch = [], []
+
+        for sample in samples:
+            token, position_id, attention_mask = self._pad_batch(
+                sample["input_ids"], sample["position_ids"], sample["attention_mask"], length_to_pad
+            )
+            token_batch.append(token)
+            position_id_batch.append(position_id)
+            attention_mask_batch.append(attention_mask)
+            choices_batch.append(sample["choice_ids"])
+            choice_target_ids_batch.append(sample["choice_indices"])
+        return {
+            "input_ids": torch.stack(token_batch),
+            "position_ids": torch.stack(position_id_batch),
+            "attention_mask": torch.stack(attention_mask_batch).unsqueeze(1),
+            "choice_ids": choices_batch,
+            "choice_indices": choice_target_ids_batch,
+        }
+
+    def build_inputs_for_multiple_choice(self, model_input: BatchEncoding, choices, max_length=None):
+        samples = [{key: value[i] for key, value in model_input.items()} for i in range(len(model_input["input_ids"]))]
+        samples = [self._build_input_for_multiple_choice(sample, choice) for sample, choice in
+                   zip(samples, choices)]
+        inputs = self._collate(samples)
+        return GLMBatchEncoding(inputs)
+
+    def build_inputs_for_generation(self, model_input: BatchEncoding, max_gen_length=512, targets=None, padding=False):
+        mask_ids = self.mask_token_ids
+        input_ids = model_input.input_ids
+        batch_size, seq_length = input_ids.shape[:2]
+        position_id, block_position_id = list(range(seq_length)), [0 for _ in range(seq_length)]
+        position_ids, block_position_ids = [], []
+        labels = None
+        if targets is not None:
+            is_batched = isinstance(targets, (list, tuple))
+            targets = self(targets, add_special_tokens=False, padding=False).input_ids
+            if not is_batched:
+                targets = [targets]
+            assert len(targets) == len(input_ids)
+            targets = [(target + [self.eop_token_id])[:max_gen_length] for target in targets]
+            if not padding:
+                max_gen_length = max(map(len, targets))
+            targets = [[self.sop_token_id] + target for target in targets]
+            labels = [target[1:] for target in targets]
+            targets = [target + [self.pad_token_id] * (max_gen_length + 1 - len(target)) for target in targets]
+            labels = [label + [-100] * (max_gen_length - len(label)) for label in labels]
+            targets = torch.tensor(targets, dtype=input_ids.dtype, device=input_ids.device)
+            labels = torch.tensor(labels, dtype=input_ids.dtype, device=input_ids.device)
+            labels = torch.cat((input_ids.new_full((batch_size, seq_length), -100), labels), dim=1)
+        for i in range(batch_size):
+            mask_positions = []
+            for mask_id in mask_ids:
+                mask_positions += (input_ids[i] == mask_id).nonzero(as_tuple=True)[0].tolist()
+            if not mask_positions:
+                raise ValueError("Cannot find mask token in the input")
+            mask_positions.sort()
+            mask_pos = mask_positions[0]
+            position_ids.append(position_id + [mask_pos] * max_gen_length)
+            block_position_ids.append(block_position_id + list(range(1, max_gen_length + 1)))
+        position_ids = torch.tensor(position_ids, dtype=input_ids.dtype, device=input_ids.device)
+        block_position_ids = torch.tensor(block_position_ids, dtype=input_ids.dtype, device=input_ids.device)
+        position_ids = torch.stack((position_ids, block_position_ids), dim=1)
+        attention_mask = model_input.attention_mask
+        attention_mask = attention_mask.unsqueeze(1).expand(-1, seq_length + max_gen_length, -1)
+        generation_attention_mask = torch.cat([attention_mask.new_zeros((seq_length, max_gen_length)),
+                                               torch.tril(attention_mask.new_ones((max_gen_length, max_gen_length)))],
+                                              dim=0).unsqueeze(0).expand(batch_size, -1, -1)
+        attention_mask = torch.cat((attention_mask, generation_attention_mask), dim=2)
+        attention_mask = attention_mask.unsqueeze(1)
+        if targets is None:
+            input_ids = torch.cat((input_ids, input_ids.new_full((batch_size, 1), self.sop_token_id)), dim=-1)
+        else:
+            input_ids = torch.cat((input_ids, targets[:, :-1]), dim=1)
+        batch = {"input_ids": input_ids, "position_ids": position_ids}
+        if labels is None:
+            batch["generation_attention_mask"] = attention_mask
+        else:
+            batch["attention_mask"] = attention_mask
+            batch["labels"] = labels
+        return BatchEncoding(batch)
+
+def encode_whitespaces(content):
+    for i in range(10, 1, -1):
+        content = content.replace(' '*i, f'<|blank_{i}|>')
+    return content
+
+def decode_whitespaces(content):
+    for i in range(10, 1, -1):
+        content = content.replace(f'<|blank_{i}|>', ' '*i)
+    return content
+
+
+class GLMChineseTokenizer(PreTrainedTokenizer, GLMTokenizerMixin):
+    vocab_files_names = {"vocab_file": "sp.model"}
+    truncation_side: str = "left"
+
+    def __init__(self, vocab_file, **kwargs):
+        self.vocab_file = vocab_file
+        self.sp_model = spm.SentencePieceProcessor()
+        self.sp_model.Load(vocab_file)
+        super().__init__(**kwargs)
+
+    @property
+    def vocab_size(self):
+        return len(self.sp_model)
+
+    def get_vocab(self):
+        vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)}
+        vocab.update(self.added_tokens_encoder)
+        return vocab
+
+    def _tokenize(self, text, **kwargs):
+        text = encode_whitespaces(text)
+        return self.sp_model.EncodeAsPieces(text)
+        #return self.sp_model.EncodeAsPieces(text, out_type=str)
+
+    def _convert_token_to_id(self, token):
+        """Converts a token (str) in an id using the vocab."""
+        return self.sp_model.PieceToId(token)
+
+    def _convert_id_to_token(self, index):
+        """Converts an index (integer) in a token (str) using the vocab."""
+        return self.sp_model.IdToPiece(index)
+
+    def convert_tokens_to_string(self, tokens):
+        res = self.sp_model.DecodeIds(tokens)
+        return decode_whitespaces(res)
+
+    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
+        if not os.path.isdir(save_directory):
+            logger.error(f"Vocabulary path ({save_directory}) should be a directory")
+            return
+        out_vocab_file = os.path.join(
+            save_directory, (filename_prefix + "-" if filename_prefix else "") + self.vocab_files_names["vocab_file"]
+        )
+
+        if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file) and os.path.isfile(self.vocab_file):
+            copyfile(self.vocab_file, out_vocab_file)
+        elif not os.path.isfile(self.vocab_file):
+            with open(out_vocab_file, "wb") as fi:
+                content_spiece_model = self.sp_model.serialized_model_proto()
+                fi.write(content_spiece_model)
+
+        return (out_vocab_file,)
+
+    def build_inputs_with_special_tokens(
+            self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
+    ) -> List[int]:
+        """
+        Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
+        adding special tokens. A BERT sequence has the following format:
+
+        - single sequence: ``[CLS] X [SEP]``
+        - pair of sequences: ``[CLS] A [SEP] B [SEP]``
+
+        Args:
+            token_ids_0 (:obj:`List[int]`):
+                List of IDs to which the special tokens will be added.
+            token_ids_1 (:obj:`List[int]`, `optional`):
+                Optional second list of IDs for sequence pairs.
+
+        Returns:
+            :obj:`List[int]`: List of `input IDs <../glossary.html#input-ids>`__ with the appropriate special tokens.
+        """
+        assert token_ids_1 is None
+        cls = [self.cls_token_id]
+        eos = [self.eos_token_id]
+        return cls + token_ids_0 + eos
+
+
+class GLMTokenizer:
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *inputs, **kwargs):
+        tokenizer_config = get_tokenizer_config(pretrained_model_name_or_path, **kwargs)
+        config_tokenizer_class = tokenizer_config.get("tokenizer_class")
+
+        if config_tokenizer_class == "GLMChineseTokenizer":
+            tokenizer_class = GLMChineseTokenizer
+        else:
+            raise NotImplementedError("Not implemented tokenizer type:", config_tokenizer_class)
+        return tokenizer_class.from_pretrained(pretrained_model_name_or_path, *inputs, **kwargs)

vlmo/tokenizer/tokenizer_config.json

@@ -0,0 +1,17 @@
+{
+  "name_or_path": "THUDM/glm-10b-chinese",
+  "eos_token": "<|endoftext|>",
+  "pad_token": "<|endoftext|>",
+  "cls_token": "[CLS]",
+  "mask_token": "[MASK]",
+  "unk_token": "[UNK]",
+  "add_prefix_space": false,
+  "tokenizer_class": "GLMChineseTokenizer",
+  "use_fast": false,
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenization_glm.GLMChineseTokenizer",
+      null
+      ]
+  }
+}

vlmo/torchscale/__init__.py

@@ -0,0 +1,2 @@
+# Copyright (c) 2022 Microsoft
+# Licensed under The MIT License [see LICENSE for details]

vlmo/torchscale/architecture/__init__.py

@@ -0,0 +1,2 @@
+# Copyright (c) 2022 Microsoft
+# Licensed under The MIT License [see LICENSE for details]

vlmo/torchscale/architecture/config.py

@@ -0,0 +1,197 @@
+# Copyright (c) 2022 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+
+
+class EncoderConfig(object):
+    def __init__(self, **kwargs):
+        self.encoder_embed_dim = kwargs.pop("encoder_embed_dim", 768)
+        self.encoder_attention_heads = kwargs.pop("encoder_attention_heads", 12)
+        self.encoder_ffn_embed_dim = kwargs.pop("encoder_ffn_embed_dim", 3072)
+        self.encoder_layers = kwargs.pop("encoder_layers", 12)
+        self.encoder_normalize_before = kwargs.pop("encoder_normalize_before", True)
+        self.normalize_output = kwargs.pop("normalize_output", True)
+        self.activation_fn = kwargs.pop("activation_fn", "gelu")
+        self.dropout = kwargs.pop("dropout", 0.0)
+        self.drop_path_rate = kwargs.pop("drop_path_rate", 0.0)
+        self.attention_dropout = kwargs.pop("attention_dropout", 0.0)
+        self.activation_dropout = kwargs.pop("activation_dropout", 0.0)
+        self.no_scale_embedding = kwargs.pop("no_scale_embedding", True)
+        self.layernorm_embedding = kwargs.pop("layernorm_embedding", False)
+        self.moe_freq = kwargs.pop("moe_freq", 0)
+        self.moe_top1_expert = kwargs.pop("moe_top1_expert", False)
+        self.moe_expert_count = kwargs.pop("moe_expert_count", 0)
+        self.moe_gating_use_fp32 = kwargs.pop("moe_gating_use_fp32", True)
+        self.moe_eval_capacity_token_fraction = kwargs.pop("moe_eval_capacity_token_fraction", 0.25)
+        self.moe_second_expert_policy = kwargs.pop("moe_second_expert_policy", "random")
+        self.moe_normalize_gate_prob_before_dropping = kwargs.pop("moe_normalize_gate_prob_before_dropping", False)
+        self.use_xmoe = kwargs.pop("use_xmoe", False)
+        self.rel_pos_buckets = kwargs.pop("rel_pos_buckets", 0)
+        self.max_rel_pos = kwargs.pop("max_rel_pos", 0)
+        self.deepnorm = kwargs.pop("deepnorm", False)
+        self.subln = kwargs.pop("subln", True)
+        self.bert_init = kwargs.pop("bert_init", False)
+        self.multiway = kwargs.pop("multiway", False)
+        self.share_encoder_input_output_embed = kwargs.pop("share_encoder_input_output_embed", False)
+        self.max_source_positions = kwargs.pop("max_source_positions", 1024)
+        self.no_output_layer = kwargs.pop("no_output_layer", False)
+        self.layernorm_eps = kwargs.pop("layernorm_eps", 1e-5)
+        self.share_layer = kwargs.pop("share_layer", False)
+        self.share_attn = kwargs.pop("share_attn", False)
+        self.mask_ratio = kwargs.pop("mask_ratio", 0)
+        self.max_text_len = kwargs.pop("max_text_len", 52)
+        self.one_attn = kwargs.pop('one_attn', False)
+
+
+        # Text
+        self.vocab_size = kwargs.pop("vocab_size", -1)
+        # Vision
+        self.img_size = kwargs.pop("img_size", 224)
+        self.patch_size = kwargs.pop("patch_size", 16)
+        self.in_chans = kwargs.pop("in_chans", 3)
+        # Fairscale
+        self.checkpoint_activations = kwargs.pop("checkpoint_activations", False)
+        self.fsdp = kwargs.pop("fsdp", False)
+        self.ddp_rank = kwargs.pop("ddp_rank", 0)
+        self.xpos_rel_pos = kwargs.pop("xpos_rel_pos", False)
+        self.xpos_scale_base = kwargs.pop("xpos_scale_base", 512)
+
+        if self.deepnorm:
+            self.encoder_normalize_before = False
+            self.subln = False
+        if self.subln:
+            self.encoder_normalize_before = True
+            self.deepnorm = False
+        if self.use_xmoe:
+            self.moe_normalize_gate_prob_before_dropping = True
+            self.moe_second_expert_policy = "random"
+            assert self.moe_freq > 0 and self.moe_expert_count > 0
+
+    def override(self, args):
+        for hp in self.__dict__.keys():
+            if getattr(args, hp, None) is not None:
+                self.__dict__[hp] = getattr(args, hp, None)
+
+
+class DecoderConfig(object):
+    def __init__(self, **kwargs):
+        self.decoder_embed_dim = kwargs.pop("decoder_embed_dim", 768)
+        self.decoder_attention_heads = kwargs.pop("decoder_attention_heads", 12)
+        self.decoder_ffn_embed_dim = kwargs.pop("decoder_ffn_embed_dim", 3072)
+        self.decoder_layers = kwargs.pop("decoder_layers", 12)
+        self.decoder_normalize_before = kwargs.pop("decoder_normalize_before", True)
+        self.activation_fn = kwargs.pop("activation_fn", "gelu")
+        self.dropout = kwargs.pop("dropout", 0.0)
+        self.drop_path_rate = kwargs.pop("drop_path_rate", 0.0)
+        self.attention_dropout = kwargs.pop("attention_dropout", 0.0)
+        self.activation_dropout = kwargs.pop("activation_dropout", 0.0)
+        self.no_scale_embedding = kwargs.pop("no_scale_embedding", True)
+        self.layernorm_embedding = kwargs.pop("layernorm_embedding", False)
+        self.moe_freq = kwargs.pop("moe_freq", 0)
+        self.moe_top1_expert = kwargs.pop("moe_top1_expert", False)
+        self.moe_expert_count = kwargs.pop("moe_expert_count", 0)
+        self.moe_gating_use_fp32 = kwargs.pop("moe_gating_use_fp32", True)
+        self.moe_eval_capacity_token_fraction = kwargs.pop("moe_eval_capacity_token_fraction", 0.25)
+        self.moe_second_expert_policy = kwargs.pop("moe_second_expert_policy", "random")
+        self.moe_normalize_gate_prob_before_dropping = kwargs.pop("moe_normalize_gate_prob_before_dropping", False)
+        self.use_xmoe = kwargs.pop("use_xmoe", False)
+        self.rel_pos_buckets = kwargs.pop("rel_pos_buckets", 0)
+        self.max_rel_pos = kwargs.pop("max_rel_pos", 0)
+        self.deepnorm = kwargs.pop("deepnorm", False)
+        self.subln = kwargs.pop("subln", True)
+        self.bert_init = kwargs.pop("bert_init", False)
+        self.multiway = kwargs.pop("multiway", False)
+        self.share_decoder_input_output_embed = kwargs.pop("share_decoder_input_output_embed", False)
+        self.max_target_positions = kwargs.pop("max_target_positions", 1024)
+        self.no_output_layer = kwargs.pop("no_output_layer", False)
+        self.layernorm_eps = kwargs.pop("layernorm_eps", 1e-5)
+        # Text
+        self.vocab_size = kwargs.pop("vocab_size", -1)
+        # Fairscale
+        self.checkpoint_activations = kwargs.pop("checkpoint_activations", False)
+        self.fsdp = kwargs.pop("fsdp", False)
+        self.ddp_rank = kwargs.pop("ddp_rank", 0)
+        self.xpos_rel_pos = kwargs.pop("xpos_rel_pos", False)
+        self.xpos_scale_base = kwargs.pop("xpos_scale_base", 512)
+
+        if self.deepnorm:
+            self.decoder_normalize_before = False
+            self.subln = False
+        if self.subln:
+            self.decoder_normalize_before = True
+            self.deepnorm = False
+        if self.use_xmoe:
+            self.moe_normalize_gate_prob_before_dropping = True
+            self.moe_second_expert_policy = "random"
+            assert self.moe_freq > 0 and self.moe_expert_count > 0
+
+    def override(self, args):
+        for hp in self.__dict__.keys():
+            if getattr(args, hp, None) is not None:
+                self.__dict__[hp] = getattr(args, hp, None)
+
+
+class EncoderDecoderConfig(object):
+    def __init__(self, **kwargs):
+        self.encoder_embed_dim = kwargs.pop("encoder_embed_dim", 768)
+        self.encoder_attention_heads = kwargs.pop("encoder_attention_heads", 12)
+        self.encoder_ffn_embed_dim = kwargs.pop("encoder_ffn_embed_dim", 3072)
+        self.encoder_layers = kwargs.pop("encoder_layers", 12)
+        self.encoder_normalize_before = kwargs.pop("encoder_normalize_before", True)
+        self.decoder_embed_dim = kwargs.pop("decoder_embed_dim", 768)
+        self.decoder_attention_heads = kwargs.pop("decoder_attention_heads", 12)
+        self.decoder_ffn_embed_dim = kwargs.pop("decoder_ffn_embed_dim", 3072)
+        self.decoder_layers = kwargs.pop("decoder_layers", 12)
+        self.decoder_normalize_before = kwargs.pop("decoder_normalize_before", True)
+        self.activation_fn = kwargs.pop("activation_fn", "gelu")
+        self.dropout = kwargs.pop("dropout", 0.0)
+        self.drop_path_rate = kwargs.pop("drop_path_rate", 0.0)
+        self.attention_dropout = kwargs.pop("attention_dropout", 0.0)
+        self.activation_dropout = kwargs.pop("activation_dropout", 0.0)
+        self.no_scale_embedding = kwargs.pop("no_scale_embedding", True)
+        self.layernorm_embedding = kwargs.pop("layernorm_embedding", False)
+        self.moe_freq = kwargs.pop("moe_freq", 0)
+        self.moe_top1_expert = kwargs.pop("moe_top1_expert", False)
+        self.moe_expert_count = kwargs.pop("moe_expert_count", 0)
+        self.moe_gating_use_fp32 = kwargs.pop("moe_gating_use_fp32", True)
+        self.moe_eval_capacity_token_fraction = kwargs.pop("moe_eval_capacity_token_fraction", 0.25)
+        self.moe_second_expert_policy = kwargs.pop("moe_second_expert_policy", "random")
+        self.moe_normalize_gate_prob_before_dropping = kwargs.pop("moe_normalize_gate_prob_before_dropping", False)
+        self.use_xmoe = kwargs.pop("use_xmoe", False)
+        self.rel_pos_buckets = kwargs.pop("rel_pos_buckets", 0)
+        self.max_rel_pos = kwargs.pop("max_rel_pos", 0)
+        self.deepnorm = kwargs.pop("deepnorm", False)
+        self.subln = kwargs.pop("subln", True)
+        self.bert_init = kwargs.pop("bert_init", False)
+        self.multiway = kwargs.pop("multiway", False)
+        self.share_all_embeddings = kwargs.pop("share_all_embeddings", False)
+        self.share_decoder_input_output_embed = kwargs.pop("share_decoder_input_output_embed", False)
+        self.max_source_positions = kwargs.pop("max_source_positions", 1024)
+        self.max_target_positions = kwargs.pop("max_target_positions", 1024)
+        self.no_output_layer = kwargs.pop("no_output_layer", False)
+        self.layernorm_eps = kwargs.pop("layernorm_eps", 1e-5)
+        # Text
+        self.vocab_size = kwargs.pop("vocab_size", -1)
+        # Fairscale
+        self.checkpoint_activations = kwargs.pop("checkpoint_activations", False)
+        self.fsdp = kwargs.pop("fsdp", False)
+        self.ddp_rank = kwargs.pop("ddp_rank", 0)
+        self.xpos_rel_pos = kwargs.pop("xpos_rel_pos", False)
+        self.xpos_scale_base = kwargs.pop("xpos_scale_base", 512)
+
+        if self.deepnorm:
+            self.encoder_normalize_before = False
+            self.decoder_normalize_before = False
+            self.subln = False
+        if self.subln:
+            self.encoder_normalize_before = True
+            self.decoder_normalize_before = True
+            self.deepnorm = False
+        if self.use_xmoe:
+            self.moe_normalize_gate_prob_before_dropping = True
+            self.moe_second_expert_policy = "random"
+            assert self.moe_freq > 0 and self.moe_expert_count > 0
+
+    def override(self, args):
+        for hp in self.__dict__.keys():
+            if getattr(args, hp, None) is not None:
+                self.__dict__[hp] = getattr(args, hp, None)

vlmo/torchscale/architecture/decoder.py

@@ -0,0 +1,428 @@
+# Copyright (c) 2022 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+
+import math
+
+import numpy as np
+import torch
+import torch.nn as nn
+from fairscale.nn import checkpoint_wrapper, wrap
+
+from vlmo.torchscale.architecture.utils import init_bert_params
+from vlmo.torchscale.component.droppath import DropPath
+from vlmo.torchscale.component.feedforward_network import FeedForwardNetwork, make_experts
+from vlmo.torchscale.component.multihead_attention import MultiheadAttention
+from vlmo.torchscale.component.relative_position_bias import RelativePositionBias
+#from vlmo.torchscale.component.xmoe.moe_layer import MOELayer
+#from vlmo.torchscale.component.xmoe.routing import Top1Gate, Top2Gate
+
+try:
+    from apex.normalization import FusedLayerNorm as LayerNorm
+except ModuleNotFoundError:
+    from torch.nn import LayerNorm
+
+
+class DecoderLayer(nn.Module):
+    def __init__(
+        self,
+        args,
+        depth,
+        is_moe_layer=False,
+        is_encoder_decoder=False,
+    ):
+        super().__init__()
+        self.args = args
+        self.embed_dim = args.decoder_embed_dim
+        self.dropout_module = torch.nn.Dropout(args.dropout)
+
+        if args.drop_path_rate > 0:
+            drop_path_prob = np.linspace(0, args.drop_path_rate, args.decoder_layers)[depth]
+            self.drop_path = DropPath(drop_path_prob)
+        else:
+            self.drop_path = None
+
+        self.self_attn = self.build_self_attention(self.embed_dim, args)
+
+        self.normalize_before = args.decoder_normalize_before
+
+        self.self_attn_layer_norm = LayerNorm(self.embed_dim, eps=args.layernorm_eps)
+
+        if not is_encoder_decoder:
+            self.encoder_attn = None
+            self.encoder_attn_layer_norm = None
+        else:
+            self.encoder_attn = self.build_encoder_attention(self.embed_dim, args)
+            self.encoder_attn_layer_norm = LayerNorm(self.embed_dim, eps=args.layernorm_eps)
+
+        self.is_moe_layer = is_moe_layer
+        self.ffn_dim = args.decoder_ffn_embed_dim
+
+        if not self.is_moe_layer:
+            self.ffn = self.build_ffn(
+                self.embed_dim,
+                self.args,
+            )
+        else:
+            if args.moe_top1_expert:
+                gate = Top1Gate(
+                    self.embed_dim,
+                    args.moe_expert_count,
+                    use_fp32=args.moe_gating_use_fp32,
+                    moe_eval_capacity_token_fraction=args.moe_eval_capacity_token_fraction,
+                    use_xmoe=args.use_xmoe,
+                )
+            else:
+                gate = Top2Gate(
+                    self.embed_dim,
+                    args.moe_expert_count,
+                    args.moe_gating_use_fp32,
+                    args.moe_second_expert_policy,
+                    args.moe_normalize_gate_prob_before_dropping,
+                    args.moe_eval_capacity_token_fraction,
+                    use_xmoe=args.use_xmoe,
+                )
+            experts = make_experts(args, self.embed_dim, self.ffn_dim)
+            self.moe_layer = MOELayer(gate, experts, args)
+
+        self.final_layer_norm = LayerNorm(self.embed_dim, eps=args.layernorm_eps)
+
+        if args.deepnorm:
+            if is_encoder_decoder:
+                self.alpha = math.pow(3.0 * args.decoder_layers, 0.25)
+            else:
+                self.alpha = math.pow(2.0 * args.decoder_layers, 0.25)
+        else:
+            self.alpha = 1.0
+
+    def build_ffn(self, embed_dim, args):
+        return FeedForwardNetwork(
+            embed_dim,
+            self.ffn_dim,
+            args.activation_fn,
+            args.dropout,
+            args.activation_dropout,
+            args.layernorm_eps,
+            args.subln,
+        )
+
+    def build_self_attention(self, embed_dim, args):
+        return MultiheadAttention(
+            args,
+            embed_dim,
+            args.decoder_attention_heads,
+            dropout=args.attention_dropout,
+            self_attention=True,
+            encoder_decoder_attention=False,
+            subln=args.subln,
+        )
+
+    def build_encoder_attention(self, embed_dim, args):
+        return MultiheadAttention(
+            args,
+            embed_dim,
+            args.decoder_attention_heads,
+            dropout=args.attention_dropout,
+            self_attention=False,
+            encoder_decoder_attention=True,
+            subln=args.subln,
+        )
+
+    def residual_connection(self, x, residual):
+        return residual * self.alpha + x
+
+    def forward(
+        self,
+        x,
+        encoder_out=None,
+        encoder_padding_mask=None,
+        incremental_state=None,
+        self_attn_mask=None,
+        self_attn_padding_mask=None,
+        self_attn_rel_pos=None,
+        cross_attn_rel_pos=None,
+    ):
+        residual = x
+        if self.normalize_before:
+            x = self.self_attn_layer_norm(x)
+
+        x, attn = self.self_attn(
+            query=x,
+            key=x,
+            value=x,
+            key_padding_mask=self_attn_padding_mask,
+            incremental_state=incremental_state,
+            attn_mask=self_attn_mask,
+            rel_pos=self_attn_rel_pos,
+        )
+        x = self.dropout_module(x)
+
+        if self.drop_path is not None:
+            x = self.drop_path(x)
+
+        x = self.residual_connection(x, residual)
+        if not self.normalize_before:
+            x = self.self_attn_layer_norm(x)
+
+        if self.encoder_attn is not None and encoder_out is not None:
+            residual = x
+            if self.normalize_before:
+                x = self.encoder_attn_layer_norm(x)
+
+            x, attn = self.encoder_attn(
+                query=x,
+                key=encoder_out,
+                value=encoder_out,
+                key_padding_mask=encoder_padding_mask,
+                incremental_state=None,
+                rel_pos=cross_attn_rel_pos,
+            )
+            x = self.dropout_module(x)
+
+            if self.drop_path is not None:
+                x = self.drop_path(x)
+
+            x = self.residual_connection(x, residual)
+            if not self.normalize_before:
+                x = self.encoder_attn_layer_norm(x)
+
+        residual = x
+        if self.normalize_before:
+            x = self.final_layer_norm(x)
+        if not self.is_moe_layer:
+            x = self.ffn(x)
+            l_aux = None
+        else:
+            x, l_aux = self.moe_layer(x)
+
+        if self.drop_path is not None:
+            x = self.drop_path(x)
+
+        x = self.residual_connection(x, residual)
+        if not self.normalize_before:
+            x = self.final_layer_norm(x)
+
+        return x, attn, None, l_aux
+
+
+class Decoder(nn.Module):
+    def __init__(
+        self, args, embed_tokens=None, embed_positions=None, output_projection=None, is_encoder_decoder=False, **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.args = args
+
+        self.dropout_module = torch.nn.Dropout(args.dropout)
+
+        embed_dim = args.decoder_embed_dim
+        self.embed_dim = embed_dim
+        self.embed_scale = 1.0 if args.no_scale_embedding else math.sqrt(embed_dim)
+
+        self.embed_tokens = embed_tokens
+        self.embed_positions = embed_positions
+
+        if output_projection is None and not args.no_output_layer and args.vocab_size > 0:
+            self.output_projection = self.build_output_projection(args)
+        else:
+            self.output_projection = output_projection
+
+        if args.layernorm_embedding:
+            self.layernorm_embedding = LayerNorm(embed_dim, eps=args.layernorm_eps)
+        else:
+            self.layernorm_embedding = None
+
+        self.layers = nn.ModuleList([])
+
+        moe_freq = args.moe_freq
+        for i in range(args.decoder_layers):
+            is_moe_layer = moe_freq != 0 and (i + 1) % moe_freq == 0
+            self.layers.append(
+                self.build_decoder_layer(
+                    args,
+                    depth=i,
+                    is_moe_layer=is_moe_layer,
+                    is_encoder_decoder=is_encoder_decoder,
+                )
+            )
+
+        self.num_layers = len(self.layers)
+
+        if args.decoder_normalize_before:
+            self.layer_norm = LayerNorm(embed_dim, eps=args.layernorm_eps)
+        else:
+            self.layer_norm = None
+
+        self.self_attn_relative_position = None
+        self.cross_attn_relative_position = None
+
+        if args.rel_pos_buckets > 0 and args.max_rel_pos > 0:
+            self.self_attn_relative_position = RelativePositionBias(
+                num_buckets=args.rel_pos_buckets,
+                max_distance=args.max_rel_pos,
+                n_heads=args.decoder_attention_heads,
+            )
+            if is_encoder_decoder:
+                self.cross_attn_relative_position = RelativePositionBias(
+                    num_buckets=args.rel_pos_buckets,
+                    max_distance=args.max_rel_pos,
+                    n_heads=args.decoder_attention_heads,
+                )
+
+        if args.bert_init:
+            self.apply(init_bert_params)
+
+        if args.deepnorm:
+            if is_encoder_decoder:
+                init_scale = math.pow(12.0 * args.decoder_layers, 0.25)
+            else:
+                init_scale = math.pow(8.0 * args.decoder_layers, 0.25)
+            for name, p in self.named_parameters():
+                if "fc1" in name or "fc2" in name or "out_proj" in name or "v_proj" in name:
+                    p.data.div_(init_scale)
+
+        if args.subln:
+            if is_encoder_decoder:
+                init_scale = math.sqrt(math.log(args.decoder_layers * 3))
+            else:
+                init_scale = math.sqrt(math.log(args.decoder_layers * 2))
+            for name, p in self.named_parameters():
+                if "encoder_attn" in name:
+                    continue
+                if "fc1" in name or "fc2" in name or "out_proj" in name or "v_proj" in name:
+                    p.data.mul_(init_scale)
+
+    def build_output_projection(
+        self,
+        args,
+    ):
+        if args.share_decoder_input_output_embed:
+            output_projection = torch.nn.Linear(
+                self.embed_tokens.weight.shape[1],
+                self.embed_tokens.weight.shape[0],
+                bias=False,
+            )
+            output_projection.weight = self.embed_tokens.weight
+        else:
+            output_projection = torch.nn.Linear(args.decoder_embed_dim, args.vocab_size, bias=False)
+            torch.nn.init.normal_(output_projection.weight, mean=0, std=args.decoder_embed_dim**-0.5)
+        return output_projection
+
+    def build_decoder_layer(self, args, depth, is_moe_layer=False, is_encoder_decoder=False):
+        layer = DecoderLayer(
+            args,
+            depth,
+            is_moe_layer=is_moe_layer,
+            is_encoder_decoder=is_encoder_decoder,
+        )
+        if args.checkpoint_activations:
+            layer = checkpoint_wrapper(layer)
+        if args.fsdp:
+            layer = wrap(layer)
+        return layer
+
+    def forward_embedding(
+        self,
+        tokens,
+        token_embedding=None,
+        incremental_state=None,
+    ):
+        positions = None
+        if self.embed_positions is not None:
+            positions = self.embed_positions(tokens, incremental_state=incremental_state)
+
+        if incremental_state is not None:
+            tokens = tokens[:, -1:]
+            if positions is not None:
+                positions = positions[:, -1:]
+
+        if token_embedding is None:
+            token_embedding = self.embed_tokens(tokens)
+
+        x = embed = self.embed_scale * token_embedding
+
+        if positions is not None:
+            x += positions
+
+        if self.layernorm_embedding is not None:
+            x = self.layernorm_embedding(x)
+
+        x = self.dropout_module(x)
+
+        return x, embed
+
+    def forward(
+        self,
+        prev_output_tokens,
+        self_attn_padding_mask=None,
+        encoder_out=None,
+        incremental_state=None,
+        features_only=False,
+        return_all_hiddens=False,
+        token_embeddings=None,
+        **kwargs
+    ):
+        # embed tokens and positions
+        x, _ = self.forward_embedding(prev_output_tokens, token_embeddings, incremental_state)
+
+        # relative position
+        self_attn_rel_pos_bias = None
+        slen = prev_output_tokens.size(1)
+        if self.self_attn_relative_position is not None:
+            self_attn_rel_pos_bias = self.self_attn_relative_position(batch_size=x.size(0), qlen=slen, klen=slen)
+            if incremental_state is not None:
+                self_attn_rel_pos_bias = self_attn_rel_pos_bias[-1:, :, :]
+        cross_attn_rel_pos_bias = None
+        if self.cross_attn_relative_position is not None:
+            cross_attn_rel_pos_bias = self.cross_attn_relative_position(
+                batch_size=x.size(0),
+                qlen=slen,
+                klen=encoder_out["encoder_out"].size(1),
+            )
+            if incremental_state is not None:
+                cross_attn_rel_pos_bias = cross_attn_rel_pos_bias[-1:, :, :]
+
+        # decoder layers
+        inner_states = [x]
+
+        if encoder_out is None:
+            l_aux = []
+        else:
+            l_aux = encoder_out["l_aux"] if "l_aux" in encoder_out else []
+
+        for idx, layer in enumerate(self.layers):
+            if incremental_state is None:
+                self_attn_mask = torch.triu(
+                    torch.zeros([x.size(1), x.size(1)]).float().fill_(float("-inf")).type_as(x),
+                    1,
+                )
+            else:
+                self_attn_mask = None
+                if idx not in incremental_state:
+                    incremental_state[idx] = {}
+
+            x, layer_attn, _, l_aux_i = layer(
+                x,
+                encoder_out["encoder_out"] if encoder_out is not None else None,
+                encoder_out["encoder_padding_mask"] if encoder_out is not None else None,
+                incremental_state[idx] if incremental_state is not None else None,
+                self_attn_mask=self_attn_mask,
+                self_attn_padding_mask=self_attn_padding_mask,
+                self_attn_rel_pos=self_attn_rel_pos_bias,
+                cross_attn_rel_pos=cross_attn_rel_pos_bias,
+            )
+            l_aux.append(l_aux_i)
+            inner_states.append(x)
+
+        if self.layer_norm is not None:
+            x = self.layer_norm(x)
+
+        if not features_only:
+            x = self.output_layer(x)
+
+        return x, {
+            "inner_states": inner_states,
+            "l_aux": l_aux,
+            "attn": None,
+        }
+
+    def output_layer(self, features):
+        return self.output_projection(features)

vlmo/torchscale/architecture/encoder.py

@@ -0,0 +1,489 @@
+# Copyright (c) 2022 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+
+import math
+
+import numpy as np
+import torch
+import torch.nn as nn
+from fairscale.nn import checkpoint_wrapper, wrap
+
+try:
+    from apex.normalization import FusedLayerNorm as LayerNorm
+except ModuleNotFoundError:
+    from torch.nn import LayerNorm
+
+from vlmo.torchscale.architecture.utils import init_bert_params
+from vlmo.torchscale.component.droppath import DropPath
+from vlmo.torchscale.component.feedforward_network import FeedForwardNetwork, make_experts
+from vlmo.torchscale.component.multihead_attention import MultiheadAttention
+from vlmo.torchscale.component.multiway_network import MultiwayWrapper, set_split_position
+from vlmo.torchscale.component.relative_position_bias import RelativePositionBias
+#from vlmo.torchscale.component.xmoe.moe_layer import MOELayer
+#from vlmo.torchscale.component.xmoe.routing import Top1Gate, Top2Gate
+# from vlmo.modules.vlmo_utils import no_sync_module_apply
+from pytorch_lightning.utilities.rank_zero import rank_zero_info
+
+def no_sync_module_apply(module, fn):
+    """FSDP module .apply will use _unshard_params_recurse which will sync params across ranks.
+    using this function when apply fn is unnecessary to sync params across ranks.
+    """
+    for child in module.children():
+        fn(child)
+        no_sync_module_apply(child, fn)
+
+class EncoderLayer(nn.Module):
+    def __init__(self, args, depth, attn=None, is_moe_layer=False, is_encoder_decoder=False):
+        super().__init__()
+        self.args = args
+        self.embed_dim = args.encoder_embed_dim
+        self.self_attn = self.build_self_attention(self.embed_dim, args) if attn is None else attn
+        self.self_attn_layer_norm = MultiwayWrapper(args, LayerNorm(self.embed_dim, eps=args.layernorm_eps))
+        self.dropout_module = torch.nn.Dropout(args.dropout)
+
+        if args.drop_path_rate > 0:
+            drop_path_prob = np.linspace(0, args.drop_path_rate, args.encoder_layers)[depth]
+            self.drop_path = DropPath(drop_path_prob)
+        else:
+            self.drop_path = None
+
+        self.normalize_before = args.encoder_normalize_before
+        self.is_moe_layer = is_moe_layer
+        self.ffn_dim = args.encoder_ffn_embed_dim
+
+        if not self.is_moe_layer:
+            self.ffn = MultiwayWrapper(
+                args,
+                self.build_ffn(
+                    self.embed_dim,
+                    self.args,
+                ),
+            )
+        else:
+            assert not self.args.multiway
+            if args.moe_top1_expert:
+                gate = Top1Gate(
+                    self.embed_dim,
+                    args.moe_expert_count,
+                    use_fp32=args.moe_gating_use_fp32,
+                    moe_eval_capacity_token_fraction=args.moe_eval_capacity_token_fraction,
+                    use_xmoe=args.use_xmoe,
+                )
+            else:
+                gate = Top2Gate(
+                    self.embed_dim,
+                    args.moe_expert_count,
+                    args.moe_gating_use_fp32,
+                    args.moe_second_expert_policy,
+                    args.moe_normalize_gate_prob_before_dropping,
+                    args.moe_eval_capacity_token_fraction,
+                    use_xmoe=args.use_xmoe,
+                )
+            experts = make_experts(args, self.embed_dim, self.ffn_dim)
+            self.moe_layer = MOELayer(gate, experts, args)
+        self.final_layer_norm = MultiwayWrapper(args, LayerNorm(self.embed_dim, eps=args.layernorm_eps))
+
+        if args.deepnorm:
+            if is_encoder_decoder:
+                self.alpha = math.pow(math.pow(args.encoder_layers, 4) * args.decoder_layers, 0.0625) * 0.81
+            else:
+                self.alpha = math.pow(2.0 * args.encoder_layers, 0.25)
+        else:
+            self.alpha = 1.0
+
+    def build_ffn(self, embed_dim, args):
+        return FeedForwardNetwork(
+            embed_dim,
+            self.ffn_dim,
+            args.activation_fn,
+            args.dropout,
+            args.activation_dropout,
+            args.layernorm_eps,
+            args.subln,
+        )
+
+    def build_self_attention(self, embed_dim, args):
+        return MultiheadAttention(
+            args,
+            embed_dim,
+            args.encoder_attention_heads,
+            dropout=args.attention_dropout,
+            self_attention=True,
+            encoder_decoder_attention=False,
+            subln=args.subln,
+            one_attn=args.one_attn,
+        )
+
+    def residual_connection(self, x, residual):
+        return residual * self.alpha + x
+
+    def forward(
+        self,
+        x,
+        encoder_padding_mask,
+        attn_mask=None,
+        rel_pos=None,
+        multiway_split_position=None,
+        incremental_state=None,
+    ):
+        if multiway_split_position is not None:
+            assert self.args.multiway
+            no_sync_module_apply(self, set_split_position(multiway_split_position))
+
+        if attn_mask is not None:
+            # float16: -1e8 equal 0
+            attn_mask = attn_mask.masked_fill(attn_mask.to(torch.bool), -1e8)
+
+        residual = x
+        if self.normalize_before:
+            x = self.self_attn_layer_norm(x)
+        x, _ = self.self_attn(
+            query=x,
+            key=x,
+            value=x,
+            key_padding_mask=encoder_padding_mask,
+            attn_mask=attn_mask,
+            rel_pos=rel_pos,
+            incremental_state=incremental_state,
+        )
+        x = self.dropout_module(x)
+
+        if self.drop_path is not None:
+            x = self.drop_path(x)
+
+        x = self.residual_connection(x, residual)
+        if not self.normalize_before:
+            x = self.self_attn_layer_norm(x)
+
+        residual = x
+        if self.normalize_before:
+            x = self.final_layer_norm(x)
+        if not self.is_moe_layer:
+            x = self.ffn(x)
+            l_aux = None
+        else:
+            x = x.transpose(0, 1)
+            x, l_aux = self.moe_layer(x)
+            x = x.transpose(0, 1)
+
+        if self.drop_path is not None:
+            x = self.drop_path(x)
+
+        x = self.residual_connection(x, residual)
+        if not self.normalize_before:
+            x = self.final_layer_norm(x)
+        return x, l_aux
+
+
+class Encoder(nn.Module):
+    def __init__(
+        self, args, embed_tokens=None, embed_positions=None, output_projection=None, is_encoder_decoder=False, **kwargs
+    ):
+        self.args = args
+        super().__init__(**kwargs)
+
+        self.dropout_module = torch.nn.Dropout(args.dropout)
+
+        embed_dim = args.encoder_embed_dim
+        self.embed_scale = 1.0 if args.no_scale_embedding else math.sqrt(embed_dim)
+        self.mask_ratio = args.mask_ratio
+        self.max_text_len = args.max_text_len
+        self.vision_len = (args.img_size // args.patch_size) * (args.img_size // args.patch_size)
+
+        self.embed_tokens = embed_tokens
+        self.embed_positions = embed_positions
+
+        if output_projection is None and not is_encoder_decoder and not args.no_output_layer and args.vocab_size > 0:
+            self.output_projection = self.build_output_projection(args)
+        else:
+            self.output_projection = output_projection
+
+        if args.layernorm_embedding:
+            self.layernorm_embedding = MultiwayWrapper(args, LayerNorm(embed_dim, eps=args.layernorm_eps), dim=1)
+        else:
+            self.layernorm_embedding = None
+
+        self.layers = nn.ModuleList([])
+        if self.args.share_layer:
+            single_layer = self.build_encoder_layer(
+                args, depth=0, is_moe_layer=False, is_encoder_decoder=is_encoder_decoder
+            )
+            for i in range(args.encoder_layers):
+                self.layers.append(single_layer)
+        elif self.args.share_attn:
+            moe_freq = args.moe_freq
+            embed_dim = args.encoder_embed_dim
+            shared_attn = self.build_self_attention(embed_dim, self.args)
+            for i in range(args.encoder_layers):
+                is_moe_layer = moe_freq != 0 and (i + 1) % moe_freq == 0
+                self.layers.append(
+                    self.build_encoder_layer(
+                        args,
+                        depth=i,
+                        attn=shared_attn,
+                        is_moe_layer=is_moe_layer,
+                        is_encoder_decoder=is_encoder_decoder,
+                    )
+                )
+
+        else:
+            moe_freq = args.moe_freq
+            for i in range(args.encoder_layers):
+                is_moe_layer = moe_freq != 0 and (i + 1) % moe_freq == 0
+                self.layers.append(
+                    self.build_encoder_layer(
+                        args,
+                        depth=i,
+                        is_moe_layer=is_moe_layer,
+                        is_encoder_decoder=is_encoder_decoder,
+                    )
+                )
+        self.num_layers = len(self.layers)
+
+        if args.encoder_normalize_before and args.normalize_output:
+            self.layer_norm = MultiwayWrapper(args, LayerNorm(embed_dim, eps=args.layernorm_eps))
+        else:
+            self.layer_norm = None
+
+        if args.rel_pos_buckets > 0 and args.max_rel_pos > 0:
+            self.relative_position = RelativePositionBias(
+                num_buckets=args.rel_pos_buckets,
+                max_distance=args.max_rel_pos,
+                n_heads=args.encoder_attention_heads,
+            )
+        else:
+            self.relative_position = None
+
+        if args.bert_init:
+            self.apply(init_bert_params)
+
+        if args.deepnorm:
+            if is_encoder_decoder:
+                init_scale = math.pow(math.pow(args.encoder_layers, 4) * args.decoder_layers, 0.0625) / 1.15
+            else:
+                init_scale = math.pow(8.0 * args.encoder_layers, 0.25)
+            for name, p in self.named_parameters():
+                if "fc1" in name or "fc2" in name or "out_proj" in name or "v_proj" in name:
+                    p.data.div_(init_scale)
+
+        if args.subln:
+            if is_encoder_decoder:
+                init_scale = math.sqrt(math.log(3 * args.decoder_layers) * math.log(2 * args.encoder_layers) / 3)
+            else:
+                init_scale = math.sqrt(math.log(args.encoder_layers * 2))
+            for name, p in self.named_parameters():
+                if "fc1" in name or "fc2" in name or "out_proj" in name or "v_proj" in name:
+                    p.data.mul_(init_scale)
+
+    def random_masking(self, x, mask_ratio):
+        N, L, D = x.shape  # batch, length, dim
+        len_keep = int(L * (1 - mask_ratio))
+
+        noise = torch.rand(N, L - 1, device=x.device)
+        ids_shuffle = torch.argsort(noise, dim=1) + torch.ones(N, L - 1, device=x.device, dtype=int)
+        ids_keep = ids_shuffle[:, :len_keep]
+
+        x_masked = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D))
+
+        x0 = x[:, 0, :]
+        x0 = x0.reshape(N, 1, D)
+        x_masked_add = torch.cat([x0, x_masked], axis=1)
+        return x_masked_add, ids_keep
+
+    def build_self_attention(self, embed_dim, args):
+        return MultiheadAttention(
+            args,
+            embed_dim,
+            args.encoder_attention_heads,
+            dropout=args.attention_dropout,
+            self_attention=True,
+            encoder_decoder_attention=False,
+            subln=args.subln,
+            one_attn=args.one_attn,
+        )
+
+    def build_output_projection(
+        self,
+        args,
+    ):
+        if args.share_encoder_input_output_embed:
+            assert args.encoder_embedding_type == "language"
+            output_projection = torch.nn.Linear(
+                self.embed_tokens.weight.shape[1],
+                self.embed_tokens.weight.shape[0],
+                bias=False,
+            )
+            output_projection.weight = self.embed_tokens.weight
+        else:
+            output_projection = torch.nn.Linear(args.encoder_embed_dim, args.vocab_size, bias=False)
+            torch.nn.init.normal_(output_projection.weight, mean=0, std=args.encoder_embed_dim**-0.5)
+        return output_projection
+
+    def checkpointing_and_params_allgather(
+        self,
+        origin_layer,
+    ):
+        origin_forward = origin_layer.forward
+
+        from deepspeed import checkpointing
+        def forward(*args, **kwargs):
+            # deepspeed checkpoint not support kwargs
+            ret = checkpointing.checkpoint(origin_forward, *args, **kwargs)
+            return ret
+
+        return forward
+
+    def build_encoder_layer(self, args, depth, attn=None, is_moe_layer=False, is_encoder_decoder=False):
+        layer = EncoderLayer(
+            args,
+            depth,
+            attn,
+            is_moe_layer=is_moe_layer,
+            is_encoder_decoder=is_encoder_decoder,
+        )
+        if args.checkpoint_activations:
+            rank_zero_info("EncoderLayer params: %s", sum(p.numel() for p in layer.parameters() if p.requires_grad))
+            layer = checkpoint_wrapper(layer)
+            # layer.ffn = checkpoint_wrapper(layer.ffn,)
+        if args.fsdp:
+            layer = wrap(layer)
+        return layer
+
+    def checkpointing_layers(self):
+        for i, layer in enumerate(self.layers):
+            rank_zero_info(f"Checkpointing wrapper EncoderLayers: {i}")
+            self.layers[i] = checkpoint_wrapper(layer)
+
+    def forward_embedding(
+        self,
+        src_tokens,
+        token_embedding=None,
+        positions=None,
+    ):
+        if token_embedding is None:
+            token_embedding = self.embed_tokens(src_tokens)
+        x = embed = self.embed_scale * token_embedding
+        if self.embed_positions is not None:
+            if src_tokens is not None:
+                x = embed + self.embed_positions(src_tokens, positions=positions)
+            else:
+                x = embed + self.embed_positions(x, positions=positions)
+        is_flip, ids_keep = 0, None
+        if self.mask_ratio > 0:
+            if x.shape[1] == self.vision_len + 1:
+                x, ids_keep = self.random_masking(x, self.mask_ratio)
+                is_flip = 1
+            elif x.shape[1] == self.vision_len + self.max_text_len + 1:
+                vision_tokens = x[:, : self.vision_len + 1, :]
+                vision_tokens, ids_keep = self.random_masking(vision_tokens, self.mask_ratio)
+                x = torch.cat(
+                    [
+                        vision_tokens,
+                        x[
+                            :,
+                            self.vision_len + 1 :,
+                        ],
+                    ],
+                    dim=1,
+                )
+                is_flip = 2
+        if self.layernorm_embedding is not None:
+            x = self.layernorm_embedding(x)
+        x = self.dropout_module(x)
+        return x, embed, ids_keep, is_flip
+
+    def forward(
+        self,
+        src_tokens,
+        encoder_padding_mask=None,
+        attn_mask=None,
+        return_all_hiddens=False,
+        token_embeddings=None,
+        multiway_split_position=None,
+        features_only=False,
+        incremental_state=None,
+        positions=None,
+        **kwargs
+    ):
+        assert src_tokens is not None or token_embeddings is not None
+
+        if encoder_padding_mask is None:
+            if src_tokens is not None:
+                encoder_padding_mask = torch.zeros_like(src_tokens, device=src_tokens.device).bool()
+            else:
+                encoder_padding_mask = torch.zeros(
+                    [token_embeddings.size(0), token_embeddings.size(1)],
+                    device=token_embeddings.device,
+                ).bool()
+
+        if multiway_split_position is not None:
+            assert self.args.multiway
+            no_sync_module_apply(self, set_split_position(multiway_split_position))
+
+        x, encoder_embedding, ids_keep, is_flip = self.forward_embedding(src_tokens, token_embeddings, positions)
+        if is_flip > 0:
+            if is_flip == 2:
+                text_ids = (
+                    torch.arange(
+                        self.vision_len + 1, self.vision_len + 1 + self.max_text_len, device=x.device, dtype=torch.int64
+                    )
+                    .unsqueeze(0)
+                    .repeat(ids_keep.shape[0], 1)
+                )
+                cls_ids = torch.zeros(ids_keep.shape[0], 1, device=x.device, dtype=torch.int64)
+                ids_keep = torch.cat([cls_ids, ids_keep, text_ids], dim=1)
+            elif is_flip == 1:
+                cls_ids = torch.zeros(ids_keep.shape[0], 1, device=x.device, dtype=torch.int64)
+                ids_keep = torch.cat([cls_ids, ids_keep], dim=1)
+            if encoder_padding_mask is not None:
+                encoder_padding_mask = torch.gather(encoder_padding_mask, dim=1, index=ids_keep)
+            if attn_mask is not None:
+                attn_mask = torch.gather(
+                    attn_mask, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, attn_mask.shape[-1])
+                )
+                attn_mask = torch.gather(attn_mask, dim=2, index=ids_keep.unsqueeze(1).repeat(1, attn_mask.shape[1], 1))
+            if multiway_split_position > 0:
+                multiway_split_position = ids_keep.shape[1] - self.max_text_len
+        x = x * (1 - encoder_padding_mask.unsqueeze(-1).type_as(x))
+
+        encoder_states = []
+
+        if return_all_hiddens:
+            encoder_states.append(x)
+
+        rel_pos_bias = None
+        if self.relative_position is not None:
+            rel_pos_bias = self.relative_position(batch_size=x.size(0), qlen=x.size(1), klen=x.size(1))
+
+        l_aux = []
+        for idx, layer in enumerate(self.layers):
+            x, l_aux_i = layer(
+                x,
+                encoder_padding_mask=encoder_padding_mask if incremental_state is None else None,
+                attn_mask=attn_mask,
+                rel_pos=rel_pos_bias,
+                multiway_split_position=multiway_split_position,
+                incremental_state=incremental_state[idx] if incremental_state is not None else None,
+            )
+            if return_all_hiddens:
+                assert encoder_states is not None
+                encoder_states.append(x)
+            l_aux.append(l_aux_i)
+
+        if multiway_split_position is not None:
+            assert self.args.multiway
+            no_sync_module_apply(self, set_split_position(multiway_split_position))
+        if self.layer_norm is not None:
+            x = self.layer_norm(x)
+
+        if not features_only and self.output_projection is not None:
+            x = self.output_projection(x)
+
+        return {
+            "encoder_out": x,
+            "encoder_embedding": encoder_embedding,
+            "encoder_padding_mask": encoder_padding_mask,
+            "encoder_states": encoder_states,
+            "l_aux": l_aux,
+            "multiway_split_position": multiway_split_position,
+        }

vlmo/torchscale/architecture/encoder_decoder.py

@@ -0,0 +1,43 @@
+# Copyright (c) 2022 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+
+import torch.nn as nn
+
+from vlmo.torchscale.architecture.decoder import Decoder
+from vlmo.torchscale.architecture.encoder import Encoder
+
+
+class EncoderDecoder(nn.Module):
+    def __init__(
+        self,
+        args,
+        encoder_embed_tokens=None,
+        encoder_embed_positions=None,
+        decoder_embed_tokens=None,
+        decoder_embed_positions=None,
+        output_projection=None,
+        **kwargs
+    ):
+        super().__init__()
+        self.args = args
+        if args.share_all_embeddings:
+            args.share_decoder_input_output_embed = True
+
+        self.encoder = Encoder(args, encoder_embed_tokens, encoder_embed_positions, is_encoder_decoder=True, **kwargs)
+
+        if args.share_all_embeddings and decoder_embed_tokens is None:
+            decoder_embed_tokens = self.encoder.embed_tokens
+
+        self.decoder = Decoder(
+            args, decoder_embed_tokens, decoder_embed_positions, output_projection, is_encoder_decoder=True, **kwargs
+        )
+
+    def forward(self, src_tokens, prev_output_tokens, return_all_hiddens=False, features_only=False, **kwargs):
+        encoder_out = self.encoder(src_tokens, return_all_hiddens=return_all_hiddens)
+        decoder_out = self.decoder(
+            prev_output_tokens,
+            encoder_out=encoder_out,
+            features_only=features_only,
+            return_all_hiddens=return_all_hiddens,
+        )
+        return decoder_out

vlmo/torchscale/architecture/utils.py

@@ -0,0 +1,33 @@
+# Copyright (c) 2022 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+
+import torch.nn as nn
+
+from vlmo.torchscale.component.multihead_attention import MultiheadAttention
+from vlmo.torchscale.component.multiway_network import MultiwayNetwork
+
+
+def init_bert_params(module):
+    def normal_(data):
+        data.copy_(data.cpu().normal_(mean=0.0, std=0.02).to(data.device))
+
+    if isinstance(module, nn.Linear):
+        normal_(module.weight.data)
+        if module.bias is not None:
+            module.bias.data.zero_()
+    if isinstance(module, nn.Embedding):
+        normal_(module.weight.data)
+        if module.padding_idx is not None:
+            module.weight.data[module.padding_idx].zero_()
+    if isinstance(module, MultiheadAttention):
+        if isinstance(module.q_proj, MultiwayNetwork):
+            normal_(module.q_proj.A.weight.data)
+            normal_(module.q_proj.B.weight.data)
+            normal_(module.k_proj.A.weight.data)
+            normal_(module.k_proj.B.weight.data)
+            normal_(module.v_proj.A.weight.data)
+            normal_(module.v_proj.B.weight.data)
+        else:
+            normal_(module.q_proj.weight.data)
+            normal_(module.k_proj.weight.data)
+            normal_(module.v_proj.weight.data)

vlmo/torchscale/component/__init__.py

@@ -0,0 +1,2 @@
+# Copyright (c) 2022 Microsoft
+# Licensed under The MIT License [see LICENSE for details]

vlmo/torchscale/component/droppath.py

@@ -0,0 +1,19 @@
+# Copyright (c) 2022 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+
+import torch.nn as nn
+from timm.models.layers import drop_path
+
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks)."""
+
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)
+
+    def extra_repr(self):
+        return "p={}".format(self.drop_prob)

vlmo/torchscale/component/embedding.py

@@ -0,0 +1,110 @@
+# Copyright (c) 2022 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class VisionLanguageEmbedding(nn.Module):
+    def __init__(self, text_embed, vision_embed):
+        super().__init__()
+        self.text_embed = text_embed
+        self.vision_embed = vision_embed
+
+    def forward(self, textual_tokens, visual_tokens, **kwargs):
+        if textual_tokens is None:
+            return self.vision_embed(visual_tokens)
+
+        if visual_tokens is None:
+            return self.text_embed(textual_tokens)
+
+        x1 = self.vision_embed(visual_tokens)
+        x2 = self.text_embed(textual_tokens)
+
+        return torch.cat([x1, x2], dim=1)
+
+
+class VisionEmbedding(nn.Module):
+    """Image to Patch Embedding"""
+
+    def __init__(
+        self,
+        img_size=224,
+        patch_size=16,
+        in_chans=3,
+        embed_dim=768,
+        contain_mask_token=False,
+        prepend_cls_token=False,
+    ):
+        super().__init__()
+        img_size = (img_size, img_size)
+        patch_size = (patch_size, patch_size)
+        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
+        self.patch_shape = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.num_patches = num_patches
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+
+        if contain_mask_token:
+            self.mask_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        else:
+            self.mask_token = None
+
+        if prepend_cls_token:
+            self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        else:
+            self.cls_token = None
+
+    def num_position_embeddings(self):
+        if self.cls_token is None:
+            return self.num_patches
+        else:
+            return self.num_patches + 1
+
+    def forward(self, x, masked_position=None, **kwargs):
+        B, C, H, W = x.shape
+        x = self.proj(x).flatten(2).transpose(1, 2)
+
+        batch_size, seq_len, _ = x.size()
+
+        if masked_position is not None:
+            assert self.mask_token is not None
+            mask_token = self.mask_token.expand(batch_size, seq_len, -1)
+            w = masked_position.unsqueeze(-1).type_as(mask_token)
+            x = x * (1 - w) + mask_token * w
+
+        if self.cls_token is not None:
+            cls_tokens = self.cls_token.expand(batch_size, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
+            x = torch.cat((cls_tokens, x), dim=1)
+
+        return x
+
+
+class TextEmbedding(nn.Embedding):
+    def reset_parameters(self):
+        nn.init.normal_(self.weight, mean=0, std=self.embedding_dim**-0.5)
+        self._fill_padding_idx_with_zero()
+
+
+class PositionalEmbedding(nn.Embedding):
+    def forward(
+        self,
+        x,
+        positions=None,
+        **kwargs,
+    ):
+        if positions is None:
+            # being consistent with Fairseq, which starts from 2.
+            positions = torch.arange(2, x.size(1) + 2, device=x.device).long().unsqueeze(0)
+        return F.embedding(
+            positions,
+            self.weight,
+            self.padding_idx,
+            self.max_norm,
+            self.norm_type,
+            self.scale_grad_by_freq,
+            self.sparse,
+        )

vlmo/torchscale/component/feedforward_network.py

@@ -0,0 +1,128 @@
+# Copyright (c) 2022 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+try:
+    from apex.normalization import FusedLayerNorm as LayerNorm
+except ModuleNotFoundError:
+    from torch.nn import LayerNorm
+
+
+class set_torch_seed(object):
+    def __init__(self, seed):
+        assert isinstance(seed, int)
+        self.rng_state = self.get_rng_state()
+
+        torch.manual_seed(seed)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed(seed)
+
+    def get_rng_state(self):
+        state = {"torch_rng_state": torch.get_rng_state()}
+        if torch.cuda.is_available():
+            state["cuda_rng_state"] = torch.cuda.get_rng_state()
+        return state
+
+    def set_rng_state(self, state):
+        torch.set_rng_state(state["torch_rng_state"])
+        if torch.cuda.is_available():
+            torch.cuda.set_rng_state(state["cuda_rng_state"])
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *exc):
+        self.set_rng_state(self.rng_state)
+
+
+def make_experts(args, embed_dim, expert_ffn_dim):
+    world_size = 1 if not torch.distributed.is_initialized() else torch.distributed.get_world_size()
+    expert_list = []
+    ddp_rank = args.ddp_rank
+    start_seed = torch.randint(1000000, (1,)).item()
+    # at least as many experts than gpus
+    if args.moe_expert_count >= world_size:
+        assert args.moe_expert_count % world_size == 0, f"{args.moe_expert_count}, {world_size}"
+        local_moe_expert_count = args.moe_expert_count // world_size
+        for i in range(local_moe_expert_count):
+            with set_torch_seed(start_seed + ddp_rank * local_moe_expert_count + i):
+                expert_list.append(
+                    FeedForwardNetwork(
+                        embed_dim,
+                        expert_ffn_dim,
+                        args.activation_fn,
+                        args.dropout,
+                        args.activation_dropout,
+                        args.layernorm_eps,
+                        args.subln,
+                    )
+                )
+    else:
+        assert world_size % args.moe_expert_count == 0, f"{world_size}, {args.moe_expert_count}"
+
+        with set_torch_seed(start_seed + ddp_rank % args.moe_expert_count):
+            expert_list.append(
+                FeedForwardNetwork(
+                    embed_dim,
+                    expert_ffn_dim,
+                    args.activation_fn,
+                    args.dropout,
+                    args.activation_dropout,
+                    args.layernorm_eps,
+                    args.subln,
+                )
+            )
+    experts = nn.ModuleList(expert_list)
+    return experts
+
+
+def get_activation_fn(activation):
+    if activation == "relu":
+        return F.relu
+    elif activation == "gelu":
+        return F.gelu
+    else:
+        raise NotImplementedError
+
+
+class FeedForwardNetwork(nn.Module):
+    def __init__(
+        self,
+        embed_dim,
+        ffn_dim,
+        activation_fn,
+        dropout,
+        activation_dropout,
+        layernorm_eps,
+        subln=False,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.activation_fn = get_activation_fn(activation=str(activation_fn))
+        self.activation_dropout_module = torch.nn.Dropout(activation_dropout)
+        self.dropout_module = torch.nn.Dropout(dropout)
+        self.fc1 = nn.Linear(self.embed_dim, ffn_dim)
+        self.fc2 = nn.Linear(ffn_dim, self.embed_dim)
+        self.ffn_layernorm = LayerNorm(ffn_dim, eps=layernorm_eps) if subln else None
+
+    def reset_parameters(self):
+        self.fc1.reset_parameters()
+        self.fc2.reset_parameters()
+        if self.ffn_layernorm is not None:
+            self.ffn_layernorm.reset_parameters()
+
+    def forward(self, x):
+        # x = x.reshape(-1, x.size(-1))
+        x = self.fc1(x)
+        # x = self.activation_fn(x.float()).type_as(x)
+        x = self.activation_fn(x)
+        x = self.activation_dropout_module(x)
+        if self.ffn_layernorm is not None:
+            x = self.ffn_layernorm(x)
+        x = self.fc2(x)
+        # x = x.view(x_shape)
+        x = self.dropout_module(x)
+        return x

vlmo/torchscale/component/multihead_attention.py

@@ -0,0 +1,154 @@
+# Copyright (c) 2022 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+
+import math
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+try:
+    from apex.normalization import FusedLayerNorm as LayerNorm
+except ModuleNotFoundError:
+    from torch.nn import LayerNorm
+
+from .multiway_network import MultiwayWrapper
+from .xpos_relative_position import XPOS
+
+
+class MultiheadAttention(nn.Module):
+    def __init__(
+        self,
+        args,
+        embed_dim,
+        num_heads,
+        dropout=0.0,
+        self_attention=False,
+        encoder_decoder_attention=False,
+        subln=False,
+        one_attn=False,
+    ):
+        super().__init__()
+        self.args = args
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.head_dim = embed_dim // num_heads
+        self.scaling = self.head_dim ** (-0.5)
+        self.self_attention = self_attention
+        self.encoder_decoder_attention = encoder_decoder_attention
+        assert self.self_attention ^ self.encoder_decoder_attention
+        if one_attn:
+            self.k_proj = nn.Linear(embed_dim, embed_dim, bias=True)
+            self.v_proj = nn.Linear(embed_dim, embed_dim, bias=True)
+            self.q_proj = nn.Linear(embed_dim, embed_dim, bias=True)
+            self.out_proj = nn.Linear(embed_dim, embed_dim, bias=True)
+        else:
+            self.k_proj = MultiwayWrapper(args, nn.Linear(embed_dim, embed_dim, bias=True))
+            self.v_proj = MultiwayWrapper(args, nn.Linear(embed_dim, embed_dim, bias=True))
+            self.q_proj = MultiwayWrapper(args, nn.Linear(embed_dim, embed_dim, bias=True))
+            # self.qkv_proj = MultiwayWrapper(args, nn.Linear(embed_dim, embed_dim*3, bias=True))
+            self.out_proj = MultiwayWrapper(args, nn.Linear(embed_dim, embed_dim, bias=True))
+        self.inner_attn_ln = (
+            MultiwayWrapper(args, LayerNorm(self.embed_dim, eps=args.layernorm_eps))
+            if subln and self.self_attention
+            else None
+        )
+        self.dropout_module = torch.nn.Dropout(dropout)
+        self.xpos = XPOS(self.head_dim, args.xpos_scale_base) if args.xpos_rel_pos and self.self_attention else None
+
+    def reset_parameters(self):
+        nn.init.xavier_uniform_(self.q_proj.weight, gain=1 / math.sqrt(2))
+        nn.init.xavier_uniform_(self.k_proj.weight, gain=1 / math.sqrt(2))
+        nn.init.xavier_uniform_(self.v_proj.weight, gain=1 / math.sqrt(2))
+        nn.init.xavier_uniform_(self.out_proj.weight)
+        nn.init.constant_(self.out_proj.bias, 0.0)
+
+    def forward(
+        self,
+        query,
+        key,
+        value,
+        incremental_state=None,
+        key_padding_mask=None,
+        attn_mask=None,
+        rel_pos=None,
+    ):
+        bsz, tgt_len, embed_dim = query.size()
+        src_len = tgt_len
+        assert embed_dim == self.embed_dim, f"query dim {embed_dim} != {self.embed_dim}"
+
+        key_bsz, src_len, _ = key.size()
+        assert key_bsz == bsz, f"{query.size(), key.size()}"
+        assert value is not None
+        assert bsz, src_len == value.shape[:2]
+        # if query is key and key is value:
+        #     qkv = self.qkv_proj(query)
+        # else:
+        #     # W*(q+k+v) = W(q) + W(k) + W(v)
+        #     qkv = self.qkv_proj(query+key+value)
+        # q,k,v = qkv.split(self.embed_dim, dim=-1)
+
+        q = self.q_proj(query)
+        k = self.k_proj(key)
+        v = self.v_proj(value)
+
+        q = (q * self.scaling).view(bsz, tgt_len, self.num_heads, self.head_dim).transpose(1, 2)
+        k = k.view(bsz, src_len, self.num_heads, self.head_dim).transpose(1, 2)
+        v = v.view(bsz, src_len, self.num_heads, self.head_dim).transpose(1, 2)
+        q = q.reshape(bsz * self.num_heads, tgt_len, self.head_dim)
+        k = k.reshape(bsz * self.num_heads, src_len, self.head_dim)
+        v = v.reshape(bsz * self.num_heads, src_len, self.head_dim)
+
+        if incremental_state is not None:
+            if "prev_key" in incremental_state:
+                prev_key = incremental_state["prev_key"].view(bsz * self.num_heads, -1, self.head_dim)
+                prev_value = incremental_state["prev_value"].view(bsz * self.num_heads, -1, self.head_dim)
+                k = torch.cat([prev_key, k], dim=1)
+                v = torch.cat([prev_value, v], dim=1)
+            incremental_state["prev_key"] = k.view(bsz, self.num_heads, -1, self.head_dim)
+            incremental_state["prev_value"] = v.view(bsz, self.num_heads, -1, self.head_dim)
+            src_len = k.size(1)
+
+        if self.xpos is not None:
+            if incremental_state is not None:
+                offset = src_len - 1
+            else:
+                offset = 0
+            k = self.xpos(k, offset=0, downscale=True)
+            q = self.xpos(q, offset=offset, downscale=False)
+
+        attn_weights = torch.bmm(q, k.transpose(1, 2))
+
+        if attn_mask is not None:
+            attn_weights = torch.nan_to_num(attn_weights)
+            if len(attn_mask.shape) != len(attn_weights.shape):
+                attn_mask = attn_mask.unsqueeze(0)
+            else:
+                attn_mask = attn_mask.repeat_interleave(self.num_heads, dim=0)
+            attn_weights += attn_mask
+
+        if key_padding_mask is not None:
+            attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
+            attn_weights = attn_weights.masked_fill(
+                key_padding_mask.unsqueeze(1).unsqueeze(2).to(torch.bool),
+                float("-inf"),
+            )
+            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
+
+        if rel_pos is not None:
+            rel_pos = rel_pos.view(attn_weights.size())
+            attn_weights = attn_weights + rel_pos
+
+        attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).type_as(attn_weights)
+        attn_probs = self.dropout_module(attn_weights)
+
+        attn = torch.bmm(attn_probs, v)
+        attn = attn.transpose(0, 1).reshape(tgt_len, bsz, embed_dim).transpose(0, 1)
+
+        if self.inner_attn_ln is not None:
+            attn = self.inner_attn_ln(attn)
+
+        attn = self.out_proj(attn)
+        attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len).transpose(1, 0)
+
+        return attn, attn_weights

vlmo/torchscale/component/multiway_network.py

@@ -0,0 +1,55 @@
+# Copyright (c) 2022 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+
+import copy
+
+import torch
+import torch.nn as nn
+
+
+def MultiwayWrapper(args, module, dim=1):
+    if args.multiway:
+        return MultiwayNetwork(module, dim=dim)
+    return module
+
+
+def set_split_position(position):
+    def apply_fn(module):
+        if hasattr(module, "split_position"):
+            module.split_position = position
+
+    return apply_fn
+
+
+class MultiwayNetwork(nn.Module):
+    def __init__(self, module, dim=1):
+        super().__init__()
+        self.dim = dim
+        self.A = module
+        self.B = copy.deepcopy(module)
+        self.B.reset_parameters()
+        self.split_position = -1
+
+    def forward(self, x, **kwargs):
+        if self.split_position == -1:
+            return self.A(x, **kwargs)
+        if self.split_position == 0:
+            return self.B(x, **kwargs)
+        x1, x2 = torch.split(
+            x,
+            [self.split_position, x.size(self.dim) - self.split_position],
+            dim=self.dim,
+        )
+        # x1, x2 = x[:self.split_position], x[self.split_position:]
+        y1, y2 = self.A(x1, **kwargs), self.B(x2, **kwargs)
+        return torch.cat([y1, y2], dim=self.dim)
+
+
+class MutliwayEmbedding(MultiwayNetwork):
+    def __init__(self, modules, dim=1):
+        super(MultiwayNetwork, self).__init__()
+        self.dim = dim
+        assert len(modules) == 2
+        self.A = modules[0]
+        self.B = modules[1]
+        self.split_position = -1

vlmo/torchscale/component/relative_position_bias.py

@@ -0,0 +1,67 @@
+# Copyright (c) 2022 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+
+import math
+
+import torch
+import torch.nn as nn
+
+
+class RelativePositionBias(nn.Module):
+    def __init__(self, bidirectional=True, num_buckets=32, max_distance=128, n_heads=12):
+        super().__init__()
+        self.bidirectional = bidirectional
+        self.num_buckets = num_buckets
+        self.max_distance = max_distance
+        self.n_heads = n_heads
+        self.relative_attention_bias = nn.Embedding(self.num_buckets, self.n_heads)
+
+    @staticmethod
+    def _relative_position_bucket(relative_position, bidirectional=True, num_buckets=32, max_distance=128):
+        ret = 0
+        n = -relative_position
+        if bidirectional:
+            num_buckets //= 2
+            ret += (n < 0).to(torch.long) * num_buckets
+            n = torch.abs(n)
+        else:
+            n = torch.max(n, torch.zeros_like(n))
+
+        max_exact = num_buckets // 2
+        is_small = n < max_exact
+
+        val_if_large = max_exact + (
+            torch.log(n.float() / max_exact) / math.log(max_distance / max_exact) * (num_buckets - max_exact)
+        ).to(torch.long)
+        val_if_large = torch.min(val_if_large, torch.full_like(val_if_large, num_buckets - 1))
+
+        ret += torch.where(is_small, n, val_if_large)
+        return ret
+
+    def compute_bias(self, qlen, klen, step=None):
+        step = 0 if step is None else step
+        context_position = torch.arange(
+            step,
+            step + qlen,
+            dtype=torch.long,
+            device=self.relative_attention_bias.weight.device,
+        )[:, None]
+        memory_position = torch.arange(klen, dtype=torch.long, device=self.relative_attention_bias.weight.device)[
+            None, :
+        ]
+        relative_position = memory_position - context_position  # shape (qlen, klen)
+
+        rp_bucket = self._relative_position_bucket(
+            relative_position,  # shape (qlen, klen)
+            bidirectional=self.bidirectional,
+            num_buckets=self.num_buckets,
+            max_distance=self.max_distance,
+        )
+        rp_bucket = rp_bucket.to(self.relative_attention_bias.weight.device)
+        values = self.relative_attention_bias(rp_bucket)  # shape (qlen, klen, num_heads)
+        values = values.permute([2, 0, 1]).unsqueeze(0)  # shape (1, num_heads, qlen, klen)
+        return values
+
+    def forward(self, batch_size, qlen, klen, step=None):
+        # shape (batch * num_heads, qlen, klen)
+        return self.compute_bias(qlen, klen, step).repeat(batch_size, 1, 1, 1).view(-1, qlen, klen)

vlmo/torchscale/component/xpos_relative_position.py

@@ -0,0 +1,62 @@
+# Copyright (c) 2022 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+
+import torch
+import torch.nn as nn
+
+
+def fixed_pos_embedding(x):
+    seq_len, dim = x.shape
+    inv_freq = 1.0 / (10000 ** (torch.arange(0, dim) / dim))
+    sinusoid_inp = torch.einsum("i , j -> i j", torch.arange(0, seq_len, dtype=torch.float), inv_freq).to(x)
+    return torch.sin(sinusoid_inp), torch.cos(sinusoid_inp)
+
+
+def rotate_every_two(x):
+    x1 = x[:, :, ::2]
+    x2 = x[:, :, 1::2]
+    x = torch.stack((-x2, x1), dim=-1)
+    return x.flatten(-2)  # in einsum notation: rearrange(x, '... d j -> ... (d j)')\
+
+
+def duplicate_interleave(m):
+    """
+    A simple version of `torch.repeat_interleave` for duplicating a matrix while interleaving the copy.
+    """
+    dim0 = m.shape[0]
+    m = m.view(-1, 1)  # flatten the matrix
+    m = m.repeat(1, 2)  # repeat all elements into the 2nd dimension
+    m = m.view(dim0, -1)  # reshape into a matrix, interleaving the copy
+    return m
+
+
+def apply_rotary_pos_emb(x, sin, cos, scale=1):
+    sin, cos = map(lambda t: duplicate_interleave(t * scale), (sin, cos))
+    # einsum notation for lambda t: repeat(t[offset:x.shape[1]+offset,:], "n d -> () n () (d j)", j=2)
+    return (x * cos) + (rotate_every_two(x) * sin)
+
+
+class XPOS(nn.Module):
+    def __init__(self, head_dim, scale_base=512):
+        super().__init__()
+        self.head_dim = head_dim
+        self.scale_base = scale_base
+        self.register_buffer("scale", (torch.arange(0, head_dim, 2) + 0.4 * head_dim) / (1.4 * head_dim))
+
+    def forward(self, x, offset=0, downscale=False):
+        length = x.shape[1]
+        min_pos = -(length + offset) // 2
+        max_pos = length + offset + min_pos
+        scale = self.scale ** torch.arange(min_pos, max_pos, 1).to(self.scale).div(self.scale_base)[:, None]
+        sin, cos = fixed_pos_embedding(scale)
+
+        if scale.shape[0] > length:
+            scale = scale[-length:]
+            sin = sin[-length:]
+            cos = cos[-length:]
+
+        if downscale:
+            scale = 1 / scale
+
+        x = apply_rotary_pos_emb(x, sin, cos, scale)
+        return x

vlmo/torchscale/model/BEiT3.py

@@ -0,0 +1,96 @@
+# Copyright (c) 2022 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+
+import torch
+import torch.nn as nn
+
+from vlmo.torchscale.architecture.encoder import Encoder
+from vlmo.torchscale.component.embedding import (
+    PositionalEmbedding,
+    TextEmbedding,
+    VisionEmbedding,
+)
+from vlmo.torchscale.component.multiway_network import MutliwayEmbedding
+
+
+class BEiT3(nn.Module):
+    def __init__(self, args, **kwargs):
+        super().__init__()
+        self.args = args
+        assert args.multiway
+        assert args.vocab_size > 0
+        assert not args.share_encoder_input_output_embed
+        self.text_embed = TextEmbedding(args.vocab_size, args.encoder_embed_dim)
+        self.vision_embed = VisionEmbedding(
+            args.img_size,
+            args.patch_size,
+            args.in_chans,
+            args.encoder_embed_dim,
+            contain_mask_token=True,
+            prepend_cls_token=True,
+        )
+        # being consistent with Fairseq, which starts from 2 for position embedding
+        embed_positions = MutliwayEmbedding(
+            modules=[
+                PositionalEmbedding(self.vision_embed.num_position_embeddings() + 2, args.encoder_embed_dim),
+                PositionalEmbedding(args.max_source_positions, args.encoder_embed_dim),
+            ],
+            dim=1,
+        )
+        self.encoder = Encoder(
+            args,
+            embed_tokens=None,
+            embed_positions=embed_positions,
+            output_projection=None,
+            is_encoder_decoder=False,
+        )
+
+    def forward(
+        self,
+        textual_tokens=None,
+        visual_tokens=None,
+        text_padding_position=None,
+        attn_mask=None,
+        vision_masked_position=None,
+        incremental_state=None,
+        positions=None,
+    ):
+        assert textual_tokens is not None or visual_tokens is not None
+
+        if textual_tokens is None:
+            x = self.vision_embed(visual_tokens, vision_masked_position)
+            encoder_padding_mask = None
+            multiway_split_position = -1
+        elif visual_tokens is None:
+            x = self.text_embed(textual_tokens)
+            encoder_padding_mask = text_padding_position
+            multiway_split_position = 0
+        else:
+            x1 = self.vision_embed(visual_tokens, vision_masked_position)
+            multiway_split_position = x1.size(1)
+            x2 = self.text_embed(textual_tokens)
+            diff = x1.shape[0] // x2.shape[0]
+            if diff != 1:
+                x2 = torch.repeat_interleave(x2, diff, dim=0)
+                text_padding_position = torch.repeat_interleave(text_padding_position, diff, dim=0)
+            x = torch.cat([x1, x2], dim=1)
+            if text_padding_position is not None:
+                encoder_padding_mask = torch.cat(
+                    [
+                        torch.zeros(x1.shape[:-1], device=x1.device, dtype=torch.bool),
+                        text_padding_position,
+                    ],
+                    dim=1,
+                )
+            else:
+                encoder_padding_mask = None
+        encoder_out = self.encoder(
+            src_tokens=None,
+            encoder_padding_mask=encoder_padding_mask,
+            attn_mask=attn_mask,
+            token_embeddings=x,
+            multiway_split_position=multiway_split_position,
+            incremental_state=incremental_state,
+            positions=positions,
+        )
+        return encoder_out

vlmo/torchscale/model/__init__.py

@@ -0,0 +1,2 @@
+# Copyright (c) 2022 Microsoft
+# Licensed under The MIT License [see LICENSE for details]

vlmo/transforms/__init__.py

@@ -0,0 +1,19 @@
+from .pixelbert import (
+    pixelbert_transform,
+    pixelbert_transform_randaug,
+)
+from .square_transform import (
+    square_transform,
+    square_transform_randaug,
+)
+
+_transforms = {
+    "pixelbert": pixelbert_transform,
+    "pixelbert_randaug": pixelbert_transform_randaug,
+    "square_transform": square_transform,
+    "square_transform_randaug": square_transform_randaug,
+}
+
+
+def keys_to_transforms(keys: list, size=224):
+    return [_transforms[key](size=size) for key in keys]

vlmo/transforms/pixelbert.py

@@ -0,0 +1,30 @@
+from .utils import (
+    inception_normalize,
+    MinMaxResize,
+)
+from torchvision import transforms
+from .randaug import RandAugment
+
+
+def pixelbert_transform(size=800):
+    longer = int((1333 / 800) * size)
+    return transforms.Compose(
+        [
+            MinMaxResize(shorter=size, longer=longer),
+            transforms.ToTensor(),
+            inception_normalize,
+        ]
+    )
+
+
+def pixelbert_transform_randaug(size=800):
+    longer = int((1333 / 800) * size)
+    trs = transforms.Compose(
+        [
+            MinMaxResize(shorter=size, longer=longer),
+            transforms.ToTensor(),
+            inception_normalize,
+        ]
+    )
+    trs.transforms.insert(0, RandAugment(2, 9))
+    return trs

vlmo/transforms/randaug.py

@@ -0,0 +1,271 @@
+# code in this file is adpated from rpmcruz/autoaugment
+# https://github.com/rpmcruz/autoaugment/blob/master/transformations.py
+import random
+
+import PIL
+
+# from PIL import ImageOps, ImageEnhance, ImageDraw
+import numpy as np
+import torch
+from PIL import Image
+
+
+def ShearX(img, v):  # [-0.3, 0.3]
+    assert -0.3 <= v <= 0.3
+    if random.random() > 0.5:
+        v = -v
+    return img.transform(img.size, PIL.Image.AFFINE, (1, v, 0, 0, 1, 0))
+
+
+def ShearY(img, v):  # [-0.3, 0.3]
+    assert -0.3 <= v <= 0.3
+    if random.random() > 0.5:
+        v = -v
+    return img.transform(img.size, PIL.Image.AFFINE, (1, 0, 0, v, 1, 0))
+
+
+def TranslateX(img, v):  # [-150, 150] => percentage: [-0.45, 0.45]
+    assert -0.45 <= v <= 0.45
+    if random.random() > 0.5:
+        v = -v
+    v = v * img.size[0]
+    return img.transform(img.size, PIL.Image.AFFINE, (1, 0, v, 0, 1, 0))
+
+
+def TranslateXabs(img, v):  # [-150, 150] => percentage: [-0.45, 0.45]
+    assert 0 <= v
+    if random.random() > 0.5:
+        v = -v
+    return img.transform(img.size, PIL.Image.AFFINE, (1, 0, v, 0, 1, 0))
+
+
+def TranslateY(img, v):  # [-150, 150] => percentage: [-0.45, 0.45]
+    assert -0.45 <= v <= 0.45
+    if random.random() > 0.5:
+        v = -v
+    v = v * img.size[1]
+    return img.transform(img.size, PIL.Image.AFFINE, (1, 0, 0, 0, 1, v))
+
+
+def TranslateYabs(img, v):  # [-150, 150] => percentage: [-0.45, 0.45]
+    assert 0 <= v
+    if random.random() > 0.5:
+        v = -v
+    return img.transform(img.size, PIL.Image.AFFINE, (1, 0, 0, 0, 1, v))
+
+
+def Rotate(img, v):  # [-30, 30]
+    assert -30 <= v <= 30
+    if random.random() > 0.5:
+        v = -v
+    return img.rotate(v)
+
+
+def AutoContrast(img, _):
+    return PIL.ImageOps.autocontrast(img)
+
+
+def Invert(img, _):
+    return PIL.ImageOps.invert(img)
+
+
+def Equalize(img, _):
+    return PIL.ImageOps.equalize(img)
+
+
+def Flip(img, _):  # not from the paper
+    return PIL.ImageOps.mirror(img)
+
+
+def Solarize(img, v):  # [0, 256]
+    assert 0 <= v <= 256
+    return PIL.ImageOps.solarize(img, v)
+
+
+def SolarizeAdd(img, addition=0, threshold=128):
+    img_np = np.array(img).astype(np.int)
+    img_np = img_np + addition
+    img_np = np.clip(img_np, 0, 255)
+    img_np = img_np.astype(np.uint8)
+    img = Image.fromarray(img_np)
+    return PIL.ImageOps.solarize(img, threshold)
+
+
+def Posterize(img, v):  # [4, 8]
+    v = int(v)
+    v = max(1, v)
+    return PIL.ImageOps.posterize(img, v)
+
+
+def Contrast(img, v):  # [0.1,1.9]
+    assert 0.1 <= v <= 1.9
+    return PIL.ImageEnhance.Contrast(img).enhance(v)
+
+
+def Color(img, v):  # [0.1,1.9]
+    assert 0.1 <= v <= 1.9
+    return PIL.ImageEnhance.Color(img).enhance(v)
+
+
+def Brightness(img, v):  # [0.1,1.9]
+    assert 0.1 <= v <= 1.9
+    return PIL.ImageEnhance.Brightness(img).enhance(v)
+
+
+def Sharpness(img, v):  # [0.1,1.9]
+    assert 0.1 <= v <= 1.9
+    return PIL.ImageEnhance.Sharpness(img).enhance(v)
+
+
+def Cutout(img, v):  # [0, 60] => percentage: [0, 0.2]
+    assert 0.0 <= v <= 0.2
+    if v <= 0.0:
+        return img
+
+    v = v * img.size[0]
+    return CutoutAbs(img, v)
+
+
+def CutoutAbs(img, v):  # [0, 60] => percentage: [0, 0.2]
+    # assert 0 <= v <= 20
+    if v < 0:
+        return img
+    w, h = img.size
+    x0 = np.random.uniform(w)
+    y0 = np.random.uniform(h)
+
+    x0 = int(max(0, x0 - v / 2.0))
+    y0 = int(max(0, y0 - v / 2.0))
+    x1 = min(w, x0 + v)
+    y1 = min(h, y0 + v)
+
+    xy = (x0, y0, x1, y1)
+    color = (125, 123, 114)
+    # color = (0, 0, 0)
+    img = img.copy()
+    PIL.ImageDraw.Draw(img).rectangle(xy, color)
+    return img
+
+
+def SamplePairing(imgs):  # [0, 0.4]
+    def f(img1, v):
+        i = np.random.choice(len(imgs))
+        img2 = PIL.Image.fromarray(imgs[i])
+        return PIL.Image.blend(img1, img2, v)
+
+    return f
+
+
+def Identity(img, v):
+    return img
+
+
+def augment_list():  # 16 oeprations and their ranges
+    # https://github.com/google-research/uda/blob/master/image/randaugment/policies.py#L57
+    # l = [
+    #     (Identity, 0., 1.0),
+    #     (ShearX, 0., 0.3),  # 0
+    #     (ShearY, 0., 0.3),  # 1
+    #     (TranslateX, 0., 0.33),  # 2
+    #     (TranslateY, 0., 0.33),  # 3
+    #     (Rotate, 0, 30),  # 4
+    #     (AutoContrast, 0, 1),  # 5
+    #     (Invert, 0, 1),  # 6
+    #     (Equalize, 0, 1),  # 7
+    #     (Solarize, 0, 110),  # 8
+    #     (Posterize, 4, 8),  # 9
+    #     # (Contrast, 0.1, 1.9),  # 10
+    #     (Color, 0.1, 1.9),  # 11
+    #     (Brightness, 0.1, 1.9),  # 12
+    #     (Sharpness, 0.1, 1.9),  # 13
+    #     # (Cutout, 0, 0.2),  # 14
+    #     # (SamplePairing(imgs), 0, 0.4),  # 15
+    # ]
+
+    # https://github.com/tensorflow/tpu/blob/8462d083dd89489a79e3200bcc8d4063bf362186/models/official/efficientnet/autoaugment.py#L505
+    l = [
+        (AutoContrast, 0, 1),
+        (Equalize, 0, 1),
+        # (Invert, 0, 1),
+        (Rotate, 0, 30),
+        (Posterize, 0, 4),
+        (Solarize, 0, 256),
+        (SolarizeAdd, 0, 110),
+        (Color, 0.1, 1.9),
+        (Contrast, 0.1, 1.9),
+        (Brightness, 0.1, 1.9),
+        (Sharpness, 0.1, 1.9),
+        (ShearX, 0.0, 0.3),
+        (ShearY, 0.0, 0.3),
+        # (CutoutAbs, 0, 40),
+        (TranslateXabs, 0.0, 100),
+        (TranslateYabs, 0.0, 100),
+    ]
+
+    return l
+
+
+class Lighting(object):
+    """Lighting noise(AlexNet - style PCA - based noise)"""
+
+    def __init__(self, alphastd, eigval, eigvec):
+        self.alphastd = alphastd
+        self.eigval = torch.Tensor(eigval)
+        self.eigvec = torch.Tensor(eigvec)
+
+    def __call__(self, img):
+        if self.alphastd == 0:
+            return img
+
+        alpha = img.new().resize_(3).normal_(0, self.alphastd)
+        rgb = (
+            self.eigvec.type_as(img)
+            .clone()
+            .mul(alpha.view(1, 3).expand(3, 3))
+            .mul(self.eigval.view(1, 3).expand(3, 3))
+            .sum(1)
+            .squeeze()
+        )
+
+        return img.add(rgb.view(3, 1, 1).expand_as(img))
+
+
+class CutoutDefault(object):
+    """
+    Reference : https://github.com/quark0/darts/blob/master/cnn/utils.py
+    """
+
+    def __init__(self, length):
+        self.length = length
+
+    def __call__(self, img):
+        h, w = img.size(1), img.size(2)
+        mask = np.ones((h, w), np.float32)
+        y = np.random.randint(h)
+        x = np.random.randint(w)
+
+        y1 = np.clip(y - self.length // 2, 0, h)
+        y2 = np.clip(y + self.length // 2, 0, h)
+        x1 = np.clip(x - self.length // 2, 0, w)
+        x2 = np.clip(x + self.length // 2, 0, w)
+
+        mask[y1:y2, x1:x2] = 0.0
+        mask = torch.from_numpy(mask)
+        mask = mask.expand_as(img)
+        img *= mask
+        return img
+
+
+class RandAugment:
+    def __init__(self, n, m):
+        self.n = n
+        self.m = m  # [0, 30]
+        self.augment_list = augment_list()
+
+    def __call__(self, img):
+        ops = random.choices(self.augment_list, k=self.n)
+        for op, minval, maxval in ops:
+            val = (float(self.m) / 30) * float(maxval - minval) + minval
+            img = op(img, val)
+
+        return img

vlmo/transforms/randaugment.py

@@ -0,0 +1,334 @@
+import cv2
+import numpy as np
+
+
+# aug functions
+def identity_func(img):
+    return img
+
+
+def autocontrast_func(img, cutoff=0):
+    """
+    same output as PIL.ImageOps.autocontrast
+    """
+    n_bins = 256
+
+    def tune_channel(ch):
+        n = ch.size
+        cut = cutoff * n // 100
+        if cut == 0:
+            high, low = ch.max(), ch.min()
+        else:
+            hist = cv2.calcHist([ch], [0], None, [n_bins], [0, n_bins])
+            low = np.argwhere(np.cumsum(hist) > cut)
+            low = 0 if low.shape[0] == 0 else low[0]
+            high = np.argwhere(np.cumsum(hist[::-1]) > cut)
+            high = n_bins - 1 if high.shape[0] == 0 else n_bins - 1 - high[0]
+        if high <= low:
+            table = np.arange(n_bins)
+        else:
+            scale = (n_bins - 1) / (high - low)
+            offset = -low * scale
+            table = np.arange(n_bins) * scale + offset
+            table[table < 0] = 0
+            table[table > n_bins - 1] = n_bins - 1
+        table = table.clip(0, 255).astype(np.uint8)
+        return table[ch]
+
+    channels = [tune_channel(ch) for ch in cv2.split(img)]
+    out = cv2.merge(channels)
+    return out
+
+
+def equalize_func(img):
+    """
+    same output as PIL.ImageOps.equalize
+    PIL's implementation is different from cv2.equalize
+    """
+    n_bins = 256
+
+    def tune_channel(ch):
+        hist = cv2.calcHist([ch], [0], None, [n_bins], [0, n_bins])
+        non_zero_hist = hist[hist != 0].reshape(-1)
+        step = np.sum(non_zero_hist[:-1]) // (n_bins - 1)
+        if step == 0:
+            return ch
+        n = np.empty_like(hist)
+        n[0] = step // 2
+        n[1:] = hist[:-1]
+        table = (np.cumsum(n) // step).clip(0, 255).astype(np.uint8)
+        return table[ch]
+
+    channels = [tune_channel(ch) for ch in cv2.split(img)]
+    out = cv2.merge(channels)
+    return out
+
+
+def rotate_func(img, degree, fill=(0, 0, 0)):
+    """
+    like PIL, rotate by degree, not radians
+    """
+    H, W = img.shape[0], img.shape[1]
+    center = W / 2, H / 2
+    M = cv2.getRotationMatrix2D(center, degree, 1)
+    out = cv2.warpAffine(img, M, (W, H), borderValue=fill)
+    return out
+
+
+def solarize_func(img, thresh=128):
+    """
+    same output as PIL.ImageOps.posterize
+    """
+    table = np.array([el if el < thresh else 255 - el for el in range(256)])
+    table = table.clip(0, 255).astype(np.uint8)
+    out = table[img]
+    return out
+
+
+def color_func(img, factor):
+    """
+    same output as PIL.ImageEnhance.Color
+    """
+    # implementation according to PIL definition, quite slow
+    #  degenerate = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)[:, :, np.newaxis]
+    #  out = blend(degenerate, img, factor)
+    #  M = (
+    #      np.eye(3) * factor
+    #      + np.float32([0.114, 0.587, 0.299]).reshape(3, 1) * (1. - factor)
+    #  )[np.newaxis, np.newaxis, :]
+    M = np.float32([[0.886, -0.114, -0.114], [-0.587, 0.413, -0.587], [-0.299, -0.299, 0.701]]) * factor + np.float32(
+        [[0.114], [0.587], [0.299]]
+    )
+    out = np.matmul(img, M).clip(0, 255).astype(np.uint8)
+    return out
+
+
+def contrast_func(img, factor):
+    """
+    same output as PIL.ImageEnhance.Contrast
+    """
+    mean = np.sum(np.mean(img, axis=(0, 1)) * np.array([0.114, 0.587, 0.299]))
+    table = np.array([(el - mean) * factor + mean for el in range(256)]).clip(0, 255).astype(np.uint8)
+    out = table[img]
+    return out
+
+
+def brightness_func(img, factor):
+    """
+    same output as PIL.ImageEnhance.Contrast
+    """
+    table = (np.arange(256, dtype=np.float32) * factor).clip(0, 255).astype(np.uint8)
+    out = table[img]
+    return out
+
+
+def sharpness_func(img, factor):
+    """
+    The differences the this result and PIL are all on the 4 boundaries, the center
+    areas are same
+    """
+    kernel = np.ones((3, 3), dtype=np.float32)
+    kernel[1][1] = 5
+    kernel /= 13
+    degenerate = cv2.filter2D(img, -1, kernel)
+    if factor == 0.0:
+        out = degenerate
+    elif factor == 1.0:
+        out = img
+    else:
+        out = img.astype(np.float32)
+        degenerate = degenerate.astype(np.float32)[1:-1, 1:-1, :]
+        out[1:-1, 1:-1, :] = degenerate + factor * (out[1:-1, 1:-1, :] - degenerate)
+        out = out.astype(np.uint8)
+    return out
+
+
+def shear_x_func(img, factor, fill=(0, 0, 0)):
+    H, W = img.shape[0], img.shape[1]
+    M = np.float32([[1, factor, 0], [0, 1, 0]])
+    out = cv2.warpAffine(img, M, (W, H), borderValue=fill, flags=cv2.INTER_LINEAR).astype(np.uint8)
+    return out
+
+
+def translate_x_func(img, offset, fill=(0, 0, 0)):
+    """
+    same output as PIL.Image.transform
+    """
+    H, W = img.shape[0], img.shape[1]
+    M = np.float32([[1, 0, -offset], [0, 1, 0]])
+    out = cv2.warpAffine(img, M, (W, H), borderValue=fill, flags=cv2.INTER_LINEAR).astype(np.uint8)
+    return out
+
+
+def translate_y_func(img, offset, fill=(0, 0, 0)):
+    """
+    same output as PIL.Image.transform
+    """
+    H, W = img.shape[0], img.shape[1]
+    M = np.float32([[1, 0, 0], [0, 1, -offset]])
+    out = cv2.warpAffine(img, M, (W, H), borderValue=fill, flags=cv2.INTER_LINEAR).astype(np.uint8)
+    return out
+
+
+def posterize_func(img, bits):
+    """
+    same output as PIL.ImageOps.posterize
+    """
+    out = np.bitwise_and(img, np.uint8(255 << (8 - bits)))
+    return out
+
+
+def shear_y_func(img, factor, fill=(0, 0, 0)):
+    H, W = img.shape[0], img.shape[1]
+    M = np.float32([[1, 0, 0], [factor, 1, 0]])
+    out = cv2.warpAffine(img, M, (W, H), borderValue=fill, flags=cv2.INTER_LINEAR).astype(np.uint8)
+    return out
+
+
+def cutout_func(img, pad_size, replace=(0, 0, 0)):
+    replace = np.array(replace, dtype=np.uint8)
+    H, W = img.shape[0], img.shape[1]
+    rh, rw = np.random.random(2)
+    pad_size = pad_size // 2
+    ch, cw = int(rh * H), int(rw * W)
+    x1, x2 = max(ch - pad_size, 0), min(ch + pad_size, H)
+    y1, y2 = max(cw - pad_size, 0), min(cw + pad_size, W)
+    out = img.copy()
+    out[x1:x2, y1:y2, :] = replace
+    return out
+
+
+# level to args
+def enhance_level_to_args(MAX_LEVEL):
+    def level_to_args(level):
+        return ((level / MAX_LEVEL) * 1.8 + 0.1,)
+
+    return level_to_args
+
+
+def shear_level_to_args(MAX_LEVEL, replace_value):
+    def level_to_args(level):
+        level = (level / MAX_LEVEL) * 0.3
+        if np.random.random() > 0.5:
+            level = -level
+        return (level, replace_value)
+
+    return level_to_args
+
+
+def translate_level_to_args(translate_const, MAX_LEVEL, replace_value):
+    def level_to_args(level):
+        level = (level / MAX_LEVEL) * float(translate_const)
+        if np.random.random() > 0.5:
+            level = -level
+        return (level, replace_value)
+
+    return level_to_args
+
+
+def cutout_level_to_args(cutout_const, MAX_LEVEL, replace_value):
+    def level_to_args(level):
+        level = int((level / MAX_LEVEL) * cutout_const)
+        return (level, replace_value)
+
+    return level_to_args
+
+
+def solarize_level_to_args(MAX_LEVEL):
+    def level_to_args(level):
+        level = int((level / MAX_LEVEL) * 256)
+        return (level,)
+
+    return level_to_args
+
+
+def none_level_to_args(level):
+    return ()
+
+
+def posterize_level_to_args(MAX_LEVEL):
+    def level_to_args(level):
+        level = int((level / MAX_LEVEL) * 4)
+        return (level,)
+
+    return level_to_args
+
+
+def rotate_level_to_args(MAX_LEVEL, replace_value):
+    def level_to_args(level):
+        level = (level / MAX_LEVEL) * 30
+        if np.random.random() < 0.5:
+            level = -level
+        return (level, replace_value)
+
+    return level_to_args
+
+
+func_dict = {
+    "Identity": identity_func,
+    "AutoContrast": autocontrast_func,
+    "Equalize": equalize_func,
+    "Rotate": rotate_func,
+    "Solarize": solarize_func,
+    "Color": color_func,
+    "Contrast": contrast_func,
+    "Brightness": brightness_func,
+    "Sharpness": sharpness_func,
+    "ShearX": shear_x_func,
+    "TranslateX": translate_x_func,
+    "TranslateY": translate_y_func,
+    "Posterize": posterize_func,
+    "ShearY": shear_y_func,
+}
+
+translate_const = 10
+MAX_LEVEL = 10
+replace_value = (128, 128, 128)
+arg_dict = {
+    "Identity": none_level_to_args,
+    "AutoContrast": none_level_to_args,
+    "Equalize": none_level_to_args,
+    "Rotate": rotate_level_to_args(MAX_LEVEL, replace_value),
+    "Solarize": solarize_level_to_args(MAX_LEVEL),
+    "Color": enhance_level_to_args(MAX_LEVEL),
+    "Contrast": enhance_level_to_args(MAX_LEVEL),
+    "Brightness": enhance_level_to_args(MAX_LEVEL),
+    "Sharpness": enhance_level_to_args(MAX_LEVEL),
+    "ShearX": shear_level_to_args(MAX_LEVEL, replace_value),
+    "TranslateX": translate_level_to_args(translate_const, MAX_LEVEL, replace_value),
+    "TranslateY": translate_level_to_args(translate_const, MAX_LEVEL, replace_value),
+    "Posterize": posterize_level_to_args(MAX_LEVEL),
+    "ShearY": shear_level_to_args(MAX_LEVEL, replace_value),
+}
+
+
+class RandomAugment(object):
+    def __init__(self, N=2, M=10, isPIL=False, augs=[]):
+        self.N = N
+        self.M = M
+        self.isPIL = isPIL
+        if augs:
+            self.augs = augs
+        else:
+            self.augs = list(arg_dict.keys())
+
+    def get_random_ops(self):
+        sampled_ops = np.random.choice(self.augs, self.N)
+        return [(op, 0.5, self.M) for op in sampled_ops]
+
+    def __call__(self, img):
+        if self.isPIL:
+            img = np.array(img)
+        ops = self.get_random_ops()
+        for name, prob, level in ops:
+            if np.random.random() > prob:
+                continue
+            args = arg_dict[name](level)
+            img = func_dict[name](img, *args)
+        return img
+
+
+if __name__ == "__main__":
+    a = RandomAugment()
+    img = np.random.randn(32, 32, 3)
+    a(img)

vlmo/transforms/square_transform.py

@@ -0,0 +1,41 @@
+# code in this file is adpated from the ALBEF repo (https://github.com/salesforce/ALBEF)
+
+from torchvision import transforms
+from .randaugment import RandomAugment
+from PIL import Image
+
+
+def square_transform(size=224):
+    return transforms.Compose(
+        [
+            transforms.Resize((size, size), interpolation=Image.BICUBIC),
+            transforms.ToTensor(),
+        ]
+    )
+
+
+def square_transform_randaug(size=224):
+    return transforms.Compose(
+        [
+            transforms.RandomResizedCrop(size, scale=(0.8, 1.0), interpolation=Image.BICUBIC),
+            transforms.RandomHorizontalFlip(),
+            RandomAugment(
+                2,
+                7,
+                isPIL=True,
+                augs=[
+                    "Identity",
+                    "AutoContrast",
+                    "Equalize",
+                    "Brightness",
+                    "Sharpness",
+                    "ShearX",
+                    "ShearY",
+                    "TranslateX",
+                    "TranslateY",
+                    "Rotate",
+                ],
+            ),
+            transforms.ToTensor(),
+        ]
+    )

vlmo/transforms/utils.py

@@ -0,0 +1,56 @@
+from torchvision import transforms
+from PIL import Image
+
+
+class MinMaxResize:
+    def __init__(self, shorter=800, longer=1333):
+        self.min = shorter
+        self.max = longer
+
+    def __call__(self, x):
+        w, h = x.size
+        scale = self.min / min(w, h)
+        if h < w:
+            newh, neww = self.min, scale * w
+        else:
+            newh, neww = scale * h, self.min
+
+        if max(newh, neww) > self.max:
+            scale = self.max / max(newh, neww)
+            newh = newh * scale
+            neww = neww * scale
+
+        newh, neww = int(newh + 0.5), int(neww + 0.5)
+        newh, neww = newh // 32 * 32, neww // 32 * 32
+
+        return x.resize((neww, newh), resample=Image.BICUBIC)
+
+
+class UnNormalize(object):
+    def __init__(self, mean, std):
+        self.mean = mean
+        self.std = std
+
+    def __call__(self, tensor):
+        """
+        Args:
+            tensor (Tensor): Tensor image of size (C, H, W) to be normalized.
+        Returns:
+            Tensor: Normalized image.
+        """
+        for t, m, s in zip(tensor, self.mean, self.std):
+            t.mul_(s).add_(m)
+            # The normalize code -> t.sub_(m).div_(s)
+        return tensor
+
+
+# This is simple maximum entropy normalization performed in Inception paper
+inception_normalize = transforms.Compose([transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])])
+
+# ViT uses simple non-biased inception normalization
+# https://github.com/google-research/vision_transformer/blob/master/vit_jax/input_pipeline.py#L132
+inception_unnormalize = transforms.Compose([UnNormalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])])
+
+cn_clip_normalize = transforms.Compose(
+    [transforms.Normalize(mean=[0.48145466, 0.4578275, 0.40821073], std=[0.26862954, 0.26130258, 0.27577711])]
+)