fix custom op bug and add custom op check

This commit is contained in:
geekun 2020-03-30 19:20:26 +08:00 committed by chang zherui
parent b18f634912
commit fb92b5b16e
2 changed files with 32 additions and 1 deletions

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) {