add huge/14 model config

rentainhe · Nov 4, 2020 · e7a4c95 · e7a4c95
1 parent 215e05a
commit e7a4c95
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 0 deletions.
diff --git a/models/configs.py b/models/configs.py
@@ -75,3 +75,19 @@ def get_l32_config():
     config = get_l16_config()
     config.patches.size = (32, 32)
     return config
+
+
+def get_h14_config():
+    """Returns the ViT-L/16 configuration."""
+    config = ml_collections.ConfigDict()
+    config.patches = ml_collections.ConfigDict({'size': (14, 14)})
+    config.hidden_size = 1280
+    config.transformer = ml_collections.ConfigDict()
+    config.transformer.mlp_dim = 5120
+    config.transformer.num_heads = 16
+    config.transformer.num_layers = 32
+    config.transformer.attention_dropout_rate = 0.0
+    config.transformer.dropout_rate = 0.1
+    config.classifier = 'token'
+    config.representation_size = None
+    return config
diff --git a/models/modeling.py b/models/modeling.py
@@ -307,5 +307,6 @@ def load_from(self, weights):
     'ViT-B_32': configs.get_b32_config(),
     'ViT-L_16': configs.get_l16_config(),
     'ViT-L_32': configs.get_l32_config(),
+    'ViT-H_14': configs.get_h14_config(),
     'testing': configs.get_testing(),
 }