From 7cadf1ac7fe55af03dcb510e6f348e32b228461e Mon Sep 17 00:00:00 2001 From: irenab Date: Mon, 9 Sep 2024 11:53:55 +0300 Subject: [PATCH 1/4] add official support up to versions: torch 2.4, onnx 1.16, onnxruntime 1.19, onnxruntime_extensions 0.12 --- .github/workflows/ci.yaml | 34 +++++++++++++------ README.md | 6 ++-- sony_custom_layers/__init__.py | 4 +-- .../pytorch/object_detection/nms.py | 2 +- .../object_detection/nms_with_indices.py | 2 +- .../object_detection/test_multiclass_nms.py | 2 +- 6 files changed, 32 insertions(+), 18 deletions(-) diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 5cf2bed..da190a9 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -92,34 +92,48 @@ jobs: fail-fast: false matrix: # Test selected versions combinations as following: - # * latest/earliest torch+torchvision X latest/earliest ort X latest/earliest ort-ext + # * latest/earliest torch+torchvision X latest/earliest ort X latest/earliest ort-ext X latest/earliest onnx # * each remaining version is covered by at least one test # !!!NOTE!!! if newer versions are added, update sony_custom_layers.__init__ pinned_requirements py_ver: [ "3.8", "3.9", "3.10", "3.11" ] - torch_ver: ["2.0", "2.1", "2.2"] - ort_ver: ["1.15", "1.17"] - ort_ext_ver: ["0.8", "0.10"] + torch_ver: ["2.0", "2.1", "2.2", "2.3", "2.4"] + ort_ver: ["1.15", "1.19"] + ort_ext_ver: ["0.8", "0.12"] + onnx_ver: ["1.14", "1.16"] exclude: - py_ver: "3.11" # incompatible versions ort_ext_ver: "0.8" include: # torchvision ver is coupled to a specific torch ver + - torch_ver: "2.4" + torchvision_ver: "0.19" + - torch_ver: "2.3" + torchvision_ver: "0.18" - torch_ver: "2.2" torchvision_ver: "0.17" - onnx_ver: "1.14" - torch_ver: "2.1" torchvision_ver: "0.16" - onnx_ver: "1.15" - torch_ver: "2.0" torchvision_ver: "0.15" - onnx_ver: "1.15" - # non-covered versions (ort, ort-ext) + # non-covered versions (ort, ort-ext, onnx) - py_ver: "3.11" - torch_ver: "2.1" - torchvision_ver: "0.16" + torch_ver: "2.4" + torchvision_ver: "0.19" ort_ver: "1.16" ort_ext_ver: "0.9" + onnx_ver: "1.15" + - py_ver: "3.11" + torch_ver: "2.3" + torchvision_ver: "0.18" + ort_ver: "1.17" + ort_ext_ver: "0.10" onnx_ver: "1.14" + - py_ver: "3.11" + torch_ver: "2.2" + torchvision_ver: "0.17" + ort_ver: "1.18" + ort_ext_ver: "0.11" + onnx_ver: "1.16" steps: - name: Checkout diff --git a/README.md b/README.md index 7d00865..9f863bf 100644 --- a/README.md +++ b/README.md @@ -47,9 +47,9 @@ pip install sony-custom-layers[torch] #### PyTorch -| **Tested FW versions** | **Tested Python version** | **Serialization** | -|--------------------------------------------------------------------------------------------------------------------------|---------------------------|---------------------------------------------------------------------------------| -| torch 2.0-2.2
torchvision 0.15-0.17
onnxruntime 1.15-1.17
onnxruntime_extensions 0.8-0.10
onnx 1.14-1.15 | 3.8-3.11 | .onnx (via torch.onnx.export)
.pt2 (via torch.export.export, torch2.2 only) | +| **Tested FW versions** | **Tested Python version** | **Serialization** | +|--------------------------------------------------------------------------------------------------------------------------|---------------------------|-------------------------------------------------------------------------------------------------------------------| +| torch 2.0-2.4
torchvision 0.15-0.19
onnxruntime 1.15-1.19
onnxruntime_extensions 0.8-0.12
onnx 1.14-1.16 | 3.8-3.11 | .onnx (via torch.onnx.export)
.pt2 (via torch.export.export, torch2.2 only - discontinued for later versions) | ## API For sony-custom-layers API see https://sony.github.io/custom_layers diff --git a/sony_custom_layers/__init__.py b/sony_custom_layers/__init__.py index ebb6a6b..cf5f332 100644 --- a/sony_custom_layers/__init__.py +++ b/sony_custom_layers/__init__.py @@ -26,6 +26,6 @@ # pinned requirements of latest tested versions for extra_requires pinned_pip_requirements = { 'tf': ['tensorflow==2.15.*'], - 'torch': ['torch==2.2.*', 'torchvision==0.17.*'], - 'torch_ort': ['onnx==1.15.*', 'onnxruntime==1.17.*', 'onnxruntime_extensions==0.10.*'] + 'torch': ['torch==2.4.*', 'torchvision==0.19.*'], + 'torch_ort': ['onnx==1.16.*', 'onnxruntime==1.19.*', 'onnxruntime_extensions==0.12.*'] } diff --git a/sony_custom_layers/pytorch/object_detection/nms.py b/sony_custom_layers/pytorch/object_detection/nms.py index d348639..ea364b6 100644 --- a/sony_custom_layers/pytorch/object_detection/nms.py +++ b/sony_custom_layers/pytorch/object_detection/nms.py @@ -121,7 +121,7 @@ def _multiclass_nms_impl(boxes: torch.Tensor, scores: torch.Tensor, score_thresh op_qualname = register_op(MULTICLASS_NMS_TORCH_OP, schema, _multiclass_nms_impl) -if is_compatible('torch>=2.2'): +if is_compatible('torch~=2.2'): @torch.library.impl_abstract(op_qualname) def _multiclass_nms_meta(boxes: torch.Tensor, scores: torch.Tensor, score_threshold: float, iou_threshold: float, diff --git a/sony_custom_layers/pytorch/object_detection/nms_with_indices.py b/sony_custom_layers/pytorch/object_detection/nms_with_indices.py index d791ce3..828928f 100644 --- a/sony_custom_layers/pytorch/object_detection/nms_with_indices.py +++ b/sony_custom_layers/pytorch/object_detection/nms_with_indices.py @@ -125,7 +125,7 @@ def _multiclass_nms_with_indices_impl(boxes: torch.Tensor, scores: torch.Tensor, op_qualname = register_op(MULTICLASS_NMS_WITH_INDICES_TORCH_OP, schema, _multiclass_nms_with_indices_impl) -if is_compatible('torch>=2.2'): +if is_compatible('torch~=2.2'): @torch.library.impl_abstract(op_qualname) def _multiclass_nms_with_indices_meta(boxes: torch.Tensor, scores: torch.Tensor, score_threshold: float, diff --git a/sony_custom_layers/pytorch/tests/object_detection/test_multiclass_nms.py b/sony_custom_layers/pytorch/tests/object_detection/test_multiclass_nms.py index b58e188..2c25e39 100644 --- a/sony_custom_layers/pytorch/tests/object_detection/test_multiclass_nms.py +++ b/sony_custom_layers/pytorch/tests/object_detection/test_multiclass_nms.py @@ -206,7 +206,7 @@ def test_ort(self, dynamic_batch, tmpdir_factory, with_indices): """ exec_in_clean_process(code, check=True) - @pytest.mark.skipif(not is_compatible('torch>=2.2'), reason='unsupported') + @pytest.mark.skipif(not is_compatible('torch~=2.2.0'), reason='unsupported') @pytest.mark.parametrize('with_indices', [True, False]) def test_pt2_export(self, tmpdir_factory, with_indices): From 19307cd81c00364c8e6d5ef0f661410fdefe0963 Mon Sep 17 00:00:00 2001 From: irenab Date: Tue, 10 Sep 2024 14:16:45 +0300 Subject: [PATCH 2/4] remove support for new torch.export --- README.md | 8 +- sony_custom_layers/pytorch/__init__.py | 74 +++++++------------ .../pytorch/object_detection/nms.py | 19 +---- .../object_detection/nms_with_indices.py | 20 +---- .../object_detection/test_multiclass_nms.py | 45 +---------- 5 files changed, 34 insertions(+), 132 deletions(-) diff --git a/README.md b/README.md index 9f863bf..71c2e39 100644 --- a/README.md +++ b/README.md @@ -47,9 +47,9 @@ pip install sony-custom-layers[torch] #### PyTorch -| **Tested FW versions** | **Tested Python version** | **Serialization** | -|--------------------------------------------------------------------------------------------------------------------------|---------------------------|-------------------------------------------------------------------------------------------------------------------| -| torch 2.0-2.4
torchvision 0.15-0.19
onnxruntime 1.15-1.19
onnxruntime_extensions 0.8-0.12
onnx 1.14-1.16 | 3.8-3.11 | .onnx (via torch.onnx.export)
.pt2 (via torch.export.export, torch2.2 only - discontinued for later versions) | +| **Tested FW versions** | **Tested Python version** | **Serialization** | +|--------------------------------------------------------------------------------------------------------------------------|---------------------------|--------------------------------| +| torch 2.0-2.4
torchvision 0.15-0.19
onnxruntime 1.15-1.19
onnxruntime_extensions 0.8-0.12
onnx 1.14-1.16 | 3.8-3.11 | .onnx (via torch.onnx.export) | ## API For sony-custom-layers API see https://sony.github.io/custom_layers @@ -66,7 +66,7 @@ For PyTorch layers see No special handling is required for torch.onnx.export and onnx.load. -For OnnxRuntime / PT2 support see [load_custom_ops](https://sony.github.io/custom_layers/sony_custom_layers/pytorch.html#load_custom_ops) +For OnnxRuntime support see [load_custom_ops](https://sony.github.io/custom_layers/sony_custom_layers/pytorch.html#load_custom_ops) ## License [Apache License 2.0](LICENSE.md). diff --git a/sony_custom_layers/pytorch/__init__.py b/sony_custom_layers/pytorch/__init__.py index 2450fe7..0c172f0 100644 --- a/sony_custom_layers/pytorch/__init__.py +++ b/sony_custom_layers/pytorch/__init__.py @@ -29,62 +29,38 @@ from .object_detection import multiclass_nms_with_indices, NMSWithIndicesResults # noqa: E402 -def load_custom_ops(load_ort: bool = False, - ort_session_ops: Optional['ort.SessionOptions'] = None) -> Optional['ort.SessionOptions']: +def load_custom_ops(ort_session_ops: Optional['ort.SessionOptions'] = None) -> 'ort.SessionOptions': """ - Note: this must be run before inferring a model with SCL in onnxruntime. - To trigger ops registration in torch any import from sony_custom_layers.pytorch is technically sufficient, - In which case this is just a dummy API to prevent unused import (e.g. when loading exported pt2 model) - - Load custom ops for torch and, optionally, for onnxruntime. - If 'load_ort' is True or 'ort_session_ops' is passed, registers the custom ops implementation for onnxruntime, and - sets up the SessionOptions object for onnxruntime session. + Registers the custom ops implementation for onnxruntime, and sets up the SessionOptions object for onnxruntime + session. Args: - load_ort: whether to register the custom ops for onnxruntime. - ort_session_ops: SessionOptions object to register the custom ops library on. If None (and 'load_ort' is True), - creates a new object. + ort_session_ops: SessionOptions object to register the custom ops library on. If None, creates a new object. Returns: - SessionOptions object if ort registration was requested, otherwise None + SessionOptions object with registered custom ops. Example: - *ONNXRuntime*: - ``` - import onnxruntime as ort - from sony_custom_layers.pytorch import load_custom_ops - - so = load_custom_ops(load_ort=True) - session = ort.InferenceSession(model_path, sess_options=so) - session.run(...) - ``` - You can also pass your own SessionOptions object upon which to register the custom ops - ``` - load_custom_ops(ort_session_options=so) - ``` - - *PT2 model*:
- If sony_custom_layers.pytorch is already imported no action is needed. Otherwise, you can use: - - ``` - from sony_custom_layers.pytorch import load_custom_ops - load_custom_ops() - - prog = torch.export.load(model_path) - y = prog.module()(x) - ``` + ``` + import onnxruntime as ort + from sony_custom_layers.pytorch import load_custom_ops + + so = load_custom_ops() + session = ort.InferenceSession(model_path, sess_options=so) + session.run(...) + ``` + You can also pass your own SessionOptions object upon which to register the custom ops + ``` + load_custom_ops(ort_session_options=so) + ``` """ - if load_ort or ort_session_ops: - validate_installed_libraries(required_libraries['torch_ort']) + validate_installed_libraries(required_libraries['torch_ort']) - # trigger onnxruntime op registration - from .object_detection import nms_ort + # trigger onnxruntime op registration + from .object_detection import nms_ort - from onnxruntime_extensions import get_library_path - from onnxruntime import SessionOptions - ort_session_ops = ort_session_ops or SessionOptions() - ort_session_ops.register_custom_ops_library(get_library_path()) - return ort_session_ops - else: - # nothing really to do after import was triggered - return None + from onnxruntime_extensions import get_library_path + from onnxruntime import SessionOptions + ort_session_ops = ort_session_ops or SessionOptions() + ort_session_ops.register_custom_ops_library(get_library_path()) + return ort_session_ops diff --git a/sony_custom_layers/pytorch/object_detection/nms.py b/sony_custom_layers/pytorch/object_detection/nms.py index ea364b6..d4f186a 100644 --- a/sony_custom_layers/pytorch/object_detection/nms.py +++ b/sony_custom_layers/pytorch/object_detection/nms.py @@ -21,7 +21,6 @@ from sony_custom_layers.pytorch.custom_lib import register_op from sony_custom_layers.pytorch.object_detection.nms_common import _batch_multiclass_nms, SCORES, LABELS -from sony_custom_layers.util.import_util import is_compatible MULTICLASS_NMS_TORCH_OP = 'multiclass_nms' @@ -119,20 +118,4 @@ def _multiclass_nms_impl(boxes: torch.Tensor, scores: torch.Tensor, score_thresh "(Tensor boxes, Tensor scores, float score_threshold, float iou_threshold, SymInt max_detections) " "-> (Tensor, Tensor, Tensor, Tensor)") -op_qualname = register_op(MULTICLASS_NMS_TORCH_OP, schema, _multiclass_nms_impl) - -if is_compatible('torch~=2.2'): - - @torch.library.impl_abstract(op_qualname) - def _multiclass_nms_meta(boxes: torch.Tensor, scores: torch.Tensor, score_threshold: float, iou_threshold: float, - max_detections: int) -> NMSResults: - """ Registers torch op's abstract implementation. It specifies the properties of the output tensors. - Needed for torch.export """ - ctx = torch.library.get_ctx() - batch = ctx.new_dynamic_size() - return NMSResults( - torch.empty((batch, max_detections, 4)), - torch.empty((batch, max_detections)), - torch.empty((batch, max_detections), dtype=torch.int64), - torch.empty((batch, 1), dtype=torch.int64) - ) # yapf: disable +register_op(MULTICLASS_NMS_TORCH_OP, schema, _multiclass_nms_impl) diff --git a/sony_custom_layers/pytorch/object_detection/nms_with_indices.py b/sony_custom_layers/pytorch/object_detection/nms_with_indices.py index 828928f..2b52b3f 100644 --- a/sony_custom_layers/pytorch/object_detection/nms_with_indices.py +++ b/sony_custom_layers/pytorch/object_detection/nms_with_indices.py @@ -18,7 +18,6 @@ import torch from torch import Tensor -from sony_custom_layers.util.import_util import is_compatible from sony_custom_layers.pytorch.custom_lib import register_op from sony_custom_layers.pytorch.object_detection.nms_common import _batch_multiclass_nms, SCORES, LABELS, INDICES @@ -123,21 +122,4 @@ def _multiclass_nms_with_indices_impl(boxes: torch.Tensor, scores: torch.Tensor, "(Tensor boxes, Tensor scores, float score_threshold, float iou_threshold, SymInt max_detections) " "-> (Tensor, Tensor, Tensor, Tensor, Tensor)") -op_qualname = register_op(MULTICLASS_NMS_WITH_INDICES_TORCH_OP, schema, _multiclass_nms_with_indices_impl) - -if is_compatible('torch~=2.2'): - - @torch.library.impl_abstract(op_qualname) - def _multiclass_nms_with_indices_meta(boxes: torch.Tensor, scores: torch.Tensor, score_threshold: float, - iou_threshold: float, max_detections: int) -> NMSWithIndicesResults: - """ Registers torch op's abstract implementation. It specifies the properties of the output tensors. - Needed for torch.export """ - ctx = torch.library.get_ctx() - batch = ctx.new_dynamic_size() - return NMSWithIndicesResults( - torch.empty((batch, max_detections, 4)), - torch.empty((batch, max_detections)), - torch.empty((batch, max_detections), dtype=torch.int64), - torch.empty((batch, max_detections), dtype=torch.int64), - torch.empty((batch, 1), dtype=torch.int64) - ) # yapf: disable +register_op(MULTICLASS_NMS_WITH_INDICES_TORCH_OP, schema, _multiclass_nms_with_indices_impl) diff --git a/sony_custom_layers/pytorch/tests/object_detection/test_multiclass_nms.py b/sony_custom_layers/pytorch/tests/object_detection/test_multiclass_nms.py index 2c25e39..0f02979 100644 --- a/sony_custom_layers/pytorch/tests/object_detection/test_multiclass_nms.py +++ b/sony_custom_layers/pytorch/tests/object_detection/test_multiclass_nms.py @@ -25,7 +25,6 @@ from sony_custom_layers.pytorch import load_custom_ops from sony_custom_layers.pytorch.object_detection.nms_common import LABELS, INDICES, SCORES from sony_custom_layers.pytorch.tests.object_detection.test_nms_common import generate_random_inputs -from sony_custom_layers.util.import_util import is_compatible from sony_custom_layers.util.test_util import exec_in_clean_process @@ -187,7 +186,7 @@ def test_ort(self, dynamic_batch, tmpdir_factory, with_indices): batch = 5 if dynamic_batch else 1 boxes, scores = generate_random_inputs(batch=batch, n_boxes=n_boxes, n_classes=n_classes, seed=42) torch_res = model(boxes, scores) - so = load_custom_ops(load_ort=True) + so = load_custom_ops() session = ort.InferenceSession(path, sess_options=so) ort_res = session.run(output_names=None, input_feed={'boxes': boxes.numpy(), 'scores': scores.numpy()}) # this is just a sanity test on random data @@ -198,7 +197,8 @@ def test_ort(self, dynamic_batch, tmpdir_factory, with_indices): import onnxruntime as ort import numpy as np from sony_custom_layers.pytorch import load_custom_ops -so = load_custom_ops(load_ort=True) +so = ort.SessionOptions() +so = load_custom_ops(so) session = ort.InferenceSession('{path}', so) boxes = np.random.rand({batch}, {n_boxes}, 4).astype(np.float32) scores = np.random.rand({batch}, {n_boxes}, {n_classes}).astype(np.float32) @@ -206,45 +206,6 @@ def test_ort(self, dynamic_batch, tmpdir_factory, with_indices): """ exec_in_clean_process(code, check=True) - @pytest.mark.skipif(not is_compatible('torch~=2.2.0'), reason='unsupported') - @pytest.mark.parametrize('with_indices', [True, False]) - def test_pt2_export(self, tmpdir_factory, with_indices): - - model = NMS(score_threshold=0.5, iou_threshold=0.3, max_detections=100, with_indices=with_indices) - prog = torch.export.export(model, args=(torch.rand(1, 10, 4), torch.rand(1, 10, 5))) - nms_node = list(prog.graph.nodes)[2] - exp_target = torch.ops.sony.multiclass_nms_with_indices if with_indices else torch.ops.sony.multiclass_nms - assert nms_node.target == exp_target.default - val = nms_node.meta['val'] - assert val[0].shape[1:] == (100, 4) - assert val[1].shape[1:] == val[2].shape[1:] == (100, ) - assert val[2].dtype == torch.int64 - if with_indices: - assert val[3].shape[1:] == (100, ) - assert val[3].dtype == torch.int64 - assert val[4].shape[1:] == (1, ) - assert val[4].dtype == torch.int64 - else: - assert val[3].shape[1:] == (1, ) - assert val[3].dtype == torch.int64 - - boxes, scores = generate_random_inputs(1, 10, 5) - torch_out = model(boxes, scores) - prog_out = prog.module()(boxes, scores) - for i in range(len(torch_out)): - assert torch.allclose(torch_out[i], prog_out[i]), i - - path = str(tmpdir_factory.mktemp('nms').join(f'nms{with_indices}.pt2')) - torch.export.save(prog, path) - # check that exported program can be loaded in a clean env - code = f""" -import torch -import sony_custom_layers.pytorch -prog = torch.export.load('{path}') -prog.module()(torch.rand(1, 10, 4), torch.rand(1, 10, 5)) - """ - exec_in_clean_process(code, check=True) - def _export_onnx(self, nms_model, n_boxes, n_classes, path, dynamic_batch: bool, with_indices: bool): input_names = ['boxes', 'scores'] output_names = ['det_boxes', 'det_scores', 'det_labels', 'valid_dets'] From 577e10603cf098a5d61a2ebe1be4fbbc67e8cb89 Mon Sep 17 00:00:00 2001 From: irenab Date: Tue, 10 Sep 2024 14:22:29 +0300 Subject: [PATCH 3/4] increment minor instead of patch when creating first dev version after release --- .github/workflows/ci.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index da190a9..862383d 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -182,7 +182,7 @@ jobs: if v.is_devrelease: print(f'{v.major}.{v.minor}.{v.micro}.dev{v.dev + 1}') else: - print(f'{v.major}.{v.minor}.{v.micro + 1}.dev0') + print(f'{v.major}.{v.minor+1}.{v.micro}.dev0') ") fi echo "new_ver=${new_version}" >> $GITHUB_ENV From 34a498a75d7d4282fbf8d598d3524787d0d125cb Mon Sep 17 00:00:00 2001 From: "GitHub Actions [Bot]" Date: Tue, 10 Sep 2024 11:29:39 +0000 Subject: [PATCH 4/4] Update docs [Created by Github action] --- docs/sony_custom_layers/pytorch.html | 515 ++++++++++++--------------- 1 file changed, 224 insertions(+), 291 deletions(-) diff --git a/docs/sony_custom_layers/pytorch.html b/docs/sony_custom_layers/pytorch.html index 2a525d1..7bcd0f3 100644 --- a/docs/sony_custom_layers/pytorch.html +++ b/docs/sony_custom_layers/pytorch.html @@ -154,65 +154,41 @@

