mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-30 20:18:08 +00:00
6203a03cd9
* Add opset and IR check. * Add test case for future opsets. https://github.com/microsoft/onnxruntime/issues/2371
131 lines
5.1 KiB
C++
131 lines
5.1 KiB
C++
// Copyright (c) Microsoft Corporation. All rights reserved.
|
|
// Licensed under the MIT License.
|
|
|
|
#pragma once
|
|
#include <list>
|
|
#include <unordered_map>
|
|
#include <memory>
|
|
#include <climits>
|
|
#include <string>
|
|
#include "core/graph/graph_viewer.h"
|
|
|
|
#include "gsl/gsl"
|
|
|
|
namespace onnxruntime {
|
|
typedef std::unordered_map<std::string, std::string> ModelMetaData;
|
|
using IOnnxRuntimeOpSchemaRegistryList = std::list<std::shared_ptr<IOnnxRuntimeOpSchemaCollection>>;
|
|
|
|
// A machine learning model representation class.
|
|
// Besides a main <Graph>, it also holds basic information, say,
|
|
// model version, model domain, model author, license etc.
|
|
class Model {
|
|
public:
|
|
static constexpr Version kNoVersion = INT64_MAX;
|
|
|
|
// Construct model from scratch.
|
|
explicit Model(const std::string& graph_name,
|
|
bool is_onnx_domain_only = false,
|
|
const ModelMetaData& model_metadata = ModelMetaData(),
|
|
const IOnnxRuntimeOpSchemaRegistryList& local_registries = IOnnxRuntimeOpSchemaRegistryList(),
|
|
const std::unordered_map<std::string, int>& domain_to_version = {},
|
|
const std::vector<ONNX_NAMESPACE::FunctionProto>& model_specific_functions = {});
|
|
|
|
// NOTE: after calling this constructor, <*this> model will
|
|
// hold a copy of <model_proto>.
|
|
explicit Model(const ONNX_NAMESPACE::ModelProto& model_proto,
|
|
const IOnnxRuntimeOpSchemaRegistryList* local_registries = nullptr);
|
|
|
|
// NOTE: after calling this constructor, <*this> model will
|
|
// own the <model_proto>.
|
|
explicit Model(std::unique_ptr<ONNX_NAMESPACE::ModelProto> model_proto,
|
|
const IOnnxRuntimeOpSchemaRegistryList* local_registries = nullptr);
|
|
|
|
// Get model's IR version.
|
|
// Return <kNoVersion> if not specified.
|
|
Version IrVersion() const;
|
|
|
|
// Get model's producer name.
|
|
// Return null pointer if not specified.
|
|
const std::string& ProducerName() const;
|
|
// Set model's producer name.
|
|
void SetProducerName(const std::string& producer_name);
|
|
|
|
// Get model's producer version.
|
|
// Return null pointer if not specified.
|
|
const std::string& ProducerVersion() const;
|
|
// Set model's producer version.
|
|
void SetProducerVersion(const std::string& producer_version);
|
|
|
|
// Get model's domain.
|
|
// Return null pointer if not specified.
|
|
const std::string& Domain() const;
|
|
// Set models' domain.
|
|
void SetDomain(const std::string& domain);
|
|
|
|
// Get model's version.
|
|
// Return null pointer if not specified.
|
|
Version ModelVersion() const;
|
|
// Set models' version.
|
|
void SetModelversion(onnxruntime::Version model_version);
|
|
|
|
// Get model's doc string.
|
|
// Return null pointer if not specified.
|
|
const std::string& DocString() const;
|
|
// Set models' doc string.
|
|
void SetDocString(const std::string& doc_string);
|
|
|
|
const ModelMetaData& MetaData() const noexcept;
|
|
|
|
// Get model's main graph.
|
|
Graph& MainGraph() noexcept;
|
|
const Graph& MainGraph() const noexcept;
|
|
|
|
// Add function proto to Model
|
|
void AddFunction(const ONNX_NAMESPACE::FunctionProto& func_proto);
|
|
|
|
// Get model's serialization proto data.
|
|
ONNX_NAMESPACE::ModelProto ToProto();
|
|
|
|
#ifdef _WIN32
|
|
static common::Status Save(Model& model, const std::wstring& file_path);
|
|
|
|
// TODO(Task:132) Use of shared_ptr<X>* in Load/Save methods is confusing.
|
|
static common::Status Load(const std::wstring& file_path, /*out*/ std::shared_ptr<Model>& p_model,
|
|
const IOnnxRuntimeOpSchemaRegistryList* local_registry = nullptr);
|
|
#endif
|
|
static common::Status Save(Model& model, const std::string& file_path);
|
|
|
|
static common::Status Save(Model& model, int fd);
|
|
|
|
static common::Status Load(std::istream& model_istream, ONNX_NAMESPACE::ModelProto* p_model_proto);
|
|
|
|
static common::Status Load(const std::string& file_path,
|
|
/*out*/ std::shared_ptr<Model>& p_model,
|
|
const IOnnxRuntimeOpSchemaRegistryList* local_registries = nullptr);
|
|
|
|
static common::Status Load(int fd, /*out*/ std::shared_ptr<Model>& p_model,
|
|
const IOnnxRuntimeOpSchemaRegistryList* local_registries = nullptr);
|
|
|
|
// 'int' rather than 'size_t' because of a protobuf design choice; let callers handle type checks
|
|
static common::Status LoadFromBytes(int count, void* pBytes, /*out*/ std::shared_ptr<Model>& p_model,
|
|
const IOnnxRuntimeOpSchemaRegistryList* local_registries = nullptr);
|
|
|
|
static common::Status Load(const ONNX_NAMESPACE::ModelProto& model_proto, /*out*/ std::shared_ptr<Model>& p_model,
|
|
const IOnnxRuntimeOpSchemaRegistryList* local_registries = nullptr);
|
|
|
|
static common::Status Load(std::unique_ptr<ONNX_NAMESPACE::ModelProto> p_model_proto,
|
|
/*out*/ std::shared_ptr<Model>& p_model,
|
|
const IOnnxRuntimeOpSchemaRegistryList* local_registries = nullptr);
|
|
|
|
private:
|
|
// Model data.
|
|
std::unique_ptr<ONNX_NAMESPACE::ModelProto> model_proto_;
|
|
|
|
// This is a duplication of <model_proto_.metadata_props()>.
|
|
// It gives better accessibility.
|
|
ModelMetaData model_metadata_;
|
|
|
|
// Main graph of the model.
|
|
std::unique_ptr<Graph> graph_;
|
|
};
|
|
} // namespace onnxruntime
|