gpt2

libai/models/utils/model_utils/bert_loader.py

+# coding=utf-8
+# Copyright 2021 The OneFlow Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+
+import oneflow as flow
+
+from .base_loader import ModelLoaderHuggerFace, ModelLoaderLiBai
+
+
+class BertLoaderHuggerFace(ModelLoaderHuggerFace):
+    def __init__(self, model, libai_cfg, pretrained_model_path, **kwargs):
+        super().__init__(model, libai_cfg, pretrained_model_path, **kwargs)
+
+        """NOTE: base_model_prefix_1 is BERT's prefix in Transformers.
+        base_model_prefix_2 is BERT's prefix in LiBai."""
+        self.base_model_prefix_1 = "bert"
+        self.base_model_prefix_2 = "bert"
+
+    def _convert_state_dict(self, flow_state_dict, cfg):
+        """Convert state_dict's keys to match model.
+
+        Args:
+            flow_state_dict (OrderedDict): model state dict.
+            cfg (dict): model's default config dict in LiBai.
+
+        Returns:
+            OrderedDict: flow state dict.
+        """
+        # The converted checkpoint.
+        oneflow_state_dict = flow_state_dict.copy()
+
+        # Get configs
+        num_heads = cfg.get("num_attention_heads")
+        hidden_size = cfg.get("hidden_size")
+        layers = cfg.get("hidden_layers")
+        head_size = int(hidden_size / num_heads)
+
+        # prefix
+        has_prefix = any(s.startswith(self.base_model_prefix_1) for s in oneflow_state_dict)
+
+        prefix = "bert." if has_prefix else ""
+        index_idx = 3 if has_prefix else 2
+        qkv_idx = 6 if has_prefix else 5
+
+        old_keys = oneflow_state_dict.keys()
+
+        for key in list(old_keys):
+
+            # Convert bert's embedding layers
+            if "embeddings" in key:
+                if "word_embeddings" in key:
+                    new_key = key.replace("word_embeddings", "vocab_embeddings")
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                elif "token_type_embeddings" in key:
+                    new_key = key.replace("token_type_embeddings", "tokentype_embeddings")
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                elif "LayerNorm.weight" in key:
+                    new_key = prefix + "encoders.0.input_layernorm.weight"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                elif "LayerNorm.bias" in key:
+                    new_key = prefix + "encoders.0.input_layernorm.bias"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                else:
+                    oneflow_state_dict[key] = oneflow_state_dict[key]
+
+            # Convert bert's attention layers
+            elif "attention" in key:
+                if "self" in key:
+                    index = key.split(".")[index_idx]
+                    if (
+                        prefix + "encoders." + index + ".self_attention.query_key_value.weight"
+                        in oneflow_state_dict.keys()
+                    ):
+                        continue
+                    q_w = key.replace(key.split(".")[qkv_idx], "query").replace(
+                        key.split(".")[qkv_idx + 1], "weight"
+                    )
+                    k_w = q_w.replace("query", "key")
+                    v_w = q_w.replace("query", "value")
+                    q_b = q_w.replace("weight", "bias")
+                    k_b = k_w.replace("weight", "bias")
+                    v_b = v_w.replace("weight", "bias")
+
+                    qkv_w = flow.cat(
+                        (
+                            oneflow_state_dict.pop(q_w),
+                            oneflow_state_dict.pop(k_w),
+                            oneflow_state_dict.pop(v_w),
+                        ),
+                        dim=0,
+                    )
+                    qkv_b = flow.cat(
+                        (
+                            oneflow_state_dict.pop(q_b),
+                            oneflow_state_dict.pop(k_b),
+                            oneflow_state_dict.pop(v_b),
+                        ),
+                        dim=-1,
+                    )
+
+                    qkv_w = self._fix_qkv_ordering(qkv_w, head_size, num_heads)
+                    qkv_b = self._fix_qkv_ordering(qkv_b, head_size, num_heads)
+
+                    new_key = (
+                        prefix + "encoders." + index + ".self_attention.query_key_value.weight"
+                    )
+                    oneflow_state_dict[new_key] = qkv_w
+
+                    new_key = prefix + "encoders." + index + ".self_attention.query_key_value.bias"
+                    oneflow_state_dict[new_key] = qkv_b
+                elif "output" in key:
+                    index = key.split(".")[index_idx]
+                    if "dense" in key:
+                        if "weight" in key:
+                            new_key = prefix + "encoders." + index + ".self_attention.dense.weight"
+                            oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                        elif "bias" in key:
+                            new_key = prefix + "encoders." + index + ".self_attention.dense.bias"
+                            oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                    elif "LayerNorm" in key:
+                        if "weight" in key:
+                            new_key = (
+                                prefix + "encoders." + index + ".post_attention_layernorm.weight"
+                            )
+                            oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                        elif "bias" in key:
+                            new_key = (
+                                prefix + "encoders." + index + ".post_attention_layernorm.bias"
+                            )
+                            oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+
+            # Convert bert's intermediate layers
+            elif "intermediate" in key:
+                index = key.split(".")[index_idx]
+                if (
+                    prefix + "encoders." + index + ".mlp.dense_h_to_4h.weight"
+                    in oneflow_state_dict.keys()
+                ):
+                    continue
+                if "weight" in key:
+                    w = key
+                    b = key.replace("weight", "bias")
+                    new_key = prefix + "encoders." + index + ".mlp.dense_h_to_4h.weight"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(w)
+                    new_key = new_key.replace("weight", "bias")
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(b)
+
+            # Convert bert's output layers
+            elif "output" in key:
+                index = key.split(".")[index_idx]
+                if "dense.weight" in key:
+                    if (
+                        prefix + "encoders." + index + ".mlp.dense_4h_to_h.weight"
+                        in oneflow_state_dict.keys()
+                    ):
+                        continue
+                    w = key
+                    b = w.replace("weight", "bias")
+                    new_key = prefix + "encoders." + index + ".mlp.dense_4h_to_h.weight"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(w)
+                    new_key = new_key.replace("weight", "bias")
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(b)
+                elif "LayerNorm.weight" in key:
+                    if (
+                        prefix + "encoders." + str(int(index) + 1) + ".input_layernorm.weight"
+                        in oneflow_state_dict.keys()
+                    ):
+                        continue
+                    w = key
+                    b = w.replace("weight", "bias")
+                    if index == str(layers - 1):
+                        new_key = prefix + "final_layernorm.weight"
+                        oneflow_state_dict[new_key] = oneflow_state_dict.pop(w)
+                        new_key = new_key.replace("weight", "bias")
+                        oneflow_state_dict[new_key] = oneflow_state_dict.pop(b)
+                        continue
+                    new_key = prefix + "encoders." + str(int(index) + 1) + ".input_layernorm.weight"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(w)
+                    new_key = new_key.replace("weight", "bias")
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(b)
+
+            # Convert bert's pooler layers
+            elif "pooler" in key:
+                if "weight" in key:
+                    new_key = prefix + "pooler.dense.weight"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                elif "bias" in key:
+                    new_key = prefix + "pooler.dense.bias"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+
+            # Convert cls_head layers
+            elif "cls" in key:
+                if "predictions.bias" in key:
+                    new_key = "cls_head.lm_logits.bias"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                elif "dense.weight" in key:
+                    new_key = "cls_head.predictions.dense.weight"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                elif "dense.bias" in key:
+                    new_key = "cls_head.predictions.dense.bias"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                elif "LayerNorm.weight" in key:
+                    new_key = "cls_head.predictions.layernorm.weight"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                elif "LayerNorm.bias" in key:
+                    new_key = "cls_head.predictions.layernorm.bias"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                elif "seq_relationship" in key:
+                    new_key = key.replace("cls", "cls_head")
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+            else:
+                oneflow_state_dict[key] = oneflow_state_dict.pop(key)
+        return oneflow_state_dict
+
+    def _load_config_from_json(self, config_file):
+        """load config from `config.json`, and update default config.
+
+        Args:
+            config_file (str): Path of config file.
+        """
+        with open(config_file, mode="r", encoding="utf-8") as f:
+            cfg_dict = json.load(f)
+
+        # update libai_cfg by config.json
+        self._update_cfg("vocab_size", cfg_dict["vocab_size"])
+        self._update_cfg("hidden_size", cfg_dict["hidden_size"])
+        self._update_cfg("hidden_layers", cfg_dict["num_hidden_layers"])
+        self._update_cfg("num_attention_heads", cfg_dict["num_attention_heads"])
+        self._update_cfg("intermediate_size", cfg_dict["intermediate_size"])
+        self._update_cfg("hidden_dropout_prob", cfg_dict["hidden_dropout_prob"])
+        self._update_cfg("attention_probs_dropout_prob", cfg_dict["attention_probs_dropout_prob"])
+        self._update_cfg("max_position_embeddings", cfg_dict["max_position_embeddings"])
+        self._update_cfg("num_tokentypes", cfg_dict["type_vocab_size"])
+        self._update_cfg("initializer_range", cfg_dict["initializer_range"])
+        self._update_cfg("layernorm_eps", cfg_dict["layer_norm_eps"])
+
+        # update libai_cfg by kwargs
+        for k, v in self.kwargs.items():
+            self._update_cfg(k, v)
+
+        # use original BERT residual connection ordering
+        self.libai_cfg.apply_residual_post_layernorm = True
+
+        self._update_cfg_log()
+
+
+class BertLoaderLiBai(ModelLoaderLiBai):
+    def __init__(self, model, libai_cfg, pretrained_model_path, **kwargs):
+        super().__init__(model, libai_cfg, pretrained_model_path, **kwargs)
+        self.base_model_prefix_2 = "bert"

