From 659a977f76827542844db544ef0491277036dc03 Mon Sep 17 00:00:00 2001 From: Thibault Date: Wed, 24 Aug 2022 09:00:20 +0200 Subject: [PATCH 1/2] Optional trt import --- alonet/torch2trt/calibrator.py | 86 +++++++++++++++++++--------------- alonet/torch2trt/utils.py | 24 ++++++---- 2 files changed, 64 insertions(+), 46 deletions(-) diff --git a/alonet/torch2trt/calibrator.py b/alonet/torch2trt/calibrator.py index 5bd7bfaf..1eea6152 100644 --- a/alonet/torch2trt/calibrator.py +++ b/alonet/torch2trt/calibrator.py @@ -3,8 +3,16 @@ import os import torch import numpy as np -import tensorrt as trt -import pycuda.driver as cuda + +try: + import tensorrt as trt + import pycuda.driver as cuda + prod_package_error = None +except Exception as e: + trt = None + cuda = None + prod_package_error = e + pass class DataBatchStreamer: @@ -74,6 +82,7 @@ def __init__( limit_batches=None, **kwargs, ): + if prod_package_error is not None: raise prod_package_error for sample in dataset[0]: if not isinstance(sample, (torch.Tensor, np.ndarray, Frame)): raise TypeError(f"unknown sample type, expected samples to be instance of {' or '.join(self.FTYPES)} got {sample.__class__.__name__} instead") @@ -195,45 +204,46 @@ def free(self): d.free() -class MinMaxCalibrator(BaseCalibrator, trt.IInt8MinMaxCalibrator): - def __init__( - self, - data_streamer, - cache_file, - **kwargs, - ): - trt.IInt8MinMaxCalibrator.__init__(self) - super(MinMaxCalibrator, self).__init__(data_streamer=data_streamer, cache_file=cache_file, **kwargs) +if trt is not None: + class MinMaxCalibrator(BaseCalibrator, trt.IInt8MinMaxCalibrator): + def __init__( + self, + data_streamer, + cache_file, + **kwargs, + ): + trt.IInt8MinMaxCalibrator.__init__(self) + super(MinMaxCalibrator, self).__init__(data_streamer=data_streamer, cache_file=cache_file, **kwargs) -class LegacyCalibrator(BaseCalibrator, trt.IInt8LegacyCalibrator): - def __init__( - self, - data_streamer, - cache_file, - **kwargs, - ): - trt.IInt8LegacyCalibrator.__init__(self) - super(LegacyCalibrator, self).__init__(data_streamer=data_streamer, cache_file=cache_file, **kwargs) + class LegacyCalibrator(BaseCalibrator, trt.IInt8LegacyCalibrator): + def __init__( + self, + data_streamer, + cache_file, + **kwargs, + ): + trt.IInt8LegacyCalibrator.__init__(self) + super(LegacyCalibrator, self).__init__(data_streamer=data_streamer, cache_file=cache_file, **kwargs) -class EntropyCalibrator(BaseCalibrator, trt.IInt8EntropyCalibrator): - def __init__( - self, - data_streamer, - cache_file, - **kwargs, - ): - trt.IInt8EntropyCalibrator.__init__(self) - super(EntropyCalibrator, self).__init__(data_streamer=data_streamer, cache_file=cache_file, **kwargs) + class EntropyCalibrator(BaseCalibrator, trt.IInt8EntropyCalibrator): + def __init__( + self, + data_streamer, + cache_file, + **kwargs, + ): + trt.IInt8EntropyCalibrator.__init__(self) + super(EntropyCalibrator, self).__init__(data_streamer=data_streamer, cache_file=cache_file, **kwargs) -class EntropyCalibrator2(BaseCalibrator, trt.IInt8EntropyCalibrator2): - def __init__( - self, - data_streamer, - cache_file, - **kwargs, - ): - trt.IInt8EntropyCalibrator2.__init__(self) - super(EntropyCalibrator2, self).__init__(data_streamer=data_streamer, cache_file=cache_file, **kwargs) + class EntropyCalibrator2(BaseCalibrator, trt.IInt8EntropyCalibrator2): + def __init__( + self, + data_streamer, + cache_file, + **kwargs, + ): + trt.IInt8EntropyCalibrator2.__init__(self) + super(EntropyCalibrator2, self).__init__(data_streamer=data_streamer, cache_file=cache_file, **kwargs) diff --git a/alonet/torch2trt/utils.py b/alonet/torch2trt/utils.py index d32fa0c2..4ce5f06d 100644 --- a/alonet/torch2trt/utils.py +++ b/alonet/torch2trt/utils.py @@ -1,10 +1,15 @@ from alonet import ALONET_ROOT -from alonet.torch2trt.calibrator import ( - LegacyCalibrator, - MinMaxCalibrator, - EntropyCalibrator, - EntropyCalibrator2, -) + +try: + from alonet.torch2trt.calibrator import ( + LegacyCalibrator, + MinMaxCalibrator, + EntropyCalibrator, + EntropyCalibrator2, + ) + calibrator_import_error = None +except Exception as e: + calibrator_import_error = e import os import ctypes @@ -13,9 +18,9 @@ import pycuda.driver as cuda import onnx_graphsurgeon as gs import tensorrt as trt - prod_package_error = None -except Exception as prod_package_error: +except Exception as e: + prod_package_error = e pass @@ -27,6 +32,9 @@ def create_calibrator(name: str, *args, **kwargs): name : str Calibrator name """ + if calibrator_import_error is not None: + raise calibrator_import_error + CALIBS = ["minmax", "entropy", "entropy2", "legacy"] if name == "entropy2": return EntropyCalibrator2(*args, **kwargs) From 7ad7477454eb1c836385cad701dd9344daac66a0 Mon Sep 17 00:00:00 2001 From: Thibault Neveu Date: Tue, 4 Oct 2022 08:04:16 +0200 Subject: [PATCH 2/2] remove only fron init --- alonet/__init__.py | 2 +- alonet/torch2trt/calibrator.py | 96 +++++++++++++++------------------- alonet/torch2trt/utils.py | 24 +++------ 3 files changed, 52 insertions(+), 70 deletions(-) diff --git a/alonet/__init__.py b/alonet/__init__.py index cb9f618a..510f1064 100644 --- a/alonet/__init__.py +++ b/alonet/__init__.py @@ -11,4 +11,4 @@ from . import detr_panoptic from . import deformable_detr_panoptic -from . import torch2trt + diff --git a/alonet/torch2trt/calibrator.py b/alonet/torch2trt/calibrator.py index 1eea6152..fe7a0a23 100644 --- a/alonet/torch2trt/calibrator.py +++ b/alonet/torch2trt/calibrator.py @@ -3,16 +3,8 @@ import os import torch import numpy as np - -try: - import tensorrt as trt - import pycuda.driver as cuda - prod_package_error = None -except Exception as e: - trt = None - cuda = None - prod_package_error = e - pass +import tensorrt as trt +import pycuda.driver as cuda class DataBatchStreamer: @@ -73,8 +65,6 @@ class DataBatchStreamer: >>> s_dataStreamer = DataBatchStreamer(dataset=s_calib) >>> m_dataStreamer = DataBatchStreamer(dataset=m_calib) """ - FTYPES = ["torch.Tensor", "ndarray", "aloscene.Frame"] - def __init__( self, dataset=None, @@ -82,10 +72,10 @@ def __init__( limit_batches=None, **kwargs, ): - if prod_package_error is not None: raise prod_package_error for sample in dataset[0]: if not isinstance(sample, (torch.Tensor, np.ndarray, Frame)): - raise TypeError(f"unknown sample type, expected samples to be instance of {' or '.join(self.FTYPES)} got {sample.__class__.__name__} instead") + ftypes = ["torch.Tensor", "ndarray", "aloscene.Frame"] + raise TypeError(f"unknown sample type, expected samples to be instance of {' or '.join(ftypes)} got {sample.__class__.__name__} instead") self.batch_idx = 0 self.dataset = dataset @@ -114,7 +104,8 @@ def convert_frame(frame): elif isinstance(frame, np.ndarray): pass else: - raise TypeError(f"Unknown sample type, expected samples to be instance of {' or '.join(self.FTYPES)} got {frame.__class__.__name__}.") + ftypes = ["torch.Tensor", "ndarray", "aloscene.Frame"] + raise TypeError(f"Unknown sample type, expected samples to be instance of {' or '.join(ftypes)} got {frame.__class__.__name__}.") return frame def next_(self): @@ -136,7 +127,7 @@ def next_(self): return None def __len__(self): - return max_batch + return self.max_batch class BaseCalibrator: @@ -204,46 +195,45 @@ def free(self): d.free() -if trt is not None: - class MinMaxCalibrator(BaseCalibrator, trt.IInt8MinMaxCalibrator): - def __init__( - self, - data_streamer, - cache_file, - **kwargs, - ): - trt.IInt8MinMaxCalibrator.__init__(self) - super(MinMaxCalibrator, self).__init__(data_streamer=data_streamer, cache_file=cache_file, **kwargs) +class MinMaxCalibrator(BaseCalibrator, trt.IInt8MinMaxCalibrator): + def __init__( + self, + data_streamer, + cache_file, + **kwargs, + ): + trt.IInt8MinMaxCalibrator.__init__(self) + super(MinMaxCalibrator, self).__init__(data_streamer=data_streamer, cache_file=cache_file, **kwargs) - class LegacyCalibrator(BaseCalibrator, trt.IInt8LegacyCalibrator): - def __init__( - self, - data_streamer, - cache_file, - **kwargs, - ): - trt.IInt8LegacyCalibrator.__init__(self) - super(LegacyCalibrator, self).__init__(data_streamer=data_streamer, cache_file=cache_file, **kwargs) +class LegacyCalibrator(BaseCalibrator, trt.IInt8LegacyCalibrator): + def __init__( + self, + data_streamer, + cache_file, + **kwargs, + ): + trt.IInt8LegacyCalibrator.__init__(self) + super(LegacyCalibrator, self).__init__(data_streamer=data_streamer, cache_file=cache_file, **kwargs) - class EntropyCalibrator(BaseCalibrator, trt.IInt8EntropyCalibrator): - def __init__( - self, - data_streamer, - cache_file, - **kwargs, - ): - trt.IInt8EntropyCalibrator.__init__(self) - super(EntropyCalibrator, self).__init__(data_streamer=data_streamer, cache_file=cache_file, **kwargs) +class EntropyCalibrator(BaseCalibrator, trt.IInt8EntropyCalibrator): + def __init__( + self, + data_streamer, + cache_file, + **kwargs, + ): + trt.IInt8EntropyCalibrator.__init__(self) + super(EntropyCalibrator, self).__init__(data_streamer=data_streamer, cache_file=cache_file, **kwargs) - class EntropyCalibrator2(BaseCalibrator, trt.IInt8EntropyCalibrator2): - def __init__( - self, - data_streamer, - cache_file, - **kwargs, - ): - trt.IInt8EntropyCalibrator2.__init__(self) - super(EntropyCalibrator2, self).__init__(data_streamer=data_streamer, cache_file=cache_file, **kwargs) +class EntropyCalibrator2(BaseCalibrator, trt.IInt8EntropyCalibrator2): + def __init__( + self, + data_streamer, + cache_file, + **kwargs, + ): + trt.IInt8EntropyCalibrator2.__init__(self) + super(EntropyCalibrator2, self).__init__(data_streamer=data_streamer, cache_file=cache_file, **kwargs) diff --git a/alonet/torch2trt/utils.py b/alonet/torch2trt/utils.py index 4ce5f06d..d32fa0c2 100644 --- a/alonet/torch2trt/utils.py +++ b/alonet/torch2trt/utils.py @@ -1,15 +1,10 @@ from alonet import ALONET_ROOT - -try: - from alonet.torch2trt.calibrator import ( - LegacyCalibrator, - MinMaxCalibrator, - EntropyCalibrator, - EntropyCalibrator2, - ) - calibrator_import_error = None -except Exception as e: - calibrator_import_error = e +from alonet.torch2trt.calibrator import ( + LegacyCalibrator, + MinMaxCalibrator, + EntropyCalibrator, + EntropyCalibrator2, +) import os import ctypes @@ -18,9 +13,9 @@ import pycuda.driver as cuda import onnx_graphsurgeon as gs import tensorrt as trt + prod_package_error = None -except Exception as e: - prod_package_error = e +except Exception as prod_package_error: pass @@ -32,9 +27,6 @@ def create_calibrator(name: str, *args, **kwargs): name : str Calibrator name """ - if calibrator_import_error is not None: - raise calibrator_import_error - CALIBS = ["minmax", "entropy", "entropy2", "legacy"] if name == "entropy2": return EntropyCalibrator2(*args, **kwargs)