| 144 | } |
| 145 | |
| 146 | void RestoreTensor(OpKernelContext* context, |
| 147 | checkpoint::TensorSliceReader::OpenTableFunction open_func, |
| 148 | int preferred_shard, bool restore_slice, int restore_index) { |
| 149 | const Tensor& file_pattern_t = context->input(0); |
| 150 | { |
| 151 | const int64 size = file_pattern_t.NumElements(); |
| 152 | OP_REQUIRES( |
| 153 | context, size == 1, |
| 154 | errors::InvalidArgument( |
| 155 | "Input 0 (file_pattern) must be a string scalar; got a tensor of ", |
| 156 | size, " elements")); |
| 157 | } |
| 158 | const string& file_pattern = file_pattern_t.flat<tstring>()(0); |
| 159 | |
| 160 | const Tensor& tensor_name_t = context->input(1); |
| 161 | { |
| 162 | const int64_t size = tensor_name_t.NumElements(); |
| 163 | OP_REQUIRES(context, size > restore_index, |
| 164 | errors::InvalidArgument( |
| 165 | "Input 1 (file_pattern) must be a have at least ", |
| 166 | restore_index + 1, " elements")); |
| 167 | } |
| 168 | const string& tensor_name = tensor_name_t.flat<tstring>()(restore_index); |
| 169 | |
| 170 | // If we cannot find a cached reader we will allocate our own. |
| 171 | std::unique_ptr<checkpoint::TensorSliceReader> allocated_reader; |
| 172 | |
| 173 | const checkpoint::TensorSliceReader* reader = nullptr; |
| 174 | |
| 175 | if (context->slice_reader_cache()) { |
| 176 | reader = context->slice_reader_cache()->GetReader(file_pattern, open_func, |
| 177 | preferred_shard); |
| 178 | } |
| 179 | if (!reader) { |
| 180 | allocated_reader.reset(new checkpoint::TensorSliceReader( |
| 181 | file_pattern, open_func, preferred_shard)); |
| 182 | reader = allocated_reader.get(); |
| 183 | } |
| 184 | OP_REQUIRES_OK(context, CHECK_NOTNULL(reader)->status()); |
| 185 | |
| 186 | // Get the shape and type from the save file. |
| 187 | DataType type; |
| 188 | TensorShape saved_shape; |
| 189 | OP_REQUIRES( |
| 190 | context, reader->HasTensor(tensor_name, &saved_shape, &type), |
| 191 | errors::NotFound("Tensor name \"", tensor_name, |
| 192 | "\" not found in checkpoint files ", file_pattern)); |
| 193 | OP_REQUIRES( |
| 194 | context, type == context->expected_output_dtype(restore_index), |
| 195 | errors::InvalidArgument("Expected to restore a tensor of type ", |
| 196 | DataTypeString(context->expected_output_dtype(0)), |
| 197 | ", got a tensor of type ", DataTypeString(type), |
| 198 | " instead: tensor_name = ", tensor_name)); |
| 199 | |
| 200 | // Shape of the output and slice to load. |
| 201 | TensorShape output_shape(saved_shape); |
| 202 | TensorSlice slice_to_load(saved_shape.dims()); |
| 203 | if (restore_slice) { |
no test coverage detected