| 30 | class TensorSliceReader; |
| 31 | |
| 32 | CheckpointReader::CheckpointReader(const string& filename, TF_Status* status) |
| 33 | : reader_(nullptr), |
| 34 | v2_reader_(nullptr), |
| 35 | var_to_shape_map_(nullptr), |
| 36 | var_to_data_type_map_(nullptr) { |
| 37 | // Depending on whether this is a V2 ckpt, initializes "reader_" or |
| 38 | // "v2_reader_". |
| 39 | std::vector<string> v2_path; |
| 40 | if (Env::Default()->GetMatchingPaths(MetaFilename(filename), &v2_path).ok() && |
| 41 | !v2_path.empty()) { |
| 42 | v2_reader_.reset( |
| 43 | new BundleReader(Env::Default(), filename /* prefix to a V2 ckpt */)); |
| 44 | if (!v2_reader_->status().ok()) { |
| 45 | Set_TF_Status_from_Status(status, v2_reader_->status()); |
| 46 | return; |
| 47 | } |
| 48 | auto result = BuildV2VarMaps(); |
| 49 | var_to_shape_map_.swap(result.first); |
| 50 | var_to_data_type_map_.swap(result.second); |
| 51 | } else { |
| 52 | reader_.reset(new TensorSliceReader(filename)); |
| 53 | if (!reader_->status().ok()) { |
| 54 | Set_TF_Status_from_Status(status, reader_->status()); |
| 55 | return; |
| 56 | } |
| 57 | var_to_shape_map_.reset( |
| 58 | new TensorSliceReader::VarToShapeMap(reader_->GetVariableToShapeMap())); |
| 59 | var_to_data_type_map_.reset(new TensorSliceReader::VarToDataTypeMap( |
| 60 | reader_->GetVariableToDataTypeMap())); |
| 61 | } |
| 62 | } |
| 63 | |
| 64 | bool CheckpointReader::HasTensor(const string& name) const { |
| 65 | if (reader_ != nullptr) { |
nothing calls this directly
no test coverage detected