Merge pull request #216 from Visual-Behavior/init_args

fix: replace max calibration batches default value with None
Visual-Behavior · Sep 9, 2022 · 712859b · 712859b
2 parents 23033db + ce01409
commit 712859b
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/alonet/torch2trt/base_exporter.py b/alonet/torch2trt/base_exporter.py
@@ -402,7 +402,7 @@ def add_argparse_args(parent_parser):
         parser.add_argument("--verbose", action="store_true", help="Helpful when debugging")
         parser.add_argument("--profiling_verbosity", default=0, type=int, help="Helpful when profiling the engine (default: %(default)s)")
         parser.add_argument("--calibration_batch_size", type=int, default=8, help="Calibration data batch size (default: %(default)s)")
-        parser.add_argument("--limit_calibration_batches", type=int, default=10, help="Limits number of batches (default: %(default)s)")
+        parser.add_argument("--limit_calibration_batches", type=int, default=None, help="Limits number of batches (default: %(default)s)")
         parser.add_argument("--cache_file", type=str, default="calib.bin", help="Path to caliaration cache file (default: %(default)s)")
         parser.add_argument(
             "--calibrator", 

diff --git a/alonet/torch2trt/calibrator.py b/alonet/torch2trt/calibrator.py
@@ -65,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,
@@ -76,7 +74,8 @@ def __init__(
             ):
         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
@@ -105,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):
@@ -127,7 +127,7 @@ def next_(self):
             return None
 
     def __len__(self):
-        return max_batch
+        return self.max_batch
 
 
 class BaseCalibrator: