| 25 | namespace checkpoint { |
| 26 | |
| 27 | void WriteRestVariables(BundleReader& reader, BundleWriter& writer, |
| 28 | const std::vector<string>& ignored_names) { |
| 29 | std::set<string> excluded_names(ignored_names.cbegin(), ignored_names.cend()); |
| 30 | std::vector<std::string> tensor_names; |
| 31 | reader.Seek(kHeaderEntryKey); |
| 32 | reader.Next(); |
| 33 | for (; reader.Valid(); reader.Next()) { |
| 34 | tensor_names.emplace_back(reader.key()); |
| 35 | } |
| 36 | for (auto& tensor_name : tensor_names) { |
| 37 | Status status; |
| 38 | DataType dtype; |
| 39 | TensorShape shape; |
| 40 | if (excluded_names.count(tensor_name)) continue; |
| 41 | status = reader.LookupDtypeAndShape(tensor_name, &dtype, &shape); |
| 42 | if (status.ok()) { |
| 43 | Tensor tensor(dtype, shape); |
| 44 | status = reader.Lookup(tensor_name, &tensor); |
| 45 | writer.Add(tensor_name, tensor); |
| 46 | } |
| 47 | } |
| 48 | } |
| 49 | |
| 50 | void WriteRestVariables(BundleReader& reader, BundleWriter& writer, |
| 51 | const std::vector<string>& ignored_names, |