diff --git a/application/backend/app/execution/training/getitune_trainer.py b/application/backend/app/execution/training/getitune_trainer.py index 055168ed648..b7720174dc1 100644 --- a/application/backend/app/execution/training/getitune_trainer.py +++ b/application/backend/app/execution/training/getitune_trainer.py @@ -1,5 +1,7 @@ # Copyright (C) 2025 Intel Corporation # SPDX-License-Identifier: Apache-2.0 +import importlib +import inspect import shutil from collections.abc import Callable from contextlib import AbstractContextManager @@ -359,14 +361,6 @@ def train_model( has_parent_revision: bool, ) -> tuple[Path, LightningEngine]: """Execute model training.""" - # TODO use weights path to initialize model from pre-downloaded weights - # after resolving https://github.com/open-edge-platform/training_extensions/issues/5100 - logger.warning( - "Argument 'weights_path' (value='{}') is not used in model training yet; " - "the weights location will be determined internally by getitune", - weights_path, - ) - # Build the DataModule logger.info("Preparing the DataModule for training (model_id={})", model_id) getitune_datamodule = DataModule.from_vision_datasets( @@ -388,6 +382,10 @@ def train_model( std=getitune_datamodule.input_std if getitune_datamodule.input_std is not None else (1.0, 1.0, 1.0), intensity_config=getitune_datamodule.input_intensity_config, ).as_dict() + if not has_parent_revision: + # Inject pre-downloaded backbone weights when training from scratch, + # but only if the model class supports the pretrained_weights_path parameter + self._inject_pretrained_weights_path(model_cfg, weights_path) model_parser = ArgumentParser() model_parser.add_subclass_arguments(LightningModel, "model", required=False, fail_untyped=False) getitune_model: LightningModel = model_parser.instantiate_classes(Namespace(model=model_cfg)).get("model") @@ -752,6 +750,35 @@ def filter_dataset(self, dm_dataset: Dataset, task: Task, training_config: Train schema=dm_dataset.schema, ) + @staticmethod + def _inject_pretrained_weights_path(model_cfg: dict, weights_path: Path) -> None: + """Inject pretrained_weights_path into model init_args if the model class supports it. + + This avoids downloading backbone weights from external servers when a local + pre-downloaded file is available. + + Args: + model_cfg: The model configuration dict with 'class_path' and 'init_args' keys. + weights_path: Path to the locally downloaded pretrained weights file. + """ + class_path = model_cfg.get("class_path", "") + try: + module_name, class_name = class_path.rsplit(".", 1) + module = importlib.import_module(module_name) + cls = getattr(module, class_name) + if "pretrained_weights_path" in inspect.signature(cls.__init__).parameters: + model_cfg["init_args"]["pretrained_weights_path"] = str(weights_path) + logger.info( + "Using pre-downloaded pretrained weights for {} from: {}", class_name, weights_path + ) + except (ImportError, AttributeError, ValueError) as exc: + logger.warning( + "Could not inject pretrained_weights_path for {}: {}. " + "The model will attempt to download weights from an external server.", + class_path, + exc, + ) + @staticmethod def __base_model_path(data_dir: Path, project_id: UUID, model_id: UUID) -> Path: return data_dir / "projects" / str(project_id) / "models" / str(model_id) diff --git a/application/backend/tests/unit/execution/training/test_getitune_trainer.py b/application/backend/tests/unit/execution/training/test_getitune_trainer.py index e9e3e93df31..49332cb890e 100644 --- a/application/backend/tests/unit/execution/training/test_getitune_trainer.py +++ b/application/backend/tests/unit/execution/training/test_getitune_trainer.py @@ -835,6 +835,126 @@ def test_train_model( assert trained_model_path == expected_checkpoint_path assert returned_engine == mock_getitune_engine + def test_train_model_without_parent_revision_injects_pretrained_weights_path( + self, + fxt_getitune_trainer: Callable[[], GetiTuneTrainer], + tmp_path: Path, + ): + """When has_parent_revision=False, pretrained_weights_path should be injected for supported model classes.""" + # Arrange + getitune_trainer = fxt_getitune_trainer() + model_id = uuid4() + + mock_dataset_info = Mock() + mock_dataset_info.getitune_training_dataset = Mock() + mock_dataset_info.getitune_validation_dataset = Mock() + mock_dataset_info.getitune_testing_dataset = Mock() + mock_dataset_info.getitune_training_subset_config = Mock() + mock_dataset_info.getitune_validation_subset_config = Mock() + mock_dataset_info.getitune_testing_subset_config = Mock() + + weights_path = tmp_path / "weights.pth" + weights_path.touch() + + training_config: dict[str, Any] = { + "model": { + "class_path": "getitune.backend.lightning.models.classification.multiclass_models.torchvision_model.TVModelMulticlassCls", + "init_args": { + "model_name": "efficientnet_b3", + }, + }, + "max_epochs": 10, + "callbacks": [], + } + + mock_datamodule = Mock() + mock_datamodule.label_info.label_names = ["cat", "dog"] + mock_datamodule.input_size = (300, 300) + mock_datamodule.input_mean = None + mock_datamodule.input_std = None + mock_datamodule.input_intensity_config = None + + mock_getitune_model = Mock() + mock_getitune_engine = Mock() + mock_getitune_engine.work_dir = str(tmp_path / f"getitune-workspace-{model_id}") + Path(mock_getitune_engine.work_dir).mkdir(parents=True) + expected_checkpoint_path = Path(mock_getitune_engine.work_dir) / "best_checkpoint.ckpt" + expected_checkpoint_path.touch() + + with patch("app.execution.training.getitune_trainer.DataModule.from_vision_datasets") as mock_dm_factory: + mock_dm_factory.return_value = mock_datamodule + with patch("app.execution.training.getitune_trainer.ArgumentParser") as mock_parser_class: + mock_parser = Mock() + mock_parser_class.return_value = mock_parser + mock_model_namespace = Mock() + mock_model_namespace.get.return_value = mock_getitune_model + mock_parser.instantiate_classes.return_value = mock_model_namespace + with patch("app.execution.training.getitune_trainer.LightningEngine") as mock_engine_class: + mock_engine_class.return_value = mock_getitune_engine + + # Act + getitune_trainer.train_model( + training_config=training_config, + dataset_info=mock_dataset_info, + weights_path=weights_path, + model_id=model_id, + device=DeviceInfo(type=DeviceType.CPU, name="Intel Core", index=None, memory=None), + has_parent_revision=False, + ) + + # Assert: pretrained_weights_path should be injected into the model init_args + assert training_config["model"]["init_args"]["pretrained_weights_path"] == str(weights_path) + # Assert: checkpoint should be None when has_parent_revision=False + engine_call_kwargs = mock_engine_class.call_args.kwargs + assert engine_call_kwargs["checkpoint"] is None + + +class TestGetiTuneTrainerInjectPretrainedWeightsPath: + """Tests for the GetiTuneTrainer._inject_pretrained_weights_path static method.""" + + def test_injects_path_for_supported_model(self, tmp_path: Path): + """When model class supports pretrained_weights_path, it should be injected.""" + weights_path = tmp_path / "weights.pth" + weights_path.touch() + + model_cfg = { + "class_path": "getitune.backend.lightning.models.classification.multiclass_models.torchvision_model.TVModelMulticlassCls", + "init_args": {"model_name": "efficientnet_b3"}, + } + + GetiTuneTrainer._inject_pretrained_weights_path(model_cfg, weights_path) + + assert model_cfg["init_args"]["pretrained_weights_path"] == str(weights_path) + + def test_does_not_inject_for_unsupported_model(self, tmp_path: Path): + """When model class does not support pretrained_weights_path, init_args should remain unchanged.""" + weights_path = tmp_path / "weights.pth" + weights_path.touch() + + model_cfg = { + "class_path": "getitune.backend.lightning.models.detection.yolox.YOLOXModel", + "init_args": {"model_name": "yolox_tiny"}, + } + + GetiTuneTrainer._inject_pretrained_weights_path(model_cfg, weights_path) + + assert "pretrained_weights_path" not in model_cfg["init_args"] + + def test_handles_invalid_class_path_gracefully(self, tmp_path: Path): + """When class_path is invalid, the method should log a warning and not raise.""" + weights_path = tmp_path / "weights.pth" + weights_path.touch() + + model_cfg = { + "class_path": "non.existent.module.SomeModel", + "init_args": {}, + } + + # Should not raise an exception + GetiTuneTrainer._inject_pretrained_weights_path(model_cfg, weights_path) + + assert "pretrained_weights_path" not in model_cfg["init_args"] + class TestGetiTuneTrainerExecuteCancellation: """Tests for the GetiTuneTrainer.execute method handling CancelledExc.""" diff --git a/library/src/getitune/backend/lightning/models/classification/backbones/torchvision.py b/library/src/getitune/backend/lightning/models/classification/backbones/torchvision.py index 59188a7ea24..d023459dafe 100644 --- a/library/src/getitune/backend/lightning/models/classification/backbones/torchvision.py +++ b/library/src/getitune/backend/lightning/models/classification/backbones/torchvision.py @@ -3,10 +3,15 @@ """Torchvison model's Backbone Class.""" +import logging +from pathlib import Path + import torch from torch import nn from torchvision.models import get_model, get_model_weights +logger = logging.getLogger(__name__) + TVModels = [ "alexnet", "convnext_base", @@ -88,6 +93,7 @@ def __init__( self, backbone: str, pretrained: bool = True, + pretrained_weights_path: str | Path | None = None, **kwargs, ): super().__init__(**kwargs) @@ -95,9 +101,17 @@ def __init__( msg = f"Backbone model name is not supported: {backbone}. Available models: {TVModels}" raise ValueError(msg) tv_model_cfg = {"name": backbone} - if pretrained: + if pretrained_weights_path is not None and Path(pretrained_weights_path).exists(): + # Load weights from local file to avoid downloading from external servers + logger.info("Loading pretrained weights for %s from local path: %s", backbone, pretrained_weights_path) + net = get_model(**tv_model_cfg) + state_dict = torch.load(pretrained_weights_path, map_location="cpu", weights_only=True) + net.load_state_dict(state_dict) + elif pretrained: tv_model_cfg["weights"] = get_model_weights(backbone) - net = get_model(**tv_model_cfg) + net = get_model(**tv_model_cfg) + else: + net = get_model(**tv_model_cfg) self.features = net.features last_layer = list(net.children())[-1] diff --git a/library/src/getitune/backend/lightning/models/classification/hlabel_models/torchvision_model.py b/library/src/getitune/backend/lightning/models/classification/hlabel_models/torchvision_model.py index 00750793a64..d2c70bab80e 100644 --- a/library/src/getitune/backend/lightning/models/classification/hlabel_models/torchvision_model.py +++ b/library/src/getitune/backend/lightning/models/classification/hlabel_models/torchvision_model.py @@ -35,6 +35,9 @@ class TVModelHLabelCls(LightningHlabelClsModel): label_info (HLabelInfo): Information about the hierarchical labels. backbone (TVModelType): The type of Torchvision backbone model. pretrained (bool, optional): Whether to use pretrained weights. Defaults to True. + pretrained_weights_path (str | None, optional): Path to a local pretrained weights file. + If provided and the file exists, weights are loaded from this path instead of downloading + from external servers. Defaults to None. optimizer (OptimizerCallable, optional): The optimizer callable. Defaults to DefaultOptimizerCallable. scheduler (LRSchedulerCallable | LRSchedulerListCallable, optional): The learning rate scheduler callable. Defaults to DefaultSchedulerCallable. @@ -49,11 +52,13 @@ def __init__( data_input_params: DataInputParams | dict | None = None, model_name: str = "efficientnet_v2_s", freeze_backbone: bool = False, + pretrained_weights_path: str | None = None, optimizer: OptimizerCallable = DefaultOptimizerCallable, scheduler: LRSchedulerCallable | LRSchedulerListCallable = DefaultSchedulerCallable, metric: MetricCallable = HLabelClsMetricCallable, torch_compile: bool = False, ) -> None: + self.pretrained_weights_path = pretrained_weights_path super().__init__( label_info=label_info, data_input_params=data_input_params, @@ -67,7 +72,7 @@ def __init__( def _create_model(self, head_config: dict | None = None) -> nn.Module: # type: ignore[override] head_config = head_config if head_config is not None else self.label_info.as_head_config_dict() - backbone = TorchvisionBackbone(backbone=self.model_name) + backbone = TorchvisionBackbone(backbone=self.model_name, pretrained_weights_path=self.pretrained_weights_path) return HLabelClassifier( backbone=backbone, neck=GlobalAveragePooling(dim=2), diff --git a/library/src/getitune/backend/lightning/models/classification/multiclass_models/torchvision_model.py b/library/src/getitune/backend/lightning/models/classification/multiclass_models/torchvision_model.py index aaa47c3657f..fc80c3c5df0 100644 --- a/library/src/getitune/backend/lightning/models/classification/multiclass_models/torchvision_model.py +++ b/library/src/getitune/backend/lightning/models/classification/multiclass_models/torchvision_model.py @@ -39,6 +39,9 @@ class TVModelMulticlassCls(LightningMulticlassClsModel): such as input size and normalization. If None is given, default parameters for the specific model will be used. model_name (str, optional): Backbone model name for feature extraction. Defaults to "efficientnet_v2_s". + pretrained_weights_path (str | None, optional): Path to a local pretrained weights file. + If provided and the file exists, weights are loaded from this path instead of downloading + from external servers. Defaults to None. optimizer (OptimizerCallable, optional): Optimizer for model training. Defaults to DefaultOptimizerCallable. scheduler (LRSchedulerCallable | LRSchedulerListCallable, optional): Learning rate scheduler. Defaults to DefaultSchedulerCallable. @@ -52,11 +55,13 @@ def __init__( data_input_params: DataInputParams | dict | None = None, model_name: str = "efficientnet_v2_s", freeze_backbone: bool = False, + pretrained_weights_path: str | None = None, optimizer: OptimizerCallable = DefaultOptimizerCallable, scheduler: LRSchedulerCallable | LRSchedulerListCallable = DefaultSchedulerCallable, metric: MetricCallable = MultiClassClsMetricCallable, torch_compile: bool = False, ) -> None: + self.pretrained_weights_path = pretrained_weights_path super().__init__( label_info=label_info, data_input_params=data_input_params, @@ -70,7 +75,7 @@ def __init__( def _create_model(self, num_classes: int | None = None) -> nn.Module: num_classes = num_classes if num_classes is not None else self.num_classes - backbone = TorchvisionBackbone(backbone=self.model_name) + backbone = TorchvisionBackbone(backbone=self.model_name, pretrained_weights_path=self.pretrained_weights_path) neck = GlobalAveragePooling(dim=2) return ImageClassifier( diff --git a/library/src/getitune/backend/lightning/models/classification/multilabel_models/torchvision_model.py b/library/src/getitune/backend/lightning/models/classification/multilabel_models/torchvision_model.py index a2af7eb739b..e8ab7013948 100644 --- a/library/src/getitune/backend/lightning/models/classification/multilabel_models/torchvision_model.py +++ b/library/src/getitune/backend/lightning/models/classification/multilabel_models/torchvision_model.py @@ -37,6 +37,9 @@ class TVModelMultilabelCls(LightningMultilabelClsModel): label_info (LabelInfoTypes): Information about the labels. backbone (TVModelType): Backbone model for feature extraction. pretrained (bool, optional): Whether to use pretrained weights. Defaults to True. + pretrained_weights_path (str | None, optional): Path to a local pretrained weights file. + If provided and the file exists, weights are loaded from this path instead of downloading + from external servers. Defaults to None. optimizer (OptimizerCallable, optional): Optimizer for model training. Defaults to DefaultOptimizerCallable. scheduler (LRSchedulerCallable | LRSchedulerListCallable, optional): Learning rate scheduler. Defaults to DefaultSchedulerCallable. @@ -51,11 +54,13 @@ def __init__( data_input_params: DataInputParams | dict | None = None, model_name: str = "efficientnet_v2_s", freeze_backbone: bool = False, + pretrained_weights_path: str | None = None, optimizer: OptimizerCallable = DefaultOptimizerCallable, scheduler: LRSchedulerCallable | LRSchedulerListCallable = DefaultSchedulerCallable, metric: MetricCallable = MultiLabelClsMetricCallable, torch_compile: bool = False, ) -> None: + self.pretrained_weights_path = pretrained_weights_path super().__init__( label_info=label_info, data_input_params=data_input_params, @@ -69,7 +74,7 @@ def __init__( def _create_model(self, num_classes: int | None = None) -> nn.Module: num_classes = num_classes if num_classes is not None else self.num_classes - backbone = TorchvisionBackbone(backbone=self.model_name) + backbone = TorchvisionBackbone(backbone=self.model_name, pretrained_weights_path=self.pretrained_weights_path) return ImageClassifier( backbone=backbone, neck=GlobalAveragePooling(dim=2),