diff --git a/pytorch_lightning/core/__init__.py b/pytorch_lightning/core/__init__.py index b737caab..56d6a258 100644 --- a/pytorch_lightning/core/__init__.py +++ b/pytorch_lightning/core/__init__.py @@ -1,146 +1,239 @@ """ -Test +A LightningModule is a strict superclass of torch.nn.Module but provides an interface to standardize +the "ingredients" for a research or production system. + +- The model/system definition (__init__) +- The model/system computations (forward) +- What happens in the training loop (training_step, training_end) +- What happens in the validation loop (validation_step, validation_end) +- What happens in the test loop (test_step, test_end) +- What optimizers to use (configure_optimizers) +- What data to use (train_dataloader, val_dataloader, test_dataloader) + +Most methods are optional. Here's a minimal example. + +.. code-block:: python + + import os + import torch + from torch.nn import functional as F + from torch.utils.data import DataLoader + from torchvision.datasets import MNIST + import torchvision.transforms as transforms + + import pytorch_lightning as pl + + class CoolModel(pl.LightningModule): + + def __init__(self): + super(CoolModel, self).__init__() + self.l1 = torch.nn.Linear(28 * 28, 10) + + def forward(self, x): + return torch.relu(self.l1(x.view(x.size(0), -1))) + + def training_step(self, batch, batch_idx): + x, y = batch + y_hat = self.forward(x) + return {'loss': F.cross_entropy(y_hat, y)} + + def validation_step(self, batch, batch_idx): + # OPTIONAL + x, y = batch + y_hat = self.forward(x) + return {'val_loss': F.cross_entropy(y_hat, y)} + + def validation_end(self, outputs): + # OPTIONAL + val_loss_mean = torch.stack([x['val_loss'] for x in outputs]).mean() + return {'val_loss': val_loss_mean} + + def test_step(self, batch, batch_idx): + # OPTIONAL + x, y = batch + y_hat = self.forward(x) + return {'test_loss': F.cross_entropy(y_hat, y)} + + def test_end(self, outputs): + # OPTIONAL + test_loss_mean = torch.stack([x['test_loss'] for x in outputs]).mean() + return {'test_loss': test_loss_mean} + + def configure_optimizers(self): + # REQUIRED + return torch.optim.Adam(self.parameters(), lr=0.02) + + @pl.data_loader + def train_dataloader(self): + return DataLoader(MNIST(os.getcwd(), train=True, download=True, + transform=transforms.ToTensor()), batch_size=32) + + @pl.data_loader + def val_dataloader(self): + # OPTIONAL + # can also return a list of val dataloaders + return DataLoader(MNIST(os.getcwd(), train=True, download=True, + transform=transforms.ToTensor()), batch_size=32) + + @pl.data_loader + def test_dataloader(self): + # OPTIONAL + # can also return a list of test dataloaders + return DataLoader(MNIST(os.getcwd(), train=False, download=True, + transform=transforms.ToTensor()), batch_size=32) + +Once you've defined the LightningModule, fit it using a trainer. + +.. code-block:: python + trainer = pl.Trainer() + model = CoolModel() + + trainer.fit(model) + +Check out this `COLAB `_ +for a live demo. + """ # """ # Lightning Module interface # ========================== -# +# # A lightning module is a strict superclass of nn.Module, it provides a standard interface # for the trainer to interact with the model. -# +# # The easiest thing to do is copy the minimal example below and modify accordingly. -# +# # Otherwise, to Define a Lightning Module, implement the following methods: -# -# +# +# # Minimal example # --------------- -# +# # .. code-block:: python -# +# # import os # import torch # from torch.nn import functional as F # from torch.utils.data import DataLoader # from torchvision.datasets import MNIST # import torchvision.transforms as transforms -# +# # import pytorch_lightning as pl -# +# # class CoolModel(pl.LightningModule): -# +# # def __init__(self): # super(CoolModel, self).__init__() # # not the best model... # self.l1 = torch.nn.Linear(28 * 28, 10) -# +# # def forward(self, x): # return torch.relu(self.l1(x.view(x.size(0), -1))) -# +# # def training_step(self, batch, batch_idx): # # REQUIRED # x, y = batch # y_hat = self.forward(x) # return {'loss': F.cross_entropy(y_hat, y)} -# +# # def validation_step(self, batch, batch_idx): # # OPTIONAL # x, y = batch # y_hat = self.forward(x) # return {'val_loss': F.cross_entropy(y_hat, y)} -# +# # def validation_end(self, outputs): # # OPTIONAL # val_loss_mean = torch.stack([x['val_loss'] for x in outputs]).mean() # return {'val_loss': val_loss_mean} -# +# # def test_step(self, batch, batch_idx): # # OPTIONAL # x, y = batch # y_hat = self.forward(x) # return {'test_loss': F.cross_entropy(y_hat, y)} -# +# # def test_end(self, outputs): # # OPTIONAL # test_loss_mean = torch.stack([x['test_loss'] for x in outputs]).mean() # return {'test_loss': test_loss_mean} -# +# # def configure_optimizers(self): # # REQUIRED # return torch.optim.Adam(self.parameters(), lr=0.02) -# +# # @pl.data_loader # def train_dataloader(self): # return DataLoader(MNIST(os.getcwd(), train=True, download=True, # transform=transforms.ToTensor()), batch_size=32) -# +# # @pl.data_loader # def val_dataloader(self): # # OPTIONAL # # can also return a list of val dataloaders # return DataLoader(MNIST(os.getcwd(), train=True, download=True, # transform=transforms.ToTensor()), batch_size=32) -# +# # @pl.data_loader # def test_dataloader(self): # # OPTIONAL # # can also return a list of test dataloaders # return DataLoader(MNIST(os.getcwd(), train=False, download=True, # transform=transforms.ToTensor()), batch_size=32) -# -# +# +# # How do these methods fit into the broader training? # --------------------------------------------------- -# +# # The LightningModule interface is on the right. Each method corresponds # to a part of a research project. Lightning automates everything not in blue. -# +# # .. figure:: docs/source/_static/images/overview_flat.jpg # :align: center -# +# # Overview. -# -# +# +# # Optional Methods # ---------------- -# +# # **add_model_specific_args** -# +# # .. code-block:: python -# +# # @staticmethod # def add_model_specific_args(parent_parser, root_dir) -# +# # Lightning has a list of default argparse commands. # This method is your chance to add or modify commands specific to your model. # The `hyperparameter argument parser # `_ # is available anywhere in your model by calling self.hparams. -# +# # **Return** # An argument parser -# +# # **Example** -# +# # .. code-block:: python -# +# # @staticmethod # def add_model_specific_args(parent_parser, root_dir): # parser = HyperOptArgumentParser(strategy=parent_parser.strategy, parents=[parent_parser]) -# +# # # param overwrites # # parser.set_defaults(gradient_clip_val=5.0) -# +# # # network params # parser.opt_list('--drop_prob', default=0.2, options=[0.2, 0.5], type=float, tunable=False) # parser.add_argument('--in_features', default=28*28) # parser.add_argument('--out_features', default=10) # # use 500 for CPU, 50000 for GPU to see speed difference # parser.add_argument('--hidden_dim', default=50000) -# +# # # data # parser.add_argument('--data_root', default=os.path.join(root_dir, 'mnist'), type=str) -# +# # # training params (opt) # parser.opt_list('--learning_rate', default=0.001, type=float, # options=[0.0001, 0.0005, 0.001, 0.005], tunable=False) @@ -149,7 +242,7 @@ Test # parser.opt_list('--optimizer_name', default='adam', type=str, # options=['adam'], tunable=False) # return parser -# +# # """ from .lightning import LightningModule diff --git a/pytorch_lightning/core/lightning.py b/pytorch_lightning/core/lightning.py index 74b17ece..7328e57b 100644 --- a/pytorch_lightning/core/lightning.py +++ b/pytorch_lightning/core/lightning.py @@ -10,7 +10,7 @@ from argparse import Namespace import torch import torch.distributed as dist -# +# from pytorch_lightning.core.decorators import data_loader from pytorch_lightning.core.grads import GradInformation from pytorch_lightning.core.hooks import ModelHooks @@ -20,81 +20,43 @@ from pytorch_lightning.overrides.data_parallel import LightningDistributedDataPa class LightningModule(ABC, GradInformation, ModelIO, ModelHooks): - # """ - # A LightningModule has the following properties which you can access at any time - # - # **logger** - # A reference to the logger you passed into trainer. - # Passing a logger is optional. If you don't pass one in, Lightning will create one - # for you automatically. This logger saves logs to `/os.getcwd()/lightning_logs`:: - # - # Trainer(logger=your_logger) - # - # - # Call it from anywhere in your LightningModule to add metrics, images, etc... - # whatever your logger supports. - # - # Here is an example using the TestTubeLogger (which is a wrapper - # on 'PyTorch SummaryWriter `_ - # with versioned folder structure). - # - # .. code-block:: python - # - # # if logger is a tensorboard logger or TestTubeLogger - # self.logger.experiment.add_embedding(...) - # self.logger.experiment.log({'val_loss': 0.9}) - # self.logger.experiment.add_scalars(...) - # - # - # **trainer** - # Last resort access to any state the trainer has. - # Changing certain properties here could affect your training run. - # - # .. code-block:: python - # - # self.trainer.optimizers - # self.trainer.current_epoch - # ... - # - # Debugging - # --------- - # - # The LightningModule also offers these tricks to help debug. - # - # **example_input_array** - # - # In the LightningModule init, you can set a dummy tensor for this property - # to get a print out of sizes coming into and out of every layer. - # - # .. code-block:: python - # - # def __init__(self): - # # put the dimensions of the first input to your system - # self.example_input_array = torch.rand(5, 28 * 28) - # """ - def __init__(self, *args, **kwargs): super(LightningModule, self).__init__(*args, **kwargs) #: Current dtype self.dtype = torch.FloatTensor + self.exp_save_path = None + #: The current epoch self.current_epoch = 0 + #: Total training batches seen across all epochs self.global_step = 0 + self.loaded_optimizer_states_dict = {} + + #: Pointer to the trainer object self.trainer = None + + #: Pointer to the logger object self.logger = None self.example_input_array = None - # track if gpu was requested for checkpointing #: True if your model is currently running on GPUs. #: Useful to set flags around the LightningModule for different CPU vs GPU behavior. self.on_gpu = False + + #: True if using dp self.use_dp = False + + #: True if using ddp self.use_ddp = False + + #: True if using ddp2 self.use_ddp2 = False + + #: True if using amp self.use_amp = False @abstractmethod