| from transformers import PretrainedConfig |
| from typing import List |
|
|
|
|
| class AugViTConfig(PretrainedConfig): |
| model_type = "augvit" |
|
|
| def __init__( |
| self, |
| image_size: int = 224, |
| patch_size: int = 32, |
| num_classes: int = 1000, |
| dim: int = 128, |
| depth: int = 2, |
| heads: int = 16, |
| mlp_dim: int = 256, |
| dropout: int = 0.1, |
| emb_dropout: int = 0.1, |
| num_channels:int=3, |
| **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 |
| self.num_channels=num_channels |
| super().__init__(**kwargs) |
