diff options
Diffstat (limited to 'tensorflow/contrib/tfprof/python/tools/tfprof/profiler_test.py')
| -rw-r--r-- | tensorflow/contrib/tfprof/python/tools/tfprof/profiler_test.py | 16 |
1 files changed, 13 insertions, 3 deletions
diff --git a/tensorflow/contrib/tfprof/python/tools/tfprof/profiler_test.py b/tensorflow/contrib/tfprof/python/tools/tfprof/profiler_test.py index 5daaafd7c8..c7113b6a57 100644 --- a/tensorflow/contrib/tfprof/python/tools/tfprof/profiler_test.py +++ b/tensorflow/contrib/tfprof/python/tools/tfprof/profiler_test.py @@ -129,7 +129,7 @@ class ProfilerTest(test.TestCase): opts = model_analyzer.PRINT_ALL_TIMING_MEMORY.copy() opts['account_type_regexes'] = ['.*'] - with session.Session() as sess, ops.device('/cpu:0'): + with session.Session() as sess: r1, r2, r3 = lib.BuildSplitableModel() sess.run(variables.global_variables_initializer()) @@ -179,8 +179,18 @@ class ProfilerTest(test.TestCase): self.assertEqual(lib.SearchTFProfNode(pb2, 'add'), None) self.assertGreater(lib.SearchTFProfNode(pb3, 'add').exec_micros, 0) - # TODO(xpan): Better test of advisor. - profiler.advise() + advice_pb = profiler.advise(model_analyzer.ALL_ADVICE) + self.assertTrue('AcceleratorUtilizationChecker' in advice_pb.checkers) + self.assertTrue('ExpensiveOperationChecker' in advice_pb.checkers) + self.assertTrue('OperationChecker' in advice_pb.checkers) + + checker = advice_pb.checkers['AcceleratorUtilizationChecker'] + if test.is_gpu_available(): + self.assertGreater(len(checker.reports), 0) + else: + self.assertEqual(len(checker.reports), 0) + checker = advice_pb.checkers['ExpensiveOperationChecker'] + self.assertGreater(len(checker.reports), 0) if __name__ == '__main__': |