experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Merge pull request #27382 from hlsyounes:master

Browse Source 
PiperOrigin-RevId: 241816998
This commit is contained in:
TensorFlower Gardener 2019-04-03 15:38:33 -07:00
commit 6009b2288a
5 changed files with 11 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
tensorflow/core/platform/env.cc
View File

@ -517,7 +517,8 @@ Status ReadBinaryProto(Env* env, const string& fname,
// respectively.
coded_stream.SetTotalBytesLimit(1024LL << 20, 512LL << 20);
if (!proto->ParseFromCodedStream(&coded_stream)) {
if (!proto->ParseFromCodedStream(&coded_stream) ||
!coded_stream.ConsumedEntireMessage()) {
TF_RETURN_IF_ERROR(stream->status());
return errors::DataLoss("Can't parse ", fname, " as binary proto");
}

3
tensorflow/core/util/memmapped_file_system.cc
View File

@ -235,8 +235,7 @@ Status MemmappedFileSystem::InitializeFromFile(Env* env,
if (!directory_
.insert(std::make_pair(
element_iter->name(),
FileRegion(element_iter->offset(),
prev_element_offset - element_iter->offset())))
FileRegion(element_iter->offset(), element_iter->length())))
.second) {
return errors::DataLoss("Corrupted memmapped model file: ", filename,
" Duplicate name of internal component ",

2
tensorflow/core/util/memmapped_file_system.proto
View File

@ -15,12 +15,14 @@ limitations under the License.
syntax = "proto3";
package tensorflow;
option cc_enable_arenas = true;
// A message that describes one region of memmapped file.
message MemmappedFileSystemDirectoryElement {
uint64 offset = 1;
string name = 2;
uint64 length = 3;
}
// A directory of regions in a memmapped file.

8
tensorflow/core/util/memmapped_file_system_writer.cc
View File

@ -47,7 +47,7 @@ Status MemmappedFileSystemWriter::SaveTensor(const Tensor& tensor,
}
// Adds pad for correct alignment after memmapping.
TF_RETURN_IF_ERROR(AdjustAlignment(Allocator::kAllocatorAlignment));
AddToDirectoryElement(element_name);
AddToDirectoryElement(element_name, tensor_data.size());
const auto result = output_file_->Append(tensor_data);
if (result.ok()) {
output_file_offset_ += tensor_data.size();
@ -69,8 +69,8 @@ Status MemmappedFileSystemWriter::SaveProtobuf(
MemmappedFileSystem::kMemmappedPackagePrefix,
" and include [A-Za-z0-9_.]");
}
AddToDirectoryElement(element_name);
const string encoded = message.SerializeAsString();
AddToDirectoryElement(element_name, encoded.size());
const auto res = output_file_->Append(encoded);
if (res.ok()) {
output_file_offset_ += encoded.size();
@ -124,11 +124,13 @@ Status MemmappedFileSystemWriter::AdjustAlignment(uint64 alignment) {
return Status::OK();
}
void MemmappedFileSystemWriter::AddToDirectoryElement(const string& name) {
void MemmappedFileSystemWriter::AddToDirectoryElement(const string& name,
uint64 length) {
MemmappedFileSystemDirectoryElement* new_directory_element =
directory_.add_element();
new_directory_element->set_offset(output_file_offset_);
new_directory_element->set_name(name);
new_directory_element->set_length(length);
}
} // namespace tensorflow

2
tensorflow/core/util/memmapped_file_system_writer.h
View File

@ -40,7 +40,7 @@ class MemmappedFileSystemWriter {
private:
Status AdjustAlignment(uint64 alignment);
void AddToDirectoryElement(const string& element_name);
void AddToDirectoryElement(const string& element_name, uint64 length);
MemmappedFileSystemDirectory directory_;
// The current offset in the file, to support alignment.
uint64 output_file_offset_ = 0;