Overhaul 2.1 Remove Dependencies / Add Full Timm Support (#3)

* remove dependencies and add full timm support

* update readme

* add unet head upsampling param

* fix flake8

* update readme

* make output_stride optional, not all timm models support an output stride arg

pyproject.toml

@@ -31,12 +31,9 @@ classifiers = [
     "Topic :: Scientific/Engineering :: Artificial Intelligence",
 ]
 dependencies = [
+    "einops>=0.7.0",
     "timm>=0.9.12",
     "torch>=1.13",
-    "torchvision>=0.14",
-    "einops>=0.3",
-    "pretrainedmodels==0.7.4",
-    "efficientnet-pytorch==0.7.1",
 ]
 dynamic = ["version"]
 
@@ -50,7 +47,6 @@ style = [
 tests = [
     "pytest>=7.3",
     "pytest-cov>=4",
-    "six>=1.16.0"
 ]
 all = [
     "torchseg[style,tests]",
@@ -71,7 +67,7 @@ exclude_lines = [
 
 [tool.isort]
 profile = "black"
-known_first_party = ["docs", "tests", "torchseg", "train"]
+known_first_party = ["tests", "torchseg", "train"]
 skip_gitignore = true
 color_output = true
 

requirements/required.txt

@@ -3,8 +3,5 @@ setuptools==69.0.0
 
 # install
 einops==0.7.0
-timm==0.9.2
-torch==2.1.2
-torchvision==0.16.2
-pretrainedmodels==0.7.4
-efficientnet-pytorch==0.7.1
+timm==0.9.12
+torch==2.1.2

requirements/tests.txt

@@ -1,5 +1,3 @@
 # tests
 pytest==7.4.4
-pytest-cov==4.1.0
-mock==5.1.0
-six==1.16.0
+pytest-cov==4.1.0

scripts/list_compatible_encoders.py

@@ -0,0 +1,40 @@
+import json
+
+import timm
+from tqdm import tqdm
+
+if __name__ == "__main__":
+    # Check for models that support `features_only=True``
+    works, fails = {}, []
+    for model in tqdm(timm.list_models()):
+        try:
+            m = timm.create_model(model, pretrained=False, features_only=True)
+            works[model] = dict(
+                indices=m.feature_info.out_indices,
+                channels=m.feature_info.channels(),
+                reduction=m.feature_info.reduction(),
+                module=m.feature_info.module_name(),
+            )
+        except RuntimeError:
+            fails.append(model)
+
+    with open("encoders_features_only_supported.json", "w") as f:
+        json.dump(works, f, indent=2)
+
+    # Check for models that support `get_intermediate_layers``
+    intermediate_layers_support = []
+    unsupported = []
+
+    for model in tqdm(fails):
+        m = timm.create_model(model, pretrained=False)
+        if hasattr(m, "get_intermediate_layers"):
+            intermediate_layers_support.append(model)
+        else:
+            unsupported.append(model)
+
+    with open("encoders_get_intermediate_layers_supported.json", "w") as f:
+        json.dump(intermediate_layers_support, f, indent=2)
+
+    # Save unsupported timm models
+    with open("encoders_unsupported.json", "w") as f:
+        json.dump(unsupported, f, indent=2)