libai/models/utils/model_utils/gpt_loader.py

+# coding=utf-8
+# Copyright 2021 The OneFlow Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+
+from .base_loader import ModelLoaderHuggerFace, ModelLoaderLiBai
+
+
+class GPT2LoaderHuggerFace(ModelLoaderHuggerFace):
+    def __init__(self, model, libai_cfg, pretrained_model_path, **kwargs):
+        super().__init__(model, libai_cfg, pretrained_model_path, **kwargs)
+
+        """NOTE: base_model_prefix_1 is GPT's prefix in Transformers.
+        base_model_prefix_2 is GPT's prefix in LiBai."""
+        self.base_model_prefix_1 = "transformer"
+        self.base_model_prefix_2 = "GPT_model"
+
+    def _convert_state_dict(self, flow_state_dict, cfg):
+        """Convert state_dict's keys to match model.
+
+        Args:
+            flow_state_dict (OrderedDict): model state dict.
+            cfg (dict): model's default config dict in LiBai.
+
+        Returns:
+            OrderedDict: flow state dict.
+        """
+        # The converted checkpoint.
+        oneflow_state_dict = flow_state_dict.copy()
+        old_keys = list(oneflow_state_dict.keys())
+
+        # Get configs
+        num_heads = cfg.get("num_attention_heads")
+        hidden_size = cfg.get("hidden_size")
+        head_size = int(hidden_size / num_heads)
+
+        # prefix
+        has_prefix = any(s.startswith(self.base_model_prefix_1) for s in oneflow_state_dict)
+        prefix1 = self.base_model_prefix_1 + "." if has_prefix else ""
+        prefix2 = "GPT_model.transformer."
+        layer_idx = 2 if has_prefix else 1
+
+        # Convert Embedding layers.
+        new_key = "GPT_model.embeddings.token_embeddings.weight"
+        old_keys.remove(prefix1 + "wte.weight")
+        oneflow_state_dict[new_key] = oneflow_state_dict.pop(prefix1 + "wte.weight")
+
+        new_key = "GPT_model.embeddings.position_embeddings.weight"
+        old_keys.remove(prefix1 + "wpe.weight")
+        oneflow_state_dict[new_key] = oneflow_state_dict.pop(prefix1 + "wpe.weight")
+
+        for key in old_keys:
+            keys = key.split(".")
+            if layer_idx >= len(keys):
+                continue
+            layer = keys[layer_idx]
+            # Convert transformer layers.
+            if "h." in key:
+                if "ln_1" in key:
+                    if "weight" in key:
+                        new_key = prefix2 + "layers." + layer + ".input_layernorm.weight"
+                    else:
+                        new_key = prefix2 + "layers." + layer + ".input_layernorm.bias"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                elif "ln_2" in key:
+                    if "weight" in key:
+                        new_key = prefix2 + "layers." + layer + ".post_attention_layernorm.weight"
+                    else:
+                        new_key = prefix2 + "layers." + layer + ".post_attention_layernorm.bias"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                elif "attn" in key:
+                    if "c_attn" in key:
+                        if "weight" in key:
+                            new_key = (
+                                prefix2
+                                + "layers."
+                                + layer
+                                + ".self_attention.query_key_value.weight"
+                            )
+                        else:
+                            new_key = (
+                                prefix2 + "layers." + layer + ".self_attention.query_key_value.bias"
+                            )
+                        qkv = oneflow_state_dict.pop(key)
+                        if qkv.ndim > 1:
+                            qkv = qkv.transpose(1, 0)
+                        qkv = self._fix_qkv_ordering(qkv, head_size, num_heads)
+                        oneflow_state_dict[new_key] = qkv
+                    elif "c_proj" in key:
+                        if "weight" in key:
+                            new_key = prefix2 + "layers." + layer + ".self_attention.dense.weight"
+                        elif "bias" in key:
+                            new_key = prefix2 + "layers." + layer + ".self_attention.dense.bias"
+                        value = oneflow_state_dict.pop(key)
+                        if value.ndim > 1:
+                            value = value.transpose(1, 0)
+                        oneflow_state_dict[new_key] = value
+                elif "mlp" in key:
+                    if "c_fc" in key:
+                        if "weight" in key:
+                            new_key = prefix2 + "layers." + layer + ".mlp.dense_h_to_4h.weight"
+                        elif "bias" in key:
+                            new_key = prefix2 + "layers." + layer + ".mlp.dense_h_to_4h.bias"
+                        value = oneflow_state_dict.pop(key)
+                        if value.ndim > 1:
+                            value = value.transpose(1, 0)
+                        oneflow_state_dict[new_key] = value
+                    elif "c_proj" in key:
+                        if "weight" in key:
+                            new_key = prefix2 + "layers." + layer + ".mlp.dense_4h_to_h.weight"
+                        elif "bias" in key:
+                            new_key = prefix2 + "layers." + layer + ".mlp.dense_4h_to_h.bias"
+                        value = oneflow_state_dict.pop(key)
+                        if value.ndim > 1:
+                            value = value.transpose(1, 0)
+                        oneflow_state_dict[new_key] = value
+            elif "ln_f" in key:
+                if "weight" in key:
+                    new_key = prefix2 + "layernorm_f.weight"
+                elif "bias" in key:
+                    new_key = prefix2 + "layernorm_f.bias"
+                oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+        return oneflow_state_dict
+
+    def _load_config_from_json(self, config_file):
+        """load config from `config.json`, and update default config.
+
+        Args:
+            config_file (str): Path of config file.
+        """
+        with open(config_file, mode="r", encoding="utf-8") as f:
+            cfg_dict = json.load(f)
+
+        # update libai_cfg by config.json
+        self._update_cfg("hidden_layers", cfg_dict["n_layer"])
+        self._update_cfg("hidden_size", cfg_dict["n_embd"])
+        self._update_cfg("num_attention_heads", cfg_dict["n_head"])
+        self._update_cfg("max_seq_length", cfg_dict["n_positions"])
+        self._update_cfg("embedding_dropout_prob", cfg_dict["embd_pdrop"])
+        self._update_cfg("attention_dropout_prob", cfg_dict["attn_pdrop"])
+        self._update_cfg("output_dropout_prob", cfg_dict["resid_pdrop"])
+        self._update_cfg("layernorm_epsilon", cfg_dict["layer_norm_epsilon"])
+        self._update_cfg("vocab_size", cfg_dict["vocab_size"])
+        self._update_cfg("initializer_range", cfg_dict["initializer_range"])
+        self._update_cfg(
+            "ffn_hidden_size",
+            cfg_dict.get("n_inner")
+            if cfg_dict.get("n_inner") is not None
+            else 4 * self.libai_cfg["hidden_size"],
+        )
+
+        # update libai_cfg by kwargs
+        for k, v in self.kwargs.items():
+            self._update_cfg(k, v)
+
+        self._update_cfg_log()
+
+
+class GPT2LoaderLiBai(ModelLoaderLiBai):
+    def __init__(self, model, libai_cfg, pretrained_model_path, **kwargs):
+        super().__init__(model, libai_cfg, pretrained_model_path, **kwargs)
+        self.base_model_prefix_2 = "GPT_model"

libai/models/utils/model_utils/roberta_loader.py

