fix custom op bug when custom op has muti-out and not use all of outputs

Merge pull request  from geekun/master
This commit is contained in:
mindspore-ci-bot 2020-04-02 21:02:25 +08:00 committed by Gitee
commit a1d6f42be8
2 changed files with 32 additions and 1 deletions
mindspore/ccsrc/transform

View File

@ -279,6 +279,31 @@ class OpAdapter : public BaseOpAdapter {
}
OutHandler getOutput(const OperatorPtr& op, int index) override {
MS_EXCEPTION_IF_NULL(op);
if (IsCustomOp(op)) {
return getCustomOutput(op, index);
}
return getNormalOutput(op, index);
}
OutHandler getCustomOutput(const OperatorPtr& op, int index) {
MS_EXCEPTION_IF_NULL(op);
auto it = cus_output_map_.find(op->GetOpType());
if (it == cus_output_map_.end()) {
MS_LOG(ERROR) << "OpAdpator(" << op->GetName() << ") has both OUTPUT is not supported!";
return OutHandler();
}
std::unordered_map<int, std::string>& output_map = it->second;
if ((output_map.find(index) != output_map.end())) {
return OutHandler(op, output_map[index]);
}
MS_LOG(ERROR) << "OpAdpator(" << op->GetName() << ") has no OUTPUT index(" << index << ")!";
return OutHandler();
}
OutHandler getNormalOutput(const OperatorPtr& op, int index) {
MS_EXCEPTION_IF_NULL(op);
if (!dyn_output_map_.empty() && !output_map_.empty()) {
MS_LOG(ERROR) << "OpAdpator(" << op->GetName() << ") has both OUTPUT and DYN_OUTPUT is not supported!";

8
mindspore/ccsrc/transform/op_adapter_util.cc Executable file → Normal file
View File

@ -223,7 +223,13 @@ bool IsCustomPrim(const PrimitivePtr& prim) {
return false;
}
return GetValue<bool>(flag);
bool is_custom_op = GetValue<bool>(flag);
if (!is_custom_op && prim->GetAttr("_custom_op_impl_config_path") != nullptr) {
MS_LOG(EXCEPTION) << "The custom op flag is false, but the op information config path is not null, non-custom op "
"can not assign the op information config path.";
}
return is_custom_op;
}
bool IsCustomCNode(const AnfNodePtr& anf) {