| 40 | namespace tensorflow { |
| 41 | |
| 42 | void SaveTensors( |
| 43 | OpKernelContext* context, |
| 44 | checkpoint::TensorSliceWriter::CreateBuilderFunction builder_func, |
| 45 | bool save_slices) { |
| 46 | const Tensor& filename_t = context->input(0); |
| 47 | { |
| 48 | const int64 size = filename_t.NumElements(); |
| 49 | OP_REQUIRES( |
| 50 | context, size == 1, |
| 51 | errors::InvalidArgument( |
| 52 | "Input 0 (filename) must be a string scalar; got a tensor of ", |
| 53 | size, "elements")); |
| 54 | } |
| 55 | |
| 56 | // Path, names, and slices if save_slices is true. |
| 57 | const int kFixedInputs = save_slices ? 3 : 2; |
| 58 | const Tensor& tensor_names_t = context->input(1); |
| 59 | OP_REQUIRES(context, |
| 60 | FastBoundsCheck(tensor_names_t.NumElements() + kFixedInputs, |
| 61 | std::numeric_limits<int>::max()), |
| 62 | errors::InvalidArgument("Too many inputs to SaveTensors")); |
| 63 | const int N = static_cast<int>(tensor_names_t.NumElements()); |
| 64 | const tstring* tensor_shapes_and_slices_ptr = nullptr; |
| 65 | if (save_slices) { |
| 66 | const Tensor& tensor_shapes_and_slices_t = context->input(2); |
| 67 | OP_REQUIRES( |
| 68 | context, |
| 69 | tensor_shapes_and_slices_t.NumElements() == static_cast<int64>(N), |
| 70 | errors::InvalidArgument("Expected ", N, |
| 71 | " elements for the tensor " |
| 72 | "shapes and slices but got ", |
| 73 | tensor_shapes_and_slices_t.NumElements())); |
| 74 | tensor_shapes_and_slices_ptr = |
| 75 | tensor_shapes_and_slices_t.flat<tstring>().data(); |
| 76 | } |
| 77 | OP_REQUIRES(context, context->num_inputs() == N + kFixedInputs, |
| 78 | errors::InvalidArgument("Expected totally ", N + kFixedInputs, |
| 79 | " inputs as input #1 (which is a string " |
| 80 | "tensor of saved names) contains ", |
| 81 | N, " names, but received ", |
| 82 | context->num_inputs(), " inputs")); |
| 83 | |
| 84 | VLOG(1) << "About to save tensors to file " << filename_t.flat<tstring>()(0) |
| 85 | << "..."; |
| 86 | checkpoint::TensorSliceWriter writer(filename_t.flat<tstring>()(0), |
| 87 | std::move(builder_func)); |
| 88 | |
| 89 | Status s; |
| 90 | auto tensor_names_flat = tensor_names_t.flat<tstring>(); |
| 91 | |
| 92 | // Process tensors in sorted name order. This allows us to avoid seeking |
| 93 | // during restoration in the common case where we are restoring a full |
| 94 | // checkpoint. |
| 95 | std::vector<size_t> sorted_name_idx(tensor_names_flat.size()); |
| 96 | std::iota(sorted_name_idx.begin(), sorted_name_idx.end(), 0); |
| 97 | std::sort(sorted_name_idx.begin(), sorted_name_idx.end(), |
| 98 | [&tensor_names_flat](size_t a, size_t b) { |
| 99 | return tensor_names_flat(a) < tensor_names_flat(b); |
no test coverage detected