def convert_adaptive_avg_pool2d(ctx):
input = ctx.method_args[0]
output_size = get_arg(ctx, 'output_size', pos=1, default=0)
output = ctx.method_return
input_trt = trt_(ctx.network, input)
if isinstance(output_size, int):
output_size = (output_size, output_size)
output_size = tuple([-1 if not o else o for o in output_size])
if output_size[0] == 1 and output_size[1] == 1:
# use reduce as max pool2d
shape_length = len(input.shape)
axes = (1 << (shape_length - 1)) + (1 << (shape_length - 2))
keepdim = True
layer = ctx.network.add_reduce(input_trt, trt.ReduceOperation.AVG,
axes, keepdim)
output._trt = layer.get_output(0)
else:
plugin = create_adaptivepool_plugin(
"adaptive_avg_pool2d_" + str(id(input)),
output_size=output_size,
pooling_type=trt.PoolingType.AVERAGE)
layer = ctx.network.add_plugin_v2(inputs=[input_trt], plugin=plugin)
output._trt = layer.get_output(0)
def convert_adaptive_max_pool2d_by_input(ctx):
input = get_arg(ctx, 'x', pos=0, default=None)
shape_wraper = get_arg(ctx, 'shape_wraper', pos=1, default=None)
output = ctx.method_return
output_size = shape_wraper.shape
input_trt = trt_(ctx.network, input)
wrapshape_trt = trt_(ctx.network, shape_wraper)
plugin = create_adaptivepool_plugin('adaptive_max_pool2d_by_input_' +
str(id(input)),
output_size=output_size,
pooling_type=trt.PoolingType.MAX)
layer = ctx.network.add_plugin_v2(inputs=[input_trt, wrapshape_trt],
plugin=plugin)
output._trt = layer.get_output(0)
def convert_adaptive_avg_pool2d(ctx):
input = ctx.method_args[0]
output_size = get_arg(ctx, 'output_size', pos=1, default=0)
output = ctx.method_return
input_trt = trt_(ctx.network, input)
if isinstance(output_size, int):
output_size = (output_size, output_size)
output_size = tuple([-1 if not o else o for o in output_size])
plugin = create_adaptivepool_plugin("adaptive_avg_pool2d_" +
str(id(input)),
output_size=output_size,
pooling_type=trt.PoolingType.AVERAGE)
layer = ctx.network.add_plugin_v2(inputs=[input_trt], plugin=plugin)
output._trt = layer.get_output(0)
def convert_adaptive_max_pool1d(ctx):
input = ctx.method_args[0]
output_size = get_arg(ctx, 'output_size', pos=1, default=0)
output = ctx.method_return
input_trt = trt_(ctx.network, input)
if output_size == 1:
# use reduce as max pool2d
shape_length = len(input.shape)
axes = (1 << (shape_length - 1))
keepdim = True
layer = ctx.network.add_reduce(input_trt, trt.ReduceOperation.MAX,
axes, keepdim)
output._trt = layer.get_output(0)
else:
output_size = (output_size, 1)
# input.unsqueeze(-1)
layer = ctx.network.add_shuffle(input_trt)
layer.reshape_dims = (0, 0, 0, 1)
input_trt = layer.get_output(0)
# adaptive pool 2d
plugin = create_adaptivepool_plugin(
'adaptive_avg_pool2d_' + str(id(input)),
output_size=output_size,
pooling_type=trt.PoolingType.MAX)
layer = ctx.network.add_plugin_v2(inputs=[input_trt], plugin=plugin)
output_trt = layer.get_output(0)
layer = ctx.network.add_shuffle(output_trt)
layer.reshape_dims = (0, 0, 0)
output_trt = layer.get_output(0)
output._trt = output_trt