| from transformers import PretrainedConfig | |
| from typing import List | |
| class AugViTConfig(PretrainedConfig): | |
| model_type = "augvit" | |
| def __init__( | |
| self, | |
| image_size: int = 32, | |
| patch_size: int = 4, | |
| num_classes: int = 10, | |
| dim: int = 128, | |
| depth: int = 6, | |
| heads: int = 16, | |
| mlp_dim: int = 256, | |
| dropout: int = 0.1, | |
| emb_dropout: int = 0.1, | |
| **kwargs, | |
| ): | |
| self.image_size = image_size | |
| self.patch_size = patch_size | |
| self.num_classes = num_classes | |
| self.dim = dim | |
| self.depth = depth | |
| self.heads = heads | |
| self.mlp_dim = mlp_dim | |
| self.dropout = dropout | |
| self.emb_dropout = emb_dropout | |
| super().__init__(**kwargs) |