Feature/upgrade timm mlflow onnxconverter (#155)

michele-milesi · lorenzomammana · web-flow · commit 24bc19f65db3 · 2026-02-13T17:02:04.000+01:00
* feat: upgraded mlflow, timm and onnxconverter-common + added new backbones

* fix: mlflow unbound properties when mlflow is not enabled

* feat: update timm import

* fix: log message and fasternet t2 backbone configuration

* feat: added cleanup step in tests

* feat: added cleanup step in tests

* feat: added cleanup step in tests

* feat: added cleanup step in tests

* feat: removed cleanup step

* fix: Fix wrong field checked in automatic batch size

---------

Co-authored-by: Lorenzo Mammana &lt;lorenzo.mammana@orobix.com&gt;
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -14,11 +14,28 @@ Starting from version 2.6.1, releases are automatically created when changes are
 
 **Note**: If a tag for the current version already exists, the workflow will skip tag and release creation to avoid duplicates.
 
+### [2.8.0]
+
+#### Added
+
+- 28 new backbone configurations: csatv2, dinov3_vitb16, dinov3_vits16, edgenext_s, efficientnet_em, efficientnetv2_b1, fasternet_t0, fasternet_t1, fasternet_t2, ghostnet_v3, hgnet_v2, mobilenet_edgetpu_v2, mobilenetv4_m, mobilenetv4_s, regnety_016, repvit, rexnet, shvit_s1, shvit_s2, shvit_s3, shvit_s4, starnet_s2, starnet_s3, starnet_s4, swiftformer_l1, swiftformer_l3, swiftformer_s, swiftformer_xs
+
+#### Updated
+
+- MLflow-skinny from ^2.3.1 to 3.8.1
+- Timm from 0.9.12 to 1.0.24
+- Onnxconverter-common from ^1.14.0 to 1.16.0
+
+#### Fixed
+
+- Avoid unbound variable errors on mlflow client properties if mlflow is not enabled.
+- Fix wrong field checked in automatic batch size corner case that could cause an infinite loop.
+
 ### [2.7.3]
 
 #### Fixed
 
-Add checks to ensure mlflow client properties can't be called if mlflow is not enabled to avoid unbound variable errors.
+- Add checks to ensure mlflow client properties can't be called if mlflow is not enabled to avoid unbound variable errors.
 
 ### [2.7.2]
 
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "quadra"
-version = "2.7.3"
+version = "2.8.0"
 description = "Deep Learning experiment orchestration library"
 authors = [
 	"Federico Belotti <federico.belotti@orobix.com>",
@@ -50,7 +50,7 @@ torchmetrics = "~0.10"
 hydra_core = "~1.3"
 hydra_colorlog = "~1.2"
 hydra_optuna_sweeper = "~1.2"
-mlflow-skinny = "^2.3.1"
+mlflow-skinny = "3.8.1"
 boto3 = "~1.26"
 minio = "~7.1"
 tensorboard = "~2.20"
@@ -70,7 +70,7 @@ label_studio_converter = "~0.0"
 scikit_multilearn = "~0.2"
 tripy = "~1.0"
 h5py = "~3.8"
-timm = "0.9.12"
+timm = "1.0.24"
 segmentation_models_pytorch = "0.5.0"
 
 anomalib-orobix = "0.7.0.dev150"
@@ -82,7 +82,7 @@ typing_extensions = { version = "4.11.0", python = "<3.10" }
 onnx = { version = "1.15.0", optional = true }
 onnxsim = { version = "0.4.28", optional = true }
 onnxruntime_gpu = { version = "1.23.2", optional = true }
-onnxconverter-common = { version = "^1.14.0", optional = true }
+onnxconverter-common = { version = "1.16.0", optional = true }
 
 [[tool.poetry.source]]
 name = "torch_cu128"
diff --git a/quadra/__init__.py b/quadra/__init__.py
@@ -1,4 +1,4 @@
-__version__ = "2.7.3"
+__version__ = "2.8.0"
 
 
 def get_version():
diff --git a/quadra/configs/backbone/csatv2.yaml b/quadra/configs/backbone/csatv2.yaml
@@ -0,0 +1,9 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: csatv2.r512_in1k
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 512
+  output_dim: 386
+  nb_heads: 8
diff --git a/quadra/configs/backbone/dinov3_vitb16.yaml b/quadra/configs/backbone/dinov3_vitb16.yaml
@@ -0,0 +1,10 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: vit_base_patch16_dinov3
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 224
+  output_dim: 768
+  patch_size: 16
+  nb_heads: 12
diff --git a/quadra/configs/backbone/dinov3_vits16.yaml b/quadra/configs/backbone/dinov3_vits16.yaml
@@ -0,0 +1,10 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: vit_small_patch16_dinov3
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 224
+  output_dim: 384
+  patch_size: 16
+  nb_heads: 6
diff --git a/quadra/configs/backbone/edgenext_s.yaml b/quadra/configs/backbone/edgenext_s.yaml
@@ -0,0 +1,9 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: edgenext_small.usi_in1k
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 256
+  output_dim: 304
+  nb_heads: 8
diff --git a/quadra/configs/backbone/efficientnet_em.yaml b/quadra/configs/backbone/efficientnet_em.yaml
@@ -0,0 +1,8 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: efficientnet_em.ra2_in1k
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 240
+  output_dim: 1280
diff --git a/quadra/configs/backbone/efficientnetv2_b1.yaml b/quadra/configs/backbone/efficientnetv2_b1.yaml
@@ -0,0 +1,8 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: tf_efficientnetv2_b1.in1k
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 191
+  output_dim: 1280
diff --git a/quadra/configs/backbone/efficientnetv2_s.yaml b/quadra/configs/backbone/efficientnetv2_s.yaml
@@ -1,6 +1,6 @@
 model:
   _target_: quadra.models.classification.TimmNetworkBuilder
-  model_name: tf_efficientnetv2_s_in21ft1k
+  model_name: tf_efficientnetv2_s.in21k_ft_in1k
   pretrained: true
   freeze: false
 metadata:
diff --git a/quadra/configs/backbone/fasternet_t0.yaml b/quadra/configs/backbone/fasternet_t0.yaml
@@ -0,0 +1,8 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: fasternet_t0.in1k
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 224
+  output_dim: 1280
diff --git a/quadra/configs/backbone/fasternet_t1.yaml b/quadra/configs/backbone/fasternet_t1.yaml
@@ -0,0 +1,8 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: fasternet_t1.in1k
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 224
+  output_dim: 1280
diff --git a/quadra/configs/backbone/fasternet_t2.yaml b/quadra/configs/backbone/fasternet_t2.yaml
@@ -0,0 +1,8 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: fasternet_t2.in1k
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 224
+  output_dim: 1280
diff --git a/quadra/configs/backbone/ghostnet_v3.yaml b/quadra/configs/backbone/ghostnet_v3.yaml
@@ -0,0 +1,8 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: ghostnetv3_100.in1k
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 224
+  output_dim: 1280
diff --git a/quadra/configs/backbone/hgnet_v2.yaml b/quadra/configs/backbone/hgnet_v2.yaml
@@ -0,0 +1,8 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: hgnetv2_b2.ssld_stage2_ft_in1k
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 224
+  output_dim: 2048
diff --git a/quadra/configs/backbone/mobilenet_edgetpu_v2.yaml b/quadra/configs/backbone/mobilenet_edgetpu_v2.yaml
@@ -0,0 +1,8 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: mobilenet_edgetpu_v2_m.ra4_e3600_r224_in1k
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 224
+  output_dim: 1344
diff --git a/quadra/configs/backbone/mobilenetv4_m.yaml b/quadra/configs/backbone/mobilenetv4_m.yaml
@@ -0,0 +1,9 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: mobilenetv4_hybrid_medium.ix_e550_r384_in1k
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 384
+  output_dim: 1280
+  nb_heads: 4
diff --git a/quadra/configs/backbone/mobilenetv4_s.yaml b/quadra/configs/backbone/mobilenetv4_s.yaml
@@ -0,0 +1,10 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: mobilenetv4_conv_small.e2400_r224_in1k
+  # model_name: mobilenetv4_conv_small.e1200_r224_in1k
+  # model_name: mobilenetv4_conv_small_050.e3000_r224_in1k
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 224
+  output_dim: 1280
diff --git a/quadra/configs/backbone/regnety_016.yaml b/quadra/configs/backbone/regnety_016.yaml
@@ -0,0 +1,8 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: regnety_016.tv2_in1k
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 224
+  output_dim: 888
diff --git a/quadra/configs/backbone/repvit.yaml b/quadra/configs/backbone/repvit.yaml
@@ -0,0 +1,10 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: repvit_m1_1.dist_450e_in1k
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 224
+  output_dim: 512
+  patch_size: 4
+  nb_heads: 8
diff --git a/quadra/configs/backbone/resnet18_ssl.yaml b/quadra/configs/backbone/resnet18_ssl.yaml
@@ -1,6 +1,6 @@
 model:
   _target_: quadra.models.classification.TimmNetworkBuilder
-  model_name: ssl_resnet18
+  model_name: resnet18.fb_ssl_yfcc100m_ft_in1k
   pretrained: true
   freeze: false
 metadata:
diff --git a/quadra/configs/backbone/rexnet.yaml b/quadra/configs/backbone/rexnet.yaml
@@ -0,0 +1,8 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: rexnet_150.nav_in1k
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 224
+  output_dim: 1920
diff --git a/quadra/configs/backbone/shvit_s1.yaml b/quadra/configs/backbone/shvit_s1.yaml
@@ -0,0 +1,8 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: shvit_s1.in1k
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 224
+  output_dim: 320
diff --git a/quadra/configs/backbone/shvit_s2.yaml b/quadra/configs/backbone/shvit_s2.yaml
@@ -0,0 +1,8 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: shvit_s2.in1k
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 224
+  output_dim: 448
diff --git a/quadra/configs/backbone/shvit_s3.yaml b/quadra/configs/backbone/shvit_s3.yaml
@@ -0,0 +1,8 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: shvit_s3.in1k
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 224
+  output_dim: 448
diff --git a/quadra/configs/backbone/shvit_s4.yaml b/quadra/configs/backbone/shvit_s4.yaml
@@ -0,0 +1,8 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: shvit_s4.in1k
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 256
+  output_dim: 448
diff --git a/quadra/configs/backbone/starnet_s2.yaml b/quadra/configs/backbone/starnet_s2.yaml
@@ -0,0 +1,8 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: starnet_s2.in1k
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 224
+  output_dim: 256
diff --git a/quadra/configs/backbone/starnet_s3.yaml b/quadra/configs/backbone/starnet_s3.yaml
@@ -0,0 +1,8 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: starnet_s3.in1k
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 224
+  output_dim: 256
diff --git a/quadra/configs/backbone/starnet_s4.yaml b/quadra/configs/backbone/starnet_s4.yaml
@@ -0,0 +1,8 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: starnet_s4.in1k
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 224
+  output_dim: 256
diff --git a/quadra/configs/backbone/swiftformer_l1.yaml b/quadra/configs/backbone/swiftformer_l1.yaml
@@ -0,0 +1,8 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: swiftformer_l1.dist_in1k
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 224
+  output_dim: 384
diff --git a/quadra/configs/backbone/swiftformer_l3.yaml b/quadra/configs/backbone/swiftformer_l3.yaml
@@ -0,0 +1,8 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: swiftformer_l3.dist_in1k
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 224
+  output_dim: 512
diff --git a/quadra/configs/backbone/swiftformer_s.yaml b/quadra/configs/backbone/swiftformer_s.yaml
@@ -0,0 +1,8 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: swiftformer_s.dist_in1k
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 224
+  output_dim: 224
diff --git a/quadra/configs/backbone/swiftformer_xs.yaml b/quadra/configs/backbone/swiftformer_xs.yaml
@@ -0,0 +1,8 @@
+model:
+  _target_: quadra.models.classification.TimmNetworkBuilder
+  model_name: swiftformer_xs.dist_in1k
+  pretrained: true
+  freeze: false
+metadata:
+  input_size: 224
+  output_dim: 220
diff --git a/quadra/models/classification/backbones.py b/quadra/models/classification/backbones.py
@@ -4,7 +4,7 @@
 
 import timm
 import torch
-from timm.models.helpers import load_checkpoint
+from timm.models import load_checkpoint
 from torch import nn
 from torchvision import models
 
diff --git a/quadra/utils/evaluation.py b/quadra/utils/evaluation.py
@@ -467,10 +467,10 @@ def wrapper(self, *args, **kwargs):
                     log.warning(
                         "The function %s went out of memory, trying to reduce the batch size to %d",
                         func.__name__,
-                        self.datamodule.batch_size,
+                        getattr(self.datamodule, batch_size_attribute_name),
                     )
 
-                    if self.datamodule.batch_size == 0:
+                    if getattr(self.datamodule, batch_size_attribute_name) == 0:
                         raise RuntimeError(
                             f"Unable to run {func.__name__} with batch size 1, the program will exit"
                         ) from e
diff --git a/quadra/utils/mlflow.py b/quadra/utils/mlflow.py
diff --git a/quadra/utils/models.py b/quadra/utils/models.py

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-__version__ = "2.7.3"`
	`1`	`+__version__ = "2.8.0"`
`2`	`2`
`3`	`3`
`4`	`4`	`def get_version():`