File size: 864 Bytes
75fcfd4 ba314eb 75fcfd4 ba314eb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 | 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() |