IVGCVSW-4545 Fix segfault parsing reshape layer
Change-Id: Ib4bbab387cec4780c8aae56fdede090fa265e4ba
Signed-off-by: Derek Lamberti <derek.lamberti@arm.com>
diff --git a/src/armnnConverter/ArmnnConverter.cpp b/src/armnnConverter/ArmnnConverter.cpp
index 70df2c3..e0a659d 100644
--- a/src/armnnConverter/ArmnnConverter.cpp
+++ b/src/armnnConverter/ArmnnConverter.cpp
@@ -420,68 +420,76 @@
ArmnnConverter converter(modelPath, inputNames, inputTensorShapes, outputNames, outputPath, isModelBinary);
- if (modelFormat.find("caffe") != std::string::npos)
+ try
{
+ if (modelFormat.find("caffe") != std::string::npos)
+ {
#if defined(ARMNN_CAFFE_PARSER)
- if (!converter.CreateNetwork<armnnCaffeParser::ICaffeParser>())
- {
- ARMNN_LOG(fatal) << "Failed to load model from file";
- return EXIT_FAILURE;
- }
+ if (!converter.CreateNetwork<armnnCaffeParser::ICaffeParser>())
+ {
+ ARMNN_LOG(fatal) << "Failed to load model from file";
+ return EXIT_FAILURE;
+ }
#else
- ARMNN_LOG(fatal) << "Not built with Caffe parser support.";
- return EXIT_FAILURE;
+ ARMNN_LOG(fatal) << "Not built with Caffe parser support.";
+ return EXIT_FAILURE;
#endif
- }
- else if (modelFormat.find("onnx") != std::string::npos)
- {
+ }
+ else if (modelFormat.find("onnx") != std::string::npos)
+ {
#if defined(ARMNN_ONNX_PARSER)
- if (!converter.CreateNetwork<armnnOnnxParser::IOnnxParser>())
- {
- ARMNN_LOG(fatal) << "Failed to load model from file";
- return EXIT_FAILURE;
- }
+ if (!converter.CreateNetwork<armnnOnnxParser::IOnnxParser>())
+ {
+ ARMNN_LOG(fatal) << "Failed to load model from file";
+ return EXIT_FAILURE;
+ }
#else
- ARMNN_LOG(fatal) << "Not built with Onnx parser support.";
- return EXIT_FAILURE;
+ ARMNN_LOG(fatal) << "Not built with Onnx parser support.";
+ return EXIT_FAILURE;
#endif
- }
- else if (modelFormat.find("tensorflow") != std::string::npos)
- {
+ }
+ else if (modelFormat.find("tensorflow") != std::string::npos)
+ {
#if defined(ARMNN_TF_PARSER)
- if (!converter.CreateNetwork<armnnTfParser::ITfParser>())
- {
- ARMNN_LOG(fatal) << "Failed to load model from file";
- return EXIT_FAILURE;
- }
+ if (!converter.CreateNetwork<armnnTfParser::ITfParser>())
+ {
+ ARMNN_LOG(fatal) << "Failed to load model from file";
+ return EXIT_FAILURE;
+ }
#else
- ARMNN_LOG(fatal) << "Not built with Tensorflow parser support.";
- return EXIT_FAILURE;
+ ARMNN_LOG(fatal) << "Not built with Tensorflow parser support.";
+ return EXIT_FAILURE;
#endif
- }
- else if (modelFormat.find("tflite") != std::string::npos)
- {
+ }
+ else if (modelFormat.find("tflite") != std::string::npos)
+ {
#if defined(ARMNN_TF_LITE_PARSER)
- if (!isModelBinary)
- {
- ARMNN_LOG(fatal) << "Unknown model format: '" << modelFormat << "'. Only 'binary' format supported \
- for tflite files";
- return EXIT_FAILURE;
- }
+ if (!isModelBinary)
+ {
+ ARMNN_LOG(fatal) << "Unknown model format: '" << modelFormat << "'. Only 'binary' format supported \
+ for tflite files";
+ return EXIT_FAILURE;
+ }
- if (!converter.CreateNetwork<armnnTfLiteParser::ITfLiteParser>())
+ if (!converter.CreateNetwork<armnnTfLiteParser::ITfLiteParser>())
+ {
+ ARMNN_LOG(fatal) << "Failed to load model from file";
+ return EXIT_FAILURE;
+ }
+#else
+ ARMNN_LOG(fatal) << "Not built with TfLite parser support.";
+ return EXIT_FAILURE;
+#endif
+ }
+ else
{
- ARMNN_LOG(fatal) << "Failed to load model from file";
+ ARMNN_LOG(fatal) << "Unknown model format: '" << modelFormat << "'";
return EXIT_FAILURE;
}
-#else
- ARMNN_LOG(fatal) << "Not built with TfLite parser support.";
- return EXIT_FAILURE;
-#endif
}
- else
+ catch(armnn::Exception& e)
{
- ARMNN_LOG(fatal) << "Unknown model format: '" << modelFormat << "'";
+ ARMNN_LOG(fatal) << "Failed to load model from file: " << e.what();
return EXIT_FAILURE;
}
diff --git a/src/armnnTfLiteParser/TfLiteParser.cpp b/src/armnnTfLiteParser/TfLiteParser.cpp
index eab9f4e..fc5041b 100644
--- a/src/armnnTfLiteParser/TfLiteParser.cpp
+++ b/src/armnnTfLiteParser/TfLiteParser.cpp
@@ -31,6 +31,14 @@
#include <algorithm>
#include <limits>
#include <numeric>
+#include <sstream>
+
+#define ARMNN_THROW_PARSE_EXCEPTION(msg) \
+ { \
+ throw armnn::ParseException( static_cast<const std::stringstream&>( std::stringstream() << msg \
+ << ": " \
+ << CHECK_LOCATION().AsString()).str()); \
+ }
using namespace armnn;
using armnn::CheckLocation;
@@ -1932,8 +1940,51 @@
armnn::TensorInfo inputTensorInfo = ToTensorInfo(inputs[0]);
armnn::TensorInfo actualOutputTensorInfo = ToTensorInfo(outputs[0]);
+
+ std::vector<int32_t> targetShape;
+ if (inputs.size() > 1 && inputs[1] != nullptr)
+ {
+ if (options != nullptr)
+ {
+ ARMNN_THROW_PARSE_EXCEPTION("Target shape defined in reshape parameters and input tensor. "
+ "Only one method expected");
+ }
+
+ if (inputs[1]->is_variable)
+ {
+ ARMNN_THROW_PARSE_EXCEPTION( "Target shapes defined in non-const input tensors is not supported");
+ }
+
+ if (inputs[1]->shape.size() != 1)
+ {
+ ARMNN_THROW_PARSE_EXCEPTION("Target 'shape' input is not a 1D tensor");
+ }
+
+ if (inputs[1]->type != tflite::TensorType_INT32)
+ {
+ ARMNN_THROW_PARSE_EXCEPTION("Target 'shape' input is not an int32 type");
+ }
+
+ auto bufferPtr = GetBuffer(m_Model, inputs[1]->buffer);
+ auto vals = reinterpret_cast<const int32_t*>(bufferPtr->data.data());
+ for (int i=0; i < inputs[1]->shape[0]; i++)
+ {
+ targetShape.push_back(vals[i]);
+ }
+ }
+ else
+ {
+ if (options == nullptr)
+ {
+ ARMNN_THROW_PARSE_EXCEPTION("Target shape not defined in reshape parameters or input tensor. "
+ "At least one method required");
+ }
+
+ targetShape = options->new_shape;
+ }
+
armnn::TensorInfo reshapeOutputTensorInfo =
- TfLiteParser::OutputShapeOfReshape(inputTensorInfo, options->new_shape);
+ TfLiteParser::OutputShapeOfReshape(inputTensorInfo, targetShape);
// Check for valid input size and that reshape parameters equal output shape
const armnn::TensorShape& reshapeOutputTensorShape = reshapeOutputTensorInfo.GetShape();
@@ -2581,10 +2632,12 @@
boost::filesystem::path pathToFile(fileName);
if (!boost::filesystem::exists(pathToFile, errorCode))
{
- throw FileNotFoundException(boost::str(boost::format("Cannot find the file (%1%) errorCode: %2% %3%") %
- fileName %
- errorCode %
- CHECK_LOCATION().AsString()));
+ std::string locationString = CHECK_LOCATION().AsString();
+ std::string msg = boost::str(boost::format("Cannot find the file (%1%) errorCode: %2% %3%") %
+ fileName %
+ errorCode %
+ locationString);
+ throw FileNotFoundException(msg);
}
std::ifstream file(fileName, std::ios::binary);
std::string fileContent((std::istreambuf_iterator<char>(file)), std::istreambuf_iterator<char>());