29from .object_detection import multiclass_nms_with_indices, NMSWithIndicesResults # noqa: E402 30 31 -32def load_custom_ops(load_ort: bool = False, -33 ort_session_ops: Optional['ort.SessionOptions'] = None) -> Optional['ort.SessionOptions']: -34 """ -35 Note: this must be run before inferring a model with SCL in onnxruntime. -36 To trigger ops registration in torch any import from sony_custom_layers.pytorch is technically sufficient, -37 In which case this is just a dummy API to prevent unused import (e.g. when loading exported pt2 model) -38 -39 Load custom ops for torch and, optionally, for onnxruntime. -40 If 'load_ort' is True or 'ort_session_ops' is passed, registers the custom ops implementation for onnxruntime, and -41 sets up the SessionOptions object for onnxruntime session. +32def load_custom_ops(ort_session_ops: Optional['ort.SessionOptions'] = None) -> 'ort.SessionOptions': +33 """ +34 Registers the custom ops implementation for onnxruntime, and sets up the SessionOptions object for onnxruntime +35 session. +36 +37 Args: +38 ort_session_ops: SessionOptions object to register the custom ops library on. If None, creates a new object. +39 +40 Returns: +41 SessionOptions object with registered custom ops. 42 -43 Args: -44 load_ort: whether to register the custom ops for onnxruntime. -45 ort_session_ops: SessionOptions object to register the custom ops library on. If None (and 'load_ort' is True), -46 creates a new object. +43 Example: +44 ``` +45 import onnxruntime as ort +46 from sony_custom_layers.pytorch import load_custom_ops 47 -48 Returns: -49 SessionOptions object if ort registration was requested, otherwise None -50 -51 Example: -52 *ONNXRuntime*: -53 ``` -54 import onnxruntime as ort -55 from sony_custom_layers.pytorch import load_custom_ops -56 -57 so = load_custom_ops(load_ort=True) -58 session = ort.InferenceSession(model_path, sess_options=so) -59 session.run(...) -60 ``` -61 You can also pass your own SessionOptions object upon which to register the custom ops -62 ``` -63 load_custom_ops(ort_session_options=so) -64 ``` -65 -66 *PT2 model*:<br> -67 If sony_custom_layers.pytorch is already imported no action is needed. Otherwise, you can use: -68 -69 ``` -70 from sony_custom_layers.pytorch import load_custom_ops -71 load_custom_ops() -72 -73 prog = torch.export.load(model_path) -74 y = prog.module()(x) -75 ``` -76 """ -77 if load_ort or ort_session_ops: -78 validate_installed_libraries(required_libraries['torch_ort']) -79 -80 # trigger onnxruntime op registration -81 from .object_detection import nms_ort -82 -83 from onnxruntime_extensions import get_library_path -84 from onnxruntime import SessionOptions -85 ort_session_ops = ort_session_ops or SessionOptions() -86 ort_session_ops.register_custom_ops_library(get_library_path()) -87 return ort_session_ops -88 else: -89 # nothing really to do after import was triggered -90 return None +48 so = load_custom_ops() +49 session = ort.InferenceSession(model_path, sess_options=so) +50 session.run(...) +51 ``` +52 You can also pass your own SessionOptions object upon which to register the custom ops +53 ``` +54 load_custom_ops(ort_session_options=so) +55 ``` +56 """ +57 validate_installed_libraries(required_libraries['torch_ort']) +58 +59 # trigger onnxruntime op registration +60 from .object_detection import nms_ort +61 +62 from onnxruntime_extensions import get_library_path +63 from onnxruntime import SessionOptions +64 ort_session_ops = ort_session_ops or SessionOptions() +65 ort_session_ops.register_custom_ops_library(get_library_path()) +66 return ort_session_ops @@ -228,49 +204,49 @@

