fix: Fix deepcopy issues of PTQ calibrators

Signed-off-by: Dheeraj Peri <[email protected]>

Author: peri044

noxfile.py

@@ -342,8 +342,6 @@ def run_l2_trt_compatibility_tests(session):
     if not USE_HOST_DEPS:
         install_deps(session)
         install_torch_trt(session)
-    download_models(session)
-    train_model(session)
     run_trt_compatibility_tests(session)
     cleanup(session)
 

py/torch_tensorrt/ptq.py

@@ -55,6 +55,11 @@ def write_calibration_cache(self, cache):
     else:
         return b""
 
+# deepcopy (which involves pickling) is performed on the compile_spec internally during compilation.
+# We register this __reduce__ function for pickler to identity the calibrator object returned by DataLoaderCalibrator during deepcopy.
+# This should be the object's local name relative to the module https://docs.python.org/3/library/pickle.html#object.__reduce__
+def __reduce__(self):
+    return self.__class__.__name__
 
 class DataLoaderCalibrator(object):
     """
@@ -114,24 +119,25 @@ def __new__(cls, *args, **kwargs):
             "get_batch": get_cache_mode_batch if use_cache else get_batch,
             "read_calibration_cache": read_calibration_cache,
             "write_calibration_cache": write_calibration_cache,
+            "__reduce__": __reduce__ # used when you deepcopy the DataLoaderCalibrator object
         }
 
         # Using type metaclass to construct calibrator class based on algorithm type
         if algo_type == CalibrationAlgo.ENTROPY_CALIBRATION:
             return type(
-                "DataLoaderCalibrator", (_C.IInt8EntropyCalibrator,), attribute_mapping
+                "Int8EntropyCalibrator", (_C.IInt8EntropyCalibrator,), attribute_mapping
             )()
         elif algo_type == CalibrationAlgo.ENTROPY_CALIBRATION_2:
             return type(
-                "DataLoaderCalibrator", (_C.IInt8MinMaxCalibrator,), attribute_mapping
+                "Int8EntropyCalibrator2", (_C.IInt8EntropyCalibrator2,), attribute_mapping
             )()
         elif algo_type == CalibrationAlgo.LEGACY_CALIBRATION:
             return type(
-                "DataLoaderCalibrator", (_C.IInt8LegacyCalibrator,), attribute_mapping
+                "Int8LegacyCalibrator", (_C.IInt8LegacyCalibrator,), attribute_mapping
             )()
         elif algo_type == CalibrationAlgo.MINMAX_CALIBRATION:
             return type(
-                "DataLoaderCalibrator", (_C.IInt8MinMaxCalibrator,), attribute_mapping
+                "Int8MinMaxCalibrator", (_C.IInt8MinMaxCalibrator,), attribute_mapping
             )()
         else:
             log(

py/torch_tensorrt/ts/_compile_spec.py

@@ -226,13 +226,7 @@ def _parse_input_signature(input_signature: Any):
 
 def _parse_compile_spec(compile_spec_: Dict[str, Any]) -> _ts_C.CompileSpec:
     # TODO: Use deepcopy to support partial compilation of collections
-    compile_spec = {}
-    for k, v in compile_spec_.items():
-        if k != "calibrator":
-            compile_spec[k] = deepcopy(v)
-        else:
-            compile_spec[k] = v
-
+    compile_spec = deepcopy(compile_spec_)
     info = _ts_C.CompileSpec()
 
     if len(compile_spec["inputs"]) > 0:

tests/py/ptq/test_ptq_dataloader_calibrator.py

@@ -81,9 +81,6 @@ def test_compile_script(self):
             device=torch.device("cuda:0"),
         )
 
-        fp32_test_acc = compute_accuracy(self.testing_dataloader, self.model)
-        log(Level.Info, "[Pyt FP32] Test Acc: {:.2f}%".format(100 * fp32_test_acc))
-
         compile_spec = {
             "inputs": [torchtrt.Input([1, 3, 32, 32])],
             "enabled_precisions": {torch.float, torch.int8},
@@ -96,8 +93,11 @@ def test_compile_script(self):
                 "allow_gpu_fallback": False,
             },
         }
-
         trt_mod = torchtrt.ts.compile(self.model, **compile_spec)
+
+        fp32_test_acc = compute_accuracy(self.testing_dataloader, self.model)
+        log(Level.Info, "[Pyt FP32] Test Acc: {:.2f}%".format(100 * fp32_test_acc))
+
         int8_test_acc = compute_accuracy(self.testing_dataloader, trt_mod)
         log(Level.Info, "[TRT INT8] Test Acc: {:.2f}%".format(100 * int8_test_acc))
         acc_diff = fp32_test_acc - int8_test_acc