Method BuildV2VarMaps

tensorflow/c/checkpoint_reader.cc:109–151

}

std::pair<std::unique_ptr<TensorSliceReader::VarToShapeMap>,
          std::unique_ptr<TensorSliceReader::VarToDataTypeMap>>
CheckpointReader::BuildV2VarMaps() {
  CHECK(v2_reader_ != nullptr);
  CHECK(v2_reader_->status().ok());

  // First pass: filters out the entries of the slices.
  std::unordered_set<string> filtered_keys;
  BundleEntryProto entry;
  v2_reader_->Seek(kHeaderEntryKey);
  for (v2_reader_->Next(); v2_reader_->Valid(); v2_reader_->Next()) {
    CHECK(entry.ParseFromArray(v2_reader_->value().data(),
                               v2_reader_->value().size()))
        << entry.InitializationErrorString();
    for (int i = 0; i < entry.slices_size(); ++i) {
      const auto& slice_proto = entry.slices(i);
      CHECK(filtered_keys
                .insert(EncodeTensorNameSlice(
                    string(v2_reader_->key()) /* full var's name */,
                    TensorSlice(slice_proto)))
                .second);
    }
  }

  // Second pass: adds the entries, ignoring the filtered keys.
  std::unique_ptr<TensorSliceReader::VarToShapeMap> var_to_shape_map(
      new TensorSliceReader::VarToShapeMap);
  std::unique_ptr<TensorSliceReader::VarToDataTypeMap> var_to_data_type_map(
      new TensorSliceReader::VarToDataTypeMap);
  v2_reader_->Seek(kHeaderEntryKey);
  for (v2_reader_->Next(); v2_reader_->Valid(); v2_reader_->Next()) {
    if (filtered_keys.count(string(v2_reader_->key())) > 0) continue;
    CHECK(entry.ParseFromArray(v2_reader_->value().data(),
                               v2_reader_->value().size()))
        << entry.InitializationErrorString();
    string key(v2_reader_->key());
    (*var_to_shape_map)[key] = TensorShape(entry.shape());
    (*var_to_data_type_map)[key] = DataType(entry.dtype());
  }
  // The returned pointers are owned by the caller.
  return std::make_pair(std::move(var_to_shape_map),
                        std::move(var_to_data_type_map));
}

}  // namespace checkpoint
}  // namespace tensorflow