+# coding=utf-8
+# Copyright 2021 The OneFlow Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import oneflow as flow
+
+from .bert_loader import BertLoaderHuggerFace, BertLoaderLiBai
+
+
+class RobertaLoaderHuggerFace(BertLoaderHuggerFace):
+    def __init__(self, model, libai_cfg, pretrained_model_path, **kwargs):
+        super().__init__(model, libai_cfg, pretrained_model_path, **kwargs)
+
+        """NOTE: base_model_prefix_1 is RoBERTa's prefix in Transformers,
+        base_model_prefix_2 is RoBERTa's prefix in LiBai."""
+        self.base_model_prefix_1 = "roberta"
+        self.base_model_prefix_2 = "roberta"
+
+    def _convert_state_dict(self, flow_state_dict, cfg):
+        """Convert state_dict's keys to match model.
+
+        Args:
+            flow_state_dict (OrderedDict): model state dict.
+            cfg (dict): model's default config dict in LiBai.
+
+        Returns:
+            OrderedDict: flow state dict.
+        """
+        # The converted checkpoint.
+        oneflow_state_dict = flow_state_dict.copy()
+
+        # Get configs
+        num_heads = cfg.get("num_attention_heads")
+        hidden_size = cfg.get("hidden_size")
+        layers = cfg.get("hidden_layers")
+        head_size = int(hidden_size / num_heads)
+
+        # prefix
+        has_prefix = any(s.startswith(self.base_model_prefix_1) for s in oneflow_state_dict)
+
+        prefix = "roberta." if has_prefix else ""
+        index_idx = 3 if has_prefix else 2
+        qkv_idx = 6 if has_prefix else 5
+
+        old_keys = oneflow_state_dict.keys()
+
+        for key in list(old_keys):
+
+            # Convert roberta's embedding layers
+            if "embeddings" in key:
+                if "word_embeddings" in key:
+                    new_key = key.replace("word_embeddings", "vocab_embeddings")
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                elif "token_type_embeddings" in key:
+                    new_key = key.replace("token_type_embeddings", "tokentype_embeddings")
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                elif "LayerNorm.weight" in key:
+                    new_key = prefix + "encoders.0.input_layernorm.weight"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                elif "LayerNorm.bias" in key:
+                    new_key = prefix + "encoders.0.input_layernorm.bias"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                else:
+                    oneflow_state_dict[key] = oneflow_state_dict[key]
+
+            # Convert roberta's attention layers
+            elif "attention" in key:
+                if "self" in key:
+                    index = key.split(".")[index_idx]
+                    if (
+                        prefix + "encoders." + index + ".self_attention.query_key_value.weight"
+                        in oneflow_state_dict.keys()
+                    ):
+                        continue
+                    q_w = key.replace(key.split(".")[qkv_idx], "query").replace(
+                        key.split(".")[qkv_idx + 1], "weight"
+                    )
+                    k_w = q_w.replace("query", "key")
+                    v_w = q_w.replace("query", "value")
+                    q_b = q_w.replace("weight", "bias")
+                    k_b = k_w.replace("weight", "bias")
+                    v_b = v_w.replace("weight", "bias")
+
+                    qkv_w = flow.cat(
+                        (
+                            oneflow_state_dict.pop(q_w),
+                            oneflow_state_dict.pop(k_w),
+                            oneflow_state_dict.pop(v_w),
+                        ),
+                        dim=0,
+                    )
+                    qkv_b = flow.cat(
+                        (
+                            oneflow_state_dict.pop(q_b),
+                            oneflow_state_dict.pop(k_b),
+                            oneflow_state_dict.pop(v_b),
+                        ),
+                        dim=-1,
+                    )
+
+                    qkv_w = self._fix_qkv_ordering(qkv_w, head_size, num_heads)
+                    qkv_b = self._fix_qkv_ordering(qkv_b, head_size, num_heads)
+
+                    new_key = (
+                        prefix + "encoders." + index + ".self_attention.query_key_value.weight"
+                    )
+                    oneflow_state_dict[new_key] = qkv_w
+
+                    new_key = prefix + "encoders." + index + ".self_attention.query_key_value.bias"
+                    oneflow_state_dict[new_key] = qkv_b
+                elif "output" in key:
+                    index = key.split(".")[index_idx]
+                    if "dense" in key:
+                        if "weight" in key:
+                            new_key = prefix + "encoders." + index + ".self_attention.dense.weight"
+                            oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                        elif "bias" in key:
+                            new_key = prefix + "encoders." + index + ".self_attention.dense.bias"
+                            oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                    elif "LayerNorm" in key:
+                        if "weight" in key:
+                            new_key = (
+                                prefix + "encoders." + index + ".post_attention_layernorm.weight"
+                            )
+                            oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                        elif "bias" in key:
+                            new_key = (
+                                prefix + "encoders." + index + ".post_attention_layernorm.bias"
+                            )
+                            oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+
+            # Convert roberta's intermediate layers
+            elif "intermediate" in key:
+                index = key.split(".")[index_idx]
+                if (
+                    prefix + "encoders." + index + ".mlp.dense_h_to_4h.weight"
+                    in oneflow_state_dict.keys()
+                ):
+                    continue
+                if "weight" in key:
+                    w = key
+                    b = key.replace("weight", "bias")
+                    new_key = prefix + "encoders." + index + ".mlp.dense_h_to_4h.weight"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(w)
+                    new_key = new_key.replace("weight", "bias")
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(b)
+
+            # Convert roberta's output layers
+            elif "output" in key:
+                index = key.split(".")[index_idx]
+                if "dense.weight" in key:
+                    if (
+                        prefix + "encoders." + index + ".mlp.dense_4h_to_h.weight"
+                        in oneflow_state_dict.keys()
+                    ):
+                        continue
+                    w = key
+                    b = w.replace("weight", "bias")
+                    new_key = prefix + "encoders." + index + ".mlp.dense_4h_to_h.weight"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(w)
+                    new_key = new_key.replace("weight", "bias")
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(b)
+                elif "LayerNorm.weight" in key:
+                    if (
+                        prefix + "encoders." + str(int(index) + 1) + ".input_layernorm.weight"
+                        in oneflow_state_dict.keys()
+                    ):
+                        continue
+                    w = key
+                    b = w.replace("weight", "bias")
+                    if index == str(layers - 1):
+                        new_key = prefix + "final_layernorm.weight"
+                        oneflow_state_dict[new_key] = oneflow_state_dict.pop(w)
+                        new_key = new_key.replace("weight", "bias")
+                        oneflow_state_dict[new_key] = oneflow_state_dict.pop(b)
+                        continue
+                    new_key = prefix + "encoders." + str(int(index) + 1) + ".input_layernorm.weight"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(w)
+                    new_key = new_key.replace("weight", "bias")
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(b)
+
+            # Convert roberta's pooler layers
+            elif "pooler" in key:
+                if "weight" in key:
+                    new_key = prefix + "pooler.dense.weight"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                elif "bias" in key:
+                    new_key = prefix + "pooler.dense.bias"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+
+            # Convert lm_head layers
+            elif "lm_head" in key:
+                if "layer_norm.weight" in key:
+                    new_key = "lm_head.layernorm.weight"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                elif "layer_norm.bias" in key:
+                    new_key = "lm_head.layernorm.bias"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                elif "seq_relationship" in key:
+                    new_key = key.replace("cls", "cls_head")
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                elif "lm_head.bias" in key:
+                    new_key = "lm_head.lm_logits.bias"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                else:
+                    oneflow_state_dict[key] = oneflow_state_dict.pop(key)
+            else:
+                oneflow_state_dict[key] = oneflow_state_dict.pop(key)
+        return oneflow_state_dict
+
+
+class RobertaLoaderLiBai(BertLoaderLiBai):
+    def __init__(self, model, libai_cfg, pretrained_model_path, **kwargs):
+        super().__init__(model, libai_cfg, pretrained_model_path, **kwargs)
+        self.base_model_prefix_2 = "roberta"

libai/models/utils/model_utils/swin_loader.py

