PyTorch Lightning 全攻略！

来源：极市平台本文约9600字，建议阅读10+分钟本文作者与你分享使用Pytorch-Lightning心得。

root-    |-data        |-__init__.py        |-data_interface.py        |-xxxdataset1.py        |-xxxdataset2.py        |-...    |-model        |-__init__.py        |-model_interface.py        |-xxxmodel1.py        |-xxxmodel2.py        |-...    |-main.py

outs = []for batch in data:    out = training_step(batch)    outs.append(out)training_epoch_end(outs)
等价Lightning代码：def training_step(self, batch, batch_idx):    prediction = ...    return prediction
def training_epoch_end(self, training_step_outputs):    for prediction in predictions:        # do something with these

def training_step(self, batch, batch_idx):    x, y, z = batch    out = self.encoder(x)    loss = self.loss(out, x)    return loss
# Multiple optimizers (e.g.: GANs)def training_step(self, batch, batch_idx, optimizer_idx):    if optimizer_idx == 0:        # do training_step with encoder    if optimizer_idx == 1:        # do training_step with decoder
# Truncated back-propagation through timedef training_step(self, batch, batch_idx, hiddens):    # hiddens are the hidden states from the previous truncated backprop step    ...    out, hiddens = self.lstm(data, hiddens)    ...    return {'loss': loss, 'hiddens': hiddens}

# most casesdef configure_optimizers(self):    opt = Adam(self.parameters(), lr=1e-3)    return opt
# multiple optimizer case (e.g.: GAN)def configure_optimizers(self):    generator_opt = Adam(self.model_gen.parameters(), lr=0.01)    disriminator_opt = Adam(self.model_disc.parameters(), lr=0.02)    return generator_opt, disriminator_opt
# example with learning rate schedulersdef configure_optimizers(self):    generator_opt = Adam(self.model_gen.parameters(), lr=0.01)    disriminator_opt = Adam(self.model_disc.parameters(), lr=0.02)    discriminator_sched = CosineAnnealing(discriminator_opt, T_max=10)    return [generator_opt, disriminator_opt], [discriminator_sched]
# example with step-based learning rate schedulersdef configure_optimizers(self):    gen_opt = Adam(self.model_gen.parameters(), lr=0.01)    dis_opt = Adam(self.model_disc.parameters(), lr=0.02)    gen_sched = {'scheduler': ExponentialLR(gen_opt, 0.99),                 'interval': 'step'}  # called after each training step    dis_sched = CosineAnnealing(discriminator_opt, T_max=10) # called every epoch    return [gen_opt, dis_opt], [gen_sched, dis_sched]
# example with optimizer frequencies# see training procedure in `Improved Training of Wasserstein GANs`, Algorithm 1# https://arxiv.org/abs/1704.00028def configure_optimizers(self):    gen_opt = Adam(self.model_gen.parameters(), lr=0.01)    dis_opt = Adam(self.model_disc.parameters(), lr=0.02)    n_critic = 5    return (        {'optimizer': dis_opt, 'frequency': n_critic},        {'optimizer': gen_opt, 'frequency': 1}    )

class LitModel(pl.LightningModule):
    def __init__(...):
    def forward(...):
    def training_step(...)
    def training_step_end(...)
    def training_epoch_end(...)
    def validation_step(...)
    def validation_step_end(...)
    def validation_epoch_end(...)
    def test_step(...)
    def test_step_end(...)
    def test_epoch_end(...)
    def configure_optimizers(...)
    def any_extra_hook(...)

model = MyLightningModule()
trainer = Trainer()trainer.fit(model, train_dataloader, val_dataloader)

def fit(...):    on_fit_start()
    if global_rank == 0:        # prepare data is called on GLOBAL_ZERO only        prepare_data()
    for gpu/tpu in gpu/tpus:        train_on_device(model.copy())
    on_fit_end()
def train_on_device(model):    # setup is called PER DEVICE    setup()    configure_optimizers()    on_pretrain_routine_start()
    for epoch in epochs:        train_loop()
    teardown()
def train_loop():    on_train_epoch_start()    train_outs = []    for train_batch in train_dataloader():        on_train_batch_start()
        # ----- train_step methods -------        out = training_step(batch)        train_outs.append(out)
        loss = out.loss
        backward()        on_after_backward()        optimizer_step()        on_before_zero_grad()        optimizer_zero_grad()
        on_train_batch_end(out)
        if should_check_val:            val_loop()
    # end training epoch    logs = training_epoch_end(outs)
