remove onnxruntime-gpu requirement

azavea · Feb 28, 2024 · f4fdad0 · f4fdad0
1 parent e525054
commit f4fdad0
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/rastervision_pytorch_learner/rastervision/pytorch_learner/utils/utils.py b/rastervision_pytorch_learner/rastervision/pytorch_learner/utils/utils.py
@@ -12,14 +12,14 @@
 from albumentations.core.transforms_interface import ImageOnlyTransform
 import cv2
 import pandas as pd
-import onnxruntime as ort
 
 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:
+    import onnxruntime as ort
     from rastervision.pytorch_learner import LearnerConfig
 
 log = logging.getLogger(__name__)
@@ -456,7 +456,7 @@ class ONNXRuntimeAdapter:
     also outputs PyTorch Tensors.
     """
 
-    def __init__(self, ort_session: ort.InferenceSession) -> None:
+    def __init__(self, ort_session: 'ort.InferenceSession') -> None:
         """Constructor.
 
         Args:
@@ -482,6 +482,8 @@ def from_file(cls, path: str, providers: Optional[List[str]] = None
         Returns:
             ONNXRuntimeAdapter: An ONNXRuntimeAdapter instance.
         """
+        import onnxruntime as ort
+
         if providers is None:
             providers = ort.get_available_providers()
             log.info(f'Using ONNX execution providers: {providers}')

diff --git a/rastervision_pytorch_learner/requirements.txt b/rastervision_pytorch_learner/requirements.txt
@@ -13,4 +13,3 @@ opencv-python-headless==4.9.0.80
 matplotlib==3.8.2
 tqdm==4.66.1
 onnx==1.15.0
-onnxruntime-gpu==1.17