def operation_multipliers_from_tensor_model(
benchmark: benchmark_module.Benchmark,
tensor_model: tensor_features_model.Model, tensor_config: Dict[Text,
Any],
settings: settings_module.Settings) -> Dict[Text, float]:
"""Runs the tensor features model to get operation weight multipliers."""
# TODO(kshi): Duplicate creation of InputValue and OutputValue objects.
inputs = [
value_module.InputValue(user_input, '')
for user_input in _user_inputs(benchmark.examples[0].inputs)
]
output = value_module.OutputValue(benchmark.examples[0].output)
example_protos = collect_tensor_data.create_tf_examples(
io_example=collect_tensor_data.IOExample(
expression='',
input_values=inputs,
output_value=output,
num_inputs=benchmark.num_inputs,
operations=[]),
max_num_inputs=tensor_config['max_num_inputs'],
permute_inputs=False)
if example_protos:
example_proto = example_protos[0]
result = tensor_features_model.eval_single_example(
tensor_model, example_proto.SerializeToString())
operation_ranks = tf.argsort(
tf.argsort(tf.squeeze(result.operation_logits), stable=True))
operation_probabilities = tf.squeeze(
tf.sigmoid(result.operation_logits))
operation_names = tensor_config['operation_names']
assert len(operation_names) == len(operation_probabilities)
multipliers = {}
for index in tf.where(operation_probabilities >=
settings.tensor_model.prioritize_threshold):
chosen_op_name = operation_names[int(index)]
if settings.printing.prioritized_operations:
print('Tensor features model prioritized {}, p={}'.format(
chosen_op_name, operation_probabilities[int(index)]))
multipliers[
chosen_op_name] = settings.tensor_model.prioritize_multiplier
for index in tf.where(operation_probabilities <=
settings.tensor_model.deprioritize_threshold):
int_index = int(index)
chosen_op_name = operation_names[int_index]
if (int(operation_ranks[int_index]) >=
settings.tensor_model.max_deprioritized):
continue
if settings.printing.deprioritized_operations:
print('Tensor features model deprioritized {}, p={}, logit={}'.
format(chosen_op_name,
operation_probabilities[int_index],
result.operation_logits[0][int_index]))
multipliers[chosen_op_name] = (
settings.tensor_model.deprioritize_multiplier)
return multipliers
else:
logging.error('Failed to create example for inputs %s and output %s',
inputs, output)
return {}
def test_extract_examples_from_value(self, max_num_inputs,
expected_expressions):
actual = collect_tensor_data.extract_examples_from_value(
self.example_value_1, max_num_inputs=max_num_inputs)
# Check that all expressions are as expected.
self.assertCountEqual([example.expression for example in actual],
expected_expressions)
# Check all elements of one IOExample namedtuple. This example is at index 1
# when max_num_inputs > 1.
if max_num_inputs > 1:
expected_index_1 = collect_tensor_data.IOExample(
expression='tf.add(tf.unique_with_counts(in1)[1], in2)',
input_values=[
value_module.InputValue([1, 1, 2, 5, 6, 5], 'in1'),
value_module.InputValue([10, 10, 20, 50, 60, 50], 'in2')
],
output_value=value_module.OutputValue([10, 10, 21, 52, 63,
52]),
num_inputs=2,
operations=[
self.add_operation, self.indexing_operation,
self.unique_with_counts_operation
])
self.assertEqual(actual[1], expected_index_1)
# Equality of Value objects is done by comparing the wrapped values. Check
# the names in input_values too.
for actual_value, expected_value in zip(
actual[1].input_values, expected_index_1.input_values):
self.assertIsInstance(actual_value, value_module.InputValue)
self.assertEqual(actual_value.name, expected_value.name)
# Check that all extracted examples actually work by eval-ing them.
for example in actual:
namespace_dict = {'tf': tf}
self.assertLen(example.input_values, example.num_inputs)
for input_value in example.input_values:
namespace_dict[input_value.name] = input_value.value
eval_output = eval(example.expression, namespace_dict) # pylint: disable=eval-used
self.assertEqual(value_module.OutputValue(eval_output),
example.output_value)
self.assertEqual(example.output_value, self.example_value_1)
def test_create_tf_examples(self):
sparse_tensor = tf.SparseTensor(values=[0, -15, 30],
indices=[[12], [34], [56]],
dense_shape=[100])
# This example does not represent a realistic tensor transformation. It uses
# variety in the input/output tensors to exercise the featurization.
io_example = collect_tensor_data.IOExample(
expression='tf.dummy_expression(in1, in2)',
input_values=[
value_module.InputValue(
[[[0.5, 2.5, 9.0], [-0.25, 0.0, 1.25]]], 'in1'),
value_module.InputValue(sparse_tensor, 'in2'),
],
output_value=value_module.OutputValue([[1.0], [0.0], [1.0],
[0.0]]),
num_inputs=2,
operations=[
self.add_operation, self.add_operation, self.gather_operation
])
with mock.patch.object(collect_tensor_data,
'COUNT_BOUNDARIES',
new=[0, 1, 3, 50, float('inf')]):
with mock.patch.object(
collect_tensor_data,
'FLOAT_BOUNDARIES',
new=[-float('inf'), -10, -1e-8, 1e-8, 10,
float('inf')]):
tf_examples = collect_tensor_data.create_tf_examples(
io_example)
operation_list = all_operations.get_operations(
include_sparse_operations=True)
expected_operations = [
2 if op.name == 'tf.add(x, y)' else
1 if op.name == 'tf.gather(params, indices)' else 0
for op in operation_list
]
expected_tf_example_1 = {
# Features from featurize_value.
'kind': [2, 2, 3, 0],
'dtype': [8, 8, 0, 0],
'rank': [2, 3, 1, 0],
'shape': [4, 1, 0, 0, 1, 2, 3, 0, 100, 0, 0, 0, 0, 0, 0, 0],
'shape_buckets': [2, 1, 0, 0, 1, 1, 2, 0, 3, 0, 0, 0, 0, 0, 0, 0],
'floats': [
1.0, 0.0, 0.5, 0.5, 9.0, -0.25, 13 / 6, 13.5 / 6, 30, -15, 5,
15, 0, 0, 0, 0
],
'float_buckets': [3, 2, 3, 3, 3, 1, 3, 3, 4, 0, 3, 4, 2, 2, 2, 2],
'counts': [
4, 4, 2, 2, 2, 0, 4, 2, 6, 6, 4, 1, 0, 1, 2, 6, 100, 3, 1, 1,
0, 1, 1, 3, 1, 1, 0, 1, 0, 0, 1, 1
],
'count_buckets': [
2, 2, 1, 1, 1, 0, 2, 1, 2, 2, 2, 1, 0, 1, 1, 2, 3, 2, 1, 1, 0,
1, 1, 2, 1, 1, 0, 1, 0, 0, 1, 1
],
'fractions': [
4 / 4, 2 / 4, 2 / 4, 2 / 4, 0 / 4, 4 / 4, 2 / 4, 6 / 6, 4 / 6,
1 / 6, 0 / 6, 1 / 6, 2 / 6, 6 / 6, 3 / 100, 1 / 3, 1 / 3,
0 / 3, 1 / 3, 1 / 3, 3 / 3, 1 / 1, 0 / 1, 1 / 1, 0 / 1, 0 / 1,
1 / 1, 1 / 1
],
'booleans': [
1, 1, 0, 1, 0, 1, 1, 0, 1, 1, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0,
0, 0, 1, 1, 1, 0, 1, 0, 1, 1, 1
],
'value_string': [
b'tf.float32:[[1.0], [0.0], [1.0], [0.0]]',
b'tf.float32:[[[0.5, 2.5, 9.0], [-0.25, 0.0, 1.25]]]',
(b'SparseTensor(indices=tf.Tensor(\n[[12]\n [34]\n [56]], '
b'shape=(3, 1), dtype=int64), '
b'values=tf.Tensor([ 0 -15 30], shape=(3,), dtype=int32), '
b'dense_shape=tf.Tensor([100], shape=(1,), dtype=int64))'),
b'0'
],
# Features from featurize_input_and_output.
'io_comparisons':
[2, 2, 2, 0, 2, 2, 1, 2, 0, 0, 2, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1],
'io_booleans': [
0, 1, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0
],
'io_counts': [1, 2, 1, 1, 2, 1, 1, 1, 1],
'io_count_buckets': [1, 1, 1, 1, 1, 1, 1, 1, 1],
'io_fractions': [
1 / 6, 2 / 4, 1 / 6, 1 / 4, 1 / 3, 2 / 4, 1 / 3, 1 / 4, 1 / 1,
1 / 1, 1 / 1, 1 / 1
],
# Features added in create_examples.
'num_inputs': [2],
'operations':
expected_operations,
'expression': [b'tf.dummy_expression(in1, in2)'],
}
print(tf_examples)
self.assertLen(tf_examples, 2)
actual_tf_example_1, actual_tf_example_2 = tf_examples
# Check the entire first example.
for key, expected in six.iteritems(expected_tf_example_1):
some_list = actual_tf_example_1.features.feature[key]
if some_list.HasField('float_list'):
actual = some_list.float_list.value
actual = [round(f, 6) for f in actual]
expected = [round(f, 6) for f in expected]
elif some_list.HasField('int64_list'):
actual = some_list.int64_list.value
elif some_list.HasField('bytes_list'):
actual = some_list.bytes_list.value
else:
self.fail('Failed to extract list from TF example.')
# Printing side-by-side like this in the test log is more helpful than the
# AssertionError message. Look at the Python3 log, which prints ints
# without the L suffix.
print('key: {}\n'
' expected: {}\n'
' got: {}'.format(key, expected, actual))
self.assertEqual(actual, expected)
# Less rigorous checks for the second example, where the two inputs have
# swapped.
self.assertEqual(
actual_tf_example_2.features.feature['rank'].int64_list.value,
[2, 1, 3, 0])
self.assertEqual(
actual_tf_example_2.features.feature['shape'].int64_list.value,
[4, 1, 0, 0, 100, 0, 0, 0, 1, 2, 3, 0, 0, 0, 0, 0])