From 9ad0072e7fbe9c5f95bc372b1ed8e26ca2af1e09 Mon Sep 17 00:00:00 2001 From: Adeel Hassan Date: Thu, 8 Feb 2024 09:37:13 -0500 Subject: [PATCH] remove LearnerPipeline and LearnerPipelineConfig --- .../rastervision/pytorch_learner/__init__.py | 5 --- .../rastervision/pytorch_learner/learner.py | 38 +++++++------------ .../pytorch_learner/learner_pipeline.py | 17 --------- .../learner_pipeline_config.py | 22 ----------- .../pytorch_learner/utils/utils.py | 33 ++++++++++++++-- 5 files changed, 42 insertions(+), 73 deletions(-) delete mode 100644 rastervision_pytorch_learner/rastervision/pytorch_learner/learner_pipeline.py delete mode 100644 rastervision_pytorch_learner/rastervision/pytorch_learner/learner_pipeline_config.py diff --git a/rastervision_pytorch_learner/rastervision/pytorch_learner/__init__.py b/rastervision_pytorch_learner/rastervision/pytorch_learner/__init__.py index 1595ffe65a..3f7943b28a 100644 --- a/rastervision_pytorch_learner/rastervision/pytorch_learner/__init__.py +++ b/rastervision_pytorch_learner/rastervision/pytorch_learner/__init__.py @@ -13,8 +13,6 @@ def register_plugin(registry: 'Registry'): import rastervision.pipeline from rastervision.pytorch_learner.learner_config import * from rastervision.pytorch_learner.learner import * -from rastervision.pytorch_learner.learner_pipeline_config import * -from rastervision.pytorch_learner.learner_pipeline import * from rastervision.pytorch_learner.classification_learner_config import * from rastervision.pytorch_learner.classification_learner import * from rastervision.pytorch_learner.regression_learner_config import * @@ -26,9 +24,6 @@ def register_plugin(registry: 'Registry'): from rastervision.pytorch_learner.dataset import * __all__ = [ - # LearnerPipeline - LearnerPipeline.__name__, - LearnerPipelineConfig.__name__, # Learner Learner.__name__, SemanticSegmentationLearner.__name__, diff --git a/rastervision_pytorch_learner/rastervision/pytorch_learner/learner.py b/rastervision_pytorch_learner/rastervision/pytorch_learner/learner.py index 6ff7d53287..48a06fc404 100644 --- a/rastervision_pytorch_learner/rastervision/pytorch_learner/learner.py +++ b/rastervision_pytorch_learner/rastervision/pytorch_learner/learner.py @@ -27,16 +27,16 @@ from rastervision.pipeline import rv_config_ as rv_config from rastervision.pipeline.utils import get_env_var from rastervision.pipeline.file_system import ( - sync_to_dir, json_to_file, file_to_json, make_dir, zipdir, - download_if_needed, download_or_copy, sync_from_dir, get_local_path, unzip, - str_to_file, is_local, get_tmp_dir) + sync_to_dir, json_to_file, make_dir, zipdir, download_if_needed, + download_or_copy, sync_from_dir, get_local_path, unzip, is_local, + get_tmp_dir) from rastervision.pipeline.file_system.utils import file_exists from rastervision.pipeline.utils import terminate_at_exit -from rastervision.pipeline.config import (build_config, upgrade_config, - save_pipeline_config) +from rastervision.pipeline.config import build_config from rastervision.pytorch_learner.utils import ( - get_hubconf_dir_from_cfg, aggregate_metrics, log_metrics_to_csv, - log_system_details, ONNXRuntimeAdapter, DDPContextManager) + aggregate_metrics, DDPContextManager, get_hubconf_dir_from_cfg, + get_learner_config_from_bundle_dir, log_metrics_to_csv, log_system_details, + ONNXRuntimeAdapter) from rastervision.pytorch_learner.dataset.visualizer import Visualizer if TYPE_CHECKING: @@ -44,8 +44,7 @@ from torch.optim.lr_scheduler import _LRScheduler from torch.utils.data import Dataset, Sampler - from rastervision.pytorch_learner import (LearnerConfig, - LearnerPipelineConfig) + from rastervision.pytorch_learner import LearnerConfig warnings.filterwarnings('ignore') @@ -305,14 +304,7 @@ def from_model_bundle(cls: Type, unzip(model_bundle_path, model_bundle_dir) if cfg is None: - config_path = join(model_bundle_dir, 'pipeline-config.json') - - config_dict = file_to_json(config_path) - config_dict = upgrade_config(config_dict) - - learner_pipeline_cfg: 'LearnerPipelineConfig' = build_config( - config_dict) - cfg = learner_pipeline_cfg.learner + cfg = get_learner_config_from_bundle_dir(model_bundle_dir) hub_dir = join(model_bundle_dir, MODULES_DIRNAME) model_def_path = None @@ -1024,8 +1016,8 @@ def setup_training(self, loss_def_path: Optional[str] = None) -> None: """ cfg = self.cfg - self.config_path = join(self.output_dir, 'learner-config.json') - str_to_file(cfg.json(), self.config_path) + self.config_path = join(self.output_dir_local, 'learner-config.json') + cfg.to_file(self.config_path) self.log_path = join(self.output_dir_local, 'log.csv') self.last_model_weights_path = join(self.output_dir_local, 'last-model.pth') @@ -1399,9 +1391,6 @@ def save_model_bundle(self, export_onnx: bool = True): This is a zip file with the model weights in .pth format and a serialized copy of the LearningConfig, which allows for making predictions in the future. """ - from rastervision.pytorch_learner.learner_pipeline_config import ( - LearnerPipelineConfig) - if self.cfg.model is None: log.warning( 'Model was not configured via ModelConfig, and therefore, ' @@ -1417,9 +1406,8 @@ def save_model_bundle(self, export_onnx: bool = True): self._bundle_modules(model_bundle_dir) self._bundle_transforms(model_bundle_dir) - pipeline_cfg = LearnerPipelineConfig(learner=self.cfg) - save_pipeline_config(pipeline_cfg, - join(model_bundle_dir, 'pipeline-config.json')) + cfg_uri = join(model_bundle_dir, 'learner-config.json') + shutil.copy(self.config_path, cfg_uri) zip_path = join(self.output_dir_local, basename(self.model_bundle_uri)) log.info(f'Saving bundle to {zip_path}.') diff --git a/rastervision_pytorch_learner/rastervision/pytorch_learner/learner_pipeline.py b/rastervision_pytorch_learner/rastervision/pytorch_learner/learner_pipeline.py deleted file mode 100644 index 2dbdfa0626..0000000000 --- a/rastervision_pytorch_learner/rastervision/pytorch_learner/learner_pipeline.py +++ /dev/null @@ -1,17 +0,0 @@ -from rastervision.pipeline.pipeline import Pipeline -from rastervision.pytorch_learner import LearnerConfig - - -class LearnerPipeline(Pipeline): - """Simple Pipeline that is a wrapper around Learner.main() - - This supports the ability to use the pytorch_learner package to train models using - the RV pipeline package and its runner functionality without the rest of RV. - """ - commands = ['train'] - gpu_commands = ['train'] - - def train(self): - learner_cfg: LearnerConfig = self.config.learner - learner = learner_cfg.build(learner_cfg, self.tmp_dir) - learner.main() diff --git a/rastervision_pytorch_learner/rastervision/pytorch_learner/learner_pipeline_config.py b/rastervision_pytorch_learner/rastervision/pytorch_learner/learner_pipeline_config.py deleted file mode 100644 index f7ed042932..0000000000 --- a/rastervision_pytorch_learner/rastervision/pytorch_learner/learner_pipeline_config.py +++ /dev/null @@ -1,22 +0,0 @@ -from rastervision.pipeline.config import register_config -from rastervision.pipeline.pipeline_config import PipelineConfig -from rastervision.pytorch_learner import LearnerConfig - - -@register_config('learner_pipeline') -class LearnerPipelineConfig(PipelineConfig): - """Configure a :class:`.LearnerPipeline`.""" - - learner: LearnerConfig - - def update(self): - super().update() - - if self.learner.output_uri is None: - self.learner.output_uri = self.root_uri - - self.learner.update() - - def build(self, tmp_dir): - from rastervision.pytorch_learner.learner_pipeline import LearnerPipeline - return LearnerPipeline(self, tmp_dir) diff --git a/rastervision_pytorch_learner/rastervision/pytorch_learner/utils/utils.py b/rastervision_pytorch_learner/rastervision/pytorch_learner/utils/utils.py index 7c1f39fa75..e6120035bc 100644 --- a/rastervision_pytorch_learner/rastervision/pytorch_learner/utils/utils.py +++ b/rastervision_pytorch_learner/rastervision/pytorch_learner/utils/utils.py @@ -1,5 +1,5 @@ -from typing import (Any, Dict, Sequence, Tuple, Optional, Union, List, - Iterable, Container) +from typing import (TYPE_CHECKING, Any, Dict, Sequence, Tuple, Optional, Union, + List, Iterable, Container) from os.path import basename, join, isfile import logging @@ -14,8 +14,13 @@ import pandas as pd import onnxruntime as ort -from rastervision.pipeline.file_system import get_tmp_dir -from rastervision.pipeline.config import ConfigError +from rastervision.pipeline.file_system.utils import (file_exists, file_to_json, + get_tmp_dir) +from rastervision.pipeline.config import (build_config, Config, ConfigError, + upgrade_config) + +if TYPE_CHECKING: + from rastervision.pytorch_learner import LearnerConfig log = logging.getLogger(__name__) @@ -487,3 +492,23 @@ def __call__(self, x: Union[torch.Tensor, np.ndarray]) -> torch.Tensor: if isinstance(out, np.ndarray): out = torch.from_numpy(out) return out + + +def get_learner_config_from_bundle_dir( + model_bundle_dir: str) -> 'LearnerConfig': + config_path = join(model_bundle_dir, 'learner-config.json') + if file_exists(config_path): + cfg = Config.from_file(config_path) + else: + # backward compatibility + config_path = join(model_bundle_dir, 'pipeline-config.json') + if not file_exists(config_path): + raise FileNotFoundError( + 'Could not find a valid config file in the bundle.') + pipeline_cfg_dict = file_to_json(config_path) + cfg_dict = pipeline_cfg_dict['learner'] + cfg_dict['plugin_versions'] = pipeline_cfg_dict['plugin_versions'] + cfg_dict = upgrade_config(cfg_dict) + cfg_dict.pop('plugin_versions', None) + cfg = build_config(cfg_dict) + return cfg