p59230764
/
InfiniTensor
forked from jiuyuan/InfiniTensor
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

refactor: 非 onnx 算子也能导出为自定义算子 

Signed-off-by: YdrMaster <ydrml@hotmail.com>
This commit is contained in:
YdrMaster 2023-09-16 14:24:29 +08:00
parent a795364612
commit d5a7713c7d
1 changed files with 4 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
src/ffi/ffi_infinitensor.cc
View File

@ -82,8 +82,9 @@ class Iterator {
auto [nodeIdx, inputs_, outputs_] = *_it++;
auto const &node = _internal->graph().nodes[nodeIdx];
auto const &name = node.name;
auto const opType = node.op->opType.name();
ASSERT(opType.substr(0, 6) == "onnx::", "Invalid opType");
auto const opType_ = node.op->opType.name();
auto opType =
opType_.substr(0, 6) == "onnx::" ? opType_.substr(6) : opType_;
auto const &attributes = node.op->attributes;
std::vector<Name> inputs(inputs_.size()), outputs(outputs_.size());
std::transform(inputs_.begin(), inputs_.end(), inputs.begin(),
@ -101,7 +102,7 @@ class Iterator {
attributes_.insert({name, attr.value});
}
}
return O{name, opType.substr(6), attributes_, inputs, outputs};
return O{name, opType, attributes_, inputs, outputs};
}
};