59 lines
1.2 KiB
Python
59 lines
1.2 KiB
Python
"""Contains preset for ViT modules.
|
|
|
|
For licensing see accompanying LICENSE file.
|
|
Copyright (C) 2025 Apple Inc. All Rights Reserved.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import dataclasses
|
|
from typing import Literal
|
|
|
|
ViTPreset = Literal["dinov2l16_384",]
|
|
|
|
MLPMode = Literal["vanilla", "glu"]
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class ViTConfig:
|
|
"""Configuration for ViT."""
|
|
|
|
in_chans: int
|
|
embed_dim: int
|
|
depth: int
|
|
num_heads: int
|
|
init_values: float
|
|
|
|
img_size: int = 384
|
|
patch_size: int = 16
|
|
|
|
num_classes: int = 21841
|
|
mlp_ratio: float = 4.0
|
|
drop_rate: float = 0.0
|
|
attn_drop_rate: float = 0.0
|
|
drop_path_rate: float = 0.0
|
|
qkv_bias: bool = True
|
|
global_pool: str = "avg"
|
|
|
|
# Properties for timm_vit.
|
|
mlp_mode: MLPMode = "vanilla"
|
|
|
|
# Properties for SPN.
|
|
intermediate_features_ids: list[int] | None = None
|
|
|
|
def asdict(self):
|
|
"""Convenience method to convert the class to a dict."""
|
|
return dataclasses.asdict(self)
|
|
|
|
|
|
VIT_CONFIG_DICT: dict[ViTPreset, ViTConfig] = {
|
|
"dinov2l16_384": ViTConfig(
|
|
in_chans=3,
|
|
embed_dim=1024,
|
|
depth=24,
|
|
num_heads=16,
|
|
init_values=1e-5,
|
|
global_pool="",
|
|
),
|
|
}
|