+# coding=utf-8
+# Copyright 2021 The OneFlow Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+
+import oneflow as flow
+
+from .base_loader import ModelLoaderHuggerFace, ModelLoaderLiBai
+
+
+class SwinLoaderHuggerFace(ModelLoaderHuggerFace):
+    def __init__(self, model, libai_cfg, pretrained_model_path, **kwargs):
+        super().__init__(model, libai_cfg, pretrained_model_path, **kwargs)
+
+        """NOTE: base_model_prefix_1 is SWIN's prefix in Transformers.
+        base_model_prefix_2 is SWIN's prefix in LiBai."""
+
+        self.base_model_prefix_1 = "swin"
+        self.base_model_prefix_2 = ""
+
+    def _convert_state_dict(self, flow_state_dict, cfg=None):
+        """Convert state_dict's keys to match model.
+
+        Args:
+            flow_state_dict (OrderedDict): model state dict.
+            cfg (dict): model's default config dict.
+
+        Returns:
+            OrderedDict: flow state dict.
+        """
+        # The converted checkpoint.
+        oneflow_state_dict = flow_state_dict.copy()
+
+        # prefix
+        has_prefix = any(s.startswith(self.base_model_prefix_1) for s in oneflow_state_dict)
+
+        index_idx_1 = 3 if has_prefix else 2
+        index_idx_2 = 5 if has_prefix else 4
+
+        old_keys = oneflow_state_dict.keys()
+
+        for key in list(old_keys):
+
+            # Convert swin's embedding layers
+            if "embeddings" in key:
+                if "patch_embeddings.projection" in key:
+                    if "weight" in key:
+                        new_key = "patch_embed.proj.weight"
+                        oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                    elif "bias" in key:
+                        new_key = "patch_embed.proj.bias"
+                        oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                elif "norm" in key:
+                    if "weight" in key:
+                        new_key = "patch_embed.norm.weight"
+                        oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                    elif "bias" in key:
+                        new_key = "patch_embed.norm.bias"
+                        oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+
+            # Convert swin's layernorm layers
+            elif "layernorm_before" in key:
+                index_layer = key.split(".")[index_idx_1]
+                index_block = key.split(".")[index_idx_2]
+                if "weight" in key:
+                    new_key = "layers." + index_layer + ".blocks." + index_block + ".norm1.weight"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                elif "bias" in key:
+                    new_key = "layers." + index_layer + ".blocks." + index_block + ".norm1.bias"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+
+            elif "layernorm_after" in key:
+                index_layer = key.split(".")[index_idx_1]
+                index_block = key.split(".")[index_idx_2]
+                if "weight" in key:
+                    new_key = "layers." + index_layer + ".blocks." + index_block + ".norm2.weight"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                elif "bias" in key:
+                    new_key = "layers." + index_layer + ".blocks." + index_block + ".norm2.bias"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+
+            # Convert swin's attention layers
+            elif "attention" in key:
+                index_layer = key.split(".")[index_idx_1]
+                index_block = key.split(".")[index_idx_2]
+                if "self" in key:
+                    if (
+                        "relative_position_bias_table" in key
+                    ):  # convert relative_position_bias_table but not index
+                        new_key = (
+                            "layers."
+                            + index_layer
+                            + ".blocks."
+                            + index_block
+                            + ".attn.relative_position_bias_table"
+                        )
+                        oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                    elif "relative_position_index" in key:
+                        new_key = (
+                            "layers."
+                            + index_layer
+                            + ".blocks."
+                            + index_block
+                            + ".attn.relative_position_index"
+                        )
+                        oneflow_state_dict.pop(key)
+                    else:
+                        if (
+                            "layers." + index_layer + ".blocks." + index_block + ".attn.qkv.weight"
+                            in oneflow_state_dict.keys()
+                        ):
+                            continue
+                        q_w = key
+                        k_w = q_w.replace("query", "key")
+                        v_w = q_w.replace("query", "value")
+                        q_b = q_w.replace("weight", "bias")
+                        k_b = k_w.replace("weight", "bias")
+                        v_b = v_w.replace("weight", "bias")
+
+                        qkv_w = flow.cat(
+                            (
+                                oneflow_state_dict.pop(q_w),
+                                oneflow_state_dict.pop(k_w),
+                                oneflow_state_dict.pop(v_w),
+                            ),
+                            dim=0,
+                        )
+                        qkv_b = flow.cat(
+                            (
+                                oneflow_state_dict.pop(q_b),
+                                oneflow_state_dict.pop(k_b),
+                                oneflow_state_dict.pop(v_b),
+                            ),
+                            dim=-1,
+                        )
+
+                        new_key = (
+                            "layers." + index_layer + ".blocks." + index_block + ".attn.qkv.weight"
+                        )
+                        oneflow_state_dict[new_key] = qkv_w
+
+                        new_key = new_key.replace("weight", "bias")
+                        oneflow_state_dict[new_key] = qkv_b
+
+                elif "output" in key:
+                    if "dense" in key:
+                        if "weight" in key:
+                            new_key = (
+                                "layers."
+                                + index_layer
+                                + ".blocks."
+                                + index_block
+                                + ".attn.proj.weight"
+                            )
+                            oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                        if "bias" in key:
+                            new_key = (
+                                "layers."
+                                + index_layer
+                                + ".blocks."
+                                + index_block
+                                + ".attn.proj.bias"
+                            )
+                            oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+
+            elif "intermediate" in key:
+                index_layer = key.split(".")[index_idx_1]
+                index_block = key.split(".")[index_idx_2]
+                if "weight" in key:
+                    if (
+                        "layers."
+                        + index_layer
+                        + ".blocks."
+                        + index_block
+                        + ".mlp.dense_h_to_4h.weight"
+                        in oneflow_state_dict.keys()
+                    ):
+                        continue
+                    w = key
+                    b = key.replace("weight", "bias")
+                    new_key = (
+                        "layers."
+                        + index_layer
+                        + ".blocks."
+                        + index_block
+                        + ".mlp.dense_h_to_4h.weight"
+                    )
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(w)
+                    new_key = new_key.replace("weight", "bias")
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(b)
+
+            elif "output" in key:
+                index_layer = key.split(".")[index_idx_1]
+                index_block = key.split(".")[index_idx_2]
+                if "dense.weight" in key:
+                    if (
+                        "layers."
+                        + index_layer
+                        + ".blocks."
+                        + index_block
+                        + ".mlp.dense_4h_to_h.weight"
+                        in oneflow_state_dict.keys()
+                    ):
+                        continue
+                    w = key
+                    b = w.replace("weight", "bias")
+                    new_key = (
+                        "layers."
+                        + index_layer
+                        + ".blocks."
+                        + index_block
+                        + ".mlp.dense_4h_to_h.weight"
+                    )
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(w)
+                    new_key = new_key.replace("weight", "bias")
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(b)
+
+            elif "downsample" in key:
+                index_layer = key.split(".")[index_idx_1]
+                if "reduction.weight" in key:
+                    new_key = "layers." + index_layer + ".downsample.reduction.weight"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                elif "norm" in key:
+                    if (
+                        "layers." + index_layer + ".downsample.norm.weight"
+                        in oneflow_state_dict.keys()
+                    ):
+                        continue
+                    w = key
+                    b = w.replace("weight", "bias")
+                    new_key = "layers." + index_layer + ".downsample.norm.weight"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(w)
+                    new_key = new_key.replace("weight", "bias")
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(b)
+
+            elif "layernorm" in key:
+                if "weight" in key:
+                    new_key = "norm.weight"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                elif "bias" in key:
+                    new_key = "norm.bias"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+
+            elif "classifier" in key:
+                if "weight" in key:
+                    new_key = "head.weight"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                elif "bias" in key:
+                    new_key = "head.bias"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+            else:
+                oneflow_state_dict[key] = oneflow_state_dict.pop(key)
+
+        return oneflow_state_dict
+
+    def _load_config_from_json(self, config_file):
+        """load config from `config.json`, and update default config.
+
+        Args:
+            config_file (str): Path of config file.
+        """
+        with open(config_file, mode="r", encoding="utf-8") as f:
+            cfg_dict = json.load(f)
+
+        # update libai_cfg by config.json
+        self._update_cfg("img_size", cfg_dict["image_size"])
+        self._update_cfg("patch_size", cfg_dict["patch_size"])
+        self._update_cfg("embed_dim", cfg_dict["embed_dim"])
+        self._update_cfg("depths", cfg_dict["depths"])
+        self._update_cfg("num_heads", cfg_dict["num_heads"])
+        self._update_cfg("window_size", cfg_dict["window_size"])
+        self._update_cfg("mlp_ratio", cfg_dict["mlp_ratio"])
+        self._update_cfg("qkv_bias", cfg_dict["qkv_bias"])
+        self._update_cfg("drop_path_rate", cfg_dict["drop_path_rate"])
+
+        # update libai_cfg by kwargs
+        for k, v in self.kwargs.items():
+            self._update_cfg(k, v)
+
+        self._update_cfg_log()
+
+
+class SwinLoaderLiBai(ModelLoaderLiBai):
+    def __init__(self, model, libai_cfg, pretrained_model_path, **kwargs):
+        super().__init__(model, libai_cfg, pretrained_model_path, **kwargs)
+        self.base_model_prefix_2 = ""

libai/models/utils/model_utils/swinv2_loader.py