def val_loop():    model.eval()    torch.set_grad_enabled(False)
    on_validation_epoch_start()    val_outs = []    for val_batch in val_dataloader():        on_validation_batch_start()
        # -------- val step methods -------        out = validation_step(val_batch)        val_outs.append(out)
        on_validation_batch_end(out)
    validation_epoch_end(val_outs)    on_validation_epoch_end()
    # set up for train    model.train()    torch.set_grad_enabled(True)

# default used by the Trainer (no scaling of batch size)trainer = Trainer(auto_scale_batch_size=None)
# run batch size scaling, result overrides hparams.batch_sizetrainer = Trainer(auto_scale_batch_size='binsearch')
# call tune to find the batch sizetrainer.tune(model)

# run learning rate finder, results override hparams.learning_ratetrainer = Trainer(auto_lr_find=True)
# run learning rate finder, results override hparams.my_lr_argtrainer = Trainer(auto_lr_find='my_lr_arg')
# call tune to find the lrtrainer.tune(model)
precision：精确度。正常是32，使用16可以减小内存消耗，增大batch。# default used by the Trainertrainer = Trainer(precision=32)
# 16-bit precisiontrainer = Trainer(precision=16, gpus=1)

# default used by the Trainertrainer = Trainer(val_check_interval=1.0)
# check validation set 4 times during a training epochtrainer = Trainer(val_check_interval=0.25)
# check validation set every 1000 training batches# use this when using iterableDataset and your dataset has no length# (ie: production cases with streaming data)trainer = Trainer(val_check_interval=1000)

# default used by the Trainer (ie: train on CPU)trainer = Trainer(gpus=None)
# equivalenttrainer = Trainer(gpus=0)
# int: train on 2 gpustrainer = Trainer(gpus=2)
# list: train on GPUs 1, 4 (by bus ordering)trainer = Trainer(gpus=[1, 4])trainer = Trainer(gpus='1, 4') # equivalent
# -1: train on all gpustrainer = Trainer(gpus=-1)trainer = Trainer(gpus='-1') # equivalent
# combine with num_nodes to train on multiple GPUs across nodes# uses 8 gpus in totaltrainer = Trainer(gpus=2, num_nodes=4)
# train only on GPUs 1 and 4 across nodestrainer = Trainer(gpus=[1, 4], num_nodes=4)

# default used by the Trainertrainer = Trainer(limit_train_batches=1.0)
# run through only 25% of the training set each epochtrainer = Trainer(limit_train_batches=0.25)
# run through only 10 batches of the training set each epochtrainer = Trainer(limit_train_batches=10)

# default used by the Trainertrainer = Trainer(fast_dev_run=False)
# runs 1 train, val, test batch and program endstrainer = Trainer(fast_dev_run=True)
# runs 7 train, val, test batches and program endstrainer = Trainer(fast_dev_run=7)

from argparse import ArgumentParser
def main(hparams):    model = LightningModule()    trainer = Trainer(gpus=hparams.gpus)    trainer.fit(model)
if __name__ == '__main__':    parser = ArgumentParser()    parser.add_argument('--gpus', default=None)    args = parser.parse_args()
    main(args)

