Add link to model weights on Hugging Face

This commit is contained in:
Beckschen 2024-05-17 06:51:35 -04:00
parent 530fb49e7e
commit 7a2ad6bce1

View File

@ -1,7 +1,7 @@
""" ViTamin
Paper: Designing Scalable Vison Models in the Vision-Language Era
Model Weights on Huggingface: https://huggingface.co/collections/jienengchen/vitamin-family-661048126b72debdaca060bf
A family of model weights on Huggingface: https://huggingface.co/collections/jienengchen/vitamin-family-661048126b72debdaca060bf
@inproceedings{chen2024vitamin,
title={ViTamin: Designing Scalable Vision Models in the Vision-language Era},
@ -529,11 +529,4 @@ def vitamin_xlarge_384(pretrained=False, **kwargs) -> VisionTransformer:
model_args = dict(img_size=384, embed_dim=1152, depth=32, num_heads=16, mlp_layer=GeGluMlp, mlp_ratio=2., class_token=False, global_pool='avg')
model = _create_vision_transformer_hybrid(
'vitamin_xlarge_384', backbone=backbone, pretrained=pretrained, **dict(model_args, **kwargs))
return model
if __name__ == "__main__":
model = timm.create_model('vitamin_large', num_classes=10).cuda()
x = torch.rand([2,3,224,224]).cuda()
y = model(x)
print(y.shape)
return model