- 
54def multiclass_nms(boxes, scores, score_threshold: float, iou_threshold: float, max_detections: int) -> NMSResults:
-55    """
-56    Multi-class non-maximum suppression.
-57    Detections are returned in descending order of their scores.
-58    The output tensors always contain a fixed number of detections, as defined by 'max_detections'.
-59    If fewer detections are selected, the output tensors are zero-padded up to 'max_detections'.
-60
-61    If you also require the input indices of the selected boxes, see `multiclass_nms_with_indices`.
-62
-63    Args:
-64        boxes (Tensor): Input boxes with shape [batch, n_boxes, 4], specified in corner coordinates
-65                        (x_min, y_min, x_max, y_max). Agnostic to the x-y axes order.
-66        scores (Tensor): Input scores with shape [batch, n_boxes, n_classes].
-67        score_threshold (float): The score threshold. Candidates with scores below the threshold are discarded.
-68        iou_threshold (float): The Intersection Over Union (IOU) threshold for boxes overlap.
-69        max_detections (int): The number of detections to return.
-70
-71    Returns:
-72        'NMSResults' named tuple:
-73        - boxes: The selected boxes with shape [batch, max_detections, 4].
-74        - scores: The corresponding scores in descending order with shape [batch, max_detections].
-75        - labels: The labels for each box with shape [batch, max_detections].
-76        - n_valid: The number of valid detections out of 'max_detections' with shape [batch, 1]
-77
-78    Raises:
-79        ValueError: If provided with invalid arguments or input tensors with unexpected or non-matching shapes.
-80
-81    Example:
-82        ```
-83        from sony_custom_layers.pytorch import multiclass_nms
-84
-85        # batch size=1, 1000 boxes, 50 classes
-86        boxes = torch.rand(1, 1000, 4)
-87        scores = torch.rand(1, 1000, 50)
-88        res = multiclass_nms(boxes,
-89                             scores,
-90                             score_threshold=0.1,
-91                             iou_threshold=0.6,
-92                             max_detections=300)
-93        # res.boxes, res.scores, res.labels, res.n_valid
-94        ```
-95    """
-96    return NMSResults(*torch.ops.sony.multiclass_nms(boxes, scores, score_threshold, iou_threshold, max_detections))
+            
53def multiclass_nms(boxes, scores, score_threshold: float, iou_threshold: float, max_detections: int) -> NMSResults:
+54    """
+55    Multi-class non-maximum suppression.
+56    Detections are returned in descending order of their scores.
+57    The output tensors always contain a fixed number of detections, as defined by 'max_detections'.
+58    If fewer detections are selected, the output tensors are zero-padded up to 'max_detections'.
+59
+60    If you also require the input indices of the selected boxes, see `multiclass_nms_with_indices`.
+61
+62    Args:
+63        boxes (Tensor): Input boxes with shape [batch, n_boxes, 4], specified in corner coordinates
+64                        (x_min, y_min, x_max, y_max). Agnostic to the x-y axes order.
+65        scores (Tensor): Input scores with shape [batch, n_boxes, n_classes].
+66        score_threshold (float): The score threshold. Candidates with scores below the threshold are discarded.
+67        iou_threshold (float): The Intersection Over Union (IOU) threshold for boxes overlap.
+68        max_detections (int): The number of detections to return.
+69
+70    Returns:
+71        'NMSResults' named tuple:
+72        - boxes: The selected boxes with shape [batch, max_detections, 4].
+73        - scores: The corresponding scores in descending order with shape [batch, max_detections].
+74        - labels: The labels for each box with shape [batch, max_detections].
+75        - n_valid: The number of valid detections out of 'max_detections' with shape [batch, 1]
+76
+77    Raises:
+78        ValueError: If provided with invalid arguments or input tensors with unexpected or non-matching shapes.
+79
+80    Example:
+81        ```
+82        from sony_custom_layers.pytorch import multiclass_nms
+83
+84        # batch size=1, 1000 boxes, 50 classes
+85        boxes = torch.rand(1, 1000, 4)
+86        scores = torch.rand(1, 1000, 50)
+87        res = multiclass_nms(boxes,
+88                             scores,
+89                             score_threshold=0.1,
+90                             iou_threshold=0.6,
+91                             max_detections=300)
+92        # res.boxes, res.scores, res.labels, res.n_valid
+93        ```
+94    """
+95    return NMSResults(*torch.ops.sony.multiclass_nms(boxes, scores, score_threshold, iou_threshold, max_detections))
 

 
 
