diff options
author | 오형석/On-Device Lab(SR)/Staff Engineer/삼성전자 <hseok82.oh@samsung.com> | 2019-10-15 20:21:54 +0900 |
---|---|---|
committer | 이한종/On-Device Lab(SR)/Engineer/삼성전자 <hanjoung.lee@samsung.com> | 2019-10-15 20:21:54 +0900 |
commit | 2b291df9e211eaef2b7e35a3eaeed021782fa60f (patch) | |
tree | b5e32e703766cbf28dfa285646c337166d92ae71 | |
parent | 2d23c9c447332a15efbf6fbe8014c724a5f3315b (diff) | |
download | nnfw-2b291df9e211eaef2b7e35a3eaeed021782fa60f.tar.gz nnfw-2b291df9e211eaef2b7e35a3eaeed021782fa60f.tar.bz2 nnfw-2b291df9e211eaef2b7e35a3eaeed021782fa60f.zip |
[m2/tool] Validation model file check in tflite_benchmark (#8150)
Pass verifier and use VerifyAndBuildFromFile in tflite_benchmark
Signed-off-by: Hyeongseok Oh <hseok82.oh@samsung.com>
-rw-r--r-- | tests/tools/tflite_benchmark/src/tflite_benchmark.cc | 21 |
1 files changed, 20 insertions, 1 deletions
diff --git a/tests/tools/tflite_benchmark/src/tflite_benchmark.cc b/tests/tools/tflite_benchmark/src/tflite_benchmark.cc index 21eee4a18..1fde0c449 100644 --- a/tests/tools/tflite_benchmark/src/tflite_benchmark.cc +++ b/tests/tools/tflite_benchmark/src/tflite_benchmark.cc @@ -70,6 +70,23 @@ bool checkParams(const int argc, char **argv) return true; } +// Verifies whether the model is a flatbuffer file. +class BMFlatBufferVerifier : public tflite::TfLiteVerifier +{ +public: + bool Verify(const char *data, int length, tflite::ErrorReporter *reporter) override + { + + flatbuffers::Verifier verifier(reinterpret_cast<const uint8_t *>(data), length); + if (!tflite::VerifyModelBuffer(verifier)) + { + reporter->Report("The model is not a valid Flatbuffer file"); + return false; + } + return true; + } +}; + int main(const int argc, char **argv) { @@ -98,7 +115,9 @@ int main(const int argc, char **argv) StderrReporter error_reporter; - auto model = FlatBufferModel::BuildFromFile(filename, &error_reporter); + std::unique_ptr<tflite::TfLiteVerifier> verifier{new BMFlatBufferVerifier}; + + auto model = FlatBufferModel::VerifyAndBuildFromFile(filename, verifier.get(), &error_reporter); if (model == nullptr) { std::cerr << "Cannot create model" << std::endl; |