[Draft] Refactor fastai metrics and Recorder to improve flexibility

[warner-benjamin]

This PR refactors fastai metrics and Recorder to allow metrics to independently compute on train, valid, or both. It allows easy logging of multiple losses, or low-compute metrics during training and low & high-compute metrics during validation without additional callbacks.

Motivation

Except for train_loss and valid_loss, Recorder is currently setup to log all metrics at train, validation, or both, with no options for fine grained control over metric logging. It is impossible to add a single metric which records during training without writing a callback and duplicating code.

Examples where this would be useful: recording individual losses in addition to total loss while performing multi-loss training or recording low-compute metrics during training and low & high-compute metrics during validation.

Overview of Changes

1. Refactor Recorder and Metric to allow any Metric to independently log on train, valid, or both, with the exception of AvgSmoothMetric.
2. Add AvgSmoothMetric, a metric version of AvgSmoothLoss.
3. Add a new func_to_metric convenience function for defining metrics from functions. This is a more generic version of the current skm_to_fastai.
4. Move the AccumMetric activation function and argmax code to Metric so all inherited classes can use it.
5. Add Custom name for metrics #3573 to all metrics.
6. Update WandbCallback, AzureMLCallback, and TensorBoardCallback to reflect changes. NeptuneCallback is unchanged as it is deprecated.
7. Update documentation, test coverage, and add code examples to reflect changes.

Defining and Using a New Metric Examples

To define new metrics via the metrics classes:

def tst_metric(out, targ): return F.mse_loss(out, targ)

def TestAvg(log_metric=LogMetric.Valid):
    return AvgMetric(tst_metric, log_metric=log_metric)

def TestAccum(log_metric=LogMetric.Valid):
    return AccumMetric(tst_metric, log_metric=log_metric)

def TestSmooth():
    return AvgSmoothMetric(tst_metric)

or with the func_to_metric convenience method:

def tst_metric(out, targ): return F.mse_loss(out, targ)

def TestAvg(log_metric=LogMetric.Valid):
    return func_to_metric(tst_metric, MetricType.Avg, False, log_metric=log_metric)

def TestAccum(log_metric=LogMetric.Valid):
    return func_to_metric(tst_metric, MetricType.Accum, False, log_metric=log_metric)

def TestSmooth():
    return func_to_metric(tst_metric, MetricType.Smooth, False)

#or more generic definition:
def TestMetric(log_metric=LogMetric.Valid, metric_type=MetricType.Accum):
    return func_to_metric(tst_metric, metric_type, False, log_metric=log_metric)

Then have TestAvg run on train and TestAccum on valid, you'd set

Learner(..., metrics=[TestAvg(LogMetric.Train), TestAccum()])
# or
Learner(..., metrics=[TestAvg('train'), TestAccum()])
# or
Learner(..., metrics=[TestMetric(log_metric=LogMetric.Train, metric_type=MetricType.Avg), 
                      TestMetric(metric_type=MetricType.Accum)])

the logger output would be this:

['epoch', 'train_loss', 'train_avg_tst_metric', 'valid_loss', 'valid_accm_tst_metric', 'time']

due to automatic metric name deduping.

Upgrade Path

For those using fastai 2.5.3 or earlier, the upgrade path is to change any functional metrics to a class initialization, eg

Learner(..., metrics=accuracy)
# to
Learner(..., metrics=Accuracy())

Outside of this one required change, this PR should be a drop-in replacement for normal use of the current metrics system.

If someone is setting train_metrics=False in Recorder.__init__, then the upgrade path is to set the log_metric argument of the new metrics to LogMetric.Train or LogMetric.Both, depending if they kept the Recorder default of valid_metrics=True or not.

Testing

In addition to writing new tests to cover the changes, I have also used this metrics refactor in a project for the past month or so without issue.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[warner-benjamin]

To do: Update all the references to and usage of existing functional metrics in tutorials, examples, etc.

[warner-benjamin]

Closed due to fastai v2 no longer accepting major changes due to imminent development of fastai v3.

Metrics refactor can be used with fastai v2 via fastxtend here.