| 312 | } |
| 313 | |
| 314 | Status SaveHashTableHelper( |
| 315 | BundleWriter* writer, const std::vector<std::pair<int64, int64>>& data, |
| 316 | const std::vector<TensibleVariable*>& tensibles, |
| 317 | const string& table_name, const std::vector<string>& tensibles_name, |
| 318 | int64 slice_beg, int64 slice_length, int64 slice_size) { |
| 319 | if (tensibles.size() != tensibles_name.size()) { |
| 320 | return errors::InvalidArgument("save tensor name error"); |
| 321 | } |
| 322 | if (data.size() == 0 && slice_beg != 0) { |
| 323 | return Status::OK(); |
| 324 | } |
| 325 | TensorSlice slice(1); |
| 326 | slice.set_start(0, slice_beg); |
| 327 | slice.set_length(0, slice_length); |
| 328 | { |
| 329 | TensorSliceProto* tensor_slice_proto; |
| 330 | TF_RETURN_IF_ERROR(writer->AddSliceHeader( |
| 331 | table_name, TensorShape({0}), DataType::DT_INT64, true, |
| 332 | &tensor_slice_proto)); |
| 333 | TensorSlice xslice(1); |
| 334 | xslice.set_start(0, 0); |
| 335 | xslice.set_length(0, data.size()); |
| 336 | xslice.AsProto(tensor_slice_proto); |
| 337 | tensor_slice_proto->set_hash_slice_begin(slice_beg); |
| 338 | tensor_slice_proto->set_hash_slice_length(slice_length); |
| 339 | tensor_slice_proto->set_hash_slice_size(slice_size); |
| 340 | string xname = checkpoint::EncodeTensorNameSlice( |
| 341 | table_name, slice); |
| 342 | SegmentBundleWriter segment_writer( |
| 343 | writer, xname, |
| 344 | TensorShape({(int64)data.size()}), DataType::DT_INT64); |
| 345 | TF_RETURN_IF_ERROR(segment_writer.Begin()); |
| 346 | for (auto&& item : data) { |
| 347 | TF_RETURN_IF_ERROR(segment_writer.WriteData(&item.first, sizeof(int64))); |
| 348 | } |
| 349 | TF_RETURN_IF_ERROR(segment_writer.End()); |
| 350 | } |
| 351 | for (size_t i = 0; i < tensibles.size(); i++) { |
| 352 | TensorShape tsx = tensibles[i]->shape(); |
| 353 | tsx.set_dim(0, 0); |
| 354 | TensorSliceProto* tensor_slice_proto; |
| 355 | TF_RETURN_IF_ERROR(writer->AddSliceHeader( |
| 356 | tensibles_name[i], tsx, tensibles[i]->dtype(), true, |
| 357 | &tensor_slice_proto)); |
| 358 | TensorSlice xslice(tsx.dims()); |
| 359 | xslice.set_start(0, 0); |
| 360 | xslice.set_length(0, data.size()); |
| 361 | xslice.AsProto(tensor_slice_proto); |
| 362 | tensor_slice_proto->set_hash_slice_begin(slice_beg); |
| 363 | tensor_slice_proto->set_hash_slice_length(slice_length); |
| 364 | tensor_slice_proto->set_hash_slice_size(slice_size); |
| 365 | string xname = checkpoint::EncodeTensorNameSlice( |
| 366 | tensibles_name[i], slice); |
| 367 | TensorShape ts = tensibles[i]->shape(); |
| 368 | ts.set_dim(0, data.size()); |
| 369 | SegmentBundleWriter segment_writer( |
| 370 | writer, xname, ts, tensibles[i]->dtype()); |
| 371 | TF_RETURN_IF_ERROR(segment_writer.Begin()); |
no test coverage detected