nhl

Author: xadupre

onnx_diagnostic/helpers/config_helper.py

@@ -1,6 +1,7 @@
 import functools
 import importlib
 import inspect
+import os
 import re
 from typing import Any, Callable, Dict, Optional, Tuple, Union
 import transformers
@@ -110,3 +111,12 @@ def config_class_from_architecture(arch: str, exc: bool = False) -> Optional[typ
     )
     cls_name = unique.pop()
     return getattr(transformers, cls_name)
+
+
+def default_num_hidden_layers():
+    """
+    Returns the default number of layers.
+    It is lower when the unit tests are running
+    when ``UNITTEST_GOING=1``.
+    """
+    return 2 if os.environ("UNITTEST_GOING", "0") == "1" else 4

onnx_diagnostic/tasks/automatic_speech_recognition.py

@@ -2,7 +2,11 @@
 import torch
 import transformers
 from ..helpers.cache_helper import make_dynamic_cache, make_encoder_decoder_cache
-from ..helpers.config_helper import update_config, check_hasattr
+from ..helpers.config_helper import (
+    update_config,
+    check_hasattr,
+    default_num_hidden_layers as nhl,
+)
 
 __TASK__ = "automatic-speech-recognition"
 
@@ -15,7 +19,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     if hasattr(config, "decoder_layers"):
         config.decoder_layers = min(config.decoder_layers, 2)
     if hasattr(config, "num_hidden_layers"):
-        config.num_hidden_layers = min(config.num_hidden_layers, 4)
+        config.num_hidden_layers = min(config.num_hidden_layers, nhl())
     update_config(config, kwargs)
     return kwargs
 

onnx_diagnostic/tasks/feature_extraction.py

@@ -1,15 +1,20 @@
 from typing import Any, Callable, Dict, Optional, Tuple
 import torch
-from ..helpers.config_helper import update_config, check_hasattr
+from ..helpers.config_helper import (
+    update_config,
+    check_hasattr,
+    default_num_hidden_layers as nhl,
+)
 from ..helpers.cache_helper import make_dynamic_cache, make_encoder_decoder_cache
 
+
 __TASK__ = "feature-extraction"
 
 
 def reduce_model_config(config: Any) -> Dict[str, Any]:
     """Reduces a model size."""
     check_hasattr(config, "num_hidden_layers")
-    kwargs = dict(num_hidden_layers=min(config.num_hidden_layers, 4))
+    kwargs = dict(num_hidden_layers=min(config.num_hidden_layers, nhl()))
     update_config(config, kwargs)
     return kwargs
 

onnx_diagnostic/tasks/fill_mask.py

@@ -1,6 +1,10 @@
 from typing import Any, Callable, Dict, Optional, Tuple
 import torch
-from ..helpers.config_helper import update_config, check_hasattr
+from ..helpers.config_helper import (
+    update_config,
+    check_hasattr,
+    default_num_hidden_layers as nhl,
+)
 
 __TASK__ = "fill-mask"
 
@@ -9,7 +13,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     """Reduces a model size."""
     check_hasattr(config, "num_attention_heads", "num_hidden_layers")
     kwargs = dict(
-        num_hidden_layers=min(config.num_hidden_layers, 4),
+        num_hidden_layers=min(config.num_hidden_layers, nhl()),
         num_attention_heads=min(config.num_attention_heads, 4),
     )
     update_config(config, kwargs)

onnx_diagnostic/tasks/image_classification.py

@@ -1,6 +1,10 @@
 from typing import Any, Callable, Dict, Optional, Tuple
 import torch
-from ..helpers.config_helper import update_config, check_hasattr
+from ..helpers.config_helper import (
+    update_config,
+    check_hasattr,
+    default_num_hidden_layers as nhl,
+)
 
 __TASK__ = "image-classification"
 
