| 18 | |
| 19 | |
| 20 | def build_model(config): |
| 21 | model_type = config.MODEL.TYPE |
| 22 | encoder_type = config.MODEL.MOBY.ENCODER |
| 23 | |
| 24 | if encoder_type == 'swin': |
| 25 | enc = partial( |
| 26 | SwinTransformer, |
| 27 | img_size=config.DATA.IMG_SIZE, |
| 28 | patch_size=config.MODEL.SWIN.PATCH_SIZE, |
| 29 | in_chans=config.MODEL.SWIN.IN_CHANS, |
| 30 | embed_dim=config.MODEL.SWIN.EMBED_DIM, |
| 31 | depths=config.MODEL.SWIN.DEPTHS, |
| 32 | num_heads=config.MODEL.SWIN.NUM_HEADS, |
| 33 | window_size=config.MODEL.SWIN.WINDOW_SIZE, |
| 34 | mlp_ratio=config.MODEL.SWIN.MLP_RATIO, |
| 35 | qkv_bias=config.MODEL.SWIN.QKV_BIAS, |
| 36 | qk_scale=config.MODEL.SWIN.QK_SCALE, |
| 37 | drop_rate=config.MODEL.DROP_RATE, |
| 38 | ape=config.MODEL.SWIN.APE, |
| 39 | patch_norm=config.MODEL.SWIN.PATCH_NORM, |
| 40 | use_checkpoint=config.TRAIN.USE_CHECKPOINT, |
| 41 | norm_befor_mlp=config.MODEL.SWIN.NORM_BEFORE_MLP, |
| 42 | ) |
| 43 | elif encoder_type.startswith('vit') or encoder_type.startswith('deit'): |
| 44 | enc = vit_models[encoder_type] |
| 45 | else: |
| 46 | raise NotImplementedError(f'--> Unknown encoder_type: {encoder_type}') |
| 47 | |
| 48 | if model_type == 'moby': |
| 49 | encoder = enc( |
| 50 | num_classes=0, |
| 51 | drop_path_rate=config.MODEL.MOBY.ONLINE_DROP_PATH_RATE, |
| 52 | ) |
| 53 | encoder_k = enc( |
| 54 | num_classes=0, |
| 55 | drop_path_rate=config.MODEL.MOBY.TARGET_DROP_PATH_RATE, |
| 56 | ) |
| 57 | model = MoBY( |
| 58 | cfg=config, |
| 59 | encoder=encoder, |
| 60 | encoder_k=encoder_k, |
| 61 | contrast_momentum=config.MODEL.MOBY.CONTRAST_MOMENTUM, |
| 62 | contrast_temperature=config.MODEL.MOBY.CONTRAST_TEMPERATURE, |
| 63 | contrast_num_negative=config.MODEL.MOBY.CONTRAST_NUM_NEGATIVE, |
| 64 | proj_num_layers=config.MODEL.MOBY.PROJ_NUM_LAYERS, |
| 65 | pred_num_layers=config.MODEL.MOBY.PRED_NUM_LAYERS, |
| 66 | ) |
| 67 | elif model_type == 'linear': |
| 68 | model = enc( |
| 69 | num_classes=config.MODEL.NUM_CLASSES, |
| 70 | drop_path_rate=config.MODEL.DROP_PATH_RATE, |
| 71 | ) |
| 72 | else: |
| 73 | raise NotImplementedError(f'--> Unknown model_type: {model_type}') |
| 74 | |
| 75 | return model |