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

Enlarge protobuf read buffer size

Browse source
This commit is contained in:
Changming Sun 2020-05-22 22:00:08 -07:00
parent eb3aaa70d6
commit 7c83118364
2 changed files with 6 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -237,7 +237,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);
google::protobuf::io::IstreamInputStream zero_copy_input(&model_istream, 1 << 20);
const bool result = p_model_proto->ParseFromZeroCopyStream(&zero_copy_input) && model_istream.eof();
if (!result) {
return Status(ONNXRUNTIME, INVALID_PROTOBUF, "Failed to load model because protobuf parsing failed.");
@ -447,7 +447,8 @@ Status Model::Load(int fd, ONNX_NAMESPACE::ModelProto& model_proto) {
}
#if GOOGLE_PROTOBUF_VERSION >= 3002000
const bool result = model_proto.ParseFromFileDescriptor(fd);
FileInputStream input(fd, 1 << 20);
const bool result = model_proto.ParseFromZeroCopyStream(&input) && input.GetErrno() == 0;
if (!result) {
return Status(ONNXRUNTIME, INVALID_PROTOBUF, "Protobuf parsing failed.");
}

4
onnxruntime/test/onnx/TestCase.cc
View file

@ -176,7 +176,9 @@ class OnnxModelInfo : public TestModelInfo {
}
ONNX_NAMESPACE::ModelProto model_pb;
if (!model_pb.ParseFromFileDescriptor(model_fd)) {
::google::protobuf::io::FileInputStream input(model_fd, 1 << 20);
const bool parse_result = model_pb.ParseFromZeroCopyStream(&input) && input.GetErrno() == 0;
if (!parse_result) {
(void)Env::Default().FileClose(model_fd);
ORT_THROW("Failed to load model because protobuf parsing failed.");
}