/usr/local/lib64/python3.6/site-packages/torch/fx/experimental/fx2trt/converters
NameSizeModeActions
__pycache__/-0755rm
acc_ops_converters.py542610644editdlrm
activation.py19880644editdlrm
adaptive_avgpool.py10970644editdlrm
add.py22610644editdlrm
batchnorm.py17370644editdlrm
convolution.py33680644editdlrm
helper_functions.py18260644editdlrm
linear.py27840644editdlrm
maxpool.py13740644editdlrm
mul.py14350644editdlrm
quantization.py19990644editdlrm
transformation.py14390644editdlrm
__init__.py4390644editdlrm
Edit: /usr/local/lib64/python3.6/site-packages/torch/fx/experimental/fx2trt/converters/maxpool.py (1374B)
import torch import tensorrt as trt from torch.fx.experimental.fx2trt.fx2trt import tensorrt_converter from .helper_functions import mark_as_int8_layer, extend_attr_to_tuple def common_maxpool(network, mod, dimension, input_val, layer_name): kernel_size = extend_attr_to_tuple(mod, "kernel_size", dimension) stride = extend_attr_to_tuple(mod, "stride", dimension) padding = extend_attr_to_tuple(mod, "padding", dimension) layer = network.add_pooling( input=input_val, type=trt.PoolingType.MAX, window_size=kernel_size) layer.stride = stride layer.padding = padding layer.name = layer_name if mod.ceil_mode: layer.padding_mode = trt.PaddingMode.EXPLICIT_ROUND_UP if input_val.dynamic_range: mark_as_int8_layer(layer, input_val.dynamic_range) return layer.get_output(0) @tensorrt_converter(torch.nn.modules.pooling.MaxPool2d) def maxpool2d(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"MaxPool2d received input {input_val} that is not part " "of the TensorRT region!") return common_maxpool(network, submod, dimension=2, input_val=input_val, layer_name=layer_name)