| | |
| | |
| | |
| | |
| | |
| |
|
| | from .intern_vit_6b import InternViT6B |
| |
|
| |
|
| | def build_model(config): |
| | model_type = config.MODEL.TYPE |
| | if model_type == 'intern_vit_6b': |
| | model = InternViT6B( |
| | num_classes=config.MODEL.NUM_CLASSES, |
| | patch_size=config.MODEL.INTERN_VIT_6B.PATCH_SIZE, |
| | img_size=config.DATA.IMG_SIZE, |
| | pretrain_size=config.MODEL.INTERN_VIT_6B.PRETRAIN_SIZE, |
| | qkv_bias=config.MODEL.INTERN_VIT_6B.QKV_BIAS, |
| | drop_path_rate=config.MODEL.DROP_PATH_RATE, |
| | embed_dim=config.MODEL.INTERN_VIT_6B.EMBED_DIM, |
| | num_heads=config.MODEL.INTERN_VIT_6B.NUM_HEADS, |
| | mlp_ratio=config.MODEL.INTERN_VIT_6B.MLP_RATIO, |
| | init_values=config.MODEL.INTERN_VIT_6B.INIT_VALUES, |
| | qk_normalization=config.MODEL.INTERN_VIT_6B.QK_NORMALIZATION, |
| | depth=config.MODEL.INTERN_VIT_6B.DEPTH, |
| | use_flash_attn=config.MODEL.INTERN_VIT_6B.USE_FLASH_ATTN, |
| | with_cp=config.TRAIN.USE_CHECKPOINT, |
| | freeze_vit=config.MODEL.INTERN_VIT_6B.FREEZE_VIT, |
| | pretrained=config.MODEL.INTERN_VIT_6B.PRETRAINED, |
| | cls_target=config.MODEL.INTERN_VIT_6B.CLS_TARGET, |
| | head_norm_type=config.MODEL.INTERN_VIT_6B.HEAD_NORM_TYPE, |
| | ) |
| | else: |
| | raise NotImplementedError(f'Unkown model: {model_type}') |
| |
|
| | return model |
| |
|
