switch to builder config
@@ -517,6 +517,7 @@ def torch2trt(module,
logger = trt.Logger(log_level)
builder = trt.Builder(logger)
+ config = builder.create_builder_config()
if isinstance(inputs, list):
inputs = tuple(inputs)
@@ -555,10 +556,16 @@ def torch2trt(module,
outputs = (outputs,)
ctx.mark_outputs(outputs, output_names)
- builder.max_workspace_size = max_workspace_size
- builder.fp16_mode = fp16_mode
+ # set max workspace size
+ config.max_workspace_size = max_workspace_size
+
+ if fp16_mode:
+ config.set_flag(trt.BuilderFlag.FP16)
+
builder.max_batch_size = max_batch_size
- builder.strict_type_constraints = strict_type_constraints
+
+ if strict_type_constraints:
+ config.set_flag(trt.BuilderFlag.STRICT_TYPES)
if int8_mode:
@@ -566,16 +573,17 @@ def torch2trt(module,
if int8_calib_dataset is None:
int8_calib_dataset = TensorBatchDataset(inputs_in)
- builder.int8_mode = True
+ config.set_flag(trt.BuilderFlag.INT8)
#Making sure not to run calibration with QAT mode on
if not 'qat_mode' in kwargs:
# @TODO(jwelsh): Should we set batch_size=max_batch_size? Need to investigate memory consumption
- builder.int8_calibrator = DatasetCalibrator(
+ calibrator = DatasetCalibrator(
inputs, int8_calib_dataset, batch_size=int8_calib_batch_size, algorithm=int8_calib_algorithm
)
+ config.int8_calibrator = calibrator
- engine = builder.build_cuda_engine(network)
+ engine = builder.build_engine(network, config)
module_trt = TRTModule(engine, input_names, output_names)
