Skip to content

GitLab

Unverified Commit d11aa330 authored by Shining Sun's avatar Shining Sun Committed by GitHub
Browse files

Add 2-GPU benchmark for NCF (#6589)

parent 8ff9eb54
Show whitespace changes
Inline Side-by-side
Showing with 10 additions and 0 deletions
official/recommendation/ncf_keras_benchmark.py
View file @ d11aa330
...@@ -113,6 +113,11 @@ class KerasNCFRealData(KerasNCFBenchmarkBase): ...@@ -113,6 +113,11 @@ class KerasNCFRealData(KerasNCFBenchmarkBase):
self._setup() self._setup()
self._run_and_report_benchmark() self._run_and_report_benchmark()
def benchmark_2_gpus(self):
self._setup()
FLAGS.num_gpus = 2
self._run_and_report_benchmark()
class KerasNCFSyntheticData(KerasNCFBenchmarkBase): class KerasNCFSyntheticData(KerasNCFBenchmarkBase):
"""Benchmark NCF model using synthetic data.""" """Benchmark NCF model using synthetic data."""
...@@ -149,3 +154,8 @@ class KerasNCFSyntheticData(KerasNCFBenchmarkBase): ...@@ -149,3 +154,8 @@ class KerasNCFSyntheticData(KerasNCFBenchmarkBase):
def benchmark_1_gpu(self): def benchmark_1_gpu(self):
self._setup() self._setup()
self._run_and_report_benchmark() self._run_and_report_benchmark()
def benchmark_2_gpus(self):
self._setup()
FLAGS.num_gpus = 2
self._run_and_report_benchmark()
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment