ArmnnConverter.cpp File Reference

#include <armnn/Logging.hpp>
#include <HeapProfiling.hpp>
#include <boost/format.hpp>
#include <boost/algorithm/string/split.hpp>
#include <boost/algorithm/string/classification.hpp>
#include <boost/program_options.hpp>
#include <cstdlib>
#include <fstream>
#include <iostream>

Go to the source code of this file.

Functions
int	main (int argc, const char *argv[])

Function Documentation

main()

int main	(	int	argc,
		const char *	argv[]
	)

Definition at line 359 of file ArmnnConverter.cpp.

References ARMNN_LOG, armnn::ConfigureLogging(), armnn::Debug, armnn::Info, and Exception::what().

{

#if (!defined(ARMNN_CAFFE_PARSER)     \
       && !defined(ARMNN_ONNX_PARSER) \
       && !defined(ARMNN_TF_PARSER)   \
       && !defined(ARMNN_TF_LITE_PARSER))
    ARMNN_LOG(fatal) << "Not built with any of the supported parsers, Caffe, Onnx, Tensorflow, or TfLite.";
    return EXIT_FAILURE;
#endif

#if !defined(ARMNN_SERIALIZER)
    ARMNN_LOG(fatal) << "Not built with Serializer support.";
    return EXIT_FAILURE;
#endif

#ifdef NDEBUG
    armnn::LogSeverity level = armnn::LogSeverity::Info;
#else
    armnn::LogSeverity level = armnn::LogSeverity::Debug;
#endif

    armnn::ConfigureLogging(true, true, level);

    std::string modelFormat;
    std::string modelPath;

    std::vector<std::string> inputNames;
    std::vector<std::string> inputTensorShapeStrs;
    std::vector<armnn::TensorShape> inputTensorShapes;

    std::vector<std::string> outputNames;
    std::string outputPath;

    bool isModelBinary = true;

    if (ParseCommandLineArgs(
        argc, argv, modelFormat, modelPath, inputNames, inputTensorShapeStrs, outputNames, outputPath, isModelBinary)
        != EXIT_SUCCESS)
    {
        return EXIT_FAILURE;
    }

    for (const std::string& shapeStr : inputTensorShapeStrs)
    {
        if (!shapeStr.empty())
        {
            std::stringstream ss(shapeStr);

            try
            {
                armnn::TensorShape shape = ParseTensorShape(ss);
                inputTensorShapes.push_back(shape);
            }
            catch (const armnn::InvalidArgumentException& e)
            {
                ARMNN_LOG(fatal) << "Cannot create tensor shape: " << e.what();
                return EXIT_FAILURE;
            }
        }
    }

    ArmnnConverter converter(modelPath, inputNames, inputTensorShapes, outputNames, outputPath, isModelBinary);

    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;
        }
#else
        ARMNN_LOG(fatal) << "Not built with Caffe parser support.";
        return EXIT_FAILURE;
#endif
    }
    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;
        }
#else
        ARMNN_LOG(fatal) << "Not built with Onnx parser support.";
        return EXIT_FAILURE;
#endif
    }
    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;
        }
#else
        ARMNN_LOG(fatal) << "Not built with Tensorflow parser support.";
        return EXIT_FAILURE;
#endif
    }
    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 (!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) << "Unknown model format: '" << modelFormat << "'";
        return EXIT_FAILURE;
    }

    if (!converter.Serialize())
    {
        ARMNN_LOG(fatal) << "Failed to serialize model";
        return EXIT_FAILURE;
    }

    return EXIT_SUCCESS;
}