| 226 | } |
| 227 | |
| 228 | Status TensorSliceReader::GetTensor( |
| 229 | const string& name, std::unique_ptr<tensorflow::Tensor>* out_tensor) const { |
| 230 | DataType type; |
| 231 | TensorShape shape; |
| 232 | TensorSlice slice; |
| 233 | { |
| 234 | mutex_lock l(mu_); |
| 235 | const TensorSliceSet* tss = gtl::FindPtrOrNull(tensors_, name); |
| 236 | if (tss == nullptr) { |
| 237 | return errors::NotFound(name, " not found in checkpoint file"); |
| 238 | } |
| 239 | |
| 240 | if (tss->Slices().size() > 1) { |
| 241 | // TODO(sherrym): Support multi-slice checkpoints. |
| 242 | return errors::Unimplemented("Sliced checkpoints are not supported"); |
| 243 | } |
| 244 | |
| 245 | type = tss->type(); |
| 246 | shape = tss->shape(); |
| 247 | slice = tss->Slices().begin()->second.slice; |
| 248 | } |
| 249 | |
| 250 | std::unique_ptr<tensorflow::Tensor> t(new tensorflow::Tensor); |
| 251 | Status s = tensorflow::Tensor::BuildTensor(type, shape, t.get()); |
| 252 | if (!s.ok()) return s; |
| 253 | bool success = false; |
| 254 | |
| 255 | #define READER_COPY(dt) \ |
| 256 | case dt: \ |
| 257 | success = CopySliceData(name, slice, \ |
| 258 | t->flat<EnumToDataType<dt>::Type>().data()); \ |
| 259 | break; |
| 260 | |
| 261 | switch (type) { |
| 262 | READER_COPY(DT_FLOAT); |
| 263 | READER_COPY(DT_DOUBLE); |
| 264 | READER_COPY(DT_INT32); |
| 265 | READER_COPY(DT_UINT8); |
| 266 | READER_COPY(DT_INT16); |
| 267 | READER_COPY(DT_INT8); |
| 268 | READER_COPY(DT_INT64); |
| 269 | READER_COPY(DT_STRING); |
| 270 | default: |
| 271 | return errors::Unimplemented("Data type not supported"); |
| 272 | } |
| 273 | #undef READER_COPY |
| 274 | |
| 275 | if (!success) { |
| 276 | return errors::NotFound(name, " not found in checkpoint file"); |
| 277 | } |
| 278 | std::swap(*out_tensor, t); |
| 279 | |
| 280 | return Status::OK(); |
| 281 | } |
| 282 | |
| 283 | TensorSliceReader::VarToShapeMap TensorSliceReader::GetVariableToShapeMap() |
| 284 | const { |