Add proper version check

lightly/__init__.py

@@ -79,6 +79,10 @@
 
 
 import os
+import packaging.version
+from packaging.version import Version as _Version
+from packaging.version import InvalidVersion
+from typing import Optional
 
 # see if torchvision vision transformer is available
 try:
@@ -97,10 +101,13 @@
     import timm
 except ImportError:
     _timm_available = False
-    _timm_version = ""
+    _timm_version: Optional[_Version] = None
 else:
     _timm_available = True
-    _timm_version = timm.__version__
+    try:
+        _timm_version = packaging.version.parse(timm.__version__)
+    except InvalidVersion:
+        _timm_version = None
 
 
 if os.getenv("LIGHTLY_DID_VERSION_CHECK", "False") == "False":

lightly/models/modules/__init__.py

@@ -27,6 +27,7 @@
     SwaVPrototypes,
 )
 from lightly.models.modules.nn_memory_bank import NNMemoryBankModule
+from packaging.version import Version as _Version
 
 if _torchvision_vit_available:
     # Requires torchvision >=0.12
@@ -35,7 +36,7 @@
         MAEDecoder,
         MAEEncoder,
     )
-if _timm_available and _timm_version >= "0.9.9":
+if _timm_available and _timm_version is not None and _timm_version >= _Version("0.9.9"):
     # Requires timm >= 0.9.9
     from lightly.models.modules.heads_timm import AIMPredictionHead
     from lightly.models.modules.masked_causal_vision_transformer import (

requirements/base.txt

@@ -9,4 +9,5 @@ six>=1.10
 tqdm>=4.44
 urllib3 >= 1.25.3
 pydantic >= 1.10.5, < 2
-aenum >= 3.1.11
+aenum >= 3.1.11
+packaging

requirements/minimal_requirements.txt

@@ -14,6 +14,7 @@ pytorch_lightning==1.7.1
 torch==1.11.0
 torchvision==0.12.0
 pillow==7.1.2
+packaging
 
 tox
 sphinx