OnnxModel.h

浏览该文件的文档.

#pragma once
#include "common/Export.h"
 
 
#include "common/Result.h"
#include "tensor/OnnxStorage.h"
 
#include <cstddef>
#include <cstdint>
#include <memory>
#include <span>
#include <string>
#include <vector>
 
namespace eve::tensor {
class OnnxCompute;
enum class OnnxElement : int { Float32 = 1, UInt8 = 2, Int8 = 3, Int32 = 6, Int64 = 7, Bool = 9 };
 
struct OnnxTensor {
    OnnxElement          element = OnnxElement::Float32;
    std::vector<int64_t> shape;
    std::vector<uint8_t> bytes;
};
struct OnnxNamedTensor {
    std::string name;
    OnnxTensor  tensor;
};
struct OnnxRunOptions {
    uint64_t seed          = 0;
    bool     requireFinite = false;
};
struct OnnxGpuResult {
    std::vector<OnnxNamedTensor> outputs;
    size_t                       dispatches = 0;
    OnnxTransferStats            transfers;
};
struct OnnxModelInfo {
    int64_t                  irVersion = 0;
    size_t                   nodeCount = 0, initializerBytes = 0;
    std::vector<std::string> inputs, outputs;
    std::vector<std::string> unsupportedNodes;
};
 
class EVENGINE_API_DOMAINS OnnxModel {
public:
    ~OnnxModel();
    [[nodiscard]] static Result<std::unique_ptr<OnnxModel>> load(std::span<const uint8_t> bytes);
    [[nodiscard]] OnnxModelInfo info() const;
    [[nodiscard]] Result<std::vector<OnnxNamedTensor>> run(std::span<const OnnxNamedTensor> feeds,
                                                           std::span<const std::string>     requested = {},
                                                           OnnxRunOptions                   options   = {}) const;
 
    [[nodiscard]] Result<OnnxGpuResult> runGpu(std::span<const OnnxNamedTensor> feeds, OnnxCompute& compute,
                                               std::span<const std::string> requested = {},
                                               OnnxRunOptions               options   = {}) const;
 
private:
    [[nodiscard]] Result<std::vector<OnnxNamedTensor>> runInternal(std::span<const OnnxNamedTensor> feeds,
                                                                   std::span<const std::string>     requested,
                                                                   OnnxCompute* compute, OnnxRunOptions options) const;
    struct Impl;
    explicit OnnxModel(std::unique_ptr<Impl> impl);
    std::unique_ptr<Impl> impl_;
};
}  // namespace eve::tensor