| # -------------------------------------------------------- | |
| # Swin Transformer | |
| # Copyright (c) 2021 Microsoft | |
| # Licensed under The MIT License [see LICENSE for details] | |
| # Written by Ze Liu | |
| # -------------------------------------------------------- | |
| from .swin_transformer import SwinTransformer | |
| def build_swin_model(model_type, freeze_at=0): | |
| if model_type == 'swin_base': | |
| model = SwinTransformer(embed_dim=128, | |
| depths=[2, 2, 18, 2], | |
| num_heads=[4, 8, 16, 32], | |
| window_size=7, | |
| drop_path_rate=0.3, | |
| out_indices=(0, 1, 2), | |
| ape=False, | |
| patch_norm=True, | |
| frozen_stages=freeze_at, | |
| use_checkpoint=False) | |
| else: | |
| raise NotImplementedError(f"Unkown model: {model_type}") | |
| return model | |