| 8 | |
| 9 | |
| 10 | def build_model(config): |
| 11 | model_type = config.MODEL.TYPE |
| 12 | if model_type == 'intern_vit_6b': |
| 13 | model = InternViT6B( |
| 14 | num_classes=config.MODEL.NUM_CLASSES, |
| 15 | patch_size=config.MODEL.INTERN_VIT_6B.PATCH_SIZE, |
| 16 | img_size=config.DATA.IMG_SIZE, |
| 17 | pretrain_size=config.MODEL.INTERN_VIT_6B.PRETRAIN_SIZE, |
| 18 | qkv_bias=config.MODEL.INTERN_VIT_6B.QKV_BIAS, |
| 19 | drop_path_rate=config.MODEL.DROP_PATH_RATE, |
| 20 | embed_dim=config.MODEL.INTERN_VIT_6B.EMBED_DIM, |
| 21 | num_heads=config.MODEL.INTERN_VIT_6B.NUM_HEADS, |
| 22 | mlp_ratio=config.MODEL.INTERN_VIT_6B.MLP_RATIO, |
| 23 | init_values=config.MODEL.INTERN_VIT_6B.INIT_VALUES, |
| 24 | qk_normalization=config.MODEL.INTERN_VIT_6B.QK_NORMALIZATION, |
| 25 | depth=config.MODEL.INTERN_VIT_6B.DEPTH, |
| 26 | use_flash_attn=config.MODEL.INTERN_VIT_6B.USE_FLASH_ATTN, |
| 27 | with_cp=config.TRAIN.USE_CHECKPOINT, |
| 28 | freeze_vit=config.MODEL.INTERN_VIT_6B.FREEZE_VIT, |
| 29 | pretrained=config.MODEL.INTERN_VIT_6B.PRETRAINED, |
| 30 | cls_target=config.MODEL.INTERN_VIT_6B.CLS_TARGET, |
| 31 | head_norm_type=config.MODEL.INTERN_VIT_6B.HEAD_NORM_TYPE, |
| 32 | ) |
| 33 | else: |
| 34 | raise NotImplementedError(f'Unkown model: {model_type}') |
| 35 | |
| 36 | return model |