class_docstr: Benchmarks for the ogbg_molpcba Flax example.
class_docstr_tok: ["Benchmarks","for","the","ogbg_molpcba","Flax","example","."]
docstr: Run training with test config for ogbg_molpcba on CPU.
docstr_tok: ["Run","training","with","test","config","for","ogbg_molpcba","on","CPU","."]
code: def test_cpu(self):
      """Run training with test config for ogbg_molpcba on CPU."""
      workdir = self.get_tmp_model_dir()
      config = test.get_config()
  
      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, _, test_accuracy = zip(*summaries["test_accuracy"])
      wall_time = np.array(wall_time)
      sec_per_epoch = np.mean(wall_time[1:] - wall_time[:-1])
      end_test_accuracy = test_accuracy[-1]
  
      _, _, test_aps = zip(*summaries["test_mean_average_precision"])
      end_test_mean_average_precision = test_aps[-1]
  
      _, _, validation_accuracy = zip(*summaries["validation_accuracy"])
      end_validation_accuracy = validation_accuracy[-1]
  
      _, _, validation_aps = zip(
          *summaries["validation_mean_average_precision"]
      )
      end_validation_mean_average_precision = validation_aps[-1]
  
      # 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,
              "test_accuracy": end_test_accuracy,
              "test_mean_average_precision": end_test_mean_average_precision,
              "validation_accuracy": end_validation_accuracy,
              "validation_mean_average_precision": (
                  end_validation_mean_average_precision
              ),
          }
      )
      self.report_extras(
          {
              "model_name": "Graph Convolutional Network",
              "description": "CPU test for ogbg_molpcba.",
              "implementation": "linen",
          }
      )
  
file_path: flax/examples/ogbg_molpcba/ogbg_molpcba_benchmark.py
imports: import time
  from absl import flags
  from absl.testing import absltest
  from flax.testing import Benchmark
  import jax
  import numpy
  import main
  from configs import default
  from configs import test
label_desc: Use image node_id 2 for calling the OgbgMolpcbaBenchmark obj's underlying member method code with example usage: obj.test_cpu() without return types