Workarond for memory issue in eager mode. (#6197)

* Workarond for memory issue in eager mode.

* Add a TODO

* Fix typo

* Address comments

* remove patch which appear hacky.

* fix typo

official/resnet/keras/keras_imagenet_benchmark.py

@@ -157,6 +157,23 @@ class Resnet50KerasBenchmarkBase(keras_benchmark.KerasBenchmark):
     FLAGS.batch_size = 128 * 8  # 8 GPUs
     self._run_and_report_benchmark()
 
+  def benchmark_8_gpu_bfc_allocator(self):
+    self._setup()
+
+    FLAGS.num_gpus = 8
+    FLAGS.enable_eager = True
+    FLAGS.distribution_strategy = 'default'
+    FLAGS.model_dir = self._get_model_dir('benchmark_8_gpu')
+    FLAGS.batch_size = 128 * 8  # 8 GPUs
+    # Use BFC allocator with limited memory on CPU as a workaround for memory
+    # spikes in eager mode.
+    # TODO(yuefengz): get rid of this test once we fix the memory issue.
+    os.environ['TF_CPU_ALLOCATOR_USE_BFC'] = 'true'
+    os.environ['TF_CPU_BFC_MEM_LIMIT_IN_MB'] = '100000'
+    self._run_and_report_benchmark()
+    del os.environ['TF_CPU_ALLOCATOR_USE_BFC']
+    del os.environ['TF_CPU_BFC_MEM_LIMIT_IN_MB']
+
   def benchmark_graph_8_gpu(self):
     self._setup()