wassname/pytorch-lightning
1
0
Fork 0
mirror of https://github.com/wassname/pytorch-lightning.git synced 2026-09-10 12:21:57 +08:00
Code Issues Packages Projects Releases Wiki Activity

cleared spaces

Browse Source
This commit is contained in:
William Falcon
2020-01-16 15:11:17 -05:00
parent 36da61eb01
commit 6b86754bbf
2 changed files with 151 additions and 96 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pytorch_lightning/core/__init__.py
+134 -41
View File
@@ -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 <https://colab.research.google.com/drive/1F_RNcHzTfFuQf-LeKvSlud6x7jXYkG31#scrollTo=HOk9c4_35FKg>`_
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
# <https://williamfalcon.github.io/test-tube/hyperparameter_optimization/HyperOptArgumentParser>`_
# 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
pytorch_lightning/core/lightning.py
+17 -55
View File
@@ -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 <https://pytorch.org/docs/stable/tensorboard.html>`_
# 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