[WIP] Add structured result output #1989

williamFalcon · 2020-05-28T16:59:31Z

This PR maintains full backward compatibility, but for those who want the option adds an optional argument to the *_step that is a structured dict with validation for managing things to return from each step

Old way (still supported)

def training_step(self, batch, batch_idx...)
    return {...}

New way

  # any loop
  def training_step(self, batch, batch_idx...)
        """
        Lightning calls this inside the training loop with the data from the training dataloader
        passed in as `batch`.
        """
        # forward pass
        x, y = batch
        y_hat = self(x)
        loss = F.cross_entropy(y_hat, y)

        # structure the return from the training loop
        return Result(loss)

Docs:

            # all options:
            def training_step(...):
                return Result(
                    minimize=loss,
                    checkpoint_on=loss,
                    early_stop_on=loss,
                    logs={'train_loss': loss},
                    progress_bar={'train_loss': loss}
                )

            # most of the time
            # will early stop and save checkpoints based on this metric by default
            return Result(loss)

            # to change what to early stop on
            return Result(loss, early_stop_on=accuracy)

            # to change what to checkpoint on
            return Result(loss, early_stop_on=accuracy, checkpoint_on=bleu_score)

            # shorthand for logging
            result = Result(loss)
            result.log('train_nce_loss', loss)

            # shorthand to put on progress bar
            result.to_bar('train_nce_loss', loss)

Additional benefits:

gets rid of 'val_loss' magic. Now user can explicitly set what to stop or save on.
gets rid of 'loss' now user can call it whatever they want (we're just minimizing it).
adds error checking if the user puts in the wrong details.
clear separation of what each return item does

pep8speaks · 2020-05-28T16:59:37Z

Hello @williamFalcon! Thanks for updating this PR.

In the file pl_examples/models/simple_template.py:

Line 86:36: E231 missing whitespace after ':'

In the file pytorch_lightning/core/step_result.py:

Line 26:120: E501 line too long (123 > 119 characters)
Line 57:120: E501 line too long (123 > 119 characters)
Line 300:38: W292 no newline at end of file

In the file pytorch_lightning/trainer/evaluation_loop.py:

Line 334:21: E303 too many blank lines (2)
Line 435:120: E501 line too long (120 > 119 characters)

In the file pytorch_lightning/trainer/training_loop.py:

Line 453:13: E265 block comment should start with '# '
Line 689:13: E731 do not assign a lambda expression, use a def

In the file pytorch_lightning/utilities/parsing.py:

Line 54:35: W292 no newline at end of file

In the file tests/base/determininistic_model.py:

Line 62:28: E226 missing whitespace around arithmetic operator
Line 62:39: E226 missing whitespace around arithmetic operator
Line 223:41: W292 no newline at end of file

Comment last updated at 2020-06-08 13:54:38 UTC

williamFalcon · 2020-05-28T16:59:53Z

@tullie is this in line with what you were thinking?
Will ping you when done

tullie · 2020-05-28T19:15:03Z

Why does StepResult need to be passed in as a training_step argument? That seems unintuitive to me.

Could we aim for something like:

def training_step(self, batch: Tensor, batch_idx: int):
        x, y = batch
        y_hat = self(x)
        loss = F.cross_entropy(y_hat, y)
        return StepResult(loss=loss, log=loss)

williamFalcon · 2020-05-28T23:17:10Z

@tullie ok, updated. check out the new API before i finish making all the deep changes

"""
        Result is an OrderedDict that gives type hints, allowed fields and validation for bad user input.

        Use as the return value for:
        - training_step
        - validation_epoch_end
        - training_epoch_end

        .. note:: Plain dictionary returns are supported but are more prone to errors

        We automatically detach anything here for you to avoid holding references to graphs

        Args:
            minimize: Metric to minimize
            logs: dictionary that will be added to your logger(s)
            early_stop_on: Metric for early stopping. If none set, will use minimize by default.
            checkpoint_on: Metric for checkpointing. If none set, will use minimize by default.
            progress_bar: dictionary of values to add to the progress bar
            hiddens: tensor of hiddens to pass to next step when using TBPTT

        .. code-block: python

            # all options:
            def training_step(...):
                return Result(
                    minimize=loss,
                    checkpoint_on=loss,
                    early_stop_on=loss,
                    logs={'train_loss': loss},
                    progress_bar={'train_loss': loss}
                )

            # most of the time
            # will early stop and save checkpoints based on this metric by default
            return Result(loss)

            # to change what to early stop on
            return Result(loss, early_stop_on=accuracy)

            # to change what to checkpoint on
            return Result(loss, early_stop_on=accuracy, checkpoint_on=bleu_score)

            # shorthand for logging
            result = Result(loss)
            result.log('train_nce_loss', loss)

            # shorthand to put on progress bar
            result.to_bar('train_nce_loss', loss)
        """

STILL SUPPORTED:

return {...}

TODO:

tests
docs
update early stopping behavior
update checkpoint behavior

jeremyjordan · 2020-05-29T01:16:43Z

related to #1256, happy to see progress here!

as i mentioned in the other issue, it might be nice to use something like pydantic for the structured output since it provides very easy data validation for free

tullie · 2020-05-29T05:09:37Z

@williamFalcon yeah this is great. Huge improvement imo! I'll let you finish the TODOs and then do a closer sweep of all the code.

Borda · 2020-05-29T07:37:44Z

as i mentioned in the other issue, it might be nice to use something like pydantic for the structured output since it provides very easy data validation for free

this looks very similar to Namespace except the validation and other features around... :]

williamFalcon · 2020-05-29T12:48:26Z

@jeremyjordan #1256 is awesome, must have missed that haha. Why don't i put together the v1 right now and you can take a stab at V2? happy to make this a joint PR. with this object we can now do whatever we want under the hood for validation :)

jeremyjordan · 2020-05-29T13:11:49Z

yeah, sounds great!

mergify · 2020-05-29T14:21:27Z