experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Update logic to read builtin op code

Browse Source 
PiperOrigin-RevId: 338017168
Change-Id: Ib5eadbf8d9d8ba4dc0c262f97728981d7d510690
This commit is contained in:
Thai Nguyen 2020-10-20 01:25:14 -07:00 committed by TensorFlower Gardener
parent c1841cf246
commit fc9652f980
2 changed files with 3 additions and 1 deletions
tensorflow/lite/tools
BUILDlist_flex_ops_no_kernel.cc

1
tensorflow/lite/tools/BUILD
View File

@ -319,6 +319,7 @@ cc_library(
hdrs = ["list_flex_ops.h"],
deps = [
"//tensorflow/lite:framework",
"//tensorflow/lite/schema:schema_utils",
"@jsoncpp_git//:jsoncpp",
],
)

3
tensorflow/lite/tools/list_flex_ops_no_kernel.cc
View File

@ -13,6 +13,7 @@ See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "json/json.h"
#include "tensorflow/lite/schema/schema_utils.h"
#include "tensorflow/lite/tools/list_flex_ops.h"
namespace tflite {
@ -40,7 +41,7 @@ void AddFlexOpsFromModel(const tflite::Model* model, OpKernelSet* flex_ops) {
for (int i = 0; i < operators->size(); ++i) {
const tflite::Operator* op = operators->Get(i);
const tflite::OperatorCode* opcode = opcodes->Get(op->opcode_index());
if (opcode->builtin_code() != tflite::BuiltinOperator_CUSTOM ||
if (tflite::GetBuiltinCode(opcode) != tflite::BuiltinOperator_CUSTOM ||
!tflite::IsFlexOp(opcode->custom_code()->c_str())) {
continue;
}