| 1225 | } |
| 1226 | |
| 1227 | Status BundleReader::GetSliceValue(StringPiece full_tensor_key, |
| 1228 | const BundleEntryProto& full_tensor_entry, |
| 1229 | const TensorSlice& slice_spec, Tensor* val) { |
| 1230 | using checkpoint::RegisterTensorSlice; |
| 1231 | using checkpoint::TensorSliceSet; |
| 1232 | DCHECK_GE(full_tensor_entry.slices_size(), 0); |
| 1233 | |
| 1234 | const TensorShape full_shape(TensorShape(full_tensor_entry.shape())); |
| 1235 | std::vector<std::pair<TensorSlice, string>> details; |
| 1236 | const string full_tensor_key_string(full_tensor_key); |
| 1237 | const TensorSliceSet* tss = |
| 1238 | gtl::FindPtrOrNull(tensor_slices_, full_tensor_key_string); |
| 1239 | |
| 1240 | // Populates the "full tensor key -> TensorSliceSet" cache. |
| 1241 | if (tss == nullptr) { |
| 1242 | if (full_tensor_entry.slices().empty()) { |
| 1243 | // Special case: a writer has saved a tensor fully, but the reader wants |
| 1244 | // to read in slices. We therefore register the full slice on-demand here |
| 1245 | // without further complicating the on-disk bundle format. |
| 1246 | TF_RETURN_IF_ERROR(RegisterTensorSlice( |
| 1247 | full_tensor_key_string, full_shape, full_tensor_entry.dtype(), |
| 1248 | /* tag */ "", |
| 1249 | /* full slice */ TensorSlice(full_shape.dims()), &tensor_slices_)); |
| 1250 | } |
| 1251 | for (const TensorSliceProto& slice : full_tensor_entry.slices()) { |
| 1252 | TF_RETURN_IF_ERROR(RegisterTensorSlice( |
| 1253 | full_tensor_key_string, full_shape, full_tensor_entry.dtype(), |
| 1254 | /* tag */ "", TensorSlice(slice), &tensor_slices_)); |
| 1255 | } |
| 1256 | tss = gtl::FindPtrOrNull(tensor_slices_, full_tensor_key_string); |
| 1257 | CHECK_NE(tss, nullptr); |
| 1258 | } |
| 1259 | if (!tss->QueryMeta(slice_spec, &details)) { |
| 1260 | return errors::InvalidArgument( |
| 1261 | "Does not have sufficient slices for partitioned tensor ", |
| 1262 | full_tensor_key, |
| 1263 | " to restore in slice_spec: ", slice_spec.DebugString()); |
| 1264 | } |
| 1265 | |
| 1266 | // The union of the slices in "details" covers "slice_spec". Performs the |
| 1267 | // copies from each. |
| 1268 | BundleEntryProto stored_slice_entry = full_tensor_entry; |
| 1269 | for (const auto& slice_tag_pair : details) { |
| 1270 | // Seeks for the stored slice. |
| 1271 | const TensorSlice& stored_slice = slice_tag_pair.first; |
| 1272 | |
| 1273 | // We already have the entry for the full tensor, so don't query again if |
| 1274 | // the slice is full. |
| 1275 | if (!stored_slice.IsFull()) { |
| 1276 | const string encoded_stored_slice_name = |
| 1277 | checkpoint::EncodeTensorNameSlice(full_tensor_key_string, |
| 1278 | stored_slice); |
| 1279 | status_ = |
| 1280 | GetBundleEntryProto(encoded_stored_slice_name, &stored_slice_entry); |
| 1281 | if (!status_.ok()) return status_; |
| 1282 | } |
| 1283 | |
| 1284 | // TODO(zongheng): should we take an OpKernelContext, so that we can call |
nothing calls this directly
no test coverage detected