From 9302fa032c107ab4a368a644e1f639daf11e9cc9 Mon Sep 17 00:00:00 2001
From: Jordan Gong <jordan.gong@protonmail.com>
Date: Sat, 20 Aug 2022 09:12:01 +0800
Subject: Add encoder option to SimCLR evaluation script

---
 simclr/evaluate.py | 4 ++++
 1 file changed, 4 insertions(+)

(limited to 'simclr')

diff --git a/simclr/evaluate.py b/simclr/evaluate.py
index 5c41b84..fff4eb2 100644
--- a/simclr/evaluate.py
+++ b/simclr/evaluate.py
@@ -40,6 +40,9 @@ def parse_args_and_config():
                         help='Path to config file (optional)')
 
     # TODO: Add model hyperparams dataclass
+    parser.add_argument('--encoder', default='resnet', type=str,
+                        choices=('resnet', 'vit'),
+                        help='Backbone of encoder')
     parser.add_argument('--hid-dim', default=2048, type=int,
                         help='Number of dimension of embedding')
     parser.add_argument('--out-dim', default=128, type=int,
@@ -261,6 +264,7 @@ if __name__ == '__main__':
         inf_mode=False,
         num_iters=args.num_iters,
         config=config,
+        encoder=args.encoder,
         hid_dim=args.hid_dim,
         out_dim=args.out_dim,
         pretrained_checkpoint=args.pretrained_checkpoint,
-- 
cgit v1.2.3