@@ -342,26 +318,26 @@ 
Example:
- 
32class NMSResults(NamedTuple):
-33    """ Container for non-maximum suppression results """
-34    boxes: Tensor
-35    scores: Tensor
-36    labels: Tensor
-37    n_valid: Tensor
-38
-39    # Note: convenience methods below are replicated in each Results container, since NamedTuple supports neither adding
-40    # new fields in derived classes nor multiple inheritance, and we want it to behave like a tuple, so no dataclasses.
-41    def detach(self) -> 'NMSResults':
-42        """ Detach all tensors and return a new object """
-43        return self.apply(lambda t: t.detach())
-44
-45    def cpu(self) -> 'NMSResults':
-46        """ Move all tensors to cpu and return a new object """
-47        return self.apply(lambda t: t.cpu())
-48
-49    def apply(self, f: Callable[[Tensor], Tensor]) -> 'NMSResults':
-50        """ Apply any function to all tensors and return a new object """
-51        return self.__class__(*[f(t) for t in self])
+            
31class NMSResults(NamedTuple):
+32    """ Container for non-maximum suppression results """
+33    boxes: Tensor
+34    scores: Tensor
+35    labels: Tensor
+36    n_valid: Tensor
+37
+38    # Note: convenience methods below are replicated in each Results container, since NamedTuple supports neither adding
+39    # new fields in derived classes nor multiple inheritance, and we want it to behave like a tuple, so no dataclasses.
+40    def detach(self) -> 'NMSResults':
+41        """ Detach all tensors and return a new object """
+42        return self.apply(lambda t: t.detach())
+43
+44    def cpu(self) -> 'NMSResults':
+45        """ Move all tensors to cpu and return a new object """
+46        return self.apply(lambda t: t.cpu())
+47
+48    def apply(self, f: Callable[[Tensor], Tensor]) -> 'NMSResults':
+49        """ Apply any function to all tensors and return a new object """
+50        return self.__class__(*[f(t) for t in self])
 

 
 
