diff --git a/include/core/graph_handler.h b/include/core/graph_handler.h index 7cc9351d..4f9ca929 100644 --- a/include/core/graph_handler.h +++ b/include/core/graph_handler.h @@ -44,6 +44,11 @@ class GraphHandlerObj { Tensor scale, Tensor bias, float momentum, float eps, bool training); + Tensor maxPool(Tensor input, Tensor output, int kh, int kw, int dh, int dw, + int ph, int pw, int sh, int sw); + Tensor avgPool(Tensor input, Tensor output, int kh, int kw, int dh, int dw, + int ph, int pw, int sh, int sw); + Tensor add(Tensor a, Tensor b, Tensor c); Tensor sub(Tensor a, Tensor b, Tensor c); Tensor mul(Tensor a, Tensor b, Tensor c); diff --git a/pyinfinitensor/src/pyinfinitensor/onnx.py b/pyinfinitensor/src/pyinfinitensor/onnx.py index 6f72dc94..32c946d0 100644 --- a/pyinfinitensor/src/pyinfinitensor/onnx.py +++ b/pyinfinitensor/src/pyinfinitensor/onnx.py @@ -46,6 +46,56 @@ def from_onnx(model: onnx.ModelProto): tensors[node.output[0]] = handler.batchNorm( input, output, mean, var, scale, bias, momentum, eps, training != 0 ) + elif node.op_type == "MaxPool": + attributes = _parse_attribute( + node, + { + "kernel_shape": None, + "dilations": [1, 1], + "pads": [0, 0], + "strides": [1, 1], + }, + ) + (k, d, p, s) = ( + attributes[name] + for name in ["kernel_shape", "dilations", "pads", "strides"] + ) + tensors[node.output[0]] = handler.maxPool( + tensors[node.input[0]], + tensors.get(node.output[0]), + k[0], + k[1], + d[0], + d[1], + p[0], + p[1], + s[0], + s[1], + ) + elif node.op_type == "AveragePool": + attributes = _parse_attribute( + node, + { + "kernel_shape": None, + "pads": [0, 0], + "strides": [1, 1], + }, + ) + (k, p, s) = ( + attributes[name] for name in ["kernel_shape", "pads", "strides"] + ) + tensors[node.output[0]] = handler.maxPool( + tensors[node.input[0]], + tensors.get(node.output[0]), + k[0], + k[1], + 1, + 1, + p[0], + p[1], + s[0], + s[1], + ) elif node.op_type == "Add": tensors[node.output[0]] = handler.add( tensors[node.input[0]], diff --git a/pyinfinitensor/tests/test_onnx.py b/pyinfinitensor/tests/test_onnx.py index 804540d7..f296e0fa 100644 --- a/pyinfinitensor/tests/test_onnx.py +++ b/pyinfinitensor/tests/test_onnx.py @@ -55,9 +55,38 @@ class TestStringMethods(unittest.TestCase): name="batchNormalization", ) make_and_import_model( - make_graph([batch_norm], "batch_norm", [x, scale, b, mean, var], [y]) + make_graph([batch_norm], "batchNorm", [x, scale, b, mean, var], [y]) ) + def test_max_pool(self): + x = make_tensor_value_info("x", TensorProto.UINT32, [1, 64, 162, 162]) + y = make_tensor_value_info("y", TensorProto.UINT32, [1, 64, 80, 80]) + pool = make_node( + "MaxPool", + ["x"], + ["y"], + kernel_shape=[3, 3], + dilations=[1, 1], + pads=[0, 0], + strides=[2, 2], + name="maxPool", + ) + make_and_import_model(make_graph([pool], "maxPool", [x], [y])) + + def test_avg_pool(self): + x = make_tensor_value_info("x", TensorProto.UINT32, [1, 64, 162, 162]) + y = make_tensor_value_info("y", TensorProto.UINT32, [1, 64, 80, 80]) + pool = make_node( + "AveragePool", + ["x"], + ["y"], + kernel_shape=[3, 3], + pads=[0, 0], + strides=[2, 2], + name="avgPool", + ) + make_and_import_model(make_graph([pool], "avgPool", [x], [y])) + def test_add(self): a = make_tensor_value_info("a", TensorProto.FLOAT, [1, 3, 5, 7]) b = make_tensor_value_info("b", TensorProto.FLOAT, [1, 3, 5, 7]) diff --git a/src/core/graph_handler.cc b/src/core/graph_handler.cc index 3840ca16..afde6690 100644 --- a/src/core/graph_handler.cc +++ b/src/core/graph_handler.cc @@ -4,9 +4,11 @@ #include "operators/element_wise.h" #include "operators/gather.h" #include "operators/matmul.h" +#include "operators/pooling.h" #include "operators/reduce_mean.h" #include "operators/reshape.h" #include "operators/unary.h" + namespace infini { static DataType dtype_repr_convert(int); @@ -46,6 +48,35 @@ Tensor GraphHandlerObj::batchNorm(Tensor input, Tensor output, Tensor mean, } } +Tensor GraphHandlerObj::maxPool(Tensor input, Tensor output, int kh, int kw, + int dh, int dw, int ph, int pw, int sh, + int sw) { + if (output) { + g->addOpWithOutputs(std::move(input), output, kh, kw, dh, + dw, ph, pw, sh, sw); + return output; + } else { + return g + ->addOp(std::move(input), output, kh, kw, dh, dw, ph, + pw, sh, sw) + ->getOutput(); + } +} +Tensor GraphHandlerObj::avgPool(Tensor input, Tensor output, int kh, int kw, + int dh, int dw, int ph, int pw, int sh, + int sw) { + if (output) { + g->addOpWithOutputs(std::move(input), output, kh, kw, dh, + dw, ph, pw, sh, sw); + return output; + } else { + return g + ->addOp(std::move(input), output, kh, kw, dh, dw, ph, + pw, sh, sw) + ->getOutput(); + } +} + // see operators/element_wise.h #define DEFINE_ELEMENT_WISE_METHOD(name, obj) \ Tensor GraphHandlerObj::name(Tensor a, Tensor b, Tensor c) { \ diff --git a/src/ffi/ffi_infinitensor.cc b/src/ffi/ffi_infinitensor.cc index d1998d6c..8d8f9ac3 100644 --- a/src/ffi/ffi_infinitensor.cc +++ b/src/ffi/ffi_infinitensor.cc @@ -50,6 +50,14 @@ void init_graph_builder(py::module &m) { py::overload_cast(&Handler::batchNorm), policy::move) + .def("maxPool", + py::overload_cast(&Handler::maxPool), + policy::move) + .def("avgPool", + py::overload_cast(&Handler::avgPool), + policy::move) .def("add", py::overload_cast(&Handler::add), policy::move) .def("sub", py::overload_cast(&Handler::sub), diff --git a/src/operators/pooling.cc b/src/operators/pooling.cc index ad4d1099..4b276db1 100644 --- a/src/operators/pooling.cc +++ b/src/operators/pooling.cc @@ -5,12 +5,12 @@ namespace infini { PoolingObj::PoolingObj(GraphObj *graph, OpType optype, Tensor input, Tensor output, int kh, int kw, int dh, int dw, int ph, int pw, int sh, int sw) - : OperatorObj(optype, {input}, {output}), kh(kh), kw(kw), dh(dh), dw(dw), - ph(ph), pw(pw), sh(sh), sw(sw) { - n = input->getDims()[0]; - c = input->getDims()[1]; - h = input->getDims()[2], w = input->getDims()[3]; + : OperatorObj(optype, {input}, {output}), + kh(kh), kw(kw), dh(dh), dw(dw), ph(ph), pw(pw), sh(sh), sw(sw), + + n(input->getDims()[0]), c(input->getDims()[1]), h(input->getDims()[2]), + w(input->getDims()[3]) { IT_ASSERT(checkValid(graph)); }