Prototype regression testing / benchmark framework for Flax examples

This provides an API/framework for users to write regression tests for their examples. Two goals were kept in mind when designing this:

• Examples code should not require any modification (or only require minimal modification) in order to be ready for regression testing.
• The code expected to be written by the user of the framework should be simple.

Using TensorBoard logs to extract regression metrics

Most examples already log the metrics of interest (e.g. time per epoch, training accuracy, perplexity, etc) to TensorBoard (TB) during training. These logs thus contain all the information required for detecting a regression. Extracting this information from the TB logs thus immediately makes most examples ready for regression testing without any modifications.

This holds true for all but the simplest examples (e.g. MNIST) that do not rely on TB. For these examples we can either add TB logging (recommended solution), or do something weird like capturing stderr output and parsing it. The latter is cumbersome, but allows for keeping the simplest examples simple.

Currently, the framework only includes a thin wrapper for reading TB scaler summaries.

API

The user API was kept minimal and similar to what unit tests would use. Specifically, it supports the following

• Extracting metrics from TB summaries
• Asserting that metric value are within pre-defined bounds
• Reporting metrics

Reporting

We also allow reporting "extras" - any additional textual information. Currently this is used for adding information about failed assertions and test descriptions.

Reporting is done via the report_metric('name', value) and report_metrics({'name': value}) methods for the metrics; and report_extra('name', 'value') and report_extras({'name': 'value'}) methods for the extras. There is also a separate report_wall_time(value) method consumed by the CI framework we plant to interface with (more on that later).

Asserting

We allow using the full spectrum of the unittest / absltest self.assert* methods for verifying benchmark metrics / correctness. This achieves two goals. First, this provides the familiar unit testing API to the end users; and second, this makes reporting metrics and checking their values completely independent, thus allowing for simple APIs for both.

The cost of this simplicity is that we need to override the default behaviour of self.assert* methods. Whereas normally a failed assertion would immediately cause the test to halt, we now detect failures as before, but defer raising them until the end of the test. This guarantees that the benchmark always produces output/report, even if the metric values were outside of the permitted value ranges. Furthermore, overriding assertion methods allows for automatically detecting and reporting benchmark failures.

Example benchmark

Together, this allows for concise benchmark code. See example below for CIFAR10.

  @flagsaver
  def test_1x_v100(self):
    """Run Wide ResNet CIFAR10 on 1x V100 GPUs for 2 epochs."""
    model_dir = tempfile.mkdtemp()
    FLAGS.num_epochs = 2
    FLAGS.arch = 'wrn26_10'
    FLAGS.model_dir = model_dir

    start_time = time.time()
    train.main([])
    benchmark_time = time.time() - start_time
    summaries = self.read_summaries(model_dir)

    # Summaries contain all the information necessary for the regression
    # metrics.
    wall_time, _, eval_error_rate = zip(*summaries['eval_error_rate'])
    wall_time = np.array(wall_time)
    sec_per_epoch = np.mean(wall_time[1:] - wall_time[:-1])
    end_error_rate = eval_error_rate[-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(sec_per_epoch, 80., 84.)
    self.assertBetween(end_error_rate, 0.30, 0.36)

    # 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,
                         'error_rate': end_error_rate})

Interfacing with monitoring / CI frameworks

The current plan is to interface with an external monitoring / CI framework by exporting benchmark results as a simple JSON file with the following format

{
      "name": <class.testMethod>
      "succeeded": true / false
      "wall_time": float (containing wall-time for the benchmark)
      "metrics": {
        "string" -> float map of other performance metrics
      }
      "extras": {
        "string" -> "string" map containing anything else of interest
      }
}

The support for this is implemented behind the scenes. So long as end uses use the self.report_* methods in their benchmark code, the framework will take care of aggregating the metrics, extras, figuring out the name of the test and the file, etc. As a bonus, if any of the assertions fail, the framework will output their error messages under extras.

JSON files (one per test method) should automatically written to disk after each test if the benchmark_output_dir is specified as an absl flag. I haven't actually tested that this works since I couldn't figure out how to get the pytest running to set flags.

What's included in this PR?

This PR implements the framework described above, and adds (for now) toy benchmarks for ImageNet (2 epochs on 8xV100), CIFAR10 (2 epochs on 1xV100) and MNIST (full training on CPU). The MNIST example was also modified to produce TB summaries.
this