1 /* Copyright 2015 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 #ifndef TENSORFLOW_C_CHECKPOINT_READER_H_ 17 #define TENSORFLOW_C_CHECKPOINT_READER_H_ 18 19 #include <memory> 20 #include <string> 21 22 #include "tensorflow/c/tf_status_helper.h" 23 #include "tensorflow/core/framework/tensor_shape.h" 24 #include "tensorflow/core/platform/status.h" 25 #include "tensorflow/core/platform/types.h" 26 #include "tensorflow/core/util/tensor_bundle/tensor_bundle.h" 27 #include "tensorflow/core/util/tensor_slice_reader.h" 28 29 namespace tensorflow { 30 namespace checkpoint { 31 32 class TensorSliceReader; 33 34 // A wrapper around BundleReader (for V2 checkpoints) and 35 // checkpoint::TensorSliceReader (for V1), that is more easily SWIG wrapped for 36 // other languages. 37 // 38 // The class currently only interacts with single-slice (i.e., non-partitioned) 39 // variables. 40 class CheckpointReader { 41 public: 42 CheckpointReader(const string& filename, TF_Status* status); 43 44 bool HasTensor(const string& name) const; 45 const string DebugString() const; 46 47 // Returns a map from variable names to their shapes. Slices of a partitioned 48 // tensor are combined into a single entry. 49 const TensorSliceReader::VarToShapeMap& GetVariableToShapeMap() const; 50 51 // Returns a map from variable names to their data types. Slices of a 52 // partitioned tensor are combined into a single entry. 53 const TensorSliceReader::VarToDataTypeMap& GetVariableToDataTypeMap() const; 54 55 // Attempts to look up the tensor named "name" and stores the found result in 56 // "out_tensor". 57 void GetTensor(const string& name, 58 std::unique_ptr<tensorflow::Tensor>* out_tensor, 59 TF_Status* out_status) const; 60 61 private: 62 // Uses "v2_reader_" to build "var name -> shape" and "var name -> data type" 63 // maps; both owned by caller. 64 // REQUIRES: "v2_reader_ != nullptr && v2_reader_.status().ok()". 65 std::pair<std::unique_ptr<TensorSliceReader::VarToShapeMap>, 66 std::unique_ptr<TensorSliceReader::VarToDataTypeMap> > 67 BuildV2VarMaps(); 68 69 // Invariant: exactly one of "reader_" and "v2_reader_" is non-null. 70 std::unique_ptr<TensorSliceReader> reader_; 71 std::unique_ptr<BundleReader> v2_reader_; 72 73 std::unique_ptr<TensorSliceReader::VarToShapeMap> var_to_shape_map_; 74 std::unique_ptr<TensorSliceReader::VarToDataTypeMap> var_to_data_type_map_; 75 76 TF_DISALLOW_COPY_AND_ASSIGN(CheckpointReader); 77 }; 78 79 } // namespace checkpoint 80 } // namespace tensorflow 81 82 #endif // TENSORFLOW_C_CHECKPOINT_READER_H_ 83