[ParseSequenceExampleOp] Optimize config construction.

Several micro-optimizations that mostly benefit small input protos:

1. Reserve vector capacity before pushing new elements in a loop.
2. Use `vector::emplace_back()` instead of `vector::push_back()` to avoid move-constructing the temporary struct literals.
3. Construct the default value tensor using the optimized scalar constructors for `Tensor`, which avoid an unnecessary allocation.
4. Evaluate `Tensor::flat<tstring>()` once for each of the key tensors, instead of re-evaluating it for each key, and similarly for the missing-assumed-empty tensor.

PiperOrigin-RevId: 309337185
Change-Id: I3d546a3927a15379f811be80c44f81d4734d0016
This commit is contained in:
Derek Murray 2020-04-30 18:11:03 -07:00 committed by TensorFlower Gardener
parent c74afd7451
commit 196b03bb55

View File

@ -569,57 +569,109 @@ class ParseSequenceExampleOp : public OpKernel {
const Tensor* dense_keys, const Tensor* sparse_keys,
const Tensor* ragged_keys,
const OpInputList& context_dense_defaults) const {
// Convert the tensors/attrs to ArraySlices once, instead of re-evaluating
// them in each loop iteration.
gtl::ArraySlice<tstring> dense_keys_slice =
dense_keys
? gtl::ArraySlice<tstring>(dense_keys->flat<tstring>().data(),
attrs_.num_context_dense)
: attrs_.context_dense_keys;
gtl::ArraySlice<tstring> sparse_keys_slice =
sparse_keys
? gtl::ArraySlice<tstring>(sparse_keys->flat<tstring>().data(),
attrs_.num_context_sparse)
: attrs_.context_sparse_keys;
gtl::ArraySlice<tstring> ragged_keys_slice(
ragged_keys->flat<tstring>().data(), attrs_.num_context_ragged);
example::FastParseExampleConfig config;
config.dense.reserve(attrs_.num_context_dense);
for (int d = 0; d < attrs_.num_context_dense; ++d) {
const tstring& key = dense_keys ? dense_keys->flat<tstring>()(d)
: attrs_.context_dense_keys[d];
config.dense.push_back({key, attrs_.context_dense_types[d],
attrs_.context_dense_shapes[d],
context_dense_defaults[d],
false /* attrs_.context_variable_length[d] */,
0 /*attrs_.context_elements_per_stride[d] */});
const tstring& key = dense_keys_slice[d];
config.dense.emplace_back(key, attrs_.context_dense_types[d],
attrs_.context_dense_shapes[d],
context_dense_defaults[d],
false /* attrs_.context_variable_length[d] */,
0 /*attrs_.context_elements_per_stride[d] */);
}
config.sparse.reserve(attrs_.num_context_sparse);
for (int d = 0; d < attrs_.num_context_sparse; ++d) {
const tstring& key = sparse_keys ? sparse_keys->flat<tstring>()(d)
: attrs_.context_sparse_keys[d];
config.sparse.push_back({key, attrs_.context_sparse_types[d]});
const tstring& key = sparse_keys_slice[d];
config.sparse.emplace_back(key, attrs_.context_sparse_types[d]);
}
config.ragged.reserve(attrs_.num_context_ragged);
for (int d = 0; d < attrs_.num_context_ragged; ++d) {
config.ragged.push_back({ragged_keys->flat<tstring>()(d),
attrs_.context_ragged_value_types[d],
attrs_.context_ragged_split_types[d]});
config.ragged.emplace_back(ragged_keys_slice[d],
attrs_.context_ragged_value_types[d],
attrs_.context_ragged_split_types[d]);
}
return config;
}
static Tensor ConstructDefaultScalar(DataType dtype) {
switch (dtype) {
case DT_INT64:
return Tensor(static_cast<int64>(0));
case DT_FLOAT:
return Tensor(static_cast<float>(0.0));
case DT_STRING:
return Tensor("");
default:
return Tensor(DT_INVALID);
}
}
example::FastParseExampleConfig MakeFeatureListConfig(
const Tensor* dense_keys, const Tensor* sparse_keys,
const Tensor* ragged_keys,
const Tensor* feature_list_dense_missing_assumed_empty) const {
// Convert the tensors/attrs to ArraySlices once, instead of re-evaluating
// them in each loop iteration.
gtl::ArraySlice<tstring> dense_keys_slice =
dense_keys
? gtl::ArraySlice<tstring>(dense_keys->flat<tstring>().data(),
attrs_.num_feature_list_dense)
: attrs_.feature_list_dense_keys;
gtl::ArraySlice<tstring> sparse_keys_slice =
sparse_keys
? gtl::ArraySlice<tstring>(sparse_keys->flat<tstring>().data(),
attrs_.num_feature_list_sparse)
: attrs_.feature_list_sparse_keys;
gtl::ArraySlice<tstring> ragged_keys_slice(
ragged_keys->flat<tstring>().data(), attrs_.num_feature_list_ragged);
// Use an empty slice to indicate that the map in attrs_ should be used
// instead.
gtl::ArraySlice<bool> feature_list_dense_missing_assumed_empty_slice =
feature_list_dense_missing_assumed_empty
? gtl::ArraySlice<bool>(
feature_list_dense_missing_assumed_empty->flat<bool>().data(),
attrs_.num_feature_list_dense)
: gtl::ArraySlice<bool>(nullptr, 0);
example::FastParseExampleConfig config;
config.dense.reserve(attrs_.num_feature_list_dense);
for (int d = 0; d < attrs_.num_feature_list_dense; ++d) {
const tstring& key = dense_keys ? dense_keys->flat<tstring>()(d)
: attrs_.feature_list_dense_keys[d];
const tstring& key = dense_keys_slice[d];
bool missing_assumed_empty =
feature_list_dense_missing_assumed_empty
? feature_list_dense_missing_assumed_empty->flat<bool>()(d)
!feature_list_dense_missing_assumed_empty_slice.empty()
? feature_list_dense_missing_assumed_empty_slice[d]
: attrs_.feature_list_dense_missing_assumed_empty.count(key) > 0;
DataType dtype = attrs_.feature_list_dense_types[d];
Tensor default_value = Tensor(dtype, TensorShape({}));
config.dense.push_back(
{key, dtype, attrs_.feature_list_dense_shapes[d], default_value,
missing_assumed_empty,
0 /*attrs_.feature_list_elements_per_stride[d] */});
config.dense.emplace_back(
key, dtype, attrs_.feature_list_dense_shapes[d],
ConstructDefaultScalar(dtype), missing_assumed_empty,
0 /*attrs_.feature_list_elements_per_stride[d] */);
}
config.sparse.reserve(attrs_.num_feature_list_sparse);
for (int d = 0; d < attrs_.num_feature_list_sparse; ++d) {
const tstring& key = sparse_keys ? sparse_keys->flat<tstring>()(d)
: attrs_.feature_list_sparse_keys[d];
config.sparse.push_back({key, attrs_.feature_list_sparse_types[d]});
const tstring& key = sparse_keys_slice[d];
config.sparse.emplace_back(key, attrs_.feature_list_sparse_types[d]);
}
config.ragged.reserve(attrs_.num_feature_list_ragged);
for (int d = 0; d < attrs_.num_feature_list_ragged; ++d) {
config.ragged.push_back({ragged_keys->flat<tstring>()(d),
attrs_.feature_list_ragged_value_types[d],
attrs_.feature_list_ragged_split_types[d]});
config.ragged.emplace_back(ragged_keys_slice[d],
attrs_.feature_list_ragged_value_types[d],
attrs_.feature_list_ragged_split_types[d]);
}
return config;
}