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. // respectively.
coded_stream.SetTotalBytesLimit(1024LL << 20, 512LL << 20); 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()); TF_RETURN_IF_ERROR(stream->status());
return errors::DataLoss("Can't parse ", fname, " as binary proto"); 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_ if (!directory_
.insert(std::make_pair( .insert(std::make_pair(
element_iter->name(), element_iter->name(),
FileRegion(element_iter->offset(), FileRegion(element_iter->offset(), element_iter->length())))
prev_element_offset - element_iter->offset())))
.second) { .second) {
return errors::DataLoss("Corrupted memmapped model file: ", filename, return errors::DataLoss("Corrupted memmapped model file: ", filename,
" Duplicate name of internal component ", " 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"; syntax = "proto3";
package tensorflow; package tensorflow;
option cc_enable_arenas = true; option cc_enable_arenas = true;
// A message that describes one region of memmapped file. // A message that describes one region of memmapped file.
message MemmappedFileSystemDirectoryElement { message MemmappedFileSystemDirectoryElement {
uint64 offset = 1; uint64 offset = 1;
string name = 2; string name = 2;
uint64 length = 3;
} }
// A directory of regions in a memmapped file. // 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. // Adds pad for correct alignment after memmapping.
TF_RETURN_IF_ERROR(AdjustAlignment(Allocator::kAllocatorAlignment)); TF_RETURN_IF_ERROR(AdjustAlignment(Allocator::kAllocatorAlignment));
AddToDirectoryElement(element_name); AddToDirectoryElement(element_name, tensor_data.size());
const auto result = output_file_->Append(tensor_data); const auto result = output_file_->Append(tensor_data);
if (result.ok()) { if (result.ok()) {
output_file_offset_ += tensor_data.size(); output_file_offset_ += tensor_data.size();
@ -69,8 +69,8 @@ Status MemmappedFileSystemWriter::SaveProtobuf(
MemmappedFileSystem::kMemmappedPackagePrefix, MemmappedFileSystem::kMemmappedPackagePrefix,
" and include [A-Za-z0-9_.]"); " and include [A-Za-z0-9_.]");
} }
AddToDirectoryElement(element_name);
const string encoded = message.SerializeAsString(); const string encoded = message.SerializeAsString();
AddToDirectoryElement(element_name, encoded.size());
const auto res = output_file_->Append(encoded); const auto res = output_file_->Append(encoded);
if (res.ok()) { if (res.ok()) {
output_file_offset_ += encoded.size(); output_file_offset_ += encoded.size();
@ -124,11 +124,13 @@ Status MemmappedFileSystemWriter::AdjustAlignment(uint64 alignment) {
return Status::OK(); return Status::OK();
} }
void MemmappedFileSystemWriter::AddToDirectoryElement(const string& name) { void MemmappedFileSystemWriter::AddToDirectoryElement(const string& name,
uint64 length) {
MemmappedFileSystemDirectoryElement* new_directory_element = MemmappedFileSystemDirectoryElement* new_directory_element =
directory_.add_element(); directory_.add_element();
new_directory_element->set_offset(output_file_offset_); new_directory_element->set_offset(output_file_offset_);
new_directory_element->set_name(name); new_directory_element->set_name(name);
new_directory_element->set_length(length);
} }
} // namespace tensorflow } // namespace tensorflow

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

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