[WIP] Add structured result output

pl_examples/models/simple_template.py

@@ -0,0 +1,147 @@
+"""
+Example template for defining a system.
+"""
+import os
+
+import torch
+from torch import Tensor
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import optim
+import torchvision.transforms as transforms
+from torchvision.datasets import MNIST
+from torch.utils.data import DataLoader
+import pytorch_lightning as pl
+
+
+class SuperLitModel(pl.LightningModule):
+    """
+    Sample model to show how to define a template.
+
+    Example:
+    """
+
+    def __init__(self,
+                 drop_prob: float = 0.2,
+                 batch_size: int = 2,
+                 in_features: int = 28 * 28,
+                 learning_rate: float = 0.001 * 8,
+                 optimizer_name: str = 'adam',
+                 out_features: int = 10,
+                 hidden_dim: int = 1000,
+                 **kwargs
+                 ):
+        # init superclass
+        super().__init__()
+        self.drop_prob = drop_prob
+        self.batch_size = batch_size
+        self.in_features = in_features
+        self.learning_rate = learning_rate
+        self.optimizer_name = optimizer_name
+        self.out_features = out_features
+        self.hidden_dim = hidden_dim
+
+        self.c_d1 = nn.Linear(in_features=self.in_features,
+                              out_features=self.hidden_dim)
+        self.c_d1_bn = nn.BatchNorm1d(self.hidden_dim)
+        self.c_d1_drop = nn.Dropout(self.drop_prob)
+
+        self.c_d2 = nn.Linear(in_features=self.hidden_dim,
+                              out_features=self.out_features)
+
+    def forward(self, x):
+        """
+        No special modification required for Lightning, define it as you normally would
+        in the `nn.Module` in vanilla PyTorch.
+        """
+        x = self.c_d1(x.view(x.size(0), -1))
+        x = torch.tanh(x)
+        x = self.c_d1_bn(x)
+        x = self.c_d1_drop(x)
+        x = self.c_d2(x)
+        return x
+
+    def training_step(self, batch: Tensor, batch_idx: int):
+        """
+        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
+        step_result = pl.Result(
+            minimize=loss,
+            checkpoint_on=loss,
+            early_stop_on=loss,
+        )
+
+        step_result.log_metric('train_loss', loss)
+        step_result.pbar_metric('pbar_loss', loss)
+
+        # return loss
+        # return step_result
+        return {'loss': loss, 'log':{'something': 1}, 'random': 'af'}
+
+    def validation_step(self, batch: Tensor, batch_idx: int):
+        # forward pass
+        x, y = batch
+        y_hat = self(x)
+        val_loss = F.cross_entropy(y_hat, y)
+
+        result = pl.Result()
+        result.log_metric('val_loss', val_loss)
+        result.pbar_metric('pbar_loss', val_loss)
+
+        return {'val_loss': val_loss, 'log': {'aa': val_loss}, 'progress_bar': {'aa': val_loss}}
+
+    def configure_optimizers(self):
+        """
+        Return whatever optimizers and learning rate schedulers you want here.
+        At least one optimizer is required.
+        """
+        return optim.Adam(self.parameters(), lr=self.learning_rate)
+
+    @staticmethod
+    def add_model_specific_args(parent_parser):  # pragma: no-cover
+        """
+        Define parameters that only apply to this model
+        """
+        parser = ArgumentParser(parents=[parent_parser], add_help=False)
+        parser.add_argument('--in_features', default=28 * 28, type=int)
+        parser.add_argument('--out_features', default=10, type=int)
+        parser.add_argument('--hidden_dim', default=5000, type=int)
+        parser.add_argument('--drop_prob', default=0.2, type=float)
+        parser.add_argument('--learning_rate', default=0.001, type=float)
+        parser.add_argument('--data_dir', default='.', type=str)
+        parser.add_argument('--batch_size', default=64, type=int)
+        return parser
+
+
+if __name__ == '__main__':
+    from argparse import ArgumentParser
+    import pytorch_lightning as pl
+
+    # add trainer args
+    parser = ArgumentParser()
+    parser = pl.Trainer.add_argparse_args(parser)
+
+    # add model args
+    parser = SuperLitModel.add_model_specific_args(parser)
+    args = parser.parse_args()
+
+    # init data, model
+    mnist_train = MNIST(args.data_dir, train=True, download=True, transform=transforms.ToTensor())
+    mnist_train = DataLoader(mnist_train, batch_size=args.batch_size, num_workers=0)
+    model = SuperLitModel(**vars(args))
+
+    # init trainer
+    trainer = pl.Trainer.from_argparse_args(args)
+
+    trainer.fit(
+        model,
+        train_dataloader=mnist_train,
+        val_dataloaders=mnist_train
+    )

pytorch_lightning/__init__.py

@@ -55,13 +55,15 @@
     from pytorch_lightning.trainer.seed import seed_everything
     from pytorch_lightning.callbacks import Callback
     from pytorch_lightning.core import data_loader
+    from pytorch_lightning.core.step_result import Result
 
     __all__ = [
         'Trainer',
         'LightningModule',
         'Callback',
         'data_loader',
         'seed_everything',
+        'Result',
     ]
 
     # necessary for regular bolts imports. Skip exception since bolts is not always installed

pytorch_lightning/callbacks/early_stopping.py

@@ -110,15 +110,20 @@ def on_train_start(self, trainer, pl_module):
         self.best = torch_inf if self.monitor_op == torch.lt else -torch_inf
 
     def on_validation_end(self, trainer, pl_module):
-        self._run_early_stopping_check(trainer, pl_module)
+        return self._run_early_stopping_check(trainer, pl_module)
 
     def _run_early_stopping_check(self, trainer, pl_module):
         logs = trainer.callback_metrics
         stop_training = False
-        if not self._validate_condition_metric(logs):
-            return stop_training
 
-        current = logs.get(self.monitor)
+        used_structured_result = 'early_stop_on' in logs
+        if not used_structured_result:
+            if not self._validate_condition_metric(logs):
+                return stop_training
+            current = logs.get(self.monitor)
+        else:
+            current = logs['early_stop_on']
+
         if not isinstance(current, torch.Tensor):
             current = torch.tensor(current)
 

pytorch_lightning/core/__init__.py

@@ -334,6 +334,7 @@ def training_step(self, batch, batch_idx):
 
 from pytorch_lightning.core.decorators import data_loader
 from pytorch_lightning.core.lightning import LightningModule
+from pytorch_lightning.core.step_result import Result
 
-__all__ = ['LightningModule', 'data_loader']
+__all__ = ['LightningModule', 'data_loader', 'Result']
 # __call__ = __all__

pytorch_lightning/core/hooks.py

@@ -115,6 +115,15 @@ def on_after_backward(self):
                                                              global_step=self.trainer.global_step)
 
         """