+# coding=utf-8
+# Copyright 2021 The OneFlow Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+
+import oneflow as flow
+
+from .base_loader import ModelLoaderHuggerFace, ModelLoaderLiBai
+
+
+class SwinV2LoaderHuggerFace(ModelLoaderHuggerFace):
+    def __init__(self, model, libai_cfg, pretrained_model_path, **kwargs):
+        super().__init__(model, libai_cfg, pretrained_model_path, **kwargs)
+
+        """NOTE: base_model_prefix_1 is SWINV2's prefix in Transformers.
+        base_model_prefix_2 is SWINV2's prefix in LiBai."""
+
+        self.base_model_prefix_1 = "swinv2"
+        self.base_model_prefix_2 = ""
+
+    def _convert_state_dict(self, flow_state_dict, cfg=None):
+        """Convert state_dict's keys to match model.
+
+        Args:
+            flow_state_dict (OrderedDict): model state dict.
+            cfg (dict): model's default config dict.
+
+        Returns:
+            OrderedDict: flow state dict.
+        """
+        # The converted checkpoint.
+        oneflow_state_dict = flow_state_dict.copy()
+
+        # prefix
+        has_prefix = any(s.startswith(self.base_model_prefix_1) for s in oneflow_state_dict)
+        index_idx_1 = 3 if has_prefix else 2
+        index_idx_2 = 5 if has_prefix else 4
+
+        old_keys = oneflow_state_dict.keys()
+
+        for key in list(old_keys):
+            # Convert swinv2's embedding layers
+            if "embeddings" in key:
+                if "patch_embeddings.projection" in key:
+                    if "weight" in key:
+                        new_key = "patch_embed.proj.weight"
+                        oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                    if "bias" in key:
+                        new_key = "patch_embed.proj.bias"
+                        oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                elif "norm" in key:
+                    if "weight" in key:
+                        new_key = "patch_embed.norm.weight"
+                        oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                    if "bias" in key:
+                        new_key = "patch_embed.norm.bias"
+                        oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+
+            # Convert swinv2's layernorm layers
+            elif "layernorm_before" in key:
+                index_layer = key.split(".")[index_idx_1]
+                index_block = key.split(".")[index_idx_2]
+                if "weight" in key:
+                    new_key = "layers." + index_layer + ".blocks." + index_block + ".norm1.weight"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                elif "bias" in key:
+                    new_key = "layers." + index_layer + ".blocks." + index_block + ".norm1.bias"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+
+            elif "layernorm_after" in key:
+                index_layer = key.split(".")[index_idx_1]
+                index_block = key.split(".")[index_idx_2]
+                if "weight" in key:
+                    new_key = "layers." + index_layer + ".blocks." + index_block + ".norm2.weight"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                elif "bias" in key:
+                    new_key = "layers." + index_layer + ".blocks." + index_block + ".norm2.bias"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+
+            # Convert swinv2's attention layers
+            elif "attention" in key:
+                index_layer = key.split(".")[index_idx_1]
+                index_block = key.split(".")[index_idx_2]
+                if "self" in key:
+                    if (
+                        "relative_position_bias_table" in key
+                    ):  # convert relative_position_bias_table but not index
+                        new_key = (
+                            "layers."
+                            + index_layer
+                            + ".blocks."
+                            + index_block
+                            + ".attn.relative_position_bias_table"
+                        )
+                        oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                    elif "relative_position_index" in key:
+                        new_key = (
+                            "layers."
+                            + index_layer
+                            + ".blocks."
+                            + index_block
+                            + ".attn.relative_position_index"
+                        )
+                        oneflow_state_dict.pop(key)
+                    elif "continuous_position_bias_mlp" in key:
+                        if (
+                            "layers."
+                            + index_layer
+                            + ".blocks."
+                            + index_block
+                            + ".attn.cpb_mlp"
+                            + ".0.weight"
+                        ) in oneflow_state_dict.keys():
+                            continue
+                        new_key = (
+                            "layers." + index_layer + ".blocks." + index_block + ".attn.cpb_mlp"
+                        )
+                        m_1_w = key
+                        m_1_b = key.replace(".0.weight", ".0.bias")
+                        m_2_w = key.replace(".0.weight", ".2.weight")
+                        oneflow_state_dict[new_key + ".0.weight"] = oneflow_state_dict.pop(m_1_w)
+                        oneflow_state_dict[new_key + ".0.bias"] = oneflow_state_dict.pop(m_1_b)
+                        oneflow_state_dict[new_key + ".2.weight"] = oneflow_state_dict.pop(m_2_w)
+                    elif "logit_scale" in key:
+                        new_key = (
+                            "layers." + index_layer + ".blocks." + index_block + ".attn.logit_scale"
+                        )
+                        oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)[None, ...]
+                    else:
+                        if (
+                            "layers." + index_layer + ".blocks." + index_block + ".attn.qkv.weight"
+                            in oneflow_state_dict.keys()
+                        ):
+                            continue
+                        q_w = key
+                        k_w = q_w.replace("query", "key")
+                        v_w = q_w.replace("query", "value")
+                        q_b = q_w.replace("weight", "bias")
+                        v_b = v_w.replace("weight", "bias")
+
+                        qkv_w = flow.cat(
+                            (
+                                oneflow_state_dict.pop(q_w),
+                                oneflow_state_dict.pop(k_w),
+                                oneflow_state_dict.pop(v_w),
+                            ),
+                            dim=0,
+                        )
+
+                        new_key = (
+                            "layers." + index_layer + ".blocks." + index_block + ".attn.qkv.weight"
+                        )
+                        oneflow_state_dict[new_key] = qkv_w
+
+                        new_key = (
+                            "layers." + index_layer + ".blocks." + index_block + ".attn.q_bias"
+                        )
+                        oneflow_state_dict[new_key] = oneflow_state_dict.pop(q_b)
+                        new_key = new_key.replace("q_bias", "v_bias")
+                        oneflow_state_dict[new_key] = oneflow_state_dict.pop(v_b)
+
+                elif "output" in key:
+                    if "dense" in key:
+                        if "weight" in key:
+                            new_key = (
+                                "layers."
+                                + index_layer
+                                + ".blocks."
+                                + index_block
+                                + ".attn.proj.weight"
+                            )
+                            oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                        if "bias" in key:
+                            new_key = (
+                                "layers."
+                                + index_layer
+                                + ".blocks."
+                                + index_block
+                                + ".attn.proj.bias"
+                            )
+                            oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+
+            elif "intermediate" in key:
+                index_layer = key.split(".")[index_idx_1]
+                index_block = key.split(".")[index_idx_2]
+                if "weight" in key:
+                    if (
+                        "layers."
+                        + index_layer
+                        + ".blocks."
+                        + index_block
+                        + ".mlp.dense_h_to_4h.weight"
+                        in oneflow_state_dict.keys()
+                    ):
+                        continue
+                    w = key
+                    b = key.replace("weight", "bias")
+                    new_key = (
+                        "layers."
+                        + index_layer
+                        + ".blocks."
+                        + index_block
+                        + ".mlp.dense_h_to_4h.weight"
+                    )
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(w)
+                    new_key = new_key.replace("weight", "bias")
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(b)
+
+            elif "output" in key:
+                index_layer = key.split(".")[index_idx_1]
+                index_block = key.split(".")[index_idx_2]
+                if "dense.weight" in key:
+                    if (
+                        "layers."
+                        + index_layer
+                        + ".blocks."
+                        + index_block
+                        + ".mlp.dense_4h_to_h.weight"
+                        in oneflow_state_dict.keys()
+                    ):
+                        continue
+                    w = key
+                    b = w.replace("weight", "bias")
+                    new_key = (
+                        "layers."
+                        + index_layer
+                        + ".blocks."
+                        + index_block
+                        + ".mlp.dense_4h_to_h.weight"
+                    )
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(w)
+                    new_key = new_key.replace("weight", "bias")
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(b)
+
+            elif "downsample" in key:
+                index_layer = key.split(".")[index_idx_1]
+                if "reduction.weight" in key:
+                    new_key = "layers." + index_layer + ".downsample.reduction.weight"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                elif "norm" in key:
+                    if (
+                        "layers." + index_layer + ".downsample.norm.weight"
+                        in oneflow_state_dict.keys()
+                    ):
+                        continue
+                    w = key
+                    b = w.replace("weight", "bias")
+                    new_key = "layers." + index_layer + ".downsample.norm.weight"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(w)
+                    new_key = new_key.replace("weight", "bias")
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(b)
+
+            elif "layernorm" in key:
+                if "weight" in key:
+                    new_key = "norm.weight"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                elif "bias" in key:
+                    new_key = "norm.bias"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+
+            elif "classifier" in key:
+                if "weight" in key:
+                    new_key = "head.weight"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                elif "bias" in key:
+                    new_key = "head.bias"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+            else:
+                oneflow_state_dict[key] = oneflow_state_dict.pop(key)
+
+        return oneflow_state_dict
+
+    def _load_config_from_json(self, config_file):
+        """load config from `config.json`, and update default config.
+
+        Args:
+            config_file (str): Path of config file.
+        """
+        with open(config_file, mode="r", encoding="utf-8") as f:
+            cfg_dict = json.load(f)
+
+        # update libai_cfg by config.json
+        self._update_cfg("img_size", cfg_dict["image_size"])
+        self._update_cfg("patch_size", cfg_dict["patch_size"])
+        self._update_cfg("embed_dim", cfg_dict["embed_dim"])
+        self._update_cfg("depths", cfg_dict["depths"])
+        self._update_cfg("num_heads", cfg_dict["num_heads"])
+        self._update_cfg("window_size", cfg_dict["window_size"])
+        self._update_cfg("mlp_ratio", cfg_dict["mlp_ratio"])
+        self._update_cfg("qkv_bias", cfg_dict["qkv_bias"])
+        self._update_cfg("drop_path_rate", cfg_dict["drop_path_rate"])
+        self._update_cfg("pretrained_window_sizes", cfg_dict["pretrained_window_sizes"])
+
+        # update libai_cfg by kwargs
+        for k, v in self.kwargs.items():
+            self._update_cfg(k, v)
+
+        self._update_cfg_log()
+
+
+class SwinV2LoaderLiBai(ModelLoaderLiBai):
+    def __init__(self, model, libai_cfg, pretrained_model_path, **kwargs):
+        super().__init__(model, libai_cfg, pretrained_model_path, **kwargs)
+        self.base_model_prefix_2 = ""

libai/models/utils/model_utils/vit_loader.py