from argparse import ArgumentParser
def main(args):    model = LightningModule()    trainer = Trainer.from_argparse_args(args)    trainer.fit(model)
if __name__ == '__main__':    parser = ArgumentParser()    parser = Trainer.add_argparse_args(        # group the Trainer arguments together        parser.add_argument_group(title="pl.Trainer args")    )    args = parser.parse_args()
    main(args)

from argparse import ArgumentParserimport pytorch_lightning as plfrom pytorch_lightning import LightningModule, Trainer
def main(args):    model = LightningModule()    trainer = Trainer.from_argparse_args(args)    trainer.fit(model)
if __name__ == '__main__':    parser = ArgumentParser()    parser.add_argument('--batch_size', default=32, type=int)    parser.add_argument('--hidden_dim', type=int, default=128)    parser = Trainer.add_argparse_args(        # group the Trainer arguments together        parser.add_argument_group(title="pl.Trainer args")    )    args = parser.parse_args()
    main(args)

class LitClassifier(pl.LightningModule):
     def __init__(self, model):         super().__init__()         self.model = model
     def training_step(self, batch, batch_idx):         x, y = batch         y_hat = self.model(x)         loss = F.cross_entropy(y_hat, y)         return loss

# put model in train modemodel.train()torch.set_grad_enabled(True)
losses = []for batch in train_dataloader:    # forward    loss = training_step(batch)    losses.append(loss.detach())
    # backward    loss.backward()
    # apply and clear grads    optimizer.step()    optimizer.zero_grad()

def training_step(self, batch, batch_idx):    x, y = batch    y_hat = self.model(x)    loss = F.cross_entropy(y_hat, y)
    # logs metrics for each training_step,    # and the average across the epoch, to the progress bar and logger    self.log('train_loss', loss, on_step=True, on_epoch=True, prog_bar=True, logger=True)    return loss

outs = []for batch in train_dataloader:    # forward    out = training_step(val_batch)
    # backward    loss.backward()
    # apply and clear grads    optimizer.step()    optimizer.zero_grad()
epoch_metric = torch.mean(torch.stack([x['train_loss'] for x in outs]))

def training_step(self, batch, batch_idx):    x, y = batch    y_hat = self.model(x)    loss = F.cross_entropy(y_hat, y)    preds = ...    return {'loss': loss, 'other_stuff': preds}
def training_epoch_end(self, training_step_outputs):   for pred in training_step_outputs:       # do something

outs = []for batch in train_dataloader:    # forward    out = training_step(val_batch)
    # backward    loss.backward()
    # apply and clear grads    optimizer.step()    optimizer.zero_grad()
training_epoch_end(outs)

class MNISTDataModule(pl.LightningDataModule):
    def __init__(self, data_dir: str = './', batch_size: int = 64, num_workers: int = 8):        super().__init__()        self.data_dir = data_dir        self.batch_size = batch_size        self.num_workers = num_workers
        self.transform = transforms.Compose([            transforms.ToTensor(),            transforms.Normalize((0.1307,), (0.3081,))        ])
        # self.dims is returned when you call dm.size()        # Setting default dims here because we know them.        # Could optionally be assigned dynamically in dm.setup()        self.dims = (1, 28, 28)        self.num_classes = 10
    def prepare_data(self):        # download        MNIST(self.data_dir, train=True, download=True)        MNIST(self.data_dir, train=False, download=True)
    def setup(self, stage=None):        # Assign train/val datasets for use in dataloaders        if stage == 'fit' or stage is None:            mnist_full = MNIST(self.data_dir, train=True, transform=self.transform)            self.mnist_train, self.mnist_val = random_split(mnist_full, [55000, 5000])
        # Assign test dataset for use in dataloader(s)        if stage == 'test' or stage is None:            self.mnist_test = MNIST(self.data_dir, train=False, transform=self.transform)
    def train_dataloader(self):        return DataLoader(self.mnist_train, batch_size=self.batch_size, num_workers=self.num_workers)
    def val_dataloader(self):        return DataLoader(self.mnist_val, batch_size=self.batch_size, num_workers=self.num_workers)
    def test_dataloader(self):        return DataLoader(self.mnist_test, batch_size=self.batch_size, num_workers=self.num_workers)

要点

from pytorch_lightning.callbacks import ModelCheckpoint
# saves a file like: my/path/sample-mnist-epoch=02-val_loss=0.32.ckptcheckpoint_callback = ModelCheckpoint(    monitor='val_loss',    filename='sample-mnist-{epoch:02d}-{val_loss:.2f}',    save_top_k=3,    mode='min',    save_last=True)
trainer = pl.Trainer(gpus=1, max_epochs=3, progress_bar_refresh_rate=20, callbacks=[checkpoint_callback])

model = MyLightingModule.load_from_checkpoint(PATH)
print(model.learning_rate)# prints the learning_rate you used in this checkpoint
model.eval()y_hat = model(x)

class LitModel(LightningModule):    def __init__(self, in_dim, out_dim):      super().__init__()      self.save_hyperparameters()      self.l1 = nn.Linear(self.hparams.in_dim, self.hparams.out_dim)
# if you train and save the model like this it will use these values when loading# the weights. But you can overwrite thisLitModel(in_dim=32, out_dim=10)
# uses in_dim=32, out_dim=10model = LitModel.load_from_checkpoint(PATH)
# uses in_dim=128, out_dim=10model = LitModel.load_from_checkpoint(PATH, in_dim=128, out_dim=10)

完全load训练状态：load包括模型的一切，以及和训练相关的一切参数，如model, epoch, step, LR schedulers, apex等

model = LitModel()trainer = Trainer(resume_from_checkpoint='some/path/to/my_checkpoint.ckpt')
# automatically restores model, epoch, step, LR schedulers, apex, etc...trainer.fit(model)

from pytorch_lightning import Trainerfrom pytorch_lightning.callbacks import EarlyStopping
early_stopping = EarlyStopping('val_loss')trainer = Trainer(callbacks=[early_stopping])

from pl_bolts.callbacks import PrintTableMetricsCallback
callback = PrintTableMetricsCallback()trainer = pl.Trainer(callbacks=[callback])trainer.fit(...)
# ------------------------------# at the end of every epoch it will print# ------------------------------
# loss│train_loss│val_loss│epoch# ──────────────────────────────# 2.2541470527648926│2.2541470527648926│2.2158432006835938│0

from pytorch_lightning import loggers as pl_loggers
# Defaulttb_logger = pl_loggers.TensorBoardLogger(    save_dir=os.getcwd(),    version=None,    name='lightning_logs')trainer = Trainer(logger=tb_logger)
# Or use the same format as otherstb_logger = pl_loggers.TensorBoardLogger('logs/')
# One Loggercomet_logger = pl_loggers.CometLogger(save_dir='logs/')trainer = Trainer(logger=comet_logger)
# Save code snapshotlogger = pl_loggers.TestTubeLogger('logs/', create_git_tag=True)
# Multiple Loggertb_logger = pl_loggers.TensorBoardLogger('logs/')comet_logger = pl_loggers.CometLogger(save_dir='logs/')trainer = Trainer(logger=[tb_logger, comet_logger])

def training_step(...):    ...    # the logger you used (in this case tensorboard)    tensorboard = self.logger.experiment    tensorboard.add_image()    tensorboard.add_histogram(...)    tensorboard.add_figure(...)

# Start tensorboard.%load_ext tensorboard%tensorboard --logdir lightning_logs/

tensorboard --logdir lightning_logs --bind_all` -> `http://SERVER-NAME:6006/

import torchvision.models as models
class ImagenetTransferLearning(LightningModule):    def __init__(self):        super().__init__()
        # init a pretrained resnet        backbone = models.resnet50(pretrained=True)        num_filters = backbone.fc.in_features        layers = list(backbone.children())[:-1]        self.feature_extractor = nn.Sequential(*layers)
        # use the pretrained model to classify cifar-10 (10 image classes)        num_target_classes = 10        self.classifier = nn.Linear(num_filters, num_target_classes)
    def forward(self, x):        self.feature_extractor.eval()        with torch.no_grad():            representations = self.feature_extractor(x).flatten(1)        x = self.classifier(representations)        ...

# badt = torch.rand(2, 2).cuda()
# good (self is LightningModule)t = torch.rand(2, 2, device=self.device)

# badself.t = torch.rand(2, 2, device=self.device)
# goodself.register_buffer("t", torch.rand(2, 2))

class RDNFuse(nn.Module):    ...    def init_norm_func(self, ref):        self.mean = torch.tensor(np.array(self.mean_sen), dtype=torch.float32).type_as(ref)
    def forward(self, x):        if not hasattr(self, 'mean'):            self.init_norm_func(x)

# Single optimizerfor epoch in epochs:    for batch in data:        loss = model.training_step(batch, batch_idx, ...)        loss.backward()        optimizer.step()        optimizer.zero_grad()
    for scheduler in schedulers:        scheduler.step()
# Multiple optimizersfor epoch in epochs:  for batch in data:     for opt in optimizers:        disable_grads_for_other_optimizers()        train_step(opt)        opt.step()
  for scheduler in schedulers:     scheduler.step()

random_split(range(10), [3, 7], generator=torch.Generator().manual_seed(42))

from torch.utils.data import DataLoader, random_splitfrom torchvision.datasets import MNIST
mnist_full = MNIST(self.data_dir, train=True, transform=self.transform)self.mnist_train, self.mnist_val = random_split(mnist_full, [55000, 5000])

Parameters：dataset (https://pytorch.org/docs/stable/data.html%23torch.utils.data.Dataset) – Dataset to be splitlengths – lengths of splits to be producedgenerator (https://pytorch.org/docs/stable/generated/torch.Generator.html%23torch.Generator) – Generator used for the random permutation.