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/activation.py (1988B)

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

from .helper_functions import mark_as_int8_layer

def common_activation(network, mod, input_val, activation_type, activation_dyn_range_fn, layer_name):
    layer = network.add_activation(
        input=input_val, type=activation_type)
    layer.name = layer_name

if input_val.dynamic_range:
        dyn_range = activation_dyn_range_fn(input_val.dynamic_range)
        mark_as_int8_layer(layer, dyn_range)

return layer.get_output(0)

@tensorrt_converter(torch.nn.functional.relu)
@tensorrt_converter(torch.nn.modules.activation.ReLU)
def relu(network, submod, args, kwargs, layer_name):
    # args/kwargs should have already been normalized to kwargs
    assert len(args) == 0
    input_val = kwargs["input"]

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

def activation_dyn_range_fn(dyn_range):
        return max(0, dyn_range[0]), max(0, dyn_range[1])

return common_activation(network, submod, input_val, trt.ActivationType.RELU, activation_dyn_range_fn, layer_name)

@tensorrt_converter(torch.nn.modules.activation.Sigmoid)
def sigmoid(network, submod, args, kwargs, layer_name):
    # args/kwargs should have already been normalized to kwargs
    assert len(args) == 0
    input_val = kwargs['input']

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

def activation_dyn_range_fn(dyn_range):
        def sigmoid_fn(x):
            return 1 / (1 + np.exp(-x))
        return sigmoid_fn(dyn_range[0]), sigmoid_fn(dyn_range[1])

return common_activation(network, submod, input_val, trt.ActivationType.SIGMOID, activation_dyn_range_fn, layer_name)

import torch import numpy as np import tensorrt as trt from torch.fx.experimental.fx2trt.fx2trt import tensorrt_converter from .helper_functions import mark_as_int8_layer def common_activation(network, mod, input_val, activation_type, activation_dyn_range_fn, layer_name): layer = network.add_activation( input=input_val, type=activation_type) layer.name = layer_name if input_val.dynamic_range: dyn_range = activation_dyn_range_fn(input_val.dynamic_range) mark_as_int8_layer(layer, dyn_range) return layer.get_output(0) @tensorrt_converter(torch.nn.functional.relu) @tensorrt_converter(torch.nn.modules.activation.ReLU) def relu(network, submod, args, kwargs, layer_name): # args/kwargs should have already been normalized to kwargs assert len(args) == 0 input_val = kwargs["input"] if not isinstance(input_val, trt.tensorrt.ITensor): raise RuntimeError(f"ReLU received input {input_val} that is not part " "of the TensorRT region!") def activation_dyn_range_fn(dyn_range): return max(0, dyn_range[0]), max(0, dyn_range[1]) return common_activation(network, submod, input_val, trt.ActivationType.RELU, activation_dyn_range_fn, layer_name) @tensorrt_converter(torch.nn.modules.activation.Sigmoid) def sigmoid(network, submod, args, kwargs, layer_name): # args/kwargs should have already been normalized to kwargs assert len(args) == 0 input_val = kwargs['input'] if not isinstance(input_val, trt.tensorrt.ITensor): raise RuntimeError(f'Sigmoid received input {input_val} that is not part ' 'of the TensorRT region!') def activation_dyn_range_fn(dyn_range): def sigmoid_fn(x): return 1 / (1 + np.exp(-x)) return sigmoid_fn(dyn_range[0]), sigmoid_fn(dyn_range[1]) return common_activation(network, submod, input_val, trt.ActivationType.SIGMOID, activation_dyn_range_fn, layer_name)