Add ViT-H (#2)
parent
18ef3c8f7e
commit
5917521c34
|
@ -405,6 +405,20 @@ def vit_large(patch_size=16, num_register_tokens=0, **kwargs):
|
|||
return model
|
||||
|
||||
|
||||
def vit_huge(patch_size=16, num_register_tokens=0, **kwargs):
|
||||
model = DinoVisionTransformer(
|
||||
patch_size=patch_size,
|
||||
embed_dim=1280,
|
||||
depth=32,
|
||||
num_heads=16,
|
||||
mlp_ratio=4,
|
||||
block_fn=partial(Block, attn_class=MemEffAttention),
|
||||
num_register_tokens=num_register_tokens,
|
||||
**kwargs,
|
||||
)
|
||||
return model
|
||||
|
||||
|
||||
def vit_giant2(patch_size=16, num_register_tokens=0, **kwargs):
|
||||
"""
|
||||
Close to ViT-giant, with embed-dim 1536 and 24 heads => embed-dim per head 64
|
||||
|
|
Loading…
Reference in New Issue