| 817 | } // namespace |
| 818 | |
| 819 | Status RestoreTensorsV2(OpKernelContext* context, const Tensor& prefix, |
| 820 | const Tensor& tensor_names, |
| 821 | const Tensor& shape_and_slices, |
| 822 | gtl::ArraySlice<DataType> dtypes) { |
| 823 | const string& prefix_string = prefix.scalar<tstring>()(); |
| 824 | |
| 825 | const auto& tensor_names_flat = tensor_names.flat<tstring>(); |
| 826 | const auto& shape_and_slices_flat = shape_and_slices.flat<tstring>(); |
| 827 | |
| 828 | // Sort lookup keys to improve locality when reading multiple tensors. |
| 829 | std::vector<size_t> sorted_name_idx(tensor_names_flat.size()); |
| 830 | std::iota(sorted_name_idx.begin(), sorted_name_idx.end(), 0); |
| 831 | std::sort(sorted_name_idx.begin(), sorted_name_idx.end(), |
| 832 | [&tensor_names_flat](size_t a, size_t b) { |
| 833 | return tensor_names_flat(a) < tensor_names_flat(b); |
| 834 | }); |
| 835 | |
| 836 | std::vector<std::unique_ptr<RestoreOp> > pool_restore_ops; |
| 837 | std::vector<std::unique_ptr<RestoreOp> > direct_restore_ops; |
| 838 | |
| 839 | BundleReader default_reader(Env::Default(), prefix_string); |
| 840 | TF_RETURN_IF_ERROR(default_reader.status()); |
| 841 | |
| 842 | std::vector<string> mismatched_errors; |
| 843 | for (const size_t i : sorted_name_idx) { |
| 844 | TensorShape restored_full_shape; |
| 845 | DataType original_dtype; |
| 846 | const string& tensor_name = tensor_names_flat(i); |
| 847 | TF_RETURN_IF_ERROR(default_reader.LookupDtypeAndShape( |
| 848 | tensor_name, &original_dtype, &restored_full_shape)); |
| 849 | if (dtypes[i] != original_dtype) { |
| 850 | string error_msg = strings::StrCat( |
| 851 | "tensor_name = ", tensor_name, "; expected dtype ", |
| 852 | DataTypeString(dtypes[i]), " does not equal original dtype ", |
| 853 | DataTypeString(original_dtype)); |
| 854 | mismatched_errors.emplace_back(error_msg); |
| 855 | } |
| 856 | } |
| 857 | if (!mismatched_errors.empty()) { |
| 858 | const string error_msg = absl::StrJoin(mismatched_errors, "\n"); |
| 859 | return errors::InvalidArgument(error_msg); |
| 860 | } |
| 861 | |
| 862 | for (auto i : sorted_name_idx) { |
| 863 | const string& tensor_name = tensor_names_flat(i); |
| 864 | const string& shape_and_slice = shape_and_slices_flat(i); |
| 865 | auto op = |
| 866 | new RestoreOp{context, i, tensor_name, shape_and_slice, prefix_string}; |
| 867 | if (op->should_run_in_pool(&default_reader)) { |
| 868 | pool_restore_ops.emplace_back(op); |
| 869 | } else { |
| 870 | direct_restore_ops.emplace_back(op); |
| 871 | } |
| 872 | } |
| 873 | |
| 874 | { |
| 875 | // Schedule any threaded operations first, skipping thread pool creation if |
| 876 | // we don't have any expensive operations. |
no test coverage detected