Create and load a backbone model given a config. Args: ---- preset: A backbone preset to load pre-defind configs. Returns: ------- A Torch module and the associated config.
(
preset: ViTPreset
)
| 47 | |
| 48 | |
| 49 | def create_backbone_model( |
| 50 | preset: ViTPreset |
| 51 | ) -> Tuple[nn.Module, ViTPreset]: |
| 52 | """Create and load a backbone model given a config. |
| 53 | |
| 54 | Args: |
| 55 | ---- |
| 56 | preset: A backbone preset to load pre-defind configs. |
| 57 | |
| 58 | Returns: |
| 59 | ------- |
| 60 | A Torch module and the associated config. |
| 61 | |
| 62 | """ |
| 63 | if preset in VIT_CONFIG_DICT: |
| 64 | config = VIT_CONFIG_DICT[preset] |
| 65 | model = create_vit(preset=preset, use_pretrained=False) |
| 66 | else: |
| 67 | raise KeyError(f"Preset {preset} not found.") |
| 68 | |
| 69 | return model, config |
| 70 | |
| 71 | |
| 72 | def create_model_and_transforms( |
no test coverage detected