Stop training on Inf/NaN loss

[mcabbott]

Closes #1981, by altering train! so that when it encounters an infinite / NaN loss, it prints a warning throws an error, and stops. It stops before updating the model, because such an update will usually make everything NaN.

Not sure it should stop, in fact. Maybe it should skip that datapoint & continue?

PR Checklist

• Tests are added
• Entry in NEWS.md
• Documentation, if applicable

[ToucheSir]

It would be good to have some way for user code to tell training has stopped prematurely. Not sure what that would look like: return value, callback, flag, etc?

[mcabbott]

One way would be to upgrade this to an error. It does depend a little on how complex/comprehensive we think train! ought to be.

I found #821, in which all options were discussed. And the idea was to write isfinite(l) || Flux.skip(), except that this got documented as something to do in a callback, which never worked, it happens too late... and nobody noticed because this was sufficiently confusing? So my opinion is that anything that complicated should be an ordinary for loop.

[ToucheSir]

Raising StopException from the inner loop doesn't seem like the worst idea. Unlike calling skip or stop in callbacks, this would be able to intercept things in time. I guess the only question is whether making StopException part of the public API for this purpose is palatable.

[mcabbott]

StopException leads directly to break, so I think there's nothing gained by implementing an auto-check that way compared to this PR's current code?

The point of the exception was, I now see, was that you could throw it from inside the loss function passed to gradient, and still get to break. But documenting it as part of the callback was a mistake.

[ToucheSir]

I missed that StopException was still being caught by the inner loop of train!. In that case, it'd have to be another exception type. Whether adding one just for this makes sense is a good question.

[mcabbott]

Ah OK, so you weren't proposing to throw and catch an error. It could just throw a DomainError, that seems the closest Base type.

[ToucheSir]

LGTM, but if @darsnack or @CarloLucibello wouldn't mind having a once-over that would be great.

[KronosTheLate]

Thanks for taking the issue I raised so seriously, it is so nice to encounter active developers who listen to user requests. Thanks for your work <3