saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Use the file size while reading onnx models. Ensure models are loaded using APIs in model.h for consistency. (#4399)

Browse source 
* Use the file size while reading onnx models. Ensure models are loaded using APIs in model.h for consistency.

* Refactor existing GetFileLength in posix.cc and address PR comments.

* Fix linux build - signed/unsigned conversion
This commit is contained in:
Pranav Sharma 2020-07-02 17:30:53 -07:00 committed by GitHub
parent d22f6fddf7
commit 4df8a1e240
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 60 additions and 22 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
onnxruntime/core/graph/model.cc
View file

@ -26,6 +26,8 @@ using namespace ONNX_NAMESPACE;
using namespace onnxruntime;
using namespace onnxruntime::common;
static constexpr int DEFAULT_PROTOBUF_BLOCK_SIZE = 4 * 1024 * 1024;
namespace onnxruntime {
Model::Model(const std::string& graph_name,
bool is_onnx_domain_only,
@ -237,6 +239,7 @@ Status Model::Load(std::istream& model_istream, ModelProto* p_model_proto) {
if (!p_model_proto) {
return Status(ONNXRUNTIME, INVALID_ARGUMENT, "Null model_proto ptr.");
}
google::protobuf::io::IstreamInputStream zero_copy_input(&model_istream);
const bool result = p_model_proto->ParseFromZeroCopyStream(&zero_copy_input) && model_istream.eof();
if (!result) {
@ -447,7 +450,13 @@ Status Model::Load(int fd, ONNX_NAMESPACE::ModelProto& model_proto) {
}
#if GOOGLE_PROTOBUF_VERSION >= 3002000
FileInputStream input(fd);
size_t file_size = 0;
int block_size = -1;
Status st = Env::Default().GetFileLength(fd, file_size);
if (st.IsOK()) {
block_size = std::min(DEFAULT_PROTOBUF_BLOCK_SIZE, static_cast<int>(file_size));
}
FileInputStream input(fd, block_size);
const bool result = model_proto.ParseFromZeroCopyStream(&input) && input.GetErrno() == 0;
if (!result) {
return Status(ONNXRUNTIME, INVALID_PROTOBUF, "Protobuf parsing failed.");

1
onnxruntime/core/platform/env.h
View file

@ -134,6 +134,7 @@ class Env {
* Gets the length of the specified file.
*/
virtual common::Status GetFileLength(_In_z_ const ORTCHAR_T* file_path, size_t& length) const = 0;
virtual common::Status GetFileLength(int fd, /*out*/ size_t& file_size) const = 0;
/**
* Copies the content of the file into the provided buffer.

27
onnxruntime/core/platform/posix/env.cc
View file

@ -231,21 +231,30 @@ class PosixEnv : public Env {
Status GetFileLength(const PathChar* file_path, size_t& length) const override {
ScopedFileDescriptor file_descriptor{open(file_path, O_RDONLY)};
if (file_descriptor.Get() < 0) {
return ReportSystemError("open", file_path);
return GetFileLength(file_descriptor.Get(), length);
}
common::Status GetFileLength(int fd, /*out*/ size_t& file_size) const override {
using namespace common;
if (fd < 0) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Invalid fd was supplied: ", fd);
}
struct stat stbuf;
if (fstat(file_descriptor.Get(), &stbuf) != 0) {
return ReportSystemError("fstat", file_path);
struct stat buf;
int rc = fstat(fd, &buf);
if (rc < 0) {
return ReportSystemError("fstat", "");
}
if (!S_ISREG(stbuf.st_mode)) {
return common::Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT,
"GetFileLength: input is not a regular file");
if (buf.st_size < 0) {
return ORT_MAKE_STATUS(SYSTEM, FAIL, "Received negative size from stat call");
}
length = static_cast<size_t>(stbuf.st_size);
if (static_cast<unsigned long long>(buf.st_size) > std::numeric_limits<size_t>::max()) {
return ORT_MAKE_STATUS(SYSTEM, FAIL, "File is too large.");
}
file_size = static_cast<size_t>(buf.st_size);
return Status::OK();
}

24
onnxruntime/core/platform/windows/env.cc
View file

@ -193,6 +193,30 @@ class WindowsEnv : public Env {
return Status::OK();
}
common::Status GetFileLength(int fd, /*out*/ size_t& file_size) const override {
using namespace common;
if (fd < 0) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Invalid fd was supplied: ", fd);
}
struct _stat buf;
int rc = _fstat(fd, &buf);
if (rc < 0) {
return Status(SYSTEM, errno);
}
if (buf.st_size < 0) {
return ORT_MAKE_STATUS(SYSTEM, FAIL, "Received negative size from stat call");
}
if (static_cast<unsigned long long>(buf.st_size) > std::numeric_limits<size_t>::max()) {
return ORT_MAKE_STATUS(SYSTEM, FAIL, "File is too large.");
}
file_size = static_cast<size_t>(buf.st_size);
return Status::OK();
}
Status ReadFileIntoBuffer(_In_z_ const ORTCHAR_T* const file_path, const FileOffsetType offset, const size_t length,
const gsl::span<char> buffer) const override {
ORT_RETURN_IF_NOT(file_path);

14
onnxruntime/core/session/inference_session.cc
View file

@ -263,9 +263,8 @@ InferenceSession::InferenceSession(const SessionOptions& session_options, const
: graph_transformation_mgr_(session_options.max_num_graph_transformation_steps),
logging_manager_(session_env.GetLoggingManager()),
insert_cast_transformer_("CastFloat16Transformer") {
google::protobuf::io::IstreamInputStream zero_copy_input(&model_istream);
const bool result = model_proto_.ParseFromZeroCopyStream(&zero_copy_input) && model_istream.eof();
ORT_ENFORCE(result, "Could not parse model successfully while constructing the inference session");
Status st = Model::Load(model_istream, &model_proto_);
ORT_ENFORCE(st.IsOK(), "Could not parse model successfully while constructing the inference session");
is_model_proto_parsed_ = true;
// Finalize session options and initialize assets of this session instance
ConstructorCommon(session_options, session_env);
@ -531,12 +530,9 @@ common::Status InferenceSession::Load(std::istream& model_istream) {
auto loader = [this, &model_istream](std::shared_ptr<onnxruntime::Model>& model) {
ModelProto model_proto;
google::protobuf::io::IstreamInputStream zero_copy_input(&model_istream);
const bool result = model_proto.ParseFromZeroCopyStream(&zero_copy_input) && model_istream.eof();
if (!result) {
return Status(common::ONNXRUNTIME, common::INVALID_PROTOBUF,
"Failed to load model because protobuf parsing failed.");
Status st = Model::Load(model_istream, &model_proto);
if (!st.IsOK()) {
return st;
}
#ifdef ENABLE_LANGUAGE_INTEROP_OPS
LoadInterOp(model_proto, interop_domains_, [&](const char* msg) { LOGS(*session_logger_, WARNING) << msg; });

5
onnxruntime/test/providers/memcpy_test.cc
View file

@ -39,9 +39,8 @@ TEST(MemcpyTest, copy1) {
ONNX_NAMESPACE::ModelProto mp;
std::ifstream model_istream("testdata/matmul_1.onnx", std::ifstream::in | std::ifstream::binary);
google::protobuf::io::IstreamInputStream zero_copy_input(&model_istream);
const bool result = mp.ParseFromZeroCopyStream(&zero_copy_input) && model_istream.eof();
ASSERT_TRUE(result);
st = Model::Load(model_istream, &mp);
ASSERT_STATUS_OK(st);
Model model(mp, nullptr, DefaultLoggingManager().DefaultLogger());
ASSERT_STATUS_OK(model.MainGraph().Resolve());