Update Select export for ONNX 1.2.

This commit is contained in:
Sergii Dymchenko 2018-06-07 16:57:26 -07:00
Родитель c98623561f
Коммит 0ec517c912
1 изменённых файлов: 3 добавлений и 9 удалений

Просмотреть файл

@ -4045,17 +4045,11 @@ ONNXIR::Node* CNTKToONNXHelper::CreateONNXNodesForSelect(const FunctionPtr &src,
oneTensor.add_float_data(1.0f);
oneNode->AddAttribute("value", oneTensor);
// ONNX can broadcast only the second input for element-wise operations,
// so instead of 1 - flag01 use -(flag01 - 1).
ONNXIR::NodeArg &subOneOutputArg = graph->CreateOwnedNodeArg(outputName + "_sub_one_out", nullptr);
ONNXIR::Node* subOneNode = graph->AddNode(outputName + "_sub_one", "Sub", "", { &ceilOutputArg, &oneOutputArg }, { &subOneOutputArg });
subOneNode->AddAttribute("broadcast", static_cast<int64_t>(1));
ONNXIR::NodeArg &negOutputArg = graph->CreateOwnedNodeArg(outputName + "_neg_out", nullptr);
graph->AddNode(outputName + "_neg", "Neg", "", { &subOneOutputArg }, { &negOutputArg });
ONNXIR::NodeArg &oneSubOutputArg = graph->CreateOwnedNodeArg(outputName + "_one_sub_out", nullptr);
graph->AddNode(outputName + "_sub_one", "Sub", "", { &oneOutputArg, &ceilOutputArg }, { &oneSubOutputArg });
ONNXIR::NodeArg &mulFalseOutputArg = graph->CreateOwnedNodeArg(outputName + "_mul_false_out", nullptr);
graph->AddNode(outputName + "_mul_false", "Mul", "", { &negOutputArg, inputs[2] }, { &mulFalseOutputArg });
graph->AddNode(outputName + "_mul_false", "Mul", "", { &oneSubOutputArg, inputs[2] }, { &mulFalseOutputArg });
ONNXIR::Node* sumNode = graph->AddNode(outputName + "_sum", "Sum", "", { &mulTrueOutputArg, &mulFalseOutputArg }, { outputs[0] });