build_model with timm

README.MD

@@ -13,7 +13,7 @@ See the available evaluations in [evaluation protocols](EVALUATION.md).
 ## commands
 
 ```bash
-vitrun train_cls.py --data_location=../data/IMNET --gin VisionTransformer.global_pool='"avg"' -w wandb:dlib/EfficientSSL/lsx2qmys 
+vitrun train_cls.py --data_location=../data/IMNET --gin build_model.model_name='"vit_tiny_patch16_224"' build_model.global_pool='"avg"'  -w wandb:dlib/EfficientSSL/ezuz0x4u --layer_decay=0.75 
 ```
 
 ## condor

vitookit/evaluation/eval_cls.py

@@ -147,9 +147,6 @@ def get_args_parser():
                         help='dataset path')
     parser.add_argument('--data_set', default='IN1K', 
                             type=str, help='ImageNet dataset path')
-    parser.add_argument('--inat_category', default='name',
-                        choices=['kingdom', 'phylum', 'class', 'order', 'supercategory', 'family', 'genus', 'name'],
-                        type=str, help='semantic granularity')
 
     parser.add_argument('--output_dir', default=None, type=str,
                         help='path where to save, empty for no saving')
@@ -336,7 +333,7 @@ def main(args):
     print(f"Model  built.")
     # load weights to evaluate
     
-    model = build_model(num_classes=args.nb_classes,drop_path_rate=args.drop_path,)
+    model = build_model(num_classes=args.nb_classes,drop_path_rate=args.drop_path)
     if args.pretrained_weights:
         load_pretrained_weights(model, args.pretrained_weights, checkpoint_key=args.checkpoint_key, prefix=args.prefix)
     if args.compile:

vitookit/models/build_model.py

 import gin
-from .vision_transformer import vit_base,vit_small, vit_tiny
+import timm
 
 @gin.configurable()
-def build_model(*args,model_fn=vit_tiny,**kwargs):
-    model = model_fn(*args,**kwargs)    
-    return model
+def build_model(model_name='vit_base_patch16_224', **kwargs):
+    model = timm.create_model(model_name, **kwargs)
+    return model
\ No newline at end of file