| from transformers import PretrainedConfig | |
| class CustomViTNanoConfig(PretrainedConfig): | |
| model_type = "custom_vit_nano" | |
| def __init__( | |
| self, | |
| image_size=224, | |
| patch_size=16, | |
| in_chans=3, | |
| num_classes=1000, | |
| embed_dim=224, | |
| depth=6, | |
| num_heads=4, | |
| mlp_hidden_dim=608, | |
| stem_channels=(32, 64, 128), | |
| dropout=0.0, | |
| **kwargs, | |
| ): | |
| super().__init__(**kwargs) | |
| self.image_size = image_size | |
| self.patch_size = patch_size | |
| self.in_chans = in_chans | |
| self.num_classes = num_classes | |
| self.embed_dim = embed_dim | |
| self.depth = depth | |
| self.num_heads = num_heads | |
| self.mlp_hidden_dim = mlp_hidden_dim | |
| self.stem_channels = stem_channels | |
| self.dropout = dropout | |
| CustomViTNanoConfig.register_for_auto_class() |