diff --git a/README.md b/README.md index c70a3be9..b7ff2f53 100644 --- a/README.md +++ b/README.md @@ -58,9 +58,12 @@ Don't worry about training on multiple gpus or speeding up your code, lightning --- ## How do I do use it? +The research code goes into a [LightningModule]((https://williamfalcon.github.io/pytorch-lightning/LightningModule/RequiredTrainerInterface/)) which you fit using a Trainer. + +Think of the LightningModule as a *system* such as seq-2-seq, GAN, etc... However, the LightningModule can ALSO just be a simple classifier such as the example below. To use lightning do 2 things: -1. [Define a LightningModel](https://williamfalcon.github.io/pytorch-lightning/LightningModule/RequiredTrainerInterface/) +1. [Define a LightningModule](https://williamfalcon.github.io/pytorch-lightning/LightningModule/RequiredTrainerInterface/) ```python import os import torch @@ -71,10 +74,10 @@ import torchvision.transforms as transforms import pytorch_lightning as pl -class CoolModel(pl.LightningModule): +class CoolSystem(pl.LightningModule): def __init__(self): - super(CoolModel, self).__init__() + super(CoolSystem, self).__init__() # not the best model... self.l1 = torch.nn.Linear(28 * 28, 10) @@ -122,7 +125,7 @@ class CoolModel(pl.LightningModule): ```python from pytorch_lightning import Trainer -model = CoolModel() +model = CoolSystem() # most basic trainer, uses good defaults trainer = Trainer()