Add 2-GPU benchmark for NCF (#6589)

d11aa330 · Shining Sun · GitHub · 8ff9eb54 · d11aa330
Unverified Commit d11aa330 authored Apr 19, 2019 by Shining Sun Committed by GitHub Apr 19, 2019
Hide whitespace changes
Inline Side-by-side

Showing with 10 additions and 0 deletions

official/recommendation/ncf_keras_benchmark.py official/recommendation/ncf_keras_benchmark.py +10 -0

No files found.
--- a/official/recommendation/ncf_keras_benchmark.py
+++ b/official/recommendation/ncf_keras_benchmark.py
@@ -113,6 +113,11 @@ class KerasNCFRealData(KerasNCFBenchmarkBase):
    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 @@ class KerasNCFSyntheticData(KerasNCFBenchmarkBase):
  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()