Add 2-GPU benchmark for NCF (tensorflow#6589)

MikeMMS · Apr 20, 2019 · d11aa33 · d11aa33
1 parent 8ff9eb5
commit d11aa33
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/official/recommendation/ncf_keras_benchmark.py b/official/recommendation/ncf_keras_benchmark.py
@@ -113,6 +113,11 @@ def benchmark_1_gpu(self):
     self._setup()
     self._run_and_report_benchmark()
 
+  def benchmark_2_gpus(self):
+    self._setup()
+    FLAGS.num_gpus = 2
+    self._run_and_report_benchmark()
+
 
 class KerasNCFSyntheticData(KerasNCFBenchmarkBase):
   """Benchmark NCF model using synthetic data."""
@@ -149,3 +154,8 @@ def _extract_benchmark_report_extras(self, stats):
   def benchmark_1_gpu(self):
     self._setup()
     self._run_and_report_benchmark()
+
+  def benchmark_2_gpus(self):
+    self._setup()
+    FLAGS.num_gpus = 2
+    self._run_and_report_benchmark()