+# coding=utf-8
+# Copyright 2021 The OneFlow Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+
+import oneflow as flow
+
+from .base_loader import ModelLoaderHuggerFace, ModelLoaderLiBai
+
+
+class ViTLoaderHuggerFace(ModelLoaderHuggerFace):
+    def __init__(self, model, libai_cfg, pretrained_model_path, **kwargs):
+        super().__init__(model, libai_cfg, pretrained_model_path, **kwargs)
+
+        """NOTE: base_model_prefix_1 is ViT's prefix in Transformers.
+        base_model_prefix_2 is ViT's prefix in LiBai."""
+
+        self.base_model_prefix_1 = "vit"
+        self.base_model_prefix_2 = ""
+
+    def _convert_state_dict(self, flow_state_dict, cfg=None):
+        """Convert state_dict's keys to match model.
+
+        Args:
+            flow_state_dict (OrderedDict): model state dict.
+            cfg (dict): model's default config dict.
+
+        Returns:
+            OrderedDict: flow state dict.
+        """
+        # The converted checkpoint.
+        oneflow_state_dict = flow_state_dict.copy()
+
+        # Get configs
+        num_heads = cfg.get("num_heads")
+        hidden_size = cfg.get("embed_dim")
+        head_size = int(hidden_size / num_heads)
+
+        # prefix
+        has_prefix = any(s.startswith(self.base_model_prefix_1) for s in oneflow_state_dict)
+
+        index_idx = 3 if has_prefix else 2
+
+        old_keys = oneflow_state_dict.keys()
+
+        for key in list(old_keys):
+
+            # Convert vit's embedding layers
+            if "embeddings" in key:
+                if "cls_token" in key:
+                    new_key = "cls_token"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                elif "position_embeddings" in key:
+                    new_key = "pos_embed"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                elif "patch_embeddings.projection" in key:
+                    if "weight" in key:
+                        new_key = "patch_embed.proj.weight"
+                        oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                    elif "bias" in key:
+                        new_key = "patch_embed.proj.bias"
+                        oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+
+            # Convert vit's layernorm layers
+            elif "layernorm_before" in key:
+                index_block = key.split(".")[index_idx]
+                if "weight" in key:
+                    new_key = "blocks." + index_block + ".input_layernorm.weight"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                elif "bias" in key:
+                    new_key = "blocks." + index_block + ".input_layernorm.bias"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+
+            elif "layernorm_after" in key:
+                index_block = key.split(".")[index_idx]
+                if "weight" in key:
+                    new_key = "blocks." + index_block + ".post_attention_layernorm.weight"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                elif "bias" in key:
+                    new_key = "blocks." + index_block + ".post_attention_layernorm.bias"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+
+            # Convert vit's attention layers
+            elif "attention" in key:
+                index_block = key.split(".")[index_idx]
+                if "attention.attention" in key:
+                    if (
+                        "blocks." + index_block + ".self_attention.query_key_value.weight"
+                        in oneflow_state_dict.keys()
+                    ):
+                        continue
+                    q_w = key
+                    k_w = q_w.replace("query", "key")
+                    v_w = q_w.replace("query", "value")
+                    q_b = q_w.replace("weight", "bias")
+                    k_b = k_w.replace("weight", "bias")
+                    v_b = v_w.replace("weight", "bias")
+
+                    qkv_w = flow.cat(
+                        (
+                            oneflow_state_dict.pop(q_w),
+                            oneflow_state_dict.pop(k_w),
+                            oneflow_state_dict.pop(v_w),
+                        ),
+                        dim=0,
+                    )
+                    qkv_b = flow.cat(
+                        (
+                            oneflow_state_dict.pop(q_b),
+                            oneflow_state_dict.pop(k_b),
+                            oneflow_state_dict.pop(v_b),
+                        ),
+                        dim=-1,
+                    )
+
+                    qkv_w = self._fix_qkv_ordering(qkv_w, head_size, num_heads)
+                    qkv_b = self._fix_qkv_ordering(qkv_b, head_size, num_heads)
+
+                    new_key = "blocks." + index_block + ".self_attention.query_key_value.weight"
+                    oneflow_state_dict[new_key] = qkv_w
+
+                    new_key = new_key.replace("weight", "bias")
+                    oneflow_state_dict[new_key] = qkv_b
+
+                elif "output" in key:
+                    if "dense" in key:
+                        if "weight" in key:
+                            new_key = "blocks." + index_block + ".self_attention.dense.weight"
+                            oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                        if "bias" in key:
+                            new_key = "blocks." + index_block + ".self_attention.dense.bias"
+                            oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+
+            elif "intermediate" in key:
+                index_block = key.split(".")[index_idx]
+                if "weight" in key:
+                    if (
+                        "blocks." + index_block + ".mlp.dense_h_to_4h.weight"
+                        in oneflow_state_dict.keys()
+                    ):
+                        continue
+                    w = key
+                    b = key.replace("weight", "bias")
+                    new_key = "blocks." + index_block + ".mlp.dense_h_to_4h.weight"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(w)
+                    new_key = new_key.replace("weight", "bias")
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(b)
+
+            elif "output" in key:
+                index_block = key.split(".")[index_idx]
+                if "dense.weight" in key:
+                    if (
+                        "blocks." + index_block + ".mlp.dense_4h_to_h.weight"
+                        in oneflow_state_dict.keys()
+                    ):
+                        continue
+                    w = key
+                    b = w.replace("weight", "bias")
+                    new_key = "blocks." + index_block + ".mlp.dense_4h_to_h.weight"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(w)
+                    new_key = new_key.replace("weight", "bias")
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(b)
+
+            elif "layernorm" in key:
+                if "weight" in key:
+                    new_key = "norm.weight"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                elif "bias" in key:
+                    new_key = "norm.bias"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+
+            elif "classifier" in key:
+                if "weight" in key:
+                    new_key = "head.weight"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                elif "bias" in key:
+                    new_key = "head.bias"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+            else:
+                oneflow_state_dict[key] = oneflow_state_dict.pop(key)
+
+        return oneflow_state_dict
+
+    def _load_config_from_json(self, config_file):
+        """load config from `config.json`, and update default config.
+
+        Args:
+            config_file (str): Path of config file.
+        """
+        with open(config_file, mode="r", encoding="utf-8") as f:
+            cfg_dict = json.load(f)
+
+        # update libai_cfg by config.json
+        self._update_cfg("img_size", cfg_dict["image_size"])
+        self._update_cfg("patch_size", cfg_dict["patch_size"])
+        self._update_cfg("in_chans", cfg_dict["num_channels"])
+        self._update_cfg("embed_dim", cfg_dict["hidden_size"])
+        self._update_cfg("depth", cfg_dict["num_hidden_layers"])
+        self._update_cfg("num_heads", cfg_dict["num_attention_heads"])
+        self._update_cfg("attn_drop_rate", cfg_dict["attention_probs_dropout_prob"])
+        self._update_cfg("drop_rate", cfg_dict["hidden_dropout_prob"])
+
+        # update libai_cfg by kwargs
+        for k, v in self.kwargs.items():
+            self._update_cfg(k, v)
+
+        self._update_cfg_log()
+
+
+class ViTLoaderLiBai(ModelLoaderLiBai):
+    def __init__(self, model, libai_cfg, pretrained_model_path, **kwargs):
+        super().__init__(model, libai_cfg, pretrained_model_path, **kwargs)
+        self.base_model_prefix_2 = ""

libai/models/utils/weight_init.py

+# coding=utf-8
+# Copyright 2021 The OneFlow Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+
+import oneflow.nn as nn
+
+
+def init_method_normal(sigma, mean=0.0):
+    """Init method based on N(0, sigma)."""
+
+    def init_(tensor):
+        return nn.init.normal_(tensor, mean=mean, std=sigma)
+
+    return init_
+
+
+def scaled_init_method_normal(sigma, num_layers, mean=0.0):
+    """Init method based on N(0, sigma/sqrt(2*num_layers)."""
+    std = sigma / math.sqrt(2.0 * num_layers)
+
+    def init_(tensor):
+        return nn.init.normal_(tensor, mean=mean, std=std)
+
+    return init_

libai/models/vision_transformer.py

+# coding=utf-8
+# Copyright 2021 The OneFlow Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import oneflow as flow
+import oneflow.nn as nn
+from flowvision.layers.weight_init import trunc_normal_
+
+import libai.utils.distributed as dist
+from libai.config.config import configurable
+from libai.layers import LayerNorm, Linear, PatchEmbedding, TransformerLayer
+
+
+class VisionTransformer(nn.Module):
+    """Vision Transformer in LiBai.
+
+    LiBai's implementation of:
+    `An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale
+    <https://arxiv.org/abs/2010.11929>`_
+
+    Args:
+        img_size (int, tuple(int)): input image size
+        patch_size (int, tuple(int)): patch size
+        in_chans (int): number of input channels
+        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
+        drop_rate (float): dropout rate
+        attn_drop_rate (float): attention dropout rate
+        drop_path_rate (float): stochastic depth rate
+        num_classes (int): number of classes for classification head
+        loss_func (callable, optional): loss function for computing the total loss
+                                        between logits and labels
+    """
+
+    @configurable
+    def __init__(
+        self,
+        img_size=224,
+        patch_size=16,
+        in_chans=3,
+        embed_dim=192,
+        depth=12,
+        num_heads=3,
+        mlp_ratio=4.0,
+        drop_rate=0.0,
+        attn_drop_rate=0.0,
+        drop_path_rate=0.0,
+        num_classes=1000,
+        loss_func=None,
+    ):
+        super().__init__()
+        self.img_size = img_size
+        self.num_classes = num_classes
+        self.patch_embed = PatchEmbedding(
+            img_size=img_size,
+            patch_size=patch_size,
+            in_chans=in_chans,
+            embed_dim=embed_dim,
+        )
+        ffn_size = int(embed_dim * mlp_ratio)
+        num_patches = self.patch_embed.num_patches
+        self.cls_token = nn.Parameter(
+            flow.zeros(
+                1,
+                1,
+                embed_dim,
+                sbp=dist.get_nd_sbp([flow.sbp.broadcast, flow.sbp.broadcast]),
+                placement=dist.get_layer_placement(0),
+            )
+        )
+        self.pos_embed = nn.Parameter(
+            flow.zeros(
+                1,
+                num_patches + 1,
+                embed_dim,
+                sbp=dist.get_nd_sbp([flow.sbp.broadcast, flow.sbp.broadcast]),
+                placement=dist.get_layer_placement(0),
+            )
+        )
+
+        self.pos_drop = nn.Dropout(p=drop_rate)
+        dpr = [
+            x.item() for x in flow.linspace(0, drop_path_rate, depth)
+        ]  # stochastic depth decay rule
+        self.blocks = nn.Sequential(
+            *[
+                TransformerLayer(
+                    hidden_size=embed_dim,
+                    ffn_hidden_size=ffn_size,
+                    num_attention_heads=num_heads,
+                    attention_dropout_prob=attn_drop_rate,
+                    output_dropout_prob=drop_rate,
+                    drop_path_prob=dpr[i],
+                    layer_idx=i,
+                )
+                for i in range(depth)
+            ]
+        )
+        self.norm = LayerNorm(embed_dim, layer_idx=-1)
+        self.head = Linear(embed_dim, num_classes, layer_idx=-1)
+
+        # loss func
+        self.loss_func = nn.CrossEntropyLoss() if loss_func is None else loss_func
+
+        # weight init
+        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, Linear):
+            trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def no_weight_decay(self):
+        return {"pos_embed", "cls_token"}
+
+    @classmethod
+    def from_config(cls, cfg):
+        return {
+            "img_size": cfg.img_size,
+            "patch_size": cfg.patch_size,
+            "in_chans": cfg.in_chans,
+            "embed_dim": cfg.embed_dim,
+            "depth": cfg.depth,
+            "num_heads": cfg.num_heads,
+            "mlp_ratio": cfg.mlp_ratio,
+            "drop_rate": cfg.drop_rate,
+            "attn_drop_rate": cfg.attn_drop_rate,
+            "drop_path_rate": cfg.drop_path_rate,
+            "num_classes": cfg.num_classes,
+            "loss_func": cfg.loss_func,
+        }
+
+    def forward_features(self, x):
+        # patch embedding
+        x = self.patch_embed(x)
+
+        cls_token = self.cls_token.expand(
+            x.shape[0], -1, -1
+        )  # stole cls_tokens impl from Phil Wang, thanks
+        cls_token = cls_token.to_global(sbp=x.sbp, placement=cls_token.placement)
+        x = flow.cat((cls_token, x), dim=1)
+
+        # position embedding
+        pos_embed = self.pos_embed.expand(x.shape[0], -1, -1)
+        pos_embed = pos_embed.to_global(sbp=x.sbp, placement=pos_embed.placement)
+        x = self.pos_drop(x + pos_embed)
+
+        # transformer block
+        x = self.blocks(x)
+        return x
+
+    def forward_head(self, x):
+        x = self.norm(x)
+        outcome = x[:, 0]
+        outcome = self.head(outcome)
+        return outcome
+
+    def forward(self, images, labels=None):
+        """
+
+        Args:
+            images (flow.Tensor): training samples.
+            labels (flow.LongTensor, optional): training targets
+
+        Returns:
+            dict:
+                A dict containing :code:`loss_value` or :code:`logits`
+                depending on training or evaluation mode.
+                :code:`{"losses": loss_value}` when training,
+                :code:`{"prediction_scores": logits}` when evaluating.
+        """
+        x = self.forward_features(images)
+        x = self.forward_head(x)
+
+        if labels is not None and self.training:
+            losses = self.loss_func(x, labels)
+            return {"losses": losses}
+        else:
+            return {"prediction_scores": x}
+
+    @staticmethod
+    def set_pipeline_stage_id(model):
+        dist_utils = dist.get_dist_util()
+
+        # Set pipeline parallelism stage_id
+        if hasattr(model.pos_embed, "config"):
+            # Old API in OneFlow 0.8
+            for module_block in model.modules():
+                if isinstance(module_block.origin, PatchEmbedding):
+                    module_block.config.set_stage(
+                        dist_utils.get_layer_stage_id(0), dist.get_layer_placement(0)
+                    )
+                elif isinstance(module_block.origin, TransformerLayer):
+                    module_block.config.set_stage(
+                        dist_utils.get_layer_stage_id(module_block.layer_idx),
+                        dist.get_layer_placement(module_block.layer_idx),
+                    )
+
+            # Set pos_embed and cls_token stage id
+            model.pos_embed.config.set_stage(
+                dist_utils.get_layer_stage_id(0), dist.get_layer_placement(0)
+            )
+            model.cls_token.config.set_stage(
+                dist_utils.get_layer_stage_id(0), dist.get_layer_placement(0)
+            )
+            model.pos_drop.config.set_stage(
+                dist_utils.get_layer_stage_id(0), dist.get_layer_placement(0)
+            )
+            model.norm.config.set_stage(
+                dist_utils.get_layer_stage_id(-1), dist.get_layer_placement(-1)
+            )
+            model.head.config.set_stage(
+                dist_utils.get_layer_stage_id(-1), dist.get_layer_placement(-1)
+            )
+            model.loss_func.config.set_stage(
+                dist_utils.get_layer_stage_id(-1), dist.get_layer_placement(-1)
+            )
+        else:
+            for module_block in model.modules():
+                if isinstance(module_block.to(nn.Module), PatchEmbedding):
+                    module_block.to(flow.nn.graph.GraphModule).set_stage(
+                        dist_utils.get_layer_stage_id(0), dist.get_layer_placement(0)
+                    )
+                elif isinstance(module_block.to(nn.Module), TransformerLayer):
+                    module_block.to(flow.nn.graph.GraphModule).set_stage(
+                        dist_utils.get_layer_stage_id(module_block.layer_idx),
+                        dist.get_layer_placement(module_block.layer_idx),
+                    )
+
+            # Set pos_embed and cls_token stage id
+            model.pos_embed.to(flow.nn.graph.GraphTensor).set_stage(
+                dist_utils.get_layer_stage_id(0), dist.get_layer_placement(0)
+            )
+            model.cls_token.to(flow.nn.graph.GraphTensor).set_stage(
+                dist_utils.get_layer_stage_id(0), dist.get_layer_placement(0)
+            )
+            model.pos_drop.to(flow.nn.graph.GraphModule).set_stage(
+                dist_utils.get_layer_stage_id(0), dist.get_layer_placement(0)
+            )
+            model.norm.to(flow.nn.graph.GraphModule).set_stage(
+                dist_utils.get_layer_stage_id(-1), dist.get_layer_placement(-1)
+            )
+            model.head.to(flow.nn.graph.GraphModule).set_stage(
+                dist_utils.get_layer_stage_id(-1), dist.get_layer_placement(-1)
+            )
+            model.loss_func.to(flow.nn.graph.GraphModule).set_stage(
+                dist_utils.get_layer_stage_id(-1), dist.get_layer_placement(-1)
+            )

