| from __future__ import annotations |
|
|
| from transformers import PretrainedConfig |
|
|
|
|
| class EfficientRAEConfig(PretrainedConfig): |
| """Configuration for the spatial compression modules around a frozen RAE.""" |
|
|
| model_type = "efficient_rae" |
|
|
| def __init__( |
| self, |
| image_size: int = 256, |
| patch_size: int = 16, |
| encoder_hidden_size: int = 1024, |
| pool_size: int = 2, |
| hidden_size: int = 1024, |
| num_attention_heads: int = 16, |
| mlp_ratio: float = 4.0, |
| pooler_num_hidden_layers: int = 16, |
| unpooler_num_hidden_layers: int = 4, |
| pooler_depth_attention: bool = True, |
| unpooler_depth_attention: bool = True, |
| depth_attention_stride: int = 4, |
| rms_norm_eps: float = 1e-6, |
| layer_norm_eps: float = 1e-6, |
| dropout: float = 0.1, |
| attention_dropout: float = 0.1, |
| semantic_num_queries: int | None = None, |
| mask_ratio_min: float = -0.1, |
| mask_ratio_max: float = 0.5, |
| noise_t_max: float = 1.0, |
| checkpoint_format_version: int = 2, |
| **kwargs, |
| ): |
| |
| |
| kwargs.pop("base_model_name_or_path", None) |
| kwargs.pop("base_model_revision", None) |
| super().__init__(**kwargs) |
| self.image_size = image_size |
| self.patch_size = patch_size |
| self.encoder_hidden_size = encoder_hidden_size |
| self.pool_size = pool_size |
| self.hidden_size = hidden_size |
| self.num_attention_heads = num_attention_heads |
| self.mlp_ratio = mlp_ratio |
| self.pooler_num_hidden_layers = pooler_num_hidden_layers |
| self.unpooler_num_hidden_layers = unpooler_num_hidden_layers |
| self.pooler_depth_attention = pooler_depth_attention |
| self.unpooler_depth_attention = unpooler_depth_attention |
| self.depth_attention_stride = depth_attention_stride |
| self.rms_norm_eps = rms_norm_eps |
| self.layer_norm_eps = layer_norm_eps |
| self.dropout = dropout |
| self.attention_dropout = attention_dropout |
| self.semantic_num_queries = ( |
| semantic_num_queries |
| if semantic_num_queries is not None |
| else self.encoder_grid_size**2 + 5 |
| ) |
| self.mask_ratio_min = mask_ratio_min |
| self.mask_ratio_max = mask_ratio_max |
| self.noise_t_max = noise_t_max |
| self.checkpoint_format_version = checkpoint_format_version |
| self._validate() |
|
|
| @property |
| def encoder_grid_size(self) -> int: |
| return self.image_size // self.patch_size |
|
|
| @property |
| def compact_grid_size(self) -> int: |
| return self.encoder_grid_size // self.pool_size |
|
|
| @property |
| def intermediate_size(self) -> int: |
| return int(self.hidden_size * self.mlp_ratio) |
|
|
| def _validate(self) -> None: |
| if self.image_size < 1 or self.patch_size < 1: |
| raise ValueError("image_size and patch_size must be positive") |
| if self.image_size % self.patch_size: |
| raise ValueError("image_size must be divisible by patch_size") |
| if self.pool_size < 1 or self.encoder_grid_size % self.pool_size: |
| raise ValueError("encoder grid size must be divisible by pool_size") |
| if self.hidden_size < 1 or self.hidden_size % self.num_attention_heads: |
| raise ValueError("hidden_size must be divisible by num_attention_heads") |
| head_size = self.hidden_size // self.num_attention_heads |
| if head_size % 4: |
| raise ValueError("attention head size must be divisible by 4 for 2D RoPE") |
| if self.mlp_ratio <= 0: |
| raise ValueError("mlp_ratio must be positive") |
| if self.pooler_num_hidden_layers < 1 or self.unpooler_num_hidden_layers < 1: |
| raise ValueError("pooler and unpooler must each have at least one layer") |
| if self.depth_attention_stride < 1: |
| raise ValueError("depth_attention_stride must be positive") |
| if self.rms_norm_eps <= 0 or self.layer_norm_eps <= 0: |
| raise ValueError("normalization epsilons must be positive") |
| if not 0.0 <= self.dropout < 1.0: |
| raise ValueError("dropout must be in [0, 1)") |
| if not 0.0 <= self.attention_dropout < 1.0: |
| raise ValueError("attention_dropout must be in [0, 1)") |
| if self.encoder_hidden_size % self.num_attention_heads: |
| raise ValueError( |
| "encoder_hidden_size must be divisible by num_attention_heads" |
| ) |
| if self.semantic_num_queries < 1: |
| raise ValueError("semantic_num_queries must be positive") |
| if self.mask_ratio_min > self.mask_ratio_max: |
| raise ValueError("mask_ratio_min must not exceed mask_ratio_max") |
| if self.mask_ratio_max > 1.0: |
| raise ValueError("mask_ratio_max must not exceed 1") |
| if not 0.0 <= self.noise_t_max <= 1.0: |
| raise ValueError("noise_t_max must be in [0, 1]") |
| if self.checkpoint_format_version != 2: |
| raise ValueError("only checkpoint_format_version=2 is supported") |
|
|
|
|
| EfficientRAEConfig.register_for_auto_class() |
|
|
