experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Switch from unordered_map to absl::flat_hash_map

Browse Source 
Showed a sizeable improvement for one large shape inference case.

PiperOrigin-RevId: 308109572
Change-Id: Ia926bd597246e399b38102a7bb09a4861b1cc7f0
This commit is contained in:
Jacques Pienaar 2020-04-23 12:46:11 -07:00 committed by TensorFlower Gardener
parent 8c607acc89
commit 2bbab6a500
1 changed files with 6 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
tensorflow/core/common_runtime/shape_refiner.h
View File

@ -17,6 +17,7 @@ limitations under the License.
#include <vector>
#include "absl/container/flat_hash_map.h"
#include "tensorflow/core/common_runtime/graph_runner.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/shape_inference.h"
@ -236,7 +237,8 @@ class ShapeRefiner {
GraphRunner graph_runner_;
// Stores a map from a node to its ExtendedInferenceContext.
std::unordered_map<const Node*, std::unique_ptr<ExtendedInferenceContext>>
absl::flat_hash_map<const Node*, std::unique_ptr<ExtendedInferenceContext>,
hash<const Node*>>
node_to_context_;
// Holds a cache from 'tensor name' to the tensor that is
@ -257,9 +259,10 @@ class ShapeRefiner {
// shape inference.
const tensorflow::FunctionLibraryDefinition* function_library_ = nullptr;
// Cache the graph corresponding to each functin definition for which shapes
// Cache the graph corresponding to each function definition for which shapes
// are refined.
std::unordered_map<const FunctionDef*, std::unique_ptr<const Graph>>
absl::flat_hash_map<const FunctionDef*, std::unique_ptr<const Graph>,
hash<const FunctionDef*>>
functions_;
TF_DISALLOW_COPY_AND_ASSIGN(ShapeRefiner);