libai/onnx_export/gpt2_to_onnx.py

+# coding=utf-8
+# Copyright 2021 The OneFlow Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import oneflow as flow
+from oneflow import nn
+from oneflow_onnx.oneflow2onnx.util import convert_to_onnx_and_check
+
+from libai.config import LazyConfig
+from libai.models.utils import GPT2LoaderLiBai
+from projects.MagicPrompt.gpt2 import GPTModel
+
+
+def get_model(config_file):
+    cfg = LazyConfig.load(config_file)
+
+    cfg.model.cfg.pretrained_model_path = None
+    cfg.dataloader = None
+    cfg.tokenization = None
+
+    print("Building model....")
+    loader = GPT2LoaderLiBai(GPTModel, cfg.cfg, "/path/to/model")
+    model = loader.load()
+    print("Build model finished.")
+
+    return model
+
+
+class gpt2Graph(nn.Graph):
+    def __init__(self, eager_model):
+        super().__init__()
+        self.model = eager_model
+
+    def build(
+        self,
+        input_ids,
+    ):
+        out = self.model(
+            input_ids,
+        )
+        return out
+
+
+if __name__ == "__main__":
+    model = get_model("projects/MagicPrompt/configs/gpt2_inference.py")
+    model.eval()
+
+    gpt2_graph = gpt2Graph(model)
+    # Build the static graph model
+    input_ids = flow.ones(
+        1, 5, dtype=flow.int64, sbp=flow.sbp.broadcast, placement=flow.placement("cuda", ranks=[0])
+    )
+
+    # check your model.forward is valid
+    # output = gpt2_graph(
+    #     input_ids
+    # )
+
+    print("Compiling the graph which may make some time, please wait for a moment....")
+
+    gpt2_graph._compile(
+        input_ids,
+    )
+
+    convert_to_onnx_and_check(
+        gpt2_graph,
+        external_data=False,
+        opset=11,
+        flow_weight_dir=None,
+        onnx_model_path="./",
+        dynamic_batch_size=False,
+        device="gpu_global",
+        input_tensor_range=[0, 10],
+    )

libai/onnx_export/onnx_inference/gpt2_onnx_infer.py

+# coding=utf-8
+# Copyright 2021 The OneFlow Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from collections import OrderedDict
+from typing import List
+
+import numpy as np
+import onnxruntime as ort
+
+
+class OnnxModel:
+    def __init__(
+        self,
+        onnx_filename,
+        providers: List[str] = None,
+        ort_optimize: bool = True,
+    ):
+        ort_sess_opt = ort.SessionOptions()
+        ort_sess_opt.graph_optimization_level = (
+            ort.GraphOptimizationLevel.ORT_ENABLE_EXTENDED
+            if ort_optimize
+            else ort.GraphOptimizationLevel.ORT_DISABLE_ALL
+        )
+        if providers is None:
+            if ort.__version__ > "1.9.0":
+                providers = [
+                    "TensorrtExecutionProvider",
+                    "CUDAExecutionProvider",
+                    "CPUExecutionProvider",
+                ]
+            else:
+                providers = ["CPUExecutionProvider"]
+        self.sess = ort.InferenceSession(
+            onnx_filename, sess_options=ort_sess_opt, providers=providers
+        )
+
+    def forward(self, input_list):
+        ipt_dict = OrderedDict()
+        for idx, ipt in enumerate(self.sess.get_inputs()):
+            ipt_dict[ipt.name] = input_list[idx]
+        onnx_res = self.sess.run([], ipt_dict)
+        return onnx_res
+
+
+if __name__ == "__main__":
+    onnx_model = OnnxModel("model.onnx")
+    input_list = [
+        np.ones((1, 5)).astype(np.int64).astype(np.int64),
+    ]
+
+    print(onnx_model.forward(input_list))

libai/onnx_export/onnx_inference/t5_onnx_infer.py

+# coding=utf-8
+# Copyright 2021 The OneFlow Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from collections import OrderedDict
+from typing import List
+
+import numpy as np
+import onnxruntime as ort
+
+
+class OnnxModel:
+    def __init__(
+        self,
+        onnx_filename,
+        providers: List[str] = None,
+        ort_optimize: bool = True,
+    ):
+        ort_sess_opt = ort.SessionOptions()
+        ort_sess_opt.graph_optimization_level = (
+            ort.GraphOptimizationLevel.ORT_ENABLE_EXTENDED
+            if ort_optimize
+            else ort.GraphOptimizationLevel.ORT_DISABLE_ALL
+        )
+        if providers is None:
+            if ort.__version__ > "1.9.0":
+                providers = [
+                    "TensorrtExecutionProvider",
+                    "CUDAExecutionProvider",
+                    "CPUExecutionProvider",
+                ]
+            else:
+                providers = ["CPUExecutionProvider"]
+        self.sess = ort.InferenceSession(
+            onnx_filename, sess_options=ort_sess_opt, providers=providers
+        )
+
+    def forward(self, input_list):
+        ipt_dict = OrderedDict()
+        for idx, ipt in enumerate(self.sess.get_inputs()):
+            ipt_dict[ipt.name] = input_list[idx]
+        onnx_res = self.sess.run([], ipt_dict)
+        return onnx_res
+
+
+if __name__ == "__main__":
+    onnx_model = OnnxModel("model.onnx")
+    input_list = [
+        np.ones((1, 5)).astype(np.int64),
+        np.ones((1, 3)).astype(np.int64),
+        np.ones((1, 5, 5)).astype(bool),
+        np.ones((1, 3, 3)).astype(bool),
+        np.ones((1, 3, 5)).astype(bool),
+    ]
+
+    print(onnx_model.forward(input_list))

