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, 3 insertions, 13 deletions
diff --git a/tensorflow/contrib/tfprof/python/tools/tfprof/profiler_test.py b/tensorflow/contrib/tfprof/python/tools/tfprof/profiler_test.py index c7113b6a57..5daaafd7c8 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: + with session.Session() as sess, ops.device('/cpu:0'): r1, r2, r3 = lib.BuildSplitableModel() sess.run(variables.global_variables_initializer()) @@ -179,18 +179,8 @@ class ProfilerTest(test.TestCase): self.assertEqual(lib.SearchTFProfNode(pb2, 'add'), None) self.assertGreater(lib.SearchTFProfNode(pb3, 'add').exec_micros, 0) - 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) + # TODO(xpan): Better test of advisor. + profiler.advise() if __name__ == '__main__': |