Spaces:
Sleeping
Sleeping
| #!/usr/bin/env python3 | |
| # Copyright (c) Facebook, Inc. and its affiliates. | |
| # | |
| # This source code is licensed under the MIT license found in the | |
| # LICENSE file in the root directory of this source tree. | |
| import argparse | |
| import torch | |
| from omegaconf import OmegaConf | |
| from fairseq.criterions.model_criterion import ModelCriterionConfig | |
| from fairseq.dataclass.configs import FairseqConfig | |
| from tasks import ImageClassificationConfig, ImagePretrainingConfig | |
| from models.data2vec_image_classification import ( | |
| Data2VecImageClassificationConfig, | |
| Data2VecImageClassificationModel, | |
| ) | |
| from models.data2vec_vision import Data2VecVisionConfig, Data2VecVisionModel | |
| def get_parser(): | |
| parser = argparse.ArgumentParser( | |
| description="convert beit checkpoint into data2vec - vision checkpoint" | |
| ) | |
| # fmt: off | |
| parser.add_argument('checkpoint', help='checkpoint to convert') | |
| parser.add_argument('--output', required=True, metavar='PATH', help='where to output converted checkpoint') | |
| parser.add_argument('--type', type=str, choices=['vision', 'image_classification'], default='image_classification', help='type of model to upgrade') | |
| parser.add_argument('--inception_norms', action='store_true', default=False) | |
| # fmt: on | |
| return parser | |
| def update_checkpoint(model_dict, prefix, is_nested): | |
| replace_paths = { | |
| "cls_token": "model.cls_emb" if is_nested else "cls_emb", | |
| "patch_embed": "model.patch_embed" if is_nested else "patch_embed", | |
| "mask_token": "mask_emb", | |
| } | |
| starts_with = { | |
| "patch_embed.proj": "model.patch_embed.conv" | |
| if is_nested | |
| else "patch_embed.conv", | |
| "lm_head": "final_proj", | |
| "fc_norm": "fc_norm", | |
| "head": "head", | |
| } | |
| partial = { | |
| "mlp.fc1": "mlp.0", | |
| "mlp.fc2": "mlp.2", | |
| } | |
| for k in list(model_dict.keys()): | |
| for sw, r in starts_with.items(): | |
| if k.startswith(sw): | |
| replace_paths[k] = k.replace(sw, r) | |
| for p, r in partial.items(): | |
| if p in k: | |
| replace_paths[k] = prefix + k.replace(p, r) | |
| if prefix != "": | |
| for k in list(model_dict.keys()): | |
| if k not in replace_paths: | |
| replace_paths[k] = prefix + k | |
| for k in list(model_dict.keys()): | |
| if k in replace_paths: | |
| model_dict[replace_paths[k]] = model_dict[k] | |
| if k != replace_paths[k]: | |
| del model_dict[k] | |
| return model_dict | |
| def main(): | |
| parser = get_parser() | |
| args = parser.parse_args() | |
| cp = torch.load(args.checkpoint, map_location="cpu") | |
| cfg = FairseqConfig( | |
| criterion=ModelCriterionConfig(_name="model", log_keys=["correct"]), | |
| ) | |
| if args.type == "image_classification": | |
| cfg.task = ImageClassificationConfig( | |
| _name="image_classification", | |
| data=".", | |
| ) | |
| if args.inception_norms: | |
| cfg.task.normalization_mean = [0.5, 0.5, 0.5] | |
| cfg.task.normalization_std = [0.5, 0.5, 0.5] | |
| cfg.model = Data2VecImageClassificationConfig( | |
| _name="data2vec_image_classification", | |
| ) | |
| cfg.model.pretrained_model_args = FairseqConfig( | |
| model=Data2VecVisionConfig( | |
| _name="data2vec_vision", shared_rel_pos_bias=False | |
| ), | |
| task=ImagePretrainingConfig( | |
| _name="image_pretraining", | |
| ), | |
| ) | |
| cfg = OmegaConf.create(cfg) | |
| state = { | |
| "cfg": OmegaConf.to_container(cfg, resolve=True, enum_to_str=True), | |
| "model": cp["module"], | |
| "best_loss": None, | |
| "optimizer": None, | |
| "extra_state": {}, | |
| } | |
| model = Data2VecImageClassificationModel(cfg.model) | |
| model.load_state_dict( | |
| update_checkpoint(state["model"], prefix="model.encoder.", is_nested=True), | |
| strict=True, | |
| ) | |
| elif args.type == "vision": | |
| cfg.task = ImagePretrainingConfig( | |
| _name="image_pretraining", | |
| data=".", | |
| ) | |
| if args.inception_norms: | |
| cfg.task.normalization_mean = [0.5, 0.5, 0.5] | |
| cfg.task.normalization_std = [0.5, 0.5, 0.5] | |
| cfg.model = Data2VecVisionConfig( | |
| _name="data2vec_vision", | |
| ) | |
| cfg = OmegaConf.create(cfg) | |
| state = { | |
| "cfg": OmegaConf.to_container(cfg, resolve=True, enum_to_str=True), | |
| "model": cp["model"], | |
| "best_loss": None, | |
| "optimizer": None, | |
| "extra_state": {}, | |
| } | |
| model = Data2VecVisionModel(cfg.model) | |
| model.load_state_dict( | |
| update_checkpoint(state["model"], prefix="encoder.", is_nested=False), | |
| strict=True, | |
| ) | |
| else: | |
| raise Exception("unsupported type " + args.type) | |
| print(state["cfg"], state.keys()) | |
| torch.save(state, args.output) | |
| if __name__ == "__main__": | |
| main() | |