class_docstr: Benchmarks for the ImageNet Flax example.
class_docstr_tok: ["Benchmarks","for","the","ImageNet","Flax","example","."]
docstr: Utility to benchmark ImageNet on 8xV100 GPUs. Use in your test func.
docstr_tok: ["Utility","to","benchmark","ImageNet","on","8xV100","GPUs",".","Use","in","your","test","func","."]
code: def _test_8x_v100_half_precision(
      self, num_epochs: int, min_accuracy, max_accuracy
  ):
      """Utility to benchmark ImageNet on 8xV100 GPUs. Use in your test func."""
      # Prepare and set flags defined in main.py.
      config = config_lib.get_config()
      config.num_epochs = num_epochs
      workdir = self.get_tmp_model_dir()
  
      FLAGS.workdir = workdir
      FLAGS.config = config
  
      start_time = time.time()
      main.main([])
      benchmark_time = time.time() - start_time
      summaries = self.read_summaries(workdir)
  
      # Summaries contain all the information necessary for the regression
      # metrics.
      wall_time, _, eval_accuracy = zip(*summaries["eval_accuracy"])
      wall_time = np.array(wall_time)
      sec_per_epoch = np.mean(wall_time[1:] - wall_time[:-1])
      end_accuracy = eval_accuracy[-1]
  
      # Assertions are deferred until the test finishes, so the metrics are
      # always reported and benchmark success is determined based on *all*
      # assertions.
      self.assertBetween(end_accuracy, min_accuracy, max_accuracy)
  
      # Use the reporting API to report single or multiple metrics/extras.
      self.report_wall_time(benchmark_time)
      self.report_metrics(
          {"sec_per_epoch": sec_per_epoch, "accuracy": end_accuracy}
      )
  
file_path: flax/examples/imagenet/imagenet_benchmark.py
imports: import time
  from absl import flags
  from absl.testing import absltest
  from absl.testing.flagsaver import flagsaver
  from flax.testing import Benchmark
  import jax
  import numpy
  import main
  from configs import v100_x8_mixed_precision
label_desc: Use image node_id 1 for calling the ImagenetBenchmark obj's underlying member method code with example usage: obj._test_8x_v100_half_precision(num_epochs, min_accuracy, max_accuracy) without return types