@@ -446,9 +422,9 @@ 
Example:
- 
41    def detach(self) -> 'NMSResults':
-42        """ Detach all tensors and return a new object """
-43        return self.apply(lambda t: t.detach())
+            
40    def detach(self) -> 'NMSResults':
+41        """ Detach all tensors and return a new object """
+42        return self.apply(lambda t: t.detach())
 

 
 
@@ -468,9 +444,9 @@ 
Example:
- 
45    def cpu(self) -> 'NMSResults':
-46        """ Move all tensors to cpu and return a new object """
-47        return self.apply(lambda t: t.cpu())
+            
44    def cpu(self) -> 'NMSResults':
+45        """ Move all tensors to cpu and return a new object """
+46        return self.apply(lambda t: t.cpu())
 

 
 
@@ -490,9 +466,9 @@ 
Example:
- 
49    def apply(self, f: Callable[[Tensor], Tensor]) -> 'NMSResults':
-50        """ Apply any function to all tensors and return a new object """
-51        return self.__class__(*[f(t) for t in self])
+            
48    def apply(self, f: Callable[[Tensor], Tensor]) -> 'NMSResults':
+49        """ Apply any function to all tensors and return a new object """
+50        return self.__class__(*[f(t) for t in self])
 

 
 
@@ -513,52 +489,52 @@ 
Example:
- 
54def multiclass_nms_with_indices(boxes, scores, score_threshold: float, iou_threshold: float,
-55                                max_detections: int) -> NMSWithIndicesResults:
-56    """
-57    Multi-class non-maximum suppression with indices.
-58    Detections are returned in descending order of their scores.
-59    The output tensors always contain a fixed number of detections, as defined by 'max_detections'.
-60    If fewer detections are selected, the output tensors are zero-padded up to 'max_detections'.
-61
-62    This operator is identical to `multiclass_nms` except that is also outputs the input indices of the selected boxes.
-63
-64    Args:
-65        boxes (Tensor): Input boxes with shape [batch, n_boxes, 4], specified in corner coordinates
-66                        (x_min, y_min, x_max, y_max). Agnostic to the x-y axes order.
-67        scores (Tensor): Input scores with shape [batch, n_boxes, n_classes].
-68        score_threshold (float): The score threshold. Candidates with scores below the threshold are discarded.
-69        iou_threshold (float): The Intersection Over Union (IOU) threshold for boxes overlap.
-70        max_detections (int): The number of detections to return.
-71
-72    Returns:
-73        'NMSWithIndicesResults' named tuple:
-74        - boxes: The selected boxes with shape [batch, max_detections, 4].
-75        - scores: The corresponding scores in descending order with shape [batch, max_detections].
-76        - labels: The labels for each box with shape [batch, max_detections].
-77        - indices: Indices of the input boxes that have been selected.
-78        - n_valid: The number of valid detections out of 'max_detections' with shape [batch, 1]
-79
-80    Raises:
-81        ValueError: If provided with invalid arguments or input tensors with unexpected or non-matching shapes.
-82
-83    Example:
-84        ```
-85        from sony_custom_layers.pytorch import multiclass_nms_with_indices
-86
-87        # batch size=1, 1000 boxes, 50 classes
-88        boxes = torch.rand(1, 1000, 4)
-89        scores = torch.rand(1, 1000, 50)
-90        res = multiclass_nms_with_indices(boxes,
-91                                          scores,
-92                                          score_threshold=0.1,
-93                                          iou_threshold=0.6,
-94                                          max_detections=300)
-95        # res.boxes, res.scores, res.labels, res.indices, res.n_valid
-96        ```
-97    """
-98    return NMSWithIndicesResults(
-99        *torch.ops.sony.multiclass_nms_with_indices(boxes, scores, score_threshold, iou_threshold, max_detections))
+            
53def multiclass_nms_with_indices(boxes, scores, score_threshold: float, iou_threshold: float,
+54                                max_detections: int) -> NMSWithIndicesResults:
+55    """
+56    Multi-class non-maximum suppression with indices.
+57    Detections are returned in descending order of their scores.
+58    The output tensors always contain a fixed number of detections, as defined by 'max_detections'.
+59    If fewer detections are selected, the output tensors are zero-padded up to 'max_detections'.
+60
+61    This operator is identical to `multiclass_nms` except that is also outputs the input indices of the selected boxes.
+62
+63    Args:
+64        boxes (Tensor): Input boxes with shape [batch, n_boxes, 4], specified in corner coordinates
+65                        (x_min, y_min, x_max, y_max). Agnostic to the x-y axes order.
+66        scores (Tensor): Input scores with shape [batch, n_boxes, n_classes].
+67        score_threshold (float): The score threshold. Candidates with scores below the threshold are discarded.
+68        iou_threshold (float): The Intersection Over Union (IOU) threshold for boxes overlap.
+69        max_detections (int): The number of detections to return.
+70
+71    Returns:
+72        'NMSWithIndicesResults' named tuple:
+73        - boxes: The selected boxes with shape [batch, max_detections, 4].
+74        - scores: The corresponding scores in descending order with shape [batch, max_detections].
+75        - labels: The labels for each box with shape [batch, max_detections].
+76        - indices: Indices of the input boxes that have been selected.
+77        - n_valid: The number of valid detections out of 'max_detections' with shape [batch, 1]
+78
+79    Raises:
+80        ValueError: If provided with invalid arguments or input tensors with unexpected or non-matching shapes.
+81
+82    Example:
+83        ```
+84        from sony_custom_layers.pytorch import multiclass_nms_with_indices
+85
+86        # batch size=1, 1000 boxes, 50 classes
+87        boxes = torch.rand(1, 1000, 4)
+88        scores = torch.rand(1, 1000, 50)
+89        res = multiclass_nms_with_indices(boxes,
+90                                          scores,
+91                                          score_threshold=0.1,
+92                                          iou_threshold=0.6,
+93                                          max_detections=300)
+94        # res.boxes, res.scores, res.labels, res.indices, res.n_valid
+95        ```
+96    """
+97    return NMSWithIndicesResults(
+98        *torch.ops.sony.multiclass_nms_with_indices(boxes, scores, score_threshold, iou_threshold, max_detections))
 

 
 
