From c6580c23423389f38c8b5ff1c34f61771f07893b Mon Sep 17 00:00:00 2001 From: =?utf8?q?=D0=A1=D0=B5=D1=80=D0=B3=D0=B5=D0=B9=20=D0=91=D0=B0=D1=80?= =?utf8?q?=D0=B0=D0=BD=D0=BD=D0=B8=D0=BA=D0=BE=D0=B2/AI=20Tools=20Lab=20/S?= =?utf8?q?RR/Engineer/=EC=82=BC=EC=84=B1=EC=A0=84=EC=9E=90?= Date: Mon, 14 Jan 2019 18:06:55 +0300 Subject: [PATCH] [nnc] Refactor ONNX importer to support operators with multiple outputs (#2831) Make operator conversion methods accept and return vectors of IODescriptors. Signed-off-by: Sergei Barannikov --- .../nnc/passes/onnx_frontend/ONNXImporterImpl.cpp | 67 +++--- .../nnc/passes/onnx_frontend/ONNXImporterImpl.h | 7 +- contrib/nnc/passes/onnx_frontend/ONNXOpCreator.cpp | 230 +++++++++++---------- contrib/nnc/passes/onnx_frontend/ONNXOpCreator.h | 92 ++++++--- 4 files changed, 217 insertions(+), 179 deletions(-) diff --git a/contrib/nnc/passes/onnx_frontend/ONNXImporterImpl.cpp b/contrib/nnc/passes/onnx_frontend/ONNXImporterImpl.cpp index e874295..d37d970 100644 --- a/contrib/nnc/passes/onnx_frontend/ONNXImporterImpl.cpp +++ b/contrib/nnc/passes/onnx_frontend/ONNXImporterImpl.cpp @@ -160,7 +160,7 @@ void ONNXImporterImpl::createGraphInputs() { const onnx::TensorProto* onnx_tensor = onnx_tensors[name]; _inputTensors.insert(std::make_pair(name, createTensor(onnx_tensor))); auto constant = _graph->create(name, _inputTensors.at(name)); - _tensorNameToPrevMirOp[name] = constant; + _tensorNameToIODescriptor[name] = constant->getOutput(0); } else { // We're dealing with graph input (assuming the picture only) auto onnx_input_shape = input.type().tensor_type().shape(); @@ -172,19 +172,20 @@ void ONNXImporterImpl::createGraphInputs() { } // TODO: Temporary solution! auto node = _graph->create(name, shape); - _tensorNameToPrevMirOp[name] = node; + _tensorNameToIODescriptor[name] = node->getOutput(0); } } } -void ONNXImporterImpl::dump(const std::vector& inputs, - const std::vector& ops, +void ONNXImporterImpl::dump(const std::vector& input_descrs, + const std::vector& out_descrs, const onnx::NodeProto& onnx_node) { - for (auto op : ops) { + for (auto out_descr : out_descrs) { + auto op = out_descr.op; std::cout << onnx_node.op_type() << " '" << op->getName() << "'"; - if (inputs[0]->getNumInputs() > 0) { + if (input_descrs[0].op->getNumInputs() > 0) { std::cout << "Input Shape: "; - dumpShape(inputs[0]->getOutputShape(0)); + dumpShape(input_descrs[0].op->getOutputShape(input_descrs[0].index)); } std::cout << " Output Shape: "; dumpShape(op->getOutputShape(0)); @@ -249,91 +250,91 @@ mir::Graph *ONNXImporterImpl::createIR() { assert(onnx_node.has_op_type()); auto op_type = onnx_node.op_type().c_str(); // Fill inputs of the given node - std::vector input_nodes(onnx_node.input_size()); + std::vector inputs(onnx_node.input_size()); for (int i = 0; i < onnx_node.input_size(); i++) { auto& name = onnx_node.input(i); - assert(_tensorNameToPrevMirOp.find(name) != _tensorNameToPrevMirOp.end()); - input_nodes[i] = _tensorNameToPrevMirOp[name]; + assert(_tensorNameToIODescriptor.find(name) != _tensorNameToIODescriptor.end()); + inputs[i] = _tensorNameToIODescriptor[name]; } - std::vector outputs; + std::vector outputs; auto* onnx_op_type = ONNXPerfectHash::getONNXOpType(op_type, onnx_node.op_type().size()); switch (onnx_op_type->opCode) { case ONNXOpCode::opConv: - outputs = _opCreator.convertConv2D(input_nodes, onnx_node); + outputs = _opCreator.convertConv2D(inputs, onnx_node); break; case ONNXOpCode::opAdd: - outputs = _opCreator.convertElementwise(input_nodes, mir::ops::ElementwiseOp::OpType::add); + outputs = _opCreator.convertElementwise(inputs, mir::ops::ElementwiseOp::OpType::add); break; case ONNXOpCode::opGather: - outputs = _opCreator.convertGather(input_nodes, onnx_node); + outputs = _opCreator.convertGather(inputs, onnx_node); break; case ONNXOpCode::opGemm: - outputs = _opCreator.convertGemm(input_nodes, onnx_node); + outputs = _opCreator.convertGemm(inputs, onnx_node); break; case ONNXOpCode::opSum: - outputs = _opCreator.convertElementwise(input_nodes, mir::ops::ElementwiseOp::OpType::add); + outputs = _opCreator.convertElementwise(inputs, mir::ops::ElementwiseOp::OpType::add); break; case ONNXOpCode::opMul: - outputs = _opCreator.convertElementwise(input_nodes, mir::ops::ElementwiseOp::OpType::mul); + outputs = _opCreator.convertElementwise(inputs, mir::ops::ElementwiseOp::OpType::mul); break; case ONNXOpCode::opMax: - outputs = _opCreator.convertElementwise(input_nodes, mir::ops::ElementwiseOp::OpType::max); + outputs = _opCreator.convertElementwise(inputs, mir::ops::ElementwiseOp::OpType::max); break; case ONNXOpCode::opGlobalAveragePool: case ONNXOpCode::opAveragePool: case ONNXOpCode::opMaxPool: - outputs = _opCreator.convertPool(input_nodes, onnx_op_type->opCode, onnx_node); + outputs = _opCreator.convertPool(inputs, onnx_op_type->opCode, onnx_node); break; case ONNXOpCode::opConcat: - outputs = _opCreator.convertConcat(input_nodes, onnx_node); + outputs = _opCreator.convertConcat(inputs, onnx_node); break; case ONNXOpCode::opReshape: - outputs = _opCreator.convertReshape(input_nodes); + outputs = _opCreator.convertReshape(inputs); break; case ONNXOpCode::opRelu: - outputs = _opCreator.convertRelu(input_nodes); + outputs = _opCreator.convertRelu(inputs); break; case ONNXOpCode::opUnsqueeze: - outputs = _opCreator.convertUnsqueeze(input_nodes[0], onnx_node); + outputs = _opCreator.convertUnsqueeze(inputs, onnx_node); break; case ONNXOpCode::opSigmoid: - outputs = _opCreator.convertSigmoid(input_nodes); + outputs = _opCreator.convertSigmoid(inputs); break; case ONNXOpCode::opSoftmax: - outputs = _opCreator.convertSoftmax(input_nodes, onnx_node); + outputs = _opCreator.convertSoftmax(inputs, onnx_node); break; case ONNXOpCode::opScale: - outputs = _opCreator.convertScale(input_nodes, onnx_node); + outputs = _opCreator.convertScale(inputs, onnx_node); break; case ONNXOpCode::opBatchNormalization: - outputs = _opCreator.convertBatchNorm(input_nodes, onnx_node, _inputTensors); + outputs = _opCreator.convertBatchNorm(inputs, onnx_node, _inputTensors); break; case ONNXOpCode::opDropout: - outputs = _opCreator.convertDropout(input_nodes, onnx_node); + outputs = _opCreator.convertDropout(inputs, onnx_node); break; default: throw PassException("Invalid ONNXOpCode" + std::to_string((int)onnx_op_type->opCode)); } // Set outputs' names for (int i = 0; i < outputs.size(); i++) { - outputs[i]->setName(onnx_node.output(i)); - auto result = _tensorNameToPrevMirOp.emplace(outputs[i]->getName(), outputs[i]); + outputs[i].op->setName(onnx_node.output(i)); + auto result = _tensorNameToIODescriptor.emplace(outputs[i].op->getName(), outputs[i]); if(!result.second) - throw PassException("Name duplication: " + outputs[i]->getName()); + throw PassException("Name duplication: " + outputs[i].op->getName()); } assert (outputs.size()); // FIXME: it should be done properly via the given graph outputs _graphOutputs.assign(outputs.begin(), outputs.end()); #if 0 - dump(input_nodes, outputs, onnx_node); + dump(inputs, outputs, onnx_node); #endif } // set graph outputs // TODO: it should be done with onnx graph outputs for (auto& output_idx : _graphOutputs) - _graph->markOutput(output_idx); + _graph->markOutput(output_idx.op); return _graph; } diff --git a/contrib/nnc/passes/onnx_frontend/ONNXImporterImpl.h b/contrib/nnc/passes/onnx_frontend/ONNXImporterImpl.h index 7fd5683..be4ff8a 100644 --- a/contrib/nnc/passes/onnx_frontend/ONNXImporterImpl.h +++ b/contrib/nnc/passes/onnx_frontend/ONNXImporterImpl.h @@ -39,7 +39,8 @@ public: void import() {}; mir::Graph *createIR() override; - void dump(const std::vector& inputs, const std::vector& op, + void dump(const std::vector& input_descrs, + const std::vector& out_descrs, const onnx::NodeProto& onnx_node); static void dumpShape(mir::Shape shape) { @@ -51,10 +52,10 @@ public: private: void createGraphInputs(); // This map maps onnx tensor names to MIR operations/nodes - std::map _tensorNameToPrevMirOp; + std::map _tensorNameToIODescriptor; // This map keeps named tensors used as graph input initializers. std::map _inputTensors; - std::vector _graphOutputs; + std::vector _graphOutputs; std::string _modelFilename; std::unique_ptr _model; mir::Graph* _graph; diff --git a/contrib/nnc/passes/onnx_frontend/ONNXOpCreator.cpp b/contrib/nnc/passes/onnx_frontend/ONNXOpCreator.cpp index 32b35f8..d3814cb 100644 --- a/contrib/nnc/passes/onnx_frontend/ONNXOpCreator.cpp +++ b/contrib/nnc/passes/onnx_frontend/ONNXOpCreator.cpp @@ -125,14 +125,15 @@ static void getKernelStridesPadding(const onnx::NodeProto &onnx_node, KernelStri } }; -std::vector ONNXOpCreator::convertConv2D(InputOps& inputs, - const onnx::NodeProto& onnx_node) { +std::vector +ONNXOpCreator::convertConv2D(const std::vector& inputs, + const onnx::NodeProto& onnx_node) { assert(inputs.size() >= 2); KernelStridesPadding cdata; getKernelStridesPadding(onnx_node, cdata); // FIXME: It can be non-constant value. - auto* in_weights = dynamic_cast(inputs[1]); + auto* in_weights = dynamic_cast(inputs[1].op); assert(in_weights && "Weights could be a constant tensor only"); const auto& in_weights_tensor = in_weights->getValue(); // We should transpose ONNX MCHW to HWOI @@ -140,105 +141,104 @@ std::vector ONNXOpCreator::convertConv2D(InputOps& inputs, mir::ops::ConstantOp* input_bias = nullptr; if (inputs.size() > 2) { - input_bias = dynamic_cast(inputs[2]); + input_bias = dynamic_cast(inputs[2].op); assert(input_bias && "1D optional bias could be a constant tensor only"); } - inputs.resize(1); - std::vector outputs; // Transpose ONNX NCHW to MIR NHWC - auto t_input = convertONNXToMIR(inputs[0]->getOutput(0)); - outputs = createOp(t_input[0]->getOutput(0), transposed, cdata.strides_shape, - cdata.padding_before, cdata.padding_after); + auto t_input = convertONNXToMIR(inputs[0]); + auto result = createOp(t_input, transposed, cdata.strides_shape, + cdata.padding_before, cdata.padding_after); if (input_bias) - outputs = createOp(outputs[0]->getOutput(0), input_bias->getValue()); + result = createOp(result->getOutput(0), input_bias->getValue()); - return convertMIRToONNX(outputs[0]->getOutput(0)); + return {convertMIRToONNX(result->getOutput(0))}; } -std::vector ONNXOpCreator::convertConcat(InputOps& inputs, - const onnx::NodeProto& onnx_node) { +std::vector +ONNXOpCreator::convertConcat(const std::vector& inputs, + const onnx::NodeProto& onnx_node) { bool found; int axis; std::tie (found, axis) = getIntAttribute(onnx_node); if (!found) throw PassException("Concat must have 'axis' attribute"); - std::vector descriptors; - for (auto input : inputs) - descriptors.push_back(input->getOutput(0)); - return createOp(descriptors, axis); + auto result = createOp(inputs, axis); + return {result->getOutput(0)}; } -std::vector -ONNXOpCreator::convertGather(ONNXOpCreator::InputOps& inputs, const onnx::NodeProto& onnx_node) { +std::vector +ONNXOpCreator::convertGather(const std::vector& inputs, + const onnx::NodeProto& onnx_node) { bool found; int value; std::tie(found, value) = getIntAttribute(onnx_node, "axis"); int axis = found ? value : 0; - return createOp(inputs[0]->getOutput(0), inputs[1]->getOutput(0), axis); + auto result = createOp(inputs[0], inputs[1], axis); + return {result->getOutput(0)}; } -std::vector ONNXOpCreator::convertPool(InputOps& inputs, ONNXOpCode op_code, - const onnx::NodeProto& onnx_node) { +std::vector +ONNXOpCreator::convertPool(const std::vector& inputs, + ONNXOpCode op_code, + const onnx::NodeProto& onnx_node) { ops::PoolOp::BorderType border_type; ops::PoolOp::PoolingType pool_type; - std::vector result; KernelStridesPadding cdata; // Transpose ONNX NCHW to MIR NHWC - auto t_input = convertONNXToMIR(inputs[0]->getOutput(0)); + auto t_input = convertONNXToMIR(inputs[0]); switch (op_code) { case ONNXOpCode::opGlobalAveragePool: { + border_type = ops::PoolOp::BorderType::ZEROFILLED; + pool_type = ops::PoolOp::PoolingType::AVG; // GlobalAveragePool is equivalent to AveragePool with kernel size equal // to the spatial dimension of input tensor - result = createOp(t_input[0]->getOutput(0), - ops::PoolOp::PoolingType::AVG, - t_input[0]->getOutputShape(0), // kernel_shape - Shape({1, 1}), // strides_shape - cdata.padding_before, cdata.padding_after, - ops::PoolOp::BorderType::ZEROFILLED, - ops::PoolOp::RoundMode::floor); - return convertMIRToONNX(result[0]->getOutput(0)); + cdata.kernel_shape = t_input.op->getOutputShape(0); + cdata.strides_shape = Shape{1, 1}; + break; } case ONNXOpCode::opAveragePool: border_type = ops::PoolOp::BorderType::ZEROFILLED; pool_type = ops::PoolOp::PoolingType::AVG; + getKernelStridesPadding(onnx_node, cdata); break; case ONNXOpCode::opMaxPool: border_type = ops::PoolOp::BorderType::EMPTY; pool_type = ops::PoolOp::PoolingType::MAX; + getKernelStridesPadding(onnx_node, cdata); break; default: assert(false); } - // Proceed with Average or Max Pool - getKernelStridesPadding(onnx_node, cdata); - result = createOp(t_input[0]->getOutput(0), pool_type, - cdata.kernel_shape, cdata.strides_shape, - cdata.padding_before, cdata.padding_after, - border_type, - ops::PoolOp::RoundMode::floor); - return convertMIRToONNX(result[0]->getOutput(0)); + auto result = createOp(t_input, pool_type, + cdata.kernel_shape, cdata.strides_shape, + cdata.padding_before, cdata.padding_after, + border_type, ops::PoolOp::RoundMode::floor); + return {convertMIRToONNX(result->getOutput(0))}; } -std::vector ONNXOpCreator::convertSoftmax(InputOps& inputs, - const onnx::NodeProto& onnx_node) { +std::vector +ONNXOpCreator::convertSoftmax(const std::vector& inputs, + const onnx::NodeProto& onnx_node) { int axis; bool found; std::tie (found, axis) = getIntAttribute(onnx_node); axis = found ? axis : 1; - return createOp(inputs[0]->getOutput(0), axis); + auto result = createOp(inputs[0], axis); + return {result->getOutput(0)}; } -std::vector ONNXOpCreator::convertReshape(InputOps& inputs) { +std::vector +ONNXOpCreator::convertReshape(const std::vector& inputs) { // The original shape - auto in_shape = inputs[0]->getInputShape(0); + auto in_shape = inputs[0].op->getOutputShape(inputs[0].index); // Input tensor describing the new shape // TODO: could it be not a constant? - auto* op = dynamic_cast(inputs[1]); + auto* op = dynamic_cast(inputs[1].op); assert(op && "We support constants only"); auto shape_tensor = op->getValue(); Shape shape_tensor_shape = (shape_tensor).getShape(); @@ -261,17 +261,18 @@ std::vector ONNXOpCreator::convertReshape(InputOps& inputs) { i++; } auto out_shape = Shape(shape_vector); - auto outputs = createOp(inputs[0]->getOutput(0), out_shape); - return outputs; + auto result = createOp(inputs[0], out_shape); + return {result->getOutput(0)}; } -std::vector -ONNXOpCreator::convertUnsqueeze(Operation* input_data, const onnx::NodeProto& onnx_node) { +std::vector +ONNXOpCreator::convertUnsqueeze(const std::vector& inputs, + const onnx::NodeProto& onnx_node) { auto* axes = findAttribute(onnx_node, "axes"); assert(axes && axes->ints_size()); - const int out_rank = input_data->getOutputShape(0).rank() + axes->ints_size(); + const Shape& input_shape = inputs[0].op->getOutputShape(inputs[0].index); + const int out_rank = input_shape.rank() + axes->ints_size(); Shape out_shape(out_rank); - const Shape& input_shape = input_data->getOutputShape(0); auto ints_iterator = axes->ints().begin(); int j = 0; for (int i = 0; i < out_rank; i++) { @@ -283,84 +284,93 @@ ONNXOpCreator::convertUnsqueeze(Operation* input_data, const onnx::NodeProto& on j++; } } - auto outputs = createOp(input_data->getOutput(0), out_shape); - return outputs; + auto result = createOp(inputs[0], out_shape); + return {result->getOutput(0)}; } -std::vector ONNXOpCreator::convertRelu(InputOps& inputs) { +std::vector +ONNXOpCreator::convertRelu(const std::vector& inputs) { assert(inputs.size() == 1); - return createOp(inputs[0]->getOutput(0)); + auto result = createOp(inputs[0]); + return {result->getOutput(0)}; } -std::vector ONNXOpCreator::convertSigmoid(InputOps& inputs) { +std::vector +ONNXOpCreator::convertSigmoid(const std::vector& inputs) { assert(inputs.size() == 1); - return createOp(inputs[0]->getOutput(0)); + auto result = createOp(inputs[0]); + return {result->getOutput(0)}; } -std::vector ONNXOpCreator::convertElementwise(InputOps& inputs, - mir::ops::ElementwiseOp::OpType op_type) { - std::vector descriptors; - for (auto input : inputs) - descriptors.push_back(input->getOutput(0)); - return createOp(descriptors, op_type); +std::vector +ONNXOpCreator::convertElementwise(const std::vector& inputs, + mir::ops::ElementwiseOp::OpType op_type) { + auto result = createOp(inputs, op_type); + return {result->getOutput(0)}; } -std::vector ONNXOpCreator::convertBatchNorm(InputOps& inputs, - const onnx::NodeProto& onnx_node, - InputTensors& input_tensors) { + +std::vector +ONNXOpCreator::convertBatchNorm(const std::vector& inputs, + const onnx::NodeProto& onnx_node, + InputTensors& input_tensors) { // overall_res = (X - mean) / sqrt(var + epsilon) * scale + bias bool found; float value; std::tie(found, value) = getFloatAttribute(onnx_node, "epsilon"); float epsilon = found ? value : 1e-05f; - const auto& scale = input_tensors.at(inputs[1]->getName()); - const auto& bias = input_tensors.at(inputs[2]->getName()); - const auto& mean = input_tensors.at(inputs[3]->getName()); - const auto& var = input_tensors.at(inputs[4]->getName()); + const auto& scale = input_tensors.at(inputs[1].op->getName()); + const auto& bias = input_tensors.at(inputs[2].op->getName()); + const auto& mean = input_tensors.at(inputs[3].op->getName()); + const auto& var = input_tensors.at(inputs[4].op->getName()); // res1 = X - mean Tensor bias_data(mean); for (auto& idx: ShapeRange(bias_data.getShape())) bias_data.at(idx) *= -1; - auto data = convertONNXToMIR(inputs[0]->getOutput(0)); - auto bias_add_1 = createOp(data[0]->getOutput(0), mean); + auto data = convertONNXToMIR(inputs[0]); + auto bias_add_1 = createOp(data, mean); // res2 = res1 * scale / (var + epsilon) Tensor multiplier(scale); Tensor var_accessor(var); for (auto& idx: ShapeRange(scale.getShape())) multiplier.at(idx) /= std::sqrt(var_accessor.at(idx) + epsilon); - auto scale_op = createOp(bias_add_1[0]->getOutput(0), scale); + auto scale_op = createOp(bias_add_1->getOutput(0), scale); // overall_res = res2 + bias - auto bias_add_2 = createOp(scale_op[0]->getOutput(0), bias); + auto bias_add_2 = createOp(scale_op->getOutput(0), bias); - return {convertMIRToONNX(bias_add_2[0]->getOutput(0))}; + return {convertMIRToONNX(bias_add_2->getOutput(0))}; } -std::vector ONNXOpCreator::convertDropout(InputOps& inputs, - const onnx::NodeProto& onnx_node) { +std::vector +ONNXOpCreator::convertDropout(const std::vector& inputs, + const onnx::NodeProto& onnx_node) { bool found; float value; std::tie(found, value) = getFloatAttribute(onnx_node, "ratio"); float ratio = found ? value : 1.0; - return createOp(inputs[0]->getOutput(0), ratio); + auto result = createOp(inputs[0], ratio); + return {result->getOutput(0)}; } -std::vector ONNXOpCreator::convertScale(InputOps& inputs, - const onnx::NodeProto& onnx_node) { +std::vector +ONNXOpCreator::convertScale(const std::vector& inputs, + const onnx::NodeProto& onnx_node) { bool found; float value; std::tie(found, value) = getFloatAttribute(onnx_node, "scale"); float scale = found ? value : 1.0; - auto outputs = createOp(inputs[0]->getOutput(0), - createTensor(scale, inputs[0]->getOutputShape(0))); - return outputs; + const auto& shape = inputs[0].op->getOutputShape(inputs[0].index); + auto result = createOp(inputs[0], createTensor(scale, shape)); + return {result->getOutput(0)}; } -std::vector ONNXOpCreator::convertGemm(InputOps& inputs, - const onnx::NodeProto& onnx_node) { +std::vector +ONNXOpCreator::convertGemm(const std::vector& inputs, + const onnx::NodeProto& onnx_node) { bool found; int ivalue; float fvalue; @@ -385,57 +395,49 @@ std::vector ONNXOpCreator::convertGemm(InputOps& inputs, // 1. Prepare input matrix A // Flatten the shape by dim(0) - mir::Shape shape0 ({inputs[0]->getOutputShape(0).dim(0), - inputs[0]->getOutputShape(0).numElements() / - inputs[0]->getOutputShape(0).dim(0)}); - auto input_a = createOp(inputs[0]->getOutput(0), shape0); + const auto& in_shape = inputs[0].op->getOutputShape(inputs[0].index); + mir::Shape shape0{in_shape.dim(0), in_shape.numElements() / in_shape.dim(0)}; + auto input_a = createOp(inputs[0], shape0); if (trans_a) - input_a = createOp(input_a[0]->getOutput(0), std::vector{1, 0}); + input_a = createOp(input_a->getOutput(0), std::vector{1, 0}); if (alpha != 1.0) - input_a = createOp(input_a[0]->getOutput(0), - createTensor(alpha, input_a[0]->getOutputShape(0))); + input_a = createOp(input_a->getOutput(0), + createTensor(alpha, input_a->getOutputShape(0))); // 2. Prepare input matrix B // - auto input_b = inputs[1]->getOutput(0); + auto input_b = inputs[1]; if (trans_b) - input_b = createOp(input_b, std::vector{1, 0})[0]->getOutput(0); + input_b = createOp(input_b, std::vector{1, 0})->getOutput(0); // Number of cols in tensor A must be equal to number of rows in tensor B - assert(input_a[0]->getOutput(0).op->getOutputShape(0).dim(1) == + assert(input_a->getOutput(0).op->getOutputShape(0).dim(1) == input_b.op->getOutputShape(0).dim(0)); - Shape mult_a_b({input_a[0]->getOutput(0).op->getOutputShape(0).dim(0), - input_b.op->getOutputShape(0).dim(1)}); + Shape mult_a_b({input_a->getOutputShape(0).dim(0), + input_b.op->getOutputShape(input_b.index).dim(1)}); // 3. Prepare input matrix C // - auto input_c = inputs[2]->getOutput(0); + auto input_c = inputs[2]; auto beta_tensor = createTensor(beta, input_c.op->getOutputShape(0)); // TODO: check 'broadcast' attribute here if ((mult_a_b.rank() == 2) && (input_c.op->getOutputShape(0).rank() == 1)) { beta_tensor = TensorVariant(beta_tensor, mult_a_b); } - auto constant = createOp(beta_tensor)[0]->getOutput(0); + auto constant = createOp(beta_tensor)->getOutput(0); std::vector descriptors = {constant, input_c}; auto c_mult = createOp(descriptors, ops::ElementwiseOp::OpType::mul); - assert(c_mult[0]->getOutput(0).op->getOutputShape(0) == mult_a_b); - return createOp(input_a[0]->getOutput(0), input_b, c_mult[0]->getOutput(0)); -} - -std::vector -ONNXOpCreator::createInput(const std::string& input_name, const mir::Shape& input_shape) { - // TODO For now we only support convolutional networks with one element per batch. - assert(input_shape.rank() == 4 && input_shape.dim(0) == 1); - auto variable = _graph->create(input_name, input_shape); - return {variable}; + assert(c_mult->getOutputShape(0) == mult_a_b); + auto result = createOp(input_a->getOutput(0), input_b, c_mult->getOutput(0)); + return {result->getOutput(0)}; } -std::vector ONNXOpCreator::convertONNXToMIR(const mir::IODescriptor& arg) { +mir::IODescriptor ONNXOpCreator::convertONNXToMIR(mir::IODescriptor arg) { // NCHW -> NHWC - return createOp(arg, std::vector{0, 2, 3, 1}); + return createOp(arg, std::vector{0, 2, 3, 1})->getOutput(0); } -std::vector ONNXOpCreator::convertMIRToONNX(const mir::IODescriptor& arg) { +mir::IODescriptor ONNXOpCreator::convertMIRToONNX(mir::IODescriptor arg) { // NHWC -> NCHW - return createOp(arg, std::vector{0, 3, 1, 2}); + return createOp(arg, std::vector{0, 3, 1, 2})->getOutput(0); } } // namespace nnc diff --git a/contrib/nnc/passes/onnx_frontend/ONNXOpCreator.h b/contrib/nnc/passes/onnx_frontend/ONNXOpCreator.h index 30999b5..7ae9144 100644 --- a/contrib/nnc/passes/onnx_frontend/ONNXOpCreator.h +++ b/contrib/nnc/passes/onnx_frontend/ONNXOpCreator.h @@ -33,46 +33,80 @@ namespace nnc { class ONNXOpCreator { public: - using InputOps = std::vector; using InputTensors = std::map; ONNXOpCreator() = default; - void setMirGraph(mir::Graph* g) {_graph = g;}; - std::vector convertConv2D(InputOps& inputs, const onnx::NodeProto& node); - std::vector convertConcat(InputOps& inputs, const onnx::NodeProto& onnx_node); - std::vector convertPool(InputOps& inputs, ONNXOpCode op_code, - const onnx::NodeProto& onnx_node); - std::vector convertSoftmax(InputOps& inputs, const onnx::NodeProto& onnx_node); - std::vector convertReshape(InputOps& inputs); - std::vector convertRelu(InputOps& inputs); - std::vector convertSigmoid(InputOps& inputs); - - std::vector - convertUnsqueeze(mir::Operation* inputs, const onnx::NodeProto& onnx_node); - std::vector convertElementwise(InputOps& inputs, - mir::ops::ElementwiseOp::OpType op_type); - std::vector convertScale(InputOps& inputs, const onnx::NodeProto& node); - std::vector convertBatchNorm(InputOps& inputs, const onnx::NodeProto& node, - InputTensors& input_tensors); - std::vector convertDropout(InputOps& inputs, const onnx::NodeProto& onnx_node); - std::vector convertGather(InputOps& inputs, const onnx::NodeProto& onnx_node); - std::vector convertGemm(InputOps& inputs, const onnx::NodeProto& onnx_node); - - std::vector createInput(const std::string&, const mir::Shape&); - std::vector convertONNXToMIR(const mir::IODescriptor& arg); - std::vector convertMIRToONNX(const mir::IODescriptor& arg); + + void setMirGraph(mir::Graph* g) { _graph = g; }; + + std::vector + convertConv2D(const std::vector& inputs, + const onnx::NodeProto& onnx_node); + + std::vector + convertConcat(const std::vector& inputs, + const onnx::NodeProto& onnx_node); + + std::vector + convertPool(const std::vector& inputs, + ONNXOpCode op_code, + const onnx::NodeProto& onnx_node); + + std::vector + convertSoftmax(const std::vector& inputs, + const onnx::NodeProto& onnx_node); + + std::vector + convertReshape(const std::vector& inputs); + + std::vector + convertRelu(const std::vector& inputs); + + std::vector + convertSigmoid(const std::vector& inputs); + + std::vector + convertUnsqueeze(const std::vector& inputs, + const onnx::NodeProto& onnx_node); + + std::vector + convertElementwise(const std::vector& inputs, + mir::ops::ElementwiseOp::OpType op_type); + + std::vector + convertScale(const std::vector& inputs, + const onnx::NodeProto& onnx_node); + + std::vector + convertBatchNorm(const std::vector& inputs, + const onnx::NodeProto& onnx_node, + InputTensors& input_tensors); + + std::vector + convertDropout(const std::vector& inputs, + const onnx::NodeProto& onnx_node); + + std::vector + convertGather(const std::vector& inputs, + const onnx::NodeProto& onnx_node); + + std::vector + convertGemm(const std::vector& inputs, + const onnx::NodeProto& onnx_node); + + mir::IODescriptor convertONNXToMIR(mir::IODescriptor arg); + mir::IODescriptor convertMIRToONNX(mir::IODescriptor arg); private: template - std::vector createOp(Types&&... args); + nnc::mir::Operation* createOp(Types&&... args); mir::Graph* _graph = nullptr; }; template -std::vector ONNXOpCreator::createOp(Types&&... args) { +nnc::mir::Operation* ONNXOpCreator::createOp(Types&&... args) { // TODO: set operation names - auto op = _graph->create("", std::forward(args)...); - return {op}; + return _graph->create("", std::forward(args)...); } } // namespace nnc #endif //NNCC_ONNX_OP_CREATOR_H -- 2.7.4