Merge pull request #211 from Visual-Behavior/trt

add create_calibrator function

Author: thibo73800

alonet/torch2trt/__init__.py

@@ -1,11 +1,13 @@
 from .TRTEngineBuilder import TRTEngineBuilder
 from .TRTExecutor import TRTExecutor
 from .base_exporter import BaseTRTExporter
-from .utils import load_trt_custom_plugins
+from .utils import load_trt_custom_plugins, create_calibrator
+from .calibrator import DataBatchStreamer
 
 from alonet import ALONET_ROOT
 import os
 
+
 MS_DEFORM_IM2COL_PLUGIN_LIB = os.path.join(
     ALONET_ROOT, "torch2trt/plugins/ms_deform_im2col/build/libms_deform_im2col_trt.so"
 )

alonet/torch2trt/calibrator.py

@@ -1,12 +1,10 @@
+from aloscene.frame import Frame
+
 import os
 import torch
 import numpy as np
 import tensorrt as trt
-
 import pycuda.driver as cuda
-import pycuda.autoinit
-
-from aloscene.frame import Frame
 
 
 class DataBatchStreamer:

alonet/torch2trt/utils.py

@@ -1,4 +1,13 @@
+from alonet import ALONET_ROOT
+from alonet.torch2trt.calibrator import (
+    LegacyCalibrator, 
+    MinMaxCalibrator,
+    EntropyCalibrator,
+    EntropyCalibrator2,
+)
+
 import os
+import ctypes
 
 try:
     import pycuda.driver as cuda
@@ -10,8 +19,25 @@
     pass
 
 
-import ctypes
-from alonet import ALONET_ROOT
+def create_calibrator(name: str, *args, **kwargs):
+    """Creates calibrator from name
+
+    Parameters
+    ----------
+        name : str
+            Calibrator name
+    """
+    CALIBS = ["minmax", "entropy", "entropy2", "legacy"]
+    if name == "entropy2":
+        return EntropyCalibrator2(*args, **kwargs)
+    elif name == "entropy":
+        return EntropyCalibrator(*args, **kwargs)
+    elif name == "minmax":
+        return MinMaxCalibrator(*args, **kwargs)
+    elif name == "legacy":
+        return LegacyCalibrator(*args, **kwargs)
+    else:
+        raise AttributeError(f"Unknown calibrator name, should be one of {' '.join(CALIBS)}")
 
 
 def load_trt_custom_plugins(lib_path: str):