aabb / config.py
toilaluan's picture
Upload folder using huggingface_hub
4e570ba verified
Raw
History Blame Contribute Delete
2.05 kB
from transformers import PretrainedConfig
class FineViTConfig(PretrainedConfig):
model_type = "finevit"
def __init__(
self,
backbone_model_name: str = "facebook/dinov2-with-registers-base",
teacher_encoder_type: str = "dinov2_with_registers",
backbone_config: dict | None = None,
num_latents: int = 64,
latent_encoder_num_layers: int = 12,
latent_encoder_head_dim: int = 64,
latent_encoder_num_attention_heads: int = 12,
latent_encoder_mlp_hidden_dim: int = 3072,
patch_decoder_num_layers: int = 8,
patch_decoder_head_dim: int = 64,
patch_decoder_num_attention_heads: int = 12,
patch_decoder_mlp_hidden_dim: int = 3072,
dropout: float = 0.0,
image_size: int = 224,
train_encoder: bool = False,
init_backbone_from_pretrained: bool = True,
initializer_range: float = 0.02,
**kwargs,
):
super().__init__(**kwargs)
self.backbone_model_name = backbone_model_name
self.teacher_encoder_type = teacher_encoder_type
self.backbone_config = backbone_config
self.num_latents = num_latents
self.latent_encoder_num_layers = latent_encoder_num_layers
self.latent_encoder_head_dim = latent_encoder_head_dim
self.latent_encoder_num_attention_heads = latent_encoder_num_attention_heads
self.latent_encoder_mlp_hidden_dim = latent_encoder_mlp_hidden_dim
self.patch_decoder_num_layers = patch_decoder_num_layers
self.patch_decoder_head_dim = patch_decoder_head_dim
self.patch_decoder_num_attention_heads = patch_decoder_num_attention_heads
self.patch_decoder_mlp_hidden_dim = patch_decoder_mlp_hidden_dim
self.dropout = dropout
self.image_size = image_size
self.train_encoder = train_encoder
self.init_backbone_from_pretrained = init_backbone_from_pretrained
self.initializer_range = initializer_range
FineViTConfig.register_for_auto_class("AutoConfig")