Internal change
PiperOrigin-RevId: 311474047 Change-Id: I2c8bcfc0c13d5bf82eaeeadc43202171eeecab8b
This commit is contained in:
parent
4afee5f519
commit
9173c6c3d3
@ -503,12 +503,10 @@ Status Reader::ReadTensors(std::vector<Tensor>* read_tensors) {
|
|||||||
size_t tensor_proto_size = tensor_proto_strs[complex_index].second;
|
size_t tensor_proto_size = tensor_proto_strs[complex_index].second;
|
||||||
TensorProto tp;
|
TensorProto tp;
|
||||||
#if defined(PLATFORM_GOOGLE)
|
#if defined(PLATFORM_GOOGLE)
|
||||||
auto tensor_proto_ptr = tensor_proto_str.release();
|
absl::string_view tensor_proto_view(tensor_proto_str.get(),
|
||||||
absl::Cord c;
|
tensor_proto_size);
|
||||||
c.AppendExternalMemory(
|
absl::Cord c = absl::MakeCordFromExternal(
|
||||||
absl::string_view(tensor_proto_ptr, tensor_proto_size),
|
tensor_proto_view, [s = std::move(tensor_proto_str)] {});
|
||||||
tensor_proto_ptr,
|
|
||||||
[](void* arg) { delete[] static_cast<char*>(arg); });
|
|
||||||
if (!tp.ParseFromCord(c)) {
|
if (!tp.ParseFromCord(c)) {
|
||||||
return errors::Internal("Could not parse TensorProto");
|
return errors::Internal("Could not parse TensorProto");
|
||||||
}
|
}
|
||||||
@ -615,11 +613,9 @@ Status Reader::ReadRecord(absl::Cord* record) {
|
|||||||
} else {
|
} else {
|
||||||
auto tmp_str = absl::make_unique<tstring>();
|
auto tmp_str = absl::make_unique<tstring>();
|
||||||
TF_RETURN_IF_ERROR(input_stream_->ReadNBytes(length, tmp_str.get()));
|
TF_RETURN_IF_ERROR(input_stream_->ReadNBytes(length, tmp_str.get()));
|
||||||
tstring* tmp_str_raw = tmp_str.release();
|
absl::string_view tmp_str_view(*tmp_str);
|
||||||
record->AppendExternalMemory(*tmp_str_raw, tmp_str_raw,
|
record->Append(
|
||||||
[](absl::string_view unused_data, void* arg) {
|
absl::MakeCordFromExternal(tmp_str_view, [s = std::move(tmp_str)] {}));
|
||||||
delete static_cast<tstring*>(arg);
|
|
||||||
});
|
|
||||||
return Status::OK();
|
return Status::OK();
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
@ -134,14 +134,7 @@ std::unique_ptr<StringListDecoder> NewStringListDecoder(const string& in) {
|
|||||||
#if defined(TENSORFLOW_PROTOBUF_USES_CORD)
|
#if defined(TENSORFLOW_PROTOBUF_USES_CORD)
|
||||||
void AssignRefCounted(StringPiece src, core::RefCounted* obj, absl::Cord* out) {
|
void AssignRefCounted(StringPiece src, core::RefCounted* obj, absl::Cord* out) {
|
||||||
obj->Ref();
|
obj->Ref();
|
||||||
out->Clear();
|
*out = absl::MakeCordFromExternal(src, [obj] { obj->Unref(); });
|
||||||
// Defines a lambda to unref "obj" when Cord deletes this piece of
|
|
||||||
// memory. +[] converts the lambda to a C style function pointer.
|
|
||||||
auto cleanup = +[](absl::string_view donotcare, void* obj) {
|
|
||||||
reinterpret_cast<core::RefCounted*>(obj)->Unref();
|
|
||||||
};
|
|
||||||
out->AppendExternalMemory(absl::string_view(src.data(), src.size()), obj,
|
|
||||||
cleanup);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
void EncodeStringList(const tstring* strings, int64 n, absl::Cord* out) {
|
void EncodeStringList(const tstring* strings, int64 n, absl::Cord* out) {
|
||||||
|
Loading…
Reference in New Issue
Block a user