@@ -631,27 +607,27 @@ 
Example:
- 
31class NMSWithIndicesResults(NamedTuple):
-32    """ Container for non-maximum suppression with indices results """
-33    boxes: Tensor
-34    scores: Tensor
-35    labels: Tensor
-36    indices: Tensor
-37    n_valid: Tensor
-38
-39    # Note: convenience methods below are replicated in each Results container, since NamedTuple supports neither adding
-40    # new fields in derived classes nor multiple inheritance, and we want it to behave like a tuple, so no dataclasses.
-41    def detach(self) -> 'NMSWithIndicesResults':
-42        """ Detach all tensors and return a new object """
-43        return self.apply(lambda t: t.detach())
-44
-45    def cpu(self) -> 'NMSWithIndicesResults':
-46        """ Move all tensors to cpu and return a new object """
-47        return self.apply(lambda t: t.cpu())
-48
-49    def apply(self, f: Callable[[Tensor], Tensor]) -> 'NMSWithIndicesResults':
-50        """ Apply any function to all tensors and return a new object """
-51        return self.__class__(*[f(t) for t in self])
+            
30class NMSWithIndicesResults(NamedTuple):
+31    """ Container for non-maximum suppression with indices results """
+32    boxes: Tensor
+33    scores: Tensor
+34    labels: Tensor
+35    indices: Tensor
+36    n_valid: Tensor
+37
+38    # Note: convenience methods below are replicated in each Results container, since NamedTuple supports neither adding
+39    # new fields in derived classes nor multiple inheritance, and we want it to behave like a tuple, so no dataclasses.
+40    def detach(self) -> 'NMSWithIndicesResults':
+41        """ Detach all tensors and return a new object """
+42        return self.apply(lambda t: t.detach())
+43
+44    def cpu(self) -> 'NMSWithIndicesResults':
+45        """ Move all tensors to cpu and return a new object """
+46        return self.apply(lambda t: t.cpu())
+47
+48    def apply(self, f: Callable[[Tensor], Tensor]) -> 'NMSWithIndicesResults':
+49        """ Apply any function to all tensors and return a new object """
+50        return self.__class__(*[f(t) for t in self])
 

 
 
