saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

handle unsqueeze change in opset13 (#8308)

Browse source 
Co-authored-by: Cheng Tang <chenta@microsoft.com>
This commit is contained in:
Tang, Cheng 2021-07-06 22:30:24 -07:00 committed by GitHub
parent 2347a0aca8
commit d7c3703371
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 36 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

41
orttraining/orttraining/core/graph/loss_func/bert_loss_legacy.cc
View file

@ -4,6 +4,7 @@
#include <vector>
#include "orttraining/core/graph/loss_func/bert_loss_legacy.h"
#include "onnx/defs/attr_proto_util.h"
#include "onnx/defs/tensor_proto_util.h"
namespace onnxruntime {
namespace training {
@ -74,12 +75,42 @@ GraphAugmenter::GraphDefs BertLossLegacy::operator()(const Graph& graph, const L
TypeProto* masked_lm_int64_type_proto = GetMaskedLMTypeProto(prediction_arg,
ONNX_NAMESPACE::TensorProto_DataType_INT64,
graph_defs);
int onnx_opset = -1;
auto onnx_domain_it = graph.DomainToVersionMap().find(kOnnxDomain);
if (onnx_domain_it != graph.DomainToVersionMap().end()) {
onnx_opset = onnx_domain_it->second;
} else {
auto onnx_domain_alias_it = graph.DomainToVersionMap().find(kOnnxDomainAlias);
if (onnx_domain_alias_it != graph.DomainToVersionMap().end())
onnx_opset = onnx_domain_alias_it->second;
else
ORT_THROW("ONNX domain not found in this model");
}
if (onnx_opset <= 12) {
new_nodes.emplace_back(NodeDef("Unsqueeze",
{ArgDef(masked_lm_positions, masked_lm_int64_type_proto)},
{ArgDef("masked_lm_positions_unsqueezed")},
{ONNX_NAMESPACE::MakeAttribute("axes", std::vector<int64_t>{static_cast<int64_t>(2)})},
"Mask_LM_Positions_Unsqueezed"));
} else {
auto t_proto = ONNX_NAMESPACE::ToTensor<int64_t>(1);
TypeProto* int64_t_proto = graph_defs.CreateTypeProto();
int64_t_proto->mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_INT64);
new_nodes.emplace_back(NodeDef("Unsqueeze",
{ArgDef(masked_lm_positions, masked_lm_int64_type_proto)},
{ArgDef("masked_lm_positions_unsqueezed")},
{ONNX_NAMESPACE::MakeAttribute("axes", std::vector<int64_t>{static_cast<int64_t>(2)})},
"Mask_LM_Positions_Unsqueezed"));
const std::string two_constant = "bert_legacy_two_constant";
new_nodes.emplace_back(NodeDef("Constant",
{},
{ArgDef(two_constant, int64_t_proto)},
{ONNX_NAMESPACE::MakeAttribute("value", t_proto)}));
new_nodes.emplace_back(NodeDef("Unsqueeze",
{ArgDef(masked_lm_positions, masked_lm_int64_type_proto),
ArgDef("two_constant", int64_t_proto)},
{ArgDef("masked_lm_positions_unsqueezed")},
NodeAttributes(),
"Mask_LM_Positions_Unsqueezed"));
}
TypeProto* gathered_prediction_type_proto = GetGatheredPredictionTypeProto(prediction_arg,
graph_defs);
new_nodes.emplace_back(NodeDef(OpDef{"GatherND", kOnnxDomain, 12},