+    def amp_scale_loss(self, unscaled_loss, optimizer, optimizer_idx):
+        if self.trainer.use_native_amp:
+            scaled_loss = self.trainer.scaler.scale(unscaled_loss)
+
+        else:
+            # TODO: remove in v0.8.0
+            scaled_loss = amp.scale_loss(unscaled_loss, optimizer)
+
+        return scaled_loss
 
     def backward(self, trainer, loss: Tensor, optimizer: Optimizer, optimizer_idx: int) -> None:
         """
@@ -134,27 +143,10 @@ def backward(self, trainer, loss: Tensor, optimizer: Optimizer, optimizer_idx: i
         Example::
 
             def backward(self, use_amp, loss, optimizer):
-                if use_amp:
-                    with amp.scale_loss(loss, optimizer) as scaled_loss:
-                        scaled_loss.backward()
-                else:
-                    loss.backward()
+                loss.backward()
 
         """
-        if trainer.precision == 16:
-            # .backward is not special on 16-bit with TPUs
-            if trainer.on_tpu:
-                return
-
-            if self.trainer.use_native_amp:
-                self.trainer.scaler.scale(loss).backward()
-
-            # TODO: remove in v0.8.0
-            else:
-                with amp.scale_loss(loss, optimizer) as scaled_loss:
-                    scaled_loss.backward()
-        else:
-            loss.backward()
+        loss.backward()
 
     def transfer_batch_to_device(self, batch: Any, device: torch.device) -> Any:
         """