@@ -749,9 +725,9 @@ 
Example:
- 
41    def detach(self) -> 'NMSWithIndicesResults':
-42        """ Detach all tensors and return a new object """
-43        return self.apply(lambda t: t.detach())
+            
40    def detach(self) -> 'NMSWithIndicesResults':
+41        """ Detach all tensors and return a new object """
+42        return self.apply(lambda t: t.detach())
 

 
 
@@ -771,9 +747,9 @@ 
Example:
- 
45    def cpu(self) -> 'NMSWithIndicesResults':
-46        """ Move all tensors to cpu and return a new object """
-47        return self.apply(lambda t: t.cpu())
+            
44    def cpu(self) -> 'NMSWithIndicesResults':
+45        """ Move all tensors to cpu and return a new object """
+46        return self.apply(lambda t: t.cpu())
 

 
 
@@ -793,9 +769,9 @@ 
Example:
- 
49    def apply(self, f: Callable[[Tensor], Tensor]) -> 'NMSWithIndicesResults':
-50        """ Apply any function to all tensors and return a new object """
-51        return self.__class__(*[f(t) for t in self])
+            
48    def apply(self, f: Callable[[Tensor], Tensor]) -> 'NMSWithIndicesResults':
+49        """ Apply any function to all tensors and return a new object """
+50        return self.__class__(*[f(t) for t in self])
 

 
 
@@ -810,123 +786,80 @@ 
Example:

 

             
         def
