Add helper function to set layer precision depending on torch2trt kwargs.
@@ -635,3 +635,25 @@ def tensorrt_converter(method, is_real=True, enabled=True, imports=[]):
return pass_converter
return register_converter
+
+
+def set_layer_precision(ctx, layer):
+ # Supported TRT precisions as given by torch2trt_kwargs.
+ INT8_MODE = "int8_mode"
+ FP16_MODE = "fp16_mode"
+
+ # Check that args exist as expected in torch2trt_kwargs.
+ trt_kwargs = ctx.torch2trt_kwargs
+ assert INT8_MODE in trt_kwargs
+ assert FP16_MODE in trt_kwargs
+
+ is_int8 = trt_kwargs.get(INT8_MODE, False)
+ is_fp16 = trt_kwargs.get(FP16_MODE, False)
+
+ if is_int8:
+ layer.precision = trt.int8
+ layer.set_output_type(0, trt.int8)
+ elif is_fp16:
+ layer.precision = trt.float16
+ layer.set_output_type(0, trt.float16)
+
