def avg_pool2d(g,
input,
kernel_size,
stride,
padding,
ceil_mode,
count_include_pad,
divisor_override=None):
if input not in sym_help._quantized_ops:
from torch.onnx.symbolic_opset9 import avg_pool2d
return avg_pool2d(g, input, kernel_size, stride, padding, ceil_mode,
count_include_pad, divisor_override)
kwargs = {
"strides_i": stride,
"pads_i": padding + padding,
"kernel_i": kernel_size[0],
"order_s": "NHWC",
"Y_scale_f": input.node()["Y_scale"],
"Y_zero_point_i": input.node()["Y_zero_point"],
}
input = nchw2nhwc(g, input)
output = g.op("_caffe2::Int8AveragePool", input, **kwargs)
output = nhwc2nchw(g, output)
sym_help._quantized_ops.add(output)
return output
def avg_pool2d(
g,
input,
kernel_size,
stride,
padding,
ceil_mode,
count_include_pad,
divisor_override=None,
):
if input not in symbolic_helper._quantized_ops:
return opset9.avg_pool2d(
g,
input,
kernel_size,
stride,
padding,
ceil_mode,
count_include_pad,
divisor_override,
)
kwargs = {
"strides_i": stride,
"pads_i": padding + padding,
"kernel_i": kernel_size[0],
"order_s": "NHWC",
"Y_scale_f": symbolic_helper._node_get(input.node(), "Y_scale"),
"Y_zero_point_i": symbolic_helper._node_get(input.node(),
"Y_zero_point"),
}
input = nchw2nhwc(g, input)
output = g.op("_caffe2::Int8AveragePool", input, **kwargs)
output = nhwc2nchw(g, output)
symbolic_helper._quantized_ops.add(output)
return output