-        load_custom_ops(	load_ort: bool = False,	ort_session_ops: Optional[onnxruntime.capi.onnxruntime_pybind11_state.SessionOptions] = None) -> Optional[onnxruntime.capi.onnxruntime_pybind11_state.SessionOptions]:
+        load_custom_ops(	ort_session_ops: Optional[onnxruntime.capi.onnxruntime_pybind11_state.SessionOptions] = None) -> onnxruntime.capi.onnxruntime_pybind11_state.SessionOptions:
 
                 
 
     

     
-            
33def load_custom_ops(load_ort: bool = False,
-34                    ort_session_ops: Optional['ort.SessionOptions'] = None) -> Optional['ort.SessionOptions']:
-35    """
-36    Note: this must be run before inferring a model with SCL in onnxruntime.
-37    To trigger ops registration in torch any import from sony_custom_layers.pytorch is technically sufficient,
-38    In which case this is just a dummy API to prevent unused import (e.g. when loading exported pt2 model)
-39
-40    Load custom ops for torch and, optionally, for onnxruntime.
-41    If 'load_ort' is True or 'ort_session_ops' is passed, registers the custom ops implementation for onnxruntime, and
-42    sets up the SessionOptions object for onnxruntime session.
+            
33def load_custom_ops(ort_session_ops: Optional['ort.SessionOptions'] = None) -> 'ort.SessionOptions':
+34    """
+35    Registers the custom ops implementation for onnxruntime, and sets up the SessionOptions object for onnxruntime
+36    session.
+37
+38    Args:
+39        ort_session_ops: SessionOptions object to register the custom ops library on. If None, creates a new object.
+40
+41    Returns:
+42        SessionOptions object with registered custom ops.
 43
-44    Args:
-45        load_ort: whether to register the custom ops for onnxruntime.
-46        ort_session_ops: SessionOptions object to register the custom ops library on. If None (and 'load_ort' is True),
-47                        creates a new object.
+44    Example:
+45        ```
+46        import onnxruntime as ort
+47        from sony_custom_layers.pytorch import load_custom_ops
 48
-49    Returns:
-50        SessionOptions object if ort registration was requested, otherwise None
-51
-52    Example:
-53        *ONNXRuntime*:
-54            ```
-55            import onnxruntime as ort
-56            from sony_custom_layers.pytorch import load_custom_ops
-57
-58            so = load_custom_ops(load_ort=True)
-59            session = ort.InferenceSession(model_path, sess_options=so)
-60            session.run(...)
-61            ```
-62            You can also pass your own SessionOptions object upon which to register the custom ops
-63            ```
-64            load_custom_ops(ort_session_options=so)
-65            ```
-66
-67        *PT2 model*:<br>
-68            If sony_custom_layers.pytorch is already imported no action is needed. Otherwise, you can use:
-69
-70            ```
-71            from sony_custom_layers.pytorch import load_custom_ops
-72            load_custom_ops()
-73
-74            prog = torch.export.load(model_path)
-75            y = prog.module()(x)
-76            ```
-77    """
-78    if load_ort or ort_session_ops:
-79        validate_installed_libraries(required_libraries['torch_ort'])
-80
-81        # trigger onnxruntime op registration
-82        from .object_detection import nms_ort
-83
-84        from onnxruntime_extensions import get_library_path
-85        from onnxruntime import SessionOptions
-86        ort_session_ops = ort_session_ops or SessionOptions()
-87        ort_session_ops.register_custom_ops_library(get_library_path())
-88        return ort_session_ops
-89    else:
-90        # nothing really to do after import was triggered
-91        return None
+49        so = load_custom_ops()
+50        session = ort.InferenceSession(model_path, sess_options=so)
+51        session.run(...)
+52        ```
+53        You can also pass your own SessionOptions object upon which to register the custom ops
+54        ```
+55        load_custom_ops(ort_session_options=so)
+56        ```
+57    """
+58    validate_installed_libraries(required_libraries['torch_ort'])
+59
+60    # trigger onnxruntime op registration
+61    from .object_detection import nms_ort
+62
+63    from onnxruntime_extensions import get_library_path
+64    from onnxruntime import SessionOptions
+65    ort_session_ops = ort_session_ops or SessionOptions()
+66    ort_session_ops.register_custom_ops_library(get_library_path())
+67    return ort_session_ops
 

 
 
-            
Note: this must be run before inferring a model with SCL in onnxruntime.
-To trigger ops registration in torch any import from sony_custom_layers.pytorch is technically sufficient,
-In which case this is just a dummy API to prevent unused import (e.g. when loading exported pt2 model)

-
-
Load custom ops for torch and, optionally, for onnxruntime.
-If 'load_ort' is True or 'ort_session_ops' is passed, registers the custom ops implementation for onnxruntime, and
-sets up the SessionOptions object for onnxruntime session.

+            
Registers the custom ops implementation for onnxruntime, and sets up the SessionOptions object for onnxruntime
+session.

 
 
Arguments:

 
 
    
-
  • load_ort:  whether to register the custom ops for onnxruntime.
    • 
-
  • ort_session_ops:  SessionOptions object to register the custom ops library on. If None (and 'load_ort' is True),
-creates a new object.
    • 
+
  • ort_session_ops:  SessionOptions object to register the custom ops library on. If None, creates a new object.
    • 
 

 
 
Returns:

 
 

-  
SessionOptions object if ort registration was requested, otherwise None

+  
SessionOptions object with registered custom ops.

 

 
 
Example:

 
 

-  
ONNXRuntime:

-  
-  
import onnxruntime as ort
+
import onnxruntime as ort
 from sony_custom_layers.pytorch import load_custom_ops
 
-so = load_custom_ops(load_ort=True)
+so = load_custom_ops()
 session = ort.InferenceSession(model_path, sess_options=so)
 session.run(...)
-

+

   
   
You can also pass your own SessionOptions object upon which to register the custom ops

-  
-  
load_custom_ops(ort_session_options=so)
-

-  
-  
PT2 model:

-      If sony_custom_layers.pytorch is already imported no action is needed. Otherwise, you can use:

-  
-  
from sony_custom_layers.pytorch import load_custom_ops
-load_custom_ops()
 
-prog = torch.export.load(model_path)
-y = prog.module()(x)
-

+
load_custom_ops(ort_session_options=so)
+