Do full inference test against test vectors for test_* models

Author: rwightman

tests/test_models.py

@@ -26,7 +26,7 @@
     has_fx_feature_extraction = False
 
 import timm
-from timm import list_models, create_model, set_scriptable, get_pretrained_cfg_value
+from timm import list_models, list_pretrained, create_model, set_scriptable, get_pretrained_cfg_value
 from timm.layers import Format, get_spatial_dim, get_channel_dim
 from timm.models import get_notrace_modules, get_notrace_functions
 
@@ -39,7 +39,8 @@
 torch_device = os.environ.get('TORCH_DEVICE', 'cpu')
 timeout = os.environ.get('TIMEOUT')
 timeout120 = int(timeout) if timeout else 120
-timeout300 = int(timeout) if timeout else 300
+timeout240 = int(timeout) if timeout else 240
+timeout360 = int(timeout) if timeout else 360
 
 if hasattr(torch._C, '_jit_set_profiling_executor'):
     # legacy executor is too slow to compile large models for unit tests
@@ -118,6 +119,50 @@ def _get_input_size(model=None, model_name='', target=None):
     return input_size
 
 
+@pytest.mark.base
+@pytest.mark.timeout(timeout240)
+@pytest.mark.parametrize('model_name', list_pretrained('test_*'))
+@pytest.mark.parametrize('batch_size', [1])
+def test_model_inference(model_name, batch_size):
+    """Run a single forward pass with each model"""
+    from PIL import Image
+    from huggingface_hub import snapshot_download
+    import tempfile
+    import safetensors
+
+    model = create_model(model_name, pretrained=True)
+    model.eval()
+    pp = timm.data.create_transform(**timm.data.resolve_data_config(model=model))
+
+    with tempfile.TemporaryDirectory()  as temp_dir:
+        snapshot_download(
+            repo_id='timm/' + model_name, repo_type='model', local_dir=temp_dir, allow_patterns='test/*'
+        )
+        rand_tensors = safetensors.torch.load_file(os.path.join(temp_dir, 'test', 'rand_tensors.safetensors'))
+        owl_tensors = safetensors.torch.load_file(os.path.join(temp_dir, 'test', 'owl_tensors.safetensors'))
+        test_owl = Image.open(os.path.join(temp_dir, 'test', 'test_owl.jpg'))
+
+    with torch.no_grad():
+        rand_output = model(rand_tensors['input'])
+        rand_features = model.forward_features(rand_tensors['input'])
+        rand_pre_logits = model.forward_head(rand_features, pre_logits=True)
+        assert torch.allclose(rand_output, rand_tensors['output'])
+        assert torch.allclose(rand_features, rand_tensors['features'])
+        assert torch.allclose(rand_pre_logits, rand_tensors['pre_logits'])
+
+        def _test_owl(owl_input):
+            owl_output = model(owl_input)
+            owl_features = model.forward_features(owl_input)
+            owl_pre_logits = model.forward_head(owl_features.clone(), pre_logits=True)
+            assert owl_output.softmax(1).argmax(1) == 24  # owl
+            assert torch.allclose(owl_output, owl_tensors['output'])
+            assert torch.allclose(owl_features, owl_tensors['features'])
+            assert torch.allclose(owl_pre_logits, owl_tensors['pre_logits'])
+
+        _test_owl(owl_tensors['input'])  # test with original pp owl tensor
+        _test_owl(pp(test_owl).unsqueeze(0))  # re-process from original jpg
+
+
 @pytest.mark.base
 @pytest.mark.timeout(timeout120)
 @pytest.mark.parametrize('model_name', list_models(exclude_filters=EXCLUDE_FILTERS))
@@ -182,7 +227,7 @@ def test_model_backward(model_name, batch_size):
 )
 
 @pytest.mark.cfg
-@pytest.mark.timeout(timeout300)
+@pytest.mark.timeout(timeout360)
 @pytest.mark.parametrize('model_name', list_models(
     exclude_filters=EXCLUDE_FILTERS + NON_STD_FILTERS, include_tags=True))
 @pytest.mark.parametrize('batch_size', [1])
@@ -260,7 +305,7 @@ def test_model_default_cfgs(model_name, batch_size):
 
 
 @pytest.mark.cfg
-@pytest.mark.timeout(timeout300)
+@pytest.mark.timeout(timeout360)
 @pytest.mark.parametrize('model_name', list_models(filter=NON_STD_FILTERS, exclude_filters=NON_STD_EXCLUDE_FILTERS, include_tags=True))
 @pytest.mark.parametrize('batch_size', [1])
 def test_model_default_cfgs_non_std(model_name, batch_size):