[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:
parent
c74afd7451
commit
196b03bb55
@ -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;
|
||||
}
|
||||
|
Loading…
x
Reference in New Issue
Block a user