def test_value_search_can_load_data(self, mock_stdout):
benchmark = all_benchmarks.find_benchmark_with_name('simple_cast')
handler = description_handler_factory.create_handler('naive_bayes')
settings = settings_module.from_dict({
'timeout':
5,
'printing.prioritized_operations':
True,
'printing.deprioritized_operations':
True,
})
tensor_config = tensor_features_model.load_config(
settings.tensor_model.config_path)
tensor_model = tensor_features_model.get_model(tensor_config)
tensor_checkpoint = tensor_features_model.create_checkpoint(
tensor_model)
tensor_checkpoint.restore(
settings.tensor_model.checkpoint_path).expect_partial()
results = value_search.run_value_search(benchmark=benchmark,
description_handler=handler,
settings=settings,
tensor_model=tensor_model,
tensor_config=tensor_config)
self.assertLen(results.solutions, 1)
self.assertIn('BOW handler prioritized tf.cast(x, dtype)',
mock_stdout.getvalue())
self.assertIn('Tensor features model prioritized tf.cast(x, dtype)',
mock_stdout.getvalue())
def test_run_value_search_works_for_small_max_weight(
self, benchmark_name, max_weight):
benchmark = all_benchmarks.find_benchmark_with_name(benchmark_name)
results = value_search.run_value_search(
benchmark=benchmark,
settings=settings_module.from_dict({'max_weight': max_weight}))
print("Note: if you're inspecting the test log, the above failure is "
"expected in test_run_value_search_works_for_small_max_weight.")
self.assertEmpty(results.solutions)
def test_run_value_handles_large_weight_constants(self):
benchmark = benchmark_module.Benchmark(examples=[
benchmark_module.Example(inputs=[[1], [2]], output=[[3]])
])
results = value_search.run_value_search(benchmark=benchmark,
settings=self.settings)
self.assertNotEmpty(results.solutions)
self.assertEqual(results.solutions[0].expression,
'tf.add(in1, tf.expand_dims(in2, 0))')
output_shape_constant = value_module.ConstantValue((1, 1))
self.assertIn(output_shape_constant, results.value_set)
# Find the element in value_set equal to output_shape_constant and assert
# that it's actually a ConstantValue, as opposed to an OperationValue.
for value in results.value_set:
if value == output_shape_constant:
self.assertIsInstance(value, value_module.ConstantValue)
def test_run_value_search_works_for_simple_benchmarks(
self, simple_benchmark):
handler = description_handler_factory.create_handler('tfidf')
results = value_search.run_value_search(
benchmark=simple_benchmark,
settings=settings_module.from_dict({
'timeout': 20,
'printing.statistics': True
}),
description_handler=handler)
self.assertLen(results.solutions, 1)
self.assertIsNotNone(results.solutions[0].expression)
self.assertGreater(results.solutions[0].weight, 0)
self.assertGreater(results.total_time, 0)
self.assertNotEmpty(results.value_set)
self.assertGreaterEqual(results.statistics.total_apply_count, 0)
def test_run_value_search_works_for_important_benchmarks(
self, benchmark_name, timeout, solution):
# These benchmarks are frequently used in demos. We should ensure that they
# are solved quickly and consistently. While the exact solution may change
# as the algorithm improves, we should ensure that the solution is always
# simple and understandable.
handler = description_handler_factory.create_handler('tfidf')
benchmark = all_benchmarks.find_benchmark_with_name(benchmark_name)
results = value_search.run_value_search(
benchmark=benchmark,
description_handler=handler,
settings=settings_module.from_dict({'timeout': timeout}))
print('Time for benchmark {}: {:.2f} sec'.format(
benchmark_name, results.total_time))
self.assertLen(results.solutions, 1)
self.assertEqual(results.solutions[0].expression, solution)
def test_run_value_search_only_minimal_solutions(self,
only_minimal_solutions,
expected_solutions):
benchmark = benchmark_module.Benchmark(examples=[
benchmark_module.Example(inputs=[[1, 4], [2, 7]], output=[3, 11])
])
results = value_search.run_value_search(
benchmark=benchmark,
settings=settings_module.from_dict({
'timeout': 20,
'max_solutions': 4,
'only_minimal_solutions': only_minimal_solutions,
'max_extra_solutions_time': 20
}))
self.assertLen(results.solutions, len(expected_solutions))
self.assertEqual(
[solution.expression for solution in results.solutions],
expected_solutions)
def test_run_value_search_works_for_simple_benchmarks(
self, simple_benchmark):
handler = description_handler_factory.create_handler('tfidf')
results = value_search.run_value_search(
benchmark=simple_benchmark,
settings=settings_module.from_dict({
'timeout': 20,
'printing.statistics': True,
'require_all_inputs_used': False,
'require_one_input_used': False,
}),
description_handler=handler)
self.assertLen(results.solutions, 1)
self.assertIsNotNone(results.solutions[0].expression)
self.assertGreater(results.solutions[0].weight, 0)
self.assertGreater(results.total_time, 0)
self.assertNotEmpty(results.value_set)
self.assertGreaterEqual(results.statistics.total_apply_count, 0)
# The simple benchmarks are so simple that we can directly write down the
# correct (minimal) target program.
self.assertEqual(results.solutions[0].expression,
simple_benchmark.target_program)
def run_on_all_benchmarks():
"""Runs value search on all benchmarks, printing results to stdout."""
benchmark_count = 0
benchmark_success = 0
unsolved_benchmarks = []
solution_times = [] # Only including successful tasks.
settings = settings_module.from_list(FLAGS.settings)
description_handler = description_handler_factory.create_handler(
settings.description_handler_name)
print('Description handler: {!r}\n'.format(description_handler))
results_json = {
'benchmark_name': FLAGS.benchmark_name,
'settings': settings.as_dict(),
'notes': FLAGS.notes,
'results': [],
}
if (settings.tensor_model.config_path and
settings.tensor_model.checkpoint_path):
tensor_config = tensor_features_model.load_config(
settings.tensor_model.config_path)
tensor_model = tensor_features_model.get_model(tensor_config)
checkpoint = tensor_features_model.create_checkpoint(tensor_model)
checkpoint.restore(settings.tensor_model.checkpoint_path).expect_partial()
# Warm up. Running the model for the first time takes an extra ~10 seconds.
print('Warming up the tensor features model...')
value_search.operation_multipliers_from_tensor_model(
all_benchmarks.find_benchmark_with_name('simple_cast'),
tensor_model, tensor_config, settings)
print('Finished warming up.')
else:
tensor_config = None
tensor_model = None
print('=' * 80)
modules = [google_benchmarks, stackoverflow_benchmarks]
for benchmark in all_benchmarks.get_chosen_benchmarks(
FLAGS.benchmark_name, modules=modules):
gc.collect()
print('Performing value search for benchmark {}.\n'
.format(benchmark.name))
benchmark_count += 1
result = value_search.run_value_search(
benchmark=benchmark,
settings=settings,
description_handler=description_handler,
tensor_model=tensor_model,
tensor_config=tensor_config)
if settings.printing.statistics:
print('\nOperation statistics:\n{}'.format(
result.statistics.statistics_as_string(
num_unique_values=len(result.value_set),
elapsed_time=result.total_time,
sort_by_time=settings.printing.statistics_sort_by_time)))
solutions = result.solutions
if solutions:
first_solution = solutions[0]
print('\nBest solution of weight {} found in {:.2f} sec:\n{}'.format(
first_solution.weight, first_solution.time,
first_solution.expression))
benchmark_success += 1
solution_times.append(first_solution.time)
else:
unsolved_benchmarks.append(benchmark)
print('=' * 80)
sys.stdout.flush()
results_json['results'].append({
'name': benchmark.name,
'solved': bool(solutions),
'solution': solutions[0].expression if solutions else None,
'solution_weight': solutions[0].weight if solutions else None,
'time': solutions[0].time if solutions else None,
})
solve_time_total = sum(solution_times)
solve_time_mean = np.mean(solution_times)
solve_time_geometric_mean = mstats.gmean(solution_times)
results_json['num_benchmarks'] = benchmark_count
results_json['num_solved'] = benchmark_success
results_json['solve_time_total'] = solve_time_total
results_json['solve_time_mean'] = solve_time_mean
results_json['solve_time_geometric_mean'] = solve_time_geometric_mean
print('Solved {} out of {} benchmarks in {:.2f} sec.'.format(
benchmark_success, benchmark_count, solve_time_total))
print('\n'
'Arithmetic mean of solve times: {:.2f} sec\n'
'Geometric mean of solve times: {:.2f} sec\n'.format(
solve_time_mean, solve_time_geometric_mean))
print('Unsolved benchmarks:')
for unsolved in unsolved_benchmarks:
print('Name: {}, target program: {}'.format(
unsolved.name, unsolved.target_program))
print()
if FLAGS.json_output and FLAGS.benchmark_name == 'ALL':
with open(FLAGS.json_output, 'w') as json_file:
json.dump(results_json, json_file,
indent=4, sort_keys=True, separators=(',', ': '))
json_file.write('\n')
print('Wrote JSON results to {}.'.format(FLAGS.json_output))
else:
print('Did not write JSON results file.')
def test_value_search_prioritizes_operations(self):
# Make sure prioritized ops get printed out.
settings = settings_module.from_dict({
'timeout':
0.5, # We don't actually care about solving the problem.
'printing.prioritized_operations':
True,
'printing.deprioritized_operations':
True,
'tensor_model.prioritize_threshold':
0.5,
'tensor_model.deprioritize_threshold':
0.2,
})
# Create data for the mocks.
all_ops = all_operations.get_operations(include_sparse_operations=True)
op_names = [op.name for op in all_ops]
abs_index = None # Prioritized by tensor features model.
argmax_index = None # Deprioritized by tensor features model.
for i, op in enumerate(all_ops):
if op.name == 'tf.abs(x)':
abs_index = i
elif op.name == 'tf.argmax(input, axis)':
argmax_index = i
self.assertIsNotNone(abs_index)
self.assertIsNotNone(argmax_index)
operation_probs = np.repeat(0.4, len(all_ops))
operation_probs[abs_index] = 0.6 # Above prioritize threshold.
operation_probs[argmax_index] = 0.1 # Below deprioritize threshold.
operation_logits = np.expand_dims(_inverse_sigmoid(operation_probs),
axis=0)
nl_data = [
# Associate "apple" with tf.zeros.
{
'docstring': ['apple pie'],
'tf_functions': ['tf.zeros'],
'comments': [],
'names': [],
'strings': []
},
] * 100
# Mock the tensor model's predictions, the NL model's data, and print().
with mock.patch('tf_coder.models.tensor_features_model.eval_single_example',
return_value=tensor_features_model.Result(
operation_logits=operation_logits)), \
mock.patch('tf_coder.datasets.github.data_loader.load_data',
return_value=nl_data), \
mock.patch('sys.stdout', new_callable=six.StringIO) as mock_stdout:
handler = bag_of_words_handlers.NaiveBayesDescriptionHandler(
max_num_prioritized=1)
benchmark = benchmark_module.Benchmark(
# I/O example doesn't matter.
examples=[
benchmark_module.Example({'my_var': [1, 2]}, [2, 1])
],
# Description contains "apple"!
description='honeycrisp apple')
value_search.run_value_search(benchmark=benchmark,
description_handler=handler,
settings=settings,
tensor_model=mock.Mock(),
tensor_config={
'max_num_inputs': 3,
'operation_names': op_names
})
self.assertIn('BOW handler prioritized tf.zeros',
mock_stdout.getvalue())
self.assertIn('Tensor features model prioritized tf.abs(x), p=0.6',
mock_stdout.getvalue())
self.assertIn(
'Tensor features model deprioritized tf.argmax(input, axis), p=0.1',
mock_stdout.getvalue())