OnnxParser.hpp

Go to the documentation of this file.

//
// Copyright © 2017 Arm Ltd. All rights reserved.
// SPDX-License-Identifier: MIT
//
#pragma once

#include "armnnOnnxParser/IOnnxParser.hpp"
#include "google/protobuf/repeated_field.h"
#include <unordered_map>

#include <onnx/onnx.pb.h>


namespace armnn
{
class TensorInfo;
enum class ActivationFunction;
}

namespace armnnOnnxParser
{

using ModelPtr = std::unique_ptr<onnx::ModelProto>;

class OnnxParser : public IOnnxParser
{

using OperationParsingFunction = void(OnnxParser::*)(const onnx::NodeProto& NodeProto);

public:

    using GraphPtr = std::unique_ptr<onnx::GraphProto>;

    virtual armnn::INetworkPtr CreateNetworkFromBinaryFile(const char* graphFile) override;

    virtual armnn::INetworkPtr CreateNetworkFromTextFile(const char* graphFile) override;

    virtual armnn::INetworkPtr CreateNetworkFromString(const std::string& protoText) override;

    virtual BindingPointInfo GetNetworkInputBindingInfo(const std::string& name) const override;

    virtual BindingPointInfo GetNetworkOutputBindingInfo(const std::string& name) const override;

public:

    OnnxParser();

    static ModelPtr LoadModelFromBinaryFile(const char * fileName);
    static ModelPtr LoadModelFromTextFile(const char * fileName);
    static ModelPtr LoadModelFromString(const std::string& inputString);

    static std::vector<std::string> GetInputs(ModelPtr& model);

    static std::vector<std::string> GetOutputs(ModelPtr& model);

private:

    armnn::INetworkPtr CreateNetworkFromModel(onnx::ModelProto& model);

    void LoadGraph();

    void SetupInfo(const google::protobuf::RepeatedPtrField<onnx::ValueInfoProto >* list);

    std::vector<armnn::TensorInfo> ComputeOutputInfo(std::vector<std::string> outNames,
                                                     const armnn::IConnectableLayer* layer,
                                                     std::vector<armnn::TensorShape> inputShapes);

    void DetectFullyConnected();

    template <typename Location>
    void GetInputAndParam(const onnx::NodeProto& node,
                          std::string* inputName,
                          std::string* constName,
                          const Location& location);

    template <typename Location>
    void To1DTensor(const std::string &name, const Location& location);

    //Broadcast Preparation functions
    std::pair<std::string, std::string> AddPrepareBroadcast(const std::string& input0, const std::string& input1);
    void PrependForBroadcast(const std::string& outputName, const std::string& input0, const std::string& input1);

    void CreateConstantLayer(const std::string& tensorName, const std::string& layerName);
    void CreateReshapeLayer(const std::string& inputName,
                            const std::string& outputName,
                            const std::string& layerName);

    void ParseBatchNormalization(const onnx::NodeProto& node);
    void ParseConstant(const onnx::NodeProto& nodeProto);

    void ParseMaxPool(const onnx::NodeProto& nodeProto);
    void ParseAveragePool(const onnx::NodeProto& nodeProto);
    void ParseGlobalAveragePool(const onnx::NodeProto& node);

    void AddPoolingLayer(const onnx::NodeProto& nodeProto, armnn::Pooling2dDescriptor& desc);

    void ParseReshape(const onnx::NodeProto& nodeProto);

    void ParseActivation(const onnx::NodeProto& nodeProto, const armnn::ActivationFunction func);
    void ParseSigmoid(const onnx::NodeProto& nodeProto);
    void ParseTanh(const onnx::NodeProto& nodeProto);
    void ParseRelu(const onnx::NodeProto& nodeProto);
    void ParseLeakyRelu(const onnx::NodeProto& nodeProto);

    void AddConvLayerWithDepthwiseConv(const onnx::NodeProto& node, const armnn::Convolution2dDescriptor& convDesc);
    void ParseConv(const onnx::NodeProto& nodeProto);

    void ParseAdd(const onnx::NodeProto& nodeProto);
    void AddFullyConnected(const onnx::NodeProto& matmulNode, const onnx::NodeProto* addNode = nullptr);

    void RegisterInputSlots(armnn::IConnectableLayer* layer, const std::vector<std::string>& tensorIndexes);
    void RegisterOutputSlots(armnn::IConnectableLayer* layer, const std::vector<std::string>& tensorIndexes);

    void SetupInputLayers();
    void SetupOutputLayers();

    void ResetParser();
    void Cleanup();

    std::pair<armnn::ConstTensor, std::unique_ptr<float[]>> CreateConstTensor(const std::string name);

    template <typename TypeList, typename Location>
    void ValidateInputs(const onnx::NodeProto& node,
                        TypeList validInputs,
                        const Location& location);

    armnn::INetworkPtr m_Network;

    GraphPtr m_Graph;

    struct OnnxTensor
    {
        std::unique_ptr<armnn::TensorInfo>          m_info;
        std::unique_ptr<const onnx::TensorProto>    m_tensor;
        onnx::TensorProto::DataType                 m_dtype;

        OnnxTensor() : m_info(nullptr), m_tensor(nullptr), m_dtype(onnx::TensorProto::FLOAT) { }
        bool isConstant() { return m_tensor != nullptr; }
    };

    std::unordered_map<std::string, OnnxTensor> m_TensorsInfo;

    static const std::map<std::string, OperationParsingFunction> m_ParserFunctions;

    struct TensorSlots
    {
        armnn::IOutputSlot* outputSlot;
        std::vector<armnn::IInputSlot*> inputSlots;

        TensorSlots() : outputSlot(nullptr) { }
    };
    std::unordered_map<std::string, TensorSlots> m_TensorConnections;

    std::unordered_map<std::string, std::pair<const onnx::NodeProto*, int>> m_OutputsMap;

    struct UsageSummary
    {
        std::vector<size_t> fusedWithNodes;
        size_t inputForNodes;

        UsageSummary() : fusedWithNodes({}), inputForNodes(0) { }

    };

    std::vector<UsageSummary> m_OutputsFusedAndUsed;
};
}