diff --git a/tests/test_models.py b/tests/test_models.py index fd09ceb2e..b5b538d9c 100644 --- a/tests/test_models.py +++ b/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):