FM / converters

/usr/local/lib64/python3.6/site-packages/torch/fx/experimental/fx2trt/converters/

/usr/local/lib64/python3.6/site-packages/torch/fx/experimental/fx2trt/converters

Name	Size	Mode	Actions
__pycache__/	-	0755	rm
acc_ops_converters.py	54261	0644	edit dl rm
activation.py	1988	0644	edit dl rm
adaptive_avgpool.py	1097	0644	edit dl rm
add.py	2261	0644	edit dl rm
batchnorm.py	1737	0644	edit dl rm
convolution.py	3368	0644	edit dl rm
helper_functions.py	1826	0644	edit dl rm
linear.py	2784	0644	edit dl rm
maxpool.py	1374	0644	edit dl rm
mul.py	1435	0644	edit dl rm
quantization.py	1999	0644	edit dl rm
transformation.py	1439	0644	edit dl rm
__init__.py	439	0644	edit dl rm

Edit: /usr/local/lib64/python3.6/site-packages/torch/fx/experimental/fx2trt/converters/quantization.py (1999B)

import torch
import tensorrt as trt
from torch.fx.experimental.fx2trt.fx2trt import tensorrt_converter

from .helper_functions import get_dyn_range, get_inputs_from_args_and_kwargs

quantize_per_tensor_inputs = ["input", "scale", "zero_point", "dtype"]

@tensorrt_converter("dequantize")
@tensorrt_converter(torch.dequantize)
@tensorrt_converter(torch.nn.quantized.modules.DeQuantize)
def dequantize(network, submod, args, kwargs, layer_name):
    input_val = args[0]

if not isinstance(input_val, trt.tensorrt.ITensor):
        raise RuntimeError(f'Dequantize received input {input_val} that is not part '
                           'of the TensorRT region!')

return input_val

@tensorrt_converter(torch.quantize_per_tensor)
@tensorrt_converter(torch.nn.quantized.modules.Quantize)
def quantize(network, submod, args, kwargs, layer_name):
    # If submod is not nn.Module then it's quantize_per_tensor
    if not isinstance(submod, torch.nn.Module):
        input_val, scale, zero_point, dtype = get_inputs_from_args_and_kwargs(args, kwargs, quantize_per_tensor_inputs)
    else:
        input_val = args[0]
        scale = submod.scale
        zero_point = submod.zero_point
        dtype = submod.dtype

if not isinstance(input_val, trt.tensorrt.ITensor):
        raise RuntimeError(f'Quantize received input {input_val} that is not part '
                           'of the TensorRT region!')

if dtype != torch.quint8:
        raise RuntimeError(f"Only support torch.quint8 quantized type for activation, get {dtype}.")

input_val.dynamic_range = get_dyn_range(scale, zero_point, dtype)
    return input_val

@tensorrt_converter(torch.nn.modules.linear.Identity)
def identity(network, submod, args, kwargs, layer_name):
    input_val = kwargs["input"]

if not isinstance(input_val, trt.tensorrt.ITensor):
        raise RuntimeError(f'Identity received input {input_val} that is not part '
                           'of the TensorRT region!')

return input_val

import torch import tensorrt as trt from torch.fx.experimental.fx2trt.fx2trt import tensorrt_converter from .helper_functions import get_dyn_range, get_inputs_from_args_and_kwargs quantize_per_tensor_inputs = ["input", "scale", "zero_point", "dtype"] @tensorrt_converter("dequantize") @tensorrt_converter(torch.dequantize) @tensorrt_converter(torch.nn.quantized.modules.DeQuantize) def dequantize(network, submod, args, kwargs, layer_name): input_val = args[0] if not isinstance(input_val, trt.tensorrt.ITensor): raise RuntimeError(f'Dequantize received input {input_val} that is not part ' 'of the TensorRT region!') return input_val @tensorrt_converter(torch.quantize_per_tensor) @tensorrt_converter(torch.nn.quantized.modules.Quantize) def quantize(network, submod, args, kwargs, layer_name): # If submod is not nn.Module then it's quantize_per_tensor if not isinstance(submod, torch.nn.Module): input_val, scale, zero_point, dtype = get_inputs_from_args_and_kwargs(args, kwargs, quantize_per_tensor_inputs) else: input_val = args[0] scale = submod.scale zero_point = submod.zero_point dtype = submod.dtype if not isinstance(input_val, trt.tensorrt.ITensor): raise RuntimeError(f'Quantize received input {input_val} that is not part ' 'of the TensorRT region!') if dtype != torch.quint8: raise RuntimeError(f"Only support torch.quint8 quantized type for activation, get {dtype}.") input_val.dynamic_range = get_dyn_range(scale, zero_point, dtype) return input_val @tensorrt_converter(torch.nn.modules.linear.Identity) def identity(network, submod, args, kwargs, layer_name): input_val = kwargs["input"] if not isinstance(input_val, trt.tensorrt.ITensor): raise RuntimeError(f'Identity received input {input_val} that is not part ' 'of the TensorRT region!') return input_val