@@ -17,7 +21,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     check_hasattr(config, ("num_hidden_layers", "hidden_sizes"))
     kwargs = dict(
         num_hidden_layers=(
-            min(config.num_hidden_layers, 4)
+            min(config.num_hidden_layers, nhl())
             if hasattr(config, "num_hidden_layers")
             else len(config.hidden_sizes)
         )

onnx_diagnostic/tasks/image_text_to_text.py

@@ -1,7 +1,12 @@
 from typing import Any, Callable, Dict, Optional, Tuple
 import torch
 from ..helpers.cache_helper import make_dynamic_cache, make_hybrid_cache
-from ..helpers.config_helper import update_config, check_hasattr, _pick
+from ..helpers.config_helper import (
+    update_config,
+    check_hasattr,
+    _pick,
+    default_num_hidden_layers as nhl,
+)
 
 __TASK__ = "image-text-to-text"
 
@@ -10,7 +15,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     """Reduces a model size."""
     kwargs: Dict[str, Any] = {}
     if hasattr(config, "num_hidden_layers"):
-        config.num_hidden_layers = min(config.num_hidden_layers, 4)
+        config.num_hidden_layers = min(config.num_hidden_layers, nhl())
     if hasattr(config, "mm_tokens_per_image"):
         config.mm_tokens_per_image = min(config.mm_tokens_per_image, 2)
     if hasattr(config, "vision_config"):

onnx_diagnostic/tasks/mask_generation.py

@@ -1,6 +1,10 @@
 from typing import Any, Callable, Dict, Optional, Tuple
 import torch
-from ..helpers.config_helper import update_config, check_hasattr
+from ..helpers.config_helper import (
+    update_config,
+    check_hasattr,
+    default_num_hidden_layers as nhl,
+)
 
 __TASK__ = "mask-generation"
 
@@ -9,7 +13,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     """Reduces a model size."""
     kwargs: Dict[str, Any] = {}
     if hasattr(config, "num_hidden_layers"):
-        config.num_hidden_layers = min(config.num_hidden_layers, 4)
+        config.num_hidden_layers = min(config.num_hidden_layers, nhl())
     if hasattr(config, "vision_config") and hasattr(config.vision_config, "num_hidden_layers"):
         config.vision_config.num_hidden_layers = min(config.vision_config.num_hidden_layers, 2)
     update_config(config, kwargs)

onnx_diagnostic/tasks/mixture_of_expert.py

@@ -2,7 +2,7 @@
 import torch
 
 # from ..helpers.cache_helper import make_dynamic_cache
-from ..helpers.config_helper import update_config  # , check_hasattr, _pick
+from ..helpers.config_helper import update_config, default_num_hidden_layers as nhl
 
 __TASK__ = "MoE"
 
@@ -11,7 +11,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     """Reduces a model size."""
     kwargs: Dict[str, Any] = {}
     if hasattr(config, "num_hidden_layers"):
-        config.num_hidden_layers = min(config.num_hidden_layers, 4)
+        config.num_hidden_layers = min(config.num_hidden_layers, nhl())
     if hasattr(config, "vision_config") and hasattr(config.vision_config, "num_hidden_layers"):
         config.vision_config.num_hidden_layers = min(config.vision_config.num_hidden_layers, 2)
     if hasattr(config, "audio_processor") and hasattr(

onnx_diagnostic/tasks/object_detection.py

@@ -1,6 +1,10 @@
 from typing import Any, Callable, Dict, Optional, Tuple
 import torch
-from ..helpers.config_helper import update_config, check_hasattr
+from ..helpers.config_helper import (
+    update_config,
+    check_hasattr,
+    default_num_hidden_layers as nhl,
+)
 
 __TASK__ = "object-detection"
 
@@ -10,7 +14,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     check_hasattr(config, ("num_hidden_layers", "hidden_sizes"))
     kwargs = dict(
         num_hidden_layers=(
-            min(config.num_hidden_layers, 4)
+            min(config.num_hidden_layers, nhl())
             if hasattr(config, "num_hidden_layers")
             else len(config.hidden_sizes)
         )

onnx_diagnostic/tasks/sentence_similarity.py

@@ -1,6 +1,10 @@
 from typing import Any, Callable, Dict, Optional, Tuple
 import torch
-from ..helpers.config_helper import update_config, check_hasattr
+from ..helpers.config_helper import (
+    update_config,
+    check_hasattr,
+    default_num_hidden_layers as nhl,
+)
 
 __TASK__ = "sentence-similarity"
 
@@ -9,7 +13,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     """Reduces a model size."""
     check_hasattr(config, "num_attention_heads", "num_hidden_layers")
     kwargs = dict(
-        num_hidden_layers=min(config.num_hidden_layers, 4),
+        num_hidden_layers=min(config.num_hidden_layers, nhl()),
         num_attention_heads=min(config.num_attention_heads, 4),
     )
     update_config(config, kwargs)