libai/onnx_export/t5_to_onnx.py

+# coding=utf-8
+# Copyright 2021 The OneFlow Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import oneflow as flow
+from oneflow import nn
+from oneflow_onnx.oneflow2onnx.util import convert_to_onnx_and_check
+
+from libai.config import LazyConfig
+from projects.MT5.mt5_model import MT5Model
+from projects.MT5.utils.mt5_loader import T5LoaderHuggerFace
+
+
+def get_model(config_file):
+    cfg = LazyConfig.load(config_file)
+
+    cfg.model.cfg.model_type = "mt5"
+    cfg.model.cfg.pretrained_model_path = None
+    cfg.dataloader = None
+    cfg.tokenization = None
+
+    print("Building model....")
+    loader = T5LoaderHuggerFace(MT5Model, cfg.model.cfg, "/path/to/model")
+    model = loader.load()
+    print("Build model finished.")
+
+    return model
+
+
+class t5Graph(nn.Graph):
+    def __init__(self, eager_model):
+        super().__init__()
+        self.model = eager_model
+
+    def build(
+        self,
+        encoder_input_ids,
+        encoder_attn_mask,
+        decoder_input_ids,
+        decoder_attn_mask,
+        encoder_decoder_attn_mask,
+    ):
+        out = self.model(
+            encoder_input_ids,
+            encoder_attn_mask,
+            decoder_input_ids,
+            decoder_attn_mask,
+            encoder_decoder_attn_mask,
+        )
+        return out
+
+
+if __name__ == "__main__":
+    model = get_model("projects/MT5/configs/mt5_pretrain.py")
+    model.eval()
+
+    t5_graph = t5Graph(model)
+    # Build the static graph model
+    encoder_input_ids = flow.ones(
+        1, 5, dtype=flow.int64, sbp=flow.sbp.broadcast, placement=flow.placement("cuda", ranks=[0])
+    )
+    encoder_attn_mask = flow.ones(
+        1, 3, dtype=flow.int64, sbp=flow.sbp.broadcast, placement=flow.placement("cuda", ranks=[0])
+    )
+    decoder_input_ids = flow.ones(
+        1,
+        5,
+        5,
+        dtype=flow.bool,
+        sbp=flow.sbp.broadcast,
+        placement=flow.placement("cuda", ranks=[0]),
+    )
+    decoder_attn_mask = flow.ones(
+        1,
+        3,
+        3,
+        dtype=flow.bool,
+        sbp=flow.sbp.broadcast,
+        placement=flow.placement("cuda", ranks=[0]),
+    )
+    encoder_decoder_attn_mask = flow.ones(
+        1,
+        3,
+        5,
+        dtype=flow.bool,
+        sbp=flow.sbp.broadcast,
+        placement=flow.placement("cuda", ranks=[0]),
+    )
+
+    # check your model.forward is valid
+    # output = t5_graph(
+    #     encoder_input_ids,
+    #     encoder_attn_mask,
+    #     decoder_input_ids,
+    #     decoder_attn_mask,
+    #     encoder_decoder_attn_mask
+    # )
+    # print(output)
+
+    print("Compiling the graph which may make some time, please wait for a moment....")
+    t5_graph._compile(
+        encoder_input_ids,
+        encoder_attn_mask,
+        decoder_input_ids,
+        decoder_attn_mask,
+        encoder_decoder_attn_mask,
+    )
+
+    convert_to_onnx_and_check(
+        t5_graph,
+        external_data=False,
+        opset=11,
+        flow_weight_dir=None,
+        onnx_model_path="./",
+        dynamic_batch_size=False,
+        device="gpu_global",
+        input_tensor_range=[0, 10],
+    )

libai/optim/__init__.py

+# coding=utf-8
+# Copyright 2021 The OneFlow Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .build import build_optimizer, get_default_optimizer_params

libai/optim/build.py

+# coding=utf-8
+# Copyright 2021 The OneFlow Authors. All rights reserved.
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+from collections import defaultdict
+from typing import Any, Dict, List
+
+import oneflow as flow
+
+from libai.config import instantiate
+from libai.layers import LayerNorm
+
+# --------------------------------------------------------
+# References:
+# https://github.com/facebookresearch/detectron2/blob/main/detectron2/solver/build.py
+# --------------------------------------------------------
+
+
+def build_optimizer(cfg, model):
+    """
+    Build an optimizer from config.
+    """
+
+    cfg.params.model = model
+    optim = instantiate(cfg)
+    return optim
+
+
+def get_default_optimizer_params(
+    model,
+    base_lr=None,
+    weight_decay=None,
+    weight_decay_norm=None,
+    weight_decay_bias=None,
+    clip_grad_max_norm=None,
+    clip_grad_norm_type=None,
+    overrides=None,
+):
+    """
+    Get default param list for optimizer, with suport for a few types of overrides.
+    If no overrides are needed, it is equivalent to `model.parameters()`.
+
+    Arguments:
+        base_lr: lr for every group by default. Can be omitted to use the one in optimizer.
+        weight_decay: weight decay for every group by default. Can be omitted to use the one
+            in optimizer.
+        weight_decay_norm: override weight decay for params in normalization layers
+        weight_decay_bias: override weight decay for bias parameters
+        overrides: if not `None`, provides values for optimizer hyperparameters
+            (LR, weight decay) for module parameters with a given name; e.g.
+            ``{"embedding": {"lr": 0.01, "weight_decay": 0.1}}`` will set the LR and
+            weight decay values for all module parameters named `embedding`.
+
+    For common transformer models, ``weight_decay_norm`` and ``weight_decay_bias``
+    are usually set to 0.
+
+    Example:
+    ::
+
+        flow.optim.AdamW(
+            get_default_optimizer_params(model, weight_decay_norm=0, weight_decay_bias=0),
+            lr=0.01,
+            weight_decay=1e-4
+        )
+    """
+    if overrides is None:
+        overrides = {}
+    defaults = {}
+    if base_lr is not None:
+        defaults["lr"] = base_lr
+    if weight_decay is not None:
+        defaults["weight_decay"] = weight_decay
+    if clip_grad_max_norm is not None and clip_grad_norm_type is not None:
+        defaults["clip_grad_max_norm"] = clip_grad_max_norm
+        defaults["clip_grad_norm_type"] = clip_grad_norm_type
+    bias_overrides = {}
+    if weight_decay_bias is not None:
+        bias_overrides["weight_decay"] = weight_decay_bias
+    if len(bias_overrides):
+        if "bias" in overrides:
+            raise ValueError("Conflicting overrides for 'bias'")
+        overrides["bias"] = bias_overrides
+
+    norm_module_types = (
+        LayerNorm,
+        flow.nn.BatchNorm1d,
+        flow.nn.BatchNorm2d,
+        flow.nn.BatchNorm3d,
+        flow.nn.GroupNorm,
+        flow.nn.InstanceNorm1d,
+        flow.nn.InstanceNorm2d,
+        flow.nn.InstanceNorm3d,
+        flow.nn.FusedBatchNorm1d,
+        flow.nn.FusedBatchNorm2d,
+        flow.nn.FusedBatchNorm3d,
+    )
+    params = []
+    memo = set()
+    for module in model.modules():
+        for model_param_name, value in module.named_parameters(recurse=False):
+            if not value.requires_grad:
+                continue
+            # Avoid duplicating parameters
+            if value in memo:
+                continue
+            memo.add(value)
+
+            hyperparams = copy.copy(defaults)
+            if isinstance(module, norm_module_types) and weight_decay_norm is not None:
+                hyperparams["weight_decay"] = weight_decay_norm
+            hyperparams.update(overrides.get(model_param_name, {}))
+            params.append({"params": [value], **hyperparams})
+    return reduce_param_groups(params)
+
+
+def _expand_param_groups(params: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """
+    Transform parameter groups into per-parameter structure.
+    Later items in `params` can overwrite parameters set in previous items.
+    """
+    ret = defaultdict(dict)
+    for item in params:
+        assert "params" in item
+        cur_params = {x: y for x, y in item.items() if x != "params"}
+        for param in item["params"]:
+            ret[param].update({"params": [param], **cur_params})
+    return list(ret.values())
+
+
+def reduce_param_groups(params: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """
+    Reorganize the parameter groups and merge duplicated groups.
+    The number of parameter groups needs to be as small as possible in order
+    to efficiently use the OneFlow multi-tensor optimizer. Therefore instead
+    of using a parameter_group per single parameter, we reorganize the
+    parameter groups and merge duplicated groups. This approach speeds
+    up multi-tensor optimizer significantly.
+    """
+    params = _expand_param_groups(params)
+    groups = defaultdict(list)  # re-group all parameter groups by their hyperparams
+    for item in params:
+        cur_params = tuple((x, y) for x, y in item.items() if x != "params")
+        groups[cur_params].extend(item["params"])
+    ret = []
+    for param_keys, param_values in groups.items():
+        cur = {kv[0]: kv[1] for kv in param_keys}
+        cur["params"] = param_values
+        ret.append(cur)
+    return ret

libai/scheduler/__init__.py

+# coding=utf-8
+# Copyright 2021 The OneFlow Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .build import build_lr_scheduler
+from .lr_scheduler import (
+    WarmupCosineAnnealingLR,
+    WarmupCosineLR,
+    WarmupExponentialLR,
+    WarmupMultiStepLR,
+    WarmupPolynomialLR,
+    WarmupStepLR,
+)