From c4572cc5aa21dfa6c81b4ef6a3479409e49561f0 Mon Sep 17 00:00:00 2001
From: Ross Wightman <rwightman@gmail.com>
Date: Mon, 24 May 2021 22:50:12 -0700
Subject: [PATCH] Add Visformer-small weighs, tweak torchscript jit test img
 size.

---
 tests/test_models.py     | 18 +++++++++++-------
 timm/models/visformer.py |  4 +++-
 2 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/tests/test_models.py b/tests/test_models.py
index c1632271..e6a73619 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -33,7 +33,11 @@ else:
 TARGET_FWD_SIZE = MAX_FWD_SIZE = 384
 TARGET_BWD_SIZE = 128
 MAX_BWD_SIZE = 320
-MAX_FWD_FEAT_SIZE = 448
+MAX_FWD_OUT_SIZE = 448
+TARGET_JIT_SIZE = 128
+MAX_JIT_SIZE = 320
+TARGET_FFEAT_SIZE = 96
+MAX_FFEAT_SIZE = 256
 
 
 def _get_input_size(model, target=None):
@@ -109,10 +113,10 @@ def test_model_default_cfgs(model_name, batch_size):
     pool_size = cfg['pool_size']
     input_size = model.default_cfg['input_size']
 
-    if all([x <= MAX_FWD_FEAT_SIZE for x in input_size]) and \
+    if all([x <= MAX_FWD_OUT_SIZE for x in input_size]) and \
             not any([fnmatch.fnmatch(model_name, x) for x in EXCLUDE_FILTERS]):
         # output sizes only checked if default res <= 448 * 448 to keep resource down
-        input_size = tuple([min(x, MAX_FWD_FEAT_SIZE) for x in input_size])
+        input_size = tuple([min(x, MAX_FWD_OUT_SIZE) for x in input_size])
         input_tensor = torch.randn((batch_size, *input_size))
 
         # test forward_features (always unpooled)
@@ -176,8 +180,8 @@ def test_model_forward_torchscript(model_name, batch_size):
         model = create_model(model_name, pretrained=False)
     model.eval()
 
-    input_size = _get_input_size(model, 128)
-    if max(input_size) > MAX_FWD_SIZE:  # NOTE using MAX_FWD_SIZE as the final limit is intentional
+    input_size = _get_input_size(model, TARGET_JIT_SIZE)
+    if max(input_size) > MAX_JIT_SIZE:  # NOTE using MAX_FWD_SIZE as the final limit is intentional
         pytest.skip("Fixed input size model > limit.")
 
     model = torch.jit.script(model)
@@ -205,8 +209,8 @@ def test_model_forward_features(model_name, batch_size):
     expected_channels = model.feature_info.channels()
     assert len(expected_channels) >= 4  # all models here should have at least 4 feature levels by default, some 5 or 6
 
-    input_size = _get_input_size(model, 96)  # jit compile is already a bit slow and we've tested normal res already...
-    if max(input_size) > MAX_FWD_SIZE:  # NOTE using MAX_FWD_SIZE as the final limit is intentional
+    input_size = _get_input_size(model, TARGET_FFEAT_SIZE)
+    if max(input_size) > MAX_FFEAT_SIZE:
         pytest.skip("Fixed input size model > limit.")
 
     outputs = model(torch.randn((batch_size, *input_size)))
diff --git a/timm/models/visformer.py b/timm/models/visformer.py
index df1d502a..936f1ddf 100644
--- a/timm/models/visformer.py
+++ b/timm/models/visformer.py
@@ -33,7 +33,9 @@ def _cfg(url='', **kwargs):
 
 default_cfgs = dict(
     visformer_tiny=_cfg(),
-    visformer_small=_cfg(),
+    visformer_small=_cfg(
+        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vt3p-weights/visformer_small-839e1f5b.pth'
+    ),
 )