Getting started

__version__ == 1.2.6

LIGHTNING IN 2 STEPS

导入下面的包：

import os
import torch
from torch import nn
import torch.nn.functional as F
from torchvision import transforms
from torchvision.datasets import MNIST
from torch.utils.data import DataLoader, random_split
import pytorch_lightning as pl

Step 1: Define LightningModule

class LitAutoEncoder(pl.LightningModule):

    def __init__(self):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(28*28, 64),
            nn.ReLU(),
            nn.Linear(64, 3)
        )
        self.decoder = nn.Sequential(
            nn.Linear(3, 64),
            nn.ReLU(),
            nn.Linear(64, 28*28)
        )

    def forward(self, x):
        # in lightning, forward defines the prediction/inference actions
        embedding = self.encoder(x)
        return embedding

    def training_step(self, batch, batch_idx):
        # training_step defined the train loop.
        # It is independent of forward
        x, y = batch
        x = x.view(x.size(0), -1)
        z = self.encoder(x)
        x_hat = self.decoder(z)
        loss = F.mse_loss(x_hat, x)
        # Logging to TensorBoard by default
        self.log('train_loss', loss)
        return loss

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
        return optimizer

SYSTEM VS MODEL

一个 lightning module 定义一个系统而不是一个模型。

系统的例子有：

Autoencoder、
BERT
DQN
GAN
Image classifier
Seq2seq
SimCLR
VAE

在 hood下 LightningModule 仍然仅仅是一个 torch.nn.Module 。将所有研究代码分组到单个文件中，使其包含：

The Train loop
The Validation loop
The Test loop
The Model or system of Models
The Optimizer

您可以通过重写 Available Callback hooks 其中的20多个hooks来定制训练的任何部分(例如反向传播)

class LitAutoEncoder(LightningModule):

    def backward(self, loss, optimizer, optimizer_idx):
        loss.backward()

FORWARD vs TRAINING_STEP

在 Lightning 中，我们把训练与推理分开。训练步骤定义了完整的训练循环。我们鼓励用户使用forward来定义推理操作。

例如，在这种情况下，我们可以定义自编码器作为 embedding 提取器：

def forward(self, x):
    embeddings = self.encoder(x)
    return embeddings

当然，没有什么可以阻止你在 training_step 中使用forward。

def training_step(self, batch, batch_idx):
    ...
    z = self(x)

这取决于你的应用。但是，我们建议您将这两种意图分开。

使用 forward 用于推理 (预测)
使用 training_step 训练

Step 2: Fit with Lightning Trainer

首先，按照您想要的方式定义数据。Lightning 需要一个 DataLoader 用于划分 train/val/test。

dataset = MNIST(os.getcwd(), download=True, transform=transforms.ToTensor())
train_loader = DataLoader(dataset)

接下来，初始化 lightning module 以及 PyTorch Lightning Trainer，然后用 data 和 model 调用 fit。

# init model
autoencoder = LitAutoEncoder()

# most basic trainer, uses good defaults (auto-tensorboard, checkpoints, logs, and more)
# trainer = pl.Trainer(gpus=8) (if you have GPUs)
trainer = pl.Trainer()
trainer.fit(autoencoder, train_loader)

Trainer 自动化：

Epoch and batch iteration
Calling of optimizer.step(), backward, zero_grad()
Calling of .eval(), enabling/disabling grads
weights loading
Tensorboard (see loggers options)
Multi-GPU support
TPU
AMP support

TIP：如果您喜欢手动管理优化器，您可以使用 Manual optimization 模式（比如： RL, GANs， etc…）

Basic features

Manual vs automatic optimization

Automatic optimization

使用Lightning，您不必担心何时启用/禁用grads、执行反向传播或更新优化器只要您从 training_step 返回一个带有graph的loss，Lightning就会自动进行优化。

def training_step(self, batch, batch_idx):
    loss = self.encoder(batch)
    return loss

Manual optimization

然而，对于某些研究，如GANs，强化学习，或有多个优化器或内部循环的东西，你可以关闭自动优化，完全控制自己的训练循环。

关闭自动优化，你就可以控制train loop ！

def __init__(self):
    self.automatic_optimization = False

def training_step(self, batch, batch_idx):
    # access your optimizers with use_pl_optimizer=False. Default is True
    opt_a, opt_b = self.optimizers(use_pl_optimizer=True)

    loss_a = self.generator(batch)
    opt_a.zero_grad()
    # use `manual_backward()` instead of `loss.backward` to automate half precision, etc...
    self.manual_backward(loss_a)
    opt_a.step()

    loss_b = self.discriminator(batch)
    opt_b.zero_grad()
    self.manual_backward(loss_b)
    opt_b.step()

Predict or Deploy

当你完成训练，你有3个选择使用你的 LightningModule 预测。

Option 1: Sub-models

取出系统内的任何模型进行预测。

# ----------------------------------
# to use as embedding extractor
# ----------------------------------
autoencoder = LitAutoEncoder.load_from_checkpoint('path/to/checkpoint_file.ckpt')
encoder_model = autoencoder.encoder
encoder_model.eval()

# ----------------------------------
# to use as image generator
# ----------------------------------
decoder_model = autoencoder.decoder
decoder_model.eval()

Option 2: Forward

您还可以添加一个forward方法来进行预测。

# ----------------------------------
# using the AE to extract embeddings
# ----------------------------------
class LitAutoEncoder(LightningModule):
    def __init__(self):
        super().__init__()
        self.encoder = nn.Sequential()

    def forward(self, x):
        embedding = self.encoder(x)
        return embedding

autoencoder = LitAutoEncoder()
autoencoder = autoencoder(torch.rand(1, 28 * 28))

# ----------------------------------
# or using the AE to generate images
# ----------------------------------
class LitAutoEncoder(LightningModule):
    def __init__(self):
        super().__init__()
        self.decoder = nn.Sequential()

    def forward(self):
        z = torch.rand(1, 3)
        image = self.decoder(z)
        image = image.view(1, 1, 28, 28)
        return image

autoencoder = LitAutoEncoder()
image_sample = autoencoder()

Option 3: Production

对于生产系统，onnx或torchscript要快得多。确保您已经添加了一个 forward 方法，或者只跟踪您需要的子模型。

# ----------------------------------
# torchscript
# ----------------------------------
autoencoder = LitAutoEncoder()
torch.jit.save(autoencoder.to_torchscript(), "model.pt")
os.path.isfile("model.pt")

# ----------------------------------
# onnx
# ----------------------------------
with tempfile.NamedTemporaryFile(suffix='.onnx', delete=False) as tmpfile:
     autoencoder = LitAutoEncoder()
     input_sample = torch.randn((1, 28 * 28))
     autoencoder.to_onnx(tmpfile.name, input_sample, export_params=True)
     os.path.isfile(tmpfile.name)

Using CPUs/GPUs/TPUs

在 Lightning 中使用cpu、gpu或tpu是很简单的。不需要改变你的代码，只需要改变 Trainer 选项。

# train on CPU
trainer = Trainer()

# train on 8 CPUs
trainer = Trainer(num_processes=8)

# train on 1024 CPUs across 128 machines
trainer = pl.Trainer(
    num_processes=8,
    num_nodes=128
)

# train on 1 GPU
trainer = pl.Trainer(gpus=1)

# train on multiple GPUs across nodes (32 gpus here)
trainer = pl.Trainer(
    gpus=4,
    num_nodes=8
)

# train on gpu 1, 3, 5 (3 gpus total)
trainer = pl.Trainer(gpus=[1, 3, 5])

# Multi GPU with mixed precision
trainer = pl.Trainer(gpus=2, precision=16)

# Train on TPUs
trainer = pl.Trainer(tpu_cores=8)

不需要修改代码中的任何一行，您现在就可以使用上面的代码执行以下操作

# train on TPUs using 16 bit precision
# using only half the training data and checking validation every quarter of a training epoch
trainer = pl.Trainer(
    tpu_cores=8,
    precision=16,
    limit_train_batches=0.5,
    val_check_interval=0.25
)

Checkpoints

Lightning 会自动保存你的模型。一旦您完成了训练，您就可以按照以下方式加载 checkpoints：

model = LitModel.load_from_checkpoint(path)

上面的 checkpoints 包含了初始化模型和设置 state dict 所需的所有参数。如果您更喜欢手动操作，以下是等效的方法

# load the ckpt
ckpt = torch.load('path/to/checkpoint.ckpt')

# equivalent to the above
model = LitModel()
model.load_state_dict(ckpt['state_dict'])

Data flow

每个循环(训练、验证、测试)都有三个可以实现的 hooks：

x_step
x_step_end
x_epoch_end

为了说明数据是如何流动的，我们将使用training loop(即x=training)

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

如果你使用DP或DDP2分布式模式(例如:在gpu上 split a batch)，使用 x_step_end 手动聚合(或者不要实现它让lightning自动聚合)。

for batch in data:
    model_copies = copy_model_per_gpu(model, num_gpus)
    batch_split = split_batch_per_gpu(batch, num_gpus)

    gpu_outs = []
    for model, batch_part in zip(model_copies, batch_split):
        # LightningModule hook
        gpu_out = model.training_step(batch_part)
        gpu_outs.append(gpu_out)

    # LightningModule hook
    out = training_step_end(gpu_outs)

在 lightning 中等价是：

def training_step(self, batch, batch_idx):
    loss = ...
    return loss

def training_step_end(self, losses):
    gpu_0_loss = losses[0]
    gpu_1_loss = losses[1]
    return (gpu_0_loss + gpu_1_loss) * 1/2

TIP： validation 和 test loops 具有相同的结构。

Logging

要记录到tensorboard，您最喜欢的记录器和/或进度条，请使用log()方法，该方法可以从LightningModule中的任何方法调用。

def training_step(self, batch, batch_idx):
    self.log('my_metric', x)

log() 方法有一些设置：

on_step (记录训练中该step的指标)
on_epoch(在时间结束时自动累积并记录)
prog_bar(记录到进度栏)
logger(记录到像Tensorboard一样的logger)

根据从何处调用log，Lightning会自动为您确定正确的模式。但是当然，您可以通过手动设置标志来覆盖默认行为

Note：设置on_epoch = True将在整个训练epoch内累积您的记录值。

def training_step(self, batch, batch_idx):
    self.log('my_loss', loss, on_step=True, on_epoch=True, prog_bar=True, logger=True)

Note：进度条中显示的损失值在最后一个值上进行了平滑（平均），因此它不同于训练/验证步骤中返回的实际损失。

您还可以直接使用logger的任何方法：

def training_step(self, batch, batch_idx):
    tensorboard = self.logger.experiment
    tensorboard.any_summary_writer_method_you_want())

训练开始后，您可以使用您喜欢的logger或启动Tensorboard日志来查看日志：

tensorboard --logdir ./lightning_logs

Note： Lightning会自动在进度栏中显示从training_step返回的损失值。因此，无需显式记录此类self.log（’loss’，loss，prog_bar = True）。

Optional extensions

Callbacks

callback是一个任意的独立程序，可以在训练循环的任意部分执行。

这是添加不太理想的学习率衰减规则的示例：

from pytorch_lightning.callbacks import Callback

class DecayLearningRate(Callback):

    def __init__(self):
        self.old_lrs = []

    def on_train_start(self, trainer, pl_module):
        # track the initial learning rates
        for opt_idx, optimizer in enumerate(trainer.optimizers):
            group = [param_group['lr'] for param_group in optimizer.param_groups]
            self.old_lrs.append(group)

    def on_train_epoch_end(self, trainer, pl_module, outputs):
        for opt_idx, optimizer in enumerate(trainer.optimizers):
            old_lr_group = self.old_lrs[opt_idx]
            new_lr_group = []
            for p_idx, param_group in enumerate(optimizer.param_groups):
                old_lr = old_lr_group[p_idx]
                new_lr = old_lr * 0.98
                new_lr_group.append(new_lr)
                param_group['lr'] = new_lr
            self.old_lrs[opt_idx] = new_lr_group

# And pass the callback to the Trainer
decay_callback = DecayLearningRate()
trainer = Trainer(callbacks=[decay_callback])

以下的事情你可以通过 callback 做到：

Send emails at some point in training
Grow the model
Update learning rates
Visualize gradients
…
You are only limited by your imagination

Learn more about custom callbacks.

LightningDataModules

DataLoader和数据处理代码往往最终散落在各处。通过将其组织到LightningDataModule中，使数据代码可重用。

class MNISTDataModule(LightningDataModule):

      def __init__(self, batch_size=32):
          super().__init__()
          self.batch_size = batch_size

      # When doing distributed training, Datamodules have two optional arguments for
      # granular control over download/prepare/splitting data:

      # OPTIONAL, called only on 1 GPU/machine
      def prepare_data(self):
          MNIST(os.getcwd(), train=True, download=True)
          MNIST(os.getcwd(), train=False, download=True)

      # OPTIONAL, called for every GPU/machine (assigning state is OK)
      def setup(self, stage: Optional[str] = None):
          # transforms
          transform=transforms.Compose([
              transforms.ToTensor(),
              transforms.Normalize((0.1307,), (0.3081,))
          ])
          # split dataset
          if stage == 'fit':
              mnist_train = MNIST(os.getcwd(), train=True, transform=transform)
              self.mnist_train, self.mnist_val = random_split(mnist_train, [55000, 5000])
          if stage == 'test':
              self.mnist_test = MNIST(os.getcwd(), train=False, transform=transform)

      # return the dataloader for each split
      def train_dataloader(self):
          mnist_train = DataLoader(self.mnist_train, batch_size=self.batch_size)
          return mnist_train

      def val_dataloader(self):
          mnist_val = DataLoader(self.mnist_val, batch_size=self.batch_size)
          return mnist_val

      def test_dataloader(self):
          mnist_test = DataLoader(self.mnist_test, batch_size=self.batch_size)
          return mnist_test

LightningDataModule旨在支持在不同项目之间共享和重用数据 splits 和 transforms。它封装了处理数据所需的所有步骤： downloading, tokenizing, processing etc。

你可以将 LightningDataModule 送入 Trainer：

# init model
model = LitModel()

# init data
dm = MNISTDataModule()

# train
trainer = pl.Trainer()
trainer.fit(model, dm)

# test
trainer.test(datamodule=dm)

DataModules对于基于数据构建模型特别有用。阅读有关 datamodules 的更多信息。

Debugging

Lightning有许多调试工具。以下是其中一些示例：

# use only 10 train batches and 3 val batches
trainer = Trainer(limit_train_batches=10, limit_val_batches=3)

# Automatically overfit the sane batch of your model for a sanity test
trainer = Trainer(overfit_batches=1)

# unit test all the code- hits every line of your code once to see if you have bugs,
# instead of waiting hours to crash on validation
trainer = Trainer(fast_dev_run=True)

# train only 20% of an epoch
trainer = Trainer(limit_train_batches=0.2)

# run validation every 25% of a training epoch
trainer = Trainer(val_check_interval=0.25)

# Profile your code to find speed/memory bottlenecks
Trainer(profiler="simple")

Other coool features

定义并训练了第一个Lightning模型后，您可能想尝试其他一些很棒的功能，例如：

Automatic early stopping
Automatic truncated-back-propagation-through-time
Automatically scale your batch size
Automatically find a good learning rate
Load checkpoints directly from S3
Scale to massive compute clusters
Use multiple dataloaders per train/val/test loop
Use multiple optimizers to do reinforcement learning or even GANs

Grid AI

Grid AI是我们的本地解决方案，可在您选择的云提供商上进行大规模训练和调整。

Best practices

Style guide

Lightning 的一个主要目标是提高可读性和再现性。想象一下查看任何GitHub repo，找到一个 lightning 模块，并知道在哪里查找您关心的事情。

本 Style guide 的目标是鼓励Lightning代码采用类似的结构。

LightningModule

这些是构建LightningModule的最佳实践。

Systems vs models

LightningModule 背后的主要原理是一个完整的系统应该是独立的。在 Lightning 中，我们区分了系统和模型。

模型类似于resnet18、RNN等。

系统定义了一组模型如何相互交互。例如：

GANs
Seq2Seq
BERT
etc

一个 LightningModule 可以定义一个系统和一个模型。

下面是定义了一个模型的 LightningModule：

class LitModel(LightningModule):
    def __init__(self, num_layers: int = 3):
        super().__init__()
        self.layer_1 = nn.Linear()
        self.layer_2 = nn.Linear()
        self.layer_3 = nn.Linear()

这是定义了一个系统的 LightningModule：

class LitModel(LightningModule):
    def __init__(self, encoder: nn.Module = None, decoder: nn.Module = None):
        super().__init__()
        self.encoder = encoder
        self.decoder = decoder

为了快速创建原型，在LightningModule中定义所有的计算通常是有用的。为了可重用性和可伸缩性，最好是传入相关的 backbones。

Self-contained

Lightning module 应该是独立的。要了解模型的 self-contained 程度，一个很好的测试方法是问自己这个问题：

“有人可以在不知道任何内部信息的情况下把这个文件放到 Trainer 那里吗? “

例如，我们将优化器与模型结合在一起，因为大多数模型需要一个特定的优化器和一个特定的学习率调度器来正常工作。

Init

LightningModules可能停止自 self-contained 的第一个地方是init。尝试在init中定义所有相关的合理默认值，这样用户就不必猜测了。

下面是一个例子，用户将不得不遍历文件来找出如何初始化这个LightningModule。

class LitModel(LightningModule):
    def __init__(self, params):
        self.lr = params.lr
        self.coef_x = params.coef_x

这样定义的模型给您留下了许多问题； coef_x 是多少? 它是一个字符串吗? 一个浮点数? 范围是什么? 等…

相反，在init中要显式

class LitModel(LightningModule):
    def __init__(self, encoder: nn.Module, coeff_x: float = 0.2, lr: float = 1e-3):
        ...

现在用户不需要猜了。进而，它们知道值的类型，并且模型有一个合理的默认值，用户可以立即看到值。

Method order

LightningModule中唯一需要的方法是：

init
training_step
configure_optimizers

但是，如果您决定实现其余的可选方法，那么推荐的顺序是

model/system definition (init)
if doing inference, define forward
training hooks
validation hooks
test hooks
configure_optimizers
any other hooks

在实践中，这段代码如下：

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(...)

Forward vs training_step

我们建议使用 forward 进行 inference/predictions 并保持 training_step 独立

def forward(...):
    embeddings = self.encoder(x)

def training_step(...):
    x, y = ...
    z = self.encoder(x)
    pred = self.decoder(z)
    ...

然而，当使用DataParallel时，您将需要手动调用forward

def training_step(...):
    x, y = ...
    z = self(x)  # < ---------- instead of self.encoder(x)
    pred = self.decoder(z)
    ...

Data

这些是处理数据的最佳实践。

Dataloaders

Lightning使用 dataloaders 来处理通过系统的所有数据流。在构建 dataloaders 时，请确保调整 number of workers 的数量以获得最大效率。

Warning：确保不要使用ddp spawn与num workers > 0一起使用，否则你的代码会出现瓶颈。

DataModules

Lightning 引入了 datamodules。dataloaders 的问题是，共享完整的数据集通常仍然具有挑战性，因为所有这些问题都需要回答：

What splits were used?
How many samples does this dataset have?
What transforms were used?
etc…

正是因为这个原因，我们建议您使用datamomodules。这在协作时特别重要，因为它也会为您的团队节省大量时间。

他们所需要做的就是将一个 datamodule 放入 lightning trainer 中，而不必担心对数据做了什么。

在学术和公司环境中都是如此，data cleaning 和 ad-hoc instructions 会减慢迭代想法的进度。

Optional extensions

LightningDataModule

Why do I need a DataModule?

在一般的 PyTorch 代码中，数据清洗/预处理通常分散在许多文件中。这使得在项目之间共享和重用精确的分割和转换是不可能的。

如果你曾经问过这些问题，Datamodules是为你准备的：

你曾用过什么splits？
你曾用过什么 transforms？
你曾用过什么 normalization？
你是如何 prepare/tokenize 这些数据的？

What is a DataModule

DataModule只是train_dataloader，val_dataloader，test_dataloader以及所需的matching transforms和数据 processing/downloads 步骤的集合。

这里是一个 PyTorch 的例子：

# regular PyTorch
test_data = MNIST(my_path, train=False, download=True)
train_data = MNIST(my_path, train=True, download=True)
train_data, val_data = random_split(train_data, [55000, 5000])

train_loader = DataLoader(train_data, batch_size=32)
val_loader = DataLoader(val_data, batch_size=32)
test_loader = DataLoader(test_data, batch_size=32)

等效的DataModule只是组织完全相同的代码，但使其可跨项目重用。

class MNISTDataModule(pl.LightningDataModule):

    def __init__(self, data_dir: str = "path/to/dir", batch_size: int = 32):
        super().__init__()
        self.data_dir = data_dir
        self.batch_size = batch_size

    def setup(self, stage: Optional[str] = None):
        self.mnist_test = MNIST(self.data_dir, train=False)
        mnist_full = MNIST(self.data_dir, train=True)
        self.mnist_train, self.mnist_val = random_split(mnist_full, [55000, 5000])

    def train_dataloader(self):
        return DataLoader(self.mnist_train, batch_size=self.batch_size)

    def val_dataloader(self):
        return DataLoader(self.mnist_val, batch_size=self.batch_size)

    def test_dataloader(self):
        return DataLoader(self.mnist_test, batch_size=self.batch_size)

    def teardown(self, stage: Optional[str] = None):
        # Used to clean-up when the run is finished
        ...

但是现在，随着你的处理复杂性的增长(transforms，多gpu训练)，你可以让Lightning处理这些细节，同时使这个数据集可重用，这样你就可以与同事分享或在不同的项目中使用。

mnist = MNISTDataModule(my_path)
model = LitClassifier()

trainer = Trainer()
trainer.fit(model, mnist)

下面是一个更现实、更复杂的数据模块，它显示了该数据模块的可重用性。

import pytorch_lightning as pl
from torch.utils.data import random_split, DataLoader

# Note - you must have torchvision installed for this example
from torchvision.datasets import MNIST
from torchvision import transforms


class MNISTDataModule(pl.LightningDataModule):

    def __init__(self, data_dir: str = './'):
        super().__init__()
        self.data_dir = data_dir
        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)

    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: Optional[str] = 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])

            # Optionally...
            # self.dims = tuple(self.mnist_train[0][0].shape)

        # 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)

            # Optionally...
            # self.dims = tuple(self.mnist_test[0][0].shape)

    def train_dataloader(self):
        return DataLoader(self.mnist_train, batch_size=32)

    def val_dataloader(self):
        return DataLoader(self.mnist_val, batch_size=32)

    def test_dataloader(self):
        return DataLoader(self.mnist_test, batch_size=32)

Note: setup 需要一个参数 stage。它用于分离 trainer.fit 和 trainer.test 的 setup 逻辑。

LightningDataModule API

定义一个 DataModule 定义 5 个方法：

prepare_data (how to download(), tokenize, etc…)
setup (how to split, etc…)
train_dataloader
val_dataloader(s)
test_dataloader(s)

prepare_data

使用此方法可以执行可能写入磁盘的操作，或者在分布式设置中只需要从单个进程完成的操作。

download
tokenize
etc…

class MNISTDataModule(pl.LightningDataModule):
    def prepare_data(self):
        # download
        MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor())
        MNIST(os.getcwd(), train=False, download=True, transform=transforms.ToTensor())

Warning： prepare_data 从一个单线程调用 (比如 GPU 0)。不要用它取设置状态 (self.x = y)。

setup

还有一些你可能想要在每个GPU上执行的数据操作。使用 setup 来做这样的事情：

count number of classes
build vocabulary
perform train/val/test splits
apply transforms (defined explicitly in your datamodule or assigned in init)
etc…

import pytorch_lightning as pl


class MNISTDataModule(pl.LightningDataModule):

    def setup(self, stage: Optional[str] = None):

        # Assign Train/val split(s) for use in Dataloaders
        if stage == 'fit' or stage is None:
            mnist_full = MNIST(
                self.data_dir,
                train=True,
                download=True,
                transform=self.transform
            )
            self.mnist_train, self.mnist_val = random_split(mnist_full, [55000, 5000])
            self.dims = self.mnist_train[0][0].shape

        # Assign Test split(s) for use in Dataloaders
        if stage == 'test' or stage is None:
            self.mnist_test = MNIST(
                self.data_dir,
                train=False,
                download=True,
                transform=self.transform
            )
            self.dims = getattr(self, 'dims', self.mnist_test[0][0].shape)

setup 在每个进程上调用。在这里设置状态是ok的。

Note： teardown 可以用来清理状态。它也会从每个进程中调用。

train_dataloader

使用此方法生成train dataloader。通常你在 setup 中包装你定义的数据集

import pytorch_lightning as pl

class MNISTDataModule(pl.LightningDataModule):
    def train_dataloader(self):
        return DataLoader(self.mnist_train, batch_size=64)

val_dataloader

使用此方法生成val dataloader。通常你在 setup 中包装你定义的数据集

import pytorch_lightning as pl


class MNISTDataModule(pl.LightningDataModule):
    def val_dataloader(self):
        return DataLoader(self.mnist_val, batch_size=64)

test_dataloader

使用此方法生成test dataloader。通常你在 setup 中包装你定义的数据集

import pytorch_lightning as pl


class MNISTDataModule(pl.LightningDataModule):
    def test_dataloader(self):
        return DataLoader(self.mnist_test, batch_size=64)

transfer_batch_to_device

重写以定义如何将任意batch 移动到 device。

class MNISTDataModule(LightningDataModule):
    def transfer_batch_to_device(self, batch, device):
        x = batch['x']
        x = CustomDataWrapper(x)
        batch['x'] = x.to(device)
        return batch

Note: 这个 hook 只在单个GPU训练和DDP上运行(no data-parallel)

on_before_batch_transfer

在将 batch 转移到 device 之前，重写可对其更改或应用数据增强。

class MNISTDataModule(LightningDataModule):
    def on_before_batch_transfer(self, batch, dataloader_idx):
        batch['x'] = transforms(batch['x'])
        return batch

Warning: 目前dataloader idx总是返回0，将来会更新以支持真实的 idx。

Note：这个 hook 只在单个GPU训练和DDP上运行(no data-parallel)

on_after_batch_transfer

在batch 被转移到 device 后，重写可对其更改或应用数据增强。

class MNISTDataModule(LightningDataModule):
    def on_after_batch_transfer(self, batch, dataloader_idx):
        batch['x'] = gpu_transforms(batch['x'])
        return batch

Warning：目前dataloader idx总是返回0，将来会更新以支持真实的 idx。

Note：这个 hook 只在单个GPU训练和DDP上运行(no data-parallel)。这个 hook 在使用CPU设备的时候也会被调用，所以在这里增加数据增强和在 on_before_batch_transfer 里意味着做同样的事情

Note：为了将数据与 transforms 分离，您可以通过 __init__ 参数化它们。

class MNISTDataModule(pl.LightningDataModule):
    def __init__(self, train_transforms, val_transforms, test_transforms):
        super().__init__()
        self.train_transforms = train_transforms
        self.val_transforms = val_transforms
        self.test_transforms = test_transforms

Using a DataModule

推荐使用 DataModule 的方式是：

dm = MNISTDataModule()
model = Model()
trainer.fit(model, dm)

trainer.test(datamodule=dm)

如果需要来自数据集的信息来构建模型，则手动运行prepare data和setup (Lightning仍然确保该方法在正确的设备上运行)。

dm = MNISTDataModule()
dm.prepare_data()
dm.setup(stage='fit')

model = Model(num_classes=dm.num_classes, width=dm.width, vocab=dm.vocab)
trainer.fit(model, dm)

dm.setup(stage='test')
trainer.test(datamodule=dm)

Datamodules without Lightning

当然，你也可以在普通的 PyTorch 代码中使用 DataModules。

# download, etc...
dm = MNISTDataModule()
dm.prepare_data()

# splits/transforms
dm.setup(stage='fit')

# use data
for batch in dm.train_dataloader():
    ...
for batch in dm.val_dataloader():
    ...

dm.teardown(stage='fit')

# lazy load test data
dm.setup(stage='test')
for batch in dm.test_dataloader():
    ...

dm.teardown(stage='test')

但是总的来说，DataModules 通过允许在统一的结构中指定数据集的所有细节来鼓励可重复性

Tutorials

Step-by-step walk-through

The research

Model

lightning module 拥有所有核心研究成分：

The model
The optimizers
The train/ val/ test steps

让我们先从模型开始。在这种情况下，我们将设计一个三层的神经网络。

import torch
from torch.nn import functional as F
from torch import nn
from pytorch_lightning.core.lightning import LightningModule

class LitMNIST(LightningModule):

  def __init__(self):
    super().__init__()

    # mnist images are (1, 28, 28) (channels, width, height)
    self.layer_1 = nn.Linear(28 * 28, 128)
    self.layer_2 = nn.Linear(128, 256)
    self.layer_3 = nn.Linear(256, 10)

  def forward(self, x):
    batch_size, channels, width, height = x.size()

    # (b, 1, 28, 28) -> (b, 1*28*28)
    x = x.view(batch_size, -1)
    x = self.layer_1(x)
    x = F.relu(x)
    x = self.layer_2(x)
    x = F.relu(x)
    x = self.layer_3(x)

    x = F.log_softmax(x, dim=1)
    return x

注意 lightning module 而不是 torch.nn.Module。 LightningModule等价于纯PyTorch Module，只是它增加了一些功能。但是，您可以像使用PyTorch Module一样使用它。

net = LitMNIST()
x = torch.randn(1, 1, 28, 28)
out = net(x)

输出：

torch.Size([1, 10])

现在我们添加包含所有训练循环逻辑的训练步骤：

class LitMNIST(LightningModule):

    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.nll_loss(logits, y)
        return loss

Data

Lightning 在纯 dataloader 上操作。这里是用于加载 MNIST 的 PyTorch 代码。

from torch.utils.data import DataLoader, random_split
from torchvision.datasets import MNIST
import os
from torchvision import datasets, transforms

# transforms
# prepare transforms standard to MNIST
transform=transforms.Compose([transforms.ToTensor(),
                              transforms.Normalize((0.1307,), (0.3081,))])

# data
mnist_train = MNIST(os.getcwd(), train=True, download=True, transform=transform)
mnist_train = DataLoader(mnist_train, batch_size=64)

你可以用三种方式使用数据加载器:

1> Pass DataLoaders to .fit()

将 dataloaders 送到 .fit() 函数

model = LitMNIST()
trainer = Trainer()
trainer.fit(model, mnist_train)

2> LightningModule DataLoaders

对于快速研究原型，将模型与数据加载器链接起来可能更容易。

class LitMNIST(pl.LightningModule):

    def train_dataloader(self):
        # transforms
        # prepare transforms standard to MNIST
        transform=transforms.Compose([transforms.ToTensor(),
                                      transforms.Normalize((0.1307,), (0.3081,))])
        # data
        mnist_train = MNIST(os.getcwd(), train=True, download=True, transform=transform)
        return DataLoader(mnist_train, batch_size=64)

    def val_dataloader(self):
        transforms = ...
        mnist_val = ...
        return DataLoader(mnist_val, batch_size=64)

    def test_dataloader(self):
        transforms = ...
        mnist_test = ...
        return DataLoader(mnist_test, batch_size=64)

数据加载器已经在模型中，不需要再使用 .fit()

model = LitMNIST()
trainer = Trainer()
trainer.fit(model)

3> DataModules(recommended)

定义 free-floating 数据加载器， splits， download instructions等等可能变得混乱。在这种情况下，最好将数据集的完整定义分组到 DataModule，它包含：

Download instructions
Processing instructions
Split instructions
Train dataloader
Val dataloader(s)
Test dataloader(s)

class MyDataModule(LightningDataModule):

    def __init__(self):
        super().__init__()
        self.train_dims = None
        self.vocab_size = 0

    def prepare_data(self):
        # called only on 1 GPU
        download_dataset()
        tokenize()
        build_vocab()

    def setup(self, stage: Optional[str] = None):
        # called on every GPU
        vocab = load_vocab()
        self.vocab_size = len(vocab)

        self.train, self.val, self.test = load_datasets()
        self.train_dims = self.train.next_batch.size()

    def train_dataloader(self):
        transforms = ...
        return DataLoader(self.train, batch_size=64)

    def val_dataloader(self):
        transforms = ...
        return DataLoader(self.val, batch_size=64)

    def test_dataloader(self):
        transforms = ...
        return DataLoader(self.test, batch_size=64)

使用 DataModules 可以更容易地共享完整数据集定义。

# use an MNIST dataset
mnist_dm = MNISTDatamodule()
model = LitModel(num_classes=mnist_dm.num_classes)
trainer.fit(model, mnist_dm)

# or other datasets with the same model
imagenet_dm = ImagenetDatamodule()
model = LitModel(num_classes=imagenet_dm.num_classes)
trainer.fit(model, imagenet_dm)

Note: prepare_data() 在分布式训练中(自动)只调用一个GPU。

Note：setup()在每个GPU上调用(自动)

Models defined by data

当您的模型需要了解数据时，最好在将数据传递给模型之前处理数据。

# init dm AND call the processing manually
dm = ImagenetDataModule()
dm.prepare_data()
dm.setup()

model = LitModel(out_features=dm.num_classes, img_width=dm.img_width, img_height=dm.img_height)
trainer.fit(model, dm)

使用prepare data()下载和处理数据集。
使用setup()进行分割，并构建模型内部构件

使用DataModule的另一种选择是将models模块的初始化延迟到LightningModule的 setup 方法，如下所示：

class LitMNIST(LightningModule):

    def __init__(self):
        self.l1 = None

    def prepare_data(self):
        download_data()
        tokenize()

    def setup(self, stage: Optional[str] = None):
        # step is either 'fit', 'validate', 'test', or 'predict'. 90% of the time not relevant
        data = load_data()
        num_classes = data.classes
        self.l1 = nn.Linear(..., num_classes)

Optimizer

接下来，我们选择使用什么优化器来训练我们的系统。在PyTorch中，我们这样做：

from torch.optim import Adam
optimizer = Adam(LitMNIST().parameters(), lr=1e-3)

在Lightning中，我们也这样做，但在configure optimizers()方法下进行组织。

class LitMNIST(LightningModule):

    def configure_optimizers(self):
        return Adam(self.parameters(), lr=1e-3)

Note: LightningModule本身有参数，所以传入self.parameters()

但是，如果有多个优化器，请使用对应的参数

class LitMNIST(LightningModule):

    def configure_optimizers(self):
        return Adam(self.generator(), lr=1e-3), Adam(self.discriminator(), lr=1e-3)

Training step

training step是在训练循环中发生的事情。

for epoch in epochs:
    for batch in data:
        # TRAINING STEP
        # ....
        # TRAINING STEP
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

在 MNIST 中，我们这样做：

for epoch in epochs:
    for batch in data:
        # ------ TRAINING STEP START ------
        x, y = batch
        logits = model(x)
        loss = F.nll_loss(logits, y)
        # ------ TRAINING STEP END ------

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

在Lightning中，训练步骤中的所有内容都在 LightningModule 中的 training_step() 函数下进行组织。

class LitMNIST(LightningModule):

    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.nll_loss(logits, y)
        return loss

同样，这是相同的PyTorch代码，只是它是由LightningModule组织的。这个代码没有限制，这意味着它可以像一个完整的 seq-2-seq, RL loop, GAN 等… 一样复杂。

The engineering

Training

到目前为止，我们在pure PyTorch中定义了4个关键成分，但使用LightningModule组织了代码。

Model.
Training data.
Optimizer.
What happens in the training loop.

为了清晰起见，我们回忆一下完整的LightningModule现在看起来是这样的。

class LitMNIST(LightningModule):
    def __init__(self):
        super().__init__()
        self.layer_1 = nn.Linear(28 * 28, 128)
        self.layer_2 = nn.Linear(128, 256)
        self.layer_3 = nn.Linear(256, 10)

    def forward(self, x):
        batch_size, channels, width, height = x.size()
        x = x.view(batch_size, -1)
        x = self.layer_1(x)
        x = F.relu(x)
        x = self.layer_2(x)
        x = F.relu(x)
        x = self.layer_3(x)
        x = F.log_softmax(x, dim=1)
        return x

    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.nll_loss(logits, y)
        return loss

同样，这是相同的PyTorch代码，只是它是由LightningModule组织的。

Logging

使用我们最喜欢的logger Tensorboard 记录，使用 log() 方法，它可以在 LightningModule 的任意方法中调用。

def training_step(self, batch, batch_idx):
    self.log('my_metric', x)

log() 有以下几个设置：

on_step(在训练的每个step记录metric)
on_epoch(在epoch结束时自动累积并记录日志)
prog_bar(记录进度条)
logger(像Tensorboard一样的日志到记录器)

根据日志的调用位置，Lightning会自动为您确定正确的模式。当然，您可以通过手动设置 flags 来覆盖默认行为。

Note：设置为epoch=True将在整个培训时间内累积您的日志值。

def training_step(self, batch, batch_idx):
    self.log('my_loss', loss, on_step=True, on_epoch=True, prog_bar=True, logger=True)

您还可以直接使用logger的任何方法。

def training_step(self, batch, batch_idx):
    tensorboard = self.logger.experiment
    tensorboard.any_summary_writer_method_you_want())

一旦你的训练开始，你可以通过使用你最喜欢的 logger 或启动 Tensorboard 日志查看日志

tensorboard --logdir ./lightning_logs

这会自动生成 tensorboard 日志 (或者你选择的logger)

但您也可以使用我们支持的任何其他记录器。

Train on CPU

from pytorch_lightning import Trainer

model = LitMNIST()
trainer = Trainer()
trainer.fit(model, train_loader)

您应该看到以下 weights summary 和 progress bar：

Train on GPU

但美妙之处在于你用trainer flag 所能做到的所有神奇之处。例如，在GPU上运行这个模型

model = LitMNIST()
trainer = Trainer(gpus=1)
trainer.fit(model, train_loader)

Train on Multi-GPU

或者你也可以在多个gpu上训练。

model = LitMNIST()
trainer = Trainer(gpus=8)
trainer.fit(model, train_loader)

或多个节点

# (32 GPUs)
model = LitMNIST()
trainer = Trainer(gpus=8, num_nodes=4, accelerator='ddp')
trainer.fit(model, train_loader)

详情参见：分布式计算手册

Hyperparameters

Lightning拥有与命令行ArgumentParser无缝交互的程序，并与您选择的超参数优化框架配合得很好。

ArgumentParser

Lightning的设计是为了增强Python内置ArgumentParser的很多功能

from argparse import ArgumentParser
parser = ArgumentParser()
parser.add_argument('--layer_1_dim', type=int, default=128)
args = parser.parse_args()

这允许你像这样调用你的程序

python trainer.py --layer_1_dim 64

Argparser Best Practices

最好的做法是把你的 arguments 分为三个部分。

Trainer args (gpus, num_nodes, etc…)
Model specific arguments (layer_dim, num_layers, learning_rate, etc…)
Program arguments (data_path, cluster_email, etc…)

我们可以这样做。首先，在LightningModule中，定义特定于该模块的参数。请记住，数据分割或数据路径也可能特定于某个模块(例如:如果你的项目有一个在Imagenet上训练的模型，另一个在CIFAR-10上训练)

class LitModel(LightningModule):

    @staticmethod
    def add_model_specific_args(parent_parser):
        parser = parent_parser.add_argument_group("LitModel")
        parser.add_argument('--encoder_layers', type=int, default=12)
        parser.add_argument('--data_path', type=str, default='/some/path')
        return parent_parser

现在在你的 main trainer 文件, 添加 Trainer 参数， program 参数，和 model 参数

# ----------------
# trainer_main.py
# ----------------
from argparse import ArgumentParser
parser = ArgumentParser()

# add PROGRAM level args
parser.add_argument('--conda_env', type=str, default='some_name')
parser.add_argument('--notification_email', type=str, default='will@email.com')

# add model specific args
parser = LitModel.add_model_specific_args(parser)

# add all the available trainer options to argparse
# ie: now --gpus --num_nodes ... --fast_dev_run all work in the cli
parser = Trainer.add_argparse_args(parser)

args = parser.parse_args()

现在你可以调用你的程序像这样运行:

python trainer_main.py --gpus 2 --num_nodes 2 --conda_env 'my_env' --encoder_layers 12

最后，确保像这样开始训练：

# init the trainer like this
trainer = Trainer.from_argparse_args(args, early_stopping_callback=...)

# NOT like this
trainer = Trainer(gpus=hparams.gpus, ...)

# init the model with Namespace directly
model = LitModel(args)

# or init the model with all the key-value pairs
dict_args = vars(args)
model = LitModel(**dict_args)

LightningModule hyperparameters

我们经常训练一个模型的多个版本。您可能会分享这个模型，或者几个月后再回到这个模型，这时了解这个模型是如何训练的就非常有用了(即:什么学习速率，神经网络等)。

Lightning有几种方法可以将这些信息保存在检查点和yaml文件中。这里的目标是提高可读性和再现性。

1> 第一种方法是要求lightning为您将 __init__ 中的所有值保存到 checkpoint。这也使得这些值可以通过self.hparams获得。

class LitMNIST(LightningModule):

    def __init__(self, layer_1_dim=128, learning_rate=1e-2, **kwargs):
        super().__init__()
        # call this to save (layer_1_dim=128, learning_rate=1e-4) to the checkpoint
        self.save_hyperparameters()

        # equivalent
        self.save_hyperparameters('layer_1_dim', 'learning_rate')

        # Now possible to access layer_1_dim from hparams
        self.hparams.layer_1_dim

2> 有时候你的init可能有你不想保存的对象或其他参数。那样的话，就选几个吧

class LitMNIST(LightningModule):

    def __init__(self, loss_fx, generator_network, layer_1_dim=128 **kwargs):
        super().__init__()
        self.layer_1_dim = layer_1_dim
        self.loss_fx = loss_fx

        # call this to save (layer_1_dim=128) to the checkpoint
        self.save_hyperparameters('layer_1_dim')

# to load specify the other args
model = LitMNIST.load_from_checkpoint(PATH, loss_fx=torch.nn.SomeOtherLoss, generator_network=MyGenerator())

3> 赋值给 self.hparams. 所有赋值给 self.hparams 将会自动保存

# using a argparse.Namespace
class LitMNIST(LightningModule):
    def __init__(self, hparams, *args, **kwargs):
        super().__init__()
        self.hparams = hparams
        self.layer_1 = nn.Linear(28 * 28, self.hparams.layer_1_dim)
        self.layer_2 = nn.Linear(self.hparams.layer_1_dim, self.hparams.layer_2_dim)
        self.layer_3 = nn.Linear(self.hparams.layer_2_dim, 10)
    def train_dataloader(self):
        return DataLoader(mnist_train, batch_size=self.hparams.batch_size)

4> 还可以将dict或Namespace等完整对象保存到 checkpoint。

# using a argparse.Namespace
class LitMNIST(LightningModule):

    def __init__(self, conf, *args, **kwargs):
        super().__init__()
        self.save_hyperparameters(conf)

        self.layer_1 = nn.Linear(28 * 28, self.hparams.layer_1_dim)
        self.layer_2 = nn.Linear(self.hparams.layer_1_dim, self.hparams.layer_2_dim)
        self.layer_3 = nn.Linear(self.hparams.layer_2_dim, 10)

conf = OmegaConf.create(...)
model = LitMNIST(conf)

# Now possible to access any stored variables from hparams
model.hparams.anything

Trainer args

回顾一下，将所有可能的训练器标记添加到argparser中，并以这种方式初始化 Trainer

parser = ArgumentParser()
parser = Trainer.add_argparse_args(parser)
hparams = parser.parse_args()

trainer = Trainer.from_argparse_args(hparams)

# or if you need to pass in callbacks
trainer = Trainer.from_argparse_args(hparams, checkpoint_callback=..., callbacks=[...])

Multiple Lightning Modules

我们经常有多个 Lightning 模块，每个模块都有不同的参数。LightningModule让您为每个定义参数，而不是污染main.py文件。

class LitMNIST(LightningModule):

    def __init__(self, layer_1_dim, **kwargs):
        super().__init__()
        self.layer_1 = nn.Linear(28 * 28, layer_1_dim)

    @staticmethod
    def add_model_specific_args(parent_parser):
        parser = parent_parser.add_argument_group("LitMNIST")
        parser.add_argument('--layer_1_dim', type=int, default=128)
        return parent_parser

class GoodGAN(LightningModule):

    def __init__(self, encoder_layers, **kwargs):
        super().__init__()
        self.encoder = Encoder(layers=encoder_layers)

    @staticmethod
    def add_model_specific_args(parent_parser):
        parser = parent_parser.add_argument_group("GoodGAN")
        parser.add_argument('--encoder_layers', type=int, default=12)
        return parent_parser

现在，我们在 main.py 里可以允许每个模型注入它需要的参数

def main(args):
    dict_args = vars(args)

    # pick model
    if args.model_name == 'gan':
        model = GoodGAN(**dict_args)
    elif args.model_name == 'mnist':
        model = LitMNIST(**dict_args)

    trainer = Trainer.from_argparse_args(args)
    trainer.fit(model)

if __name__ == '__main__':
    parser = ArgumentParser()
    parser = Trainer.add_argparse_args(parser)

    # figure out which model to use
    parser.add_argument('--model_name', type=str, default='gan', help='gan or mnist')

    # THIS LINE IS KEY TO PULL THE MODEL NAME
    temp_args, _ = parser.parse_known_args()

    # let the model add what it wants
    if temp_args.model_name == 'gan':
        parser = GoodGAN.add_model_specific_args(parser)
    elif temp_args.model_name == 'mnist':
        parser = LitMNIST.add_model_specific_args(parser)

    args = parser.parse_args()

    # train
    main(args)

现在我们可以使用命令行接口训练MNIST或GAN

$ python main.py --model_name gan --encoder_layers 24
$ python main.py --model_name mnist --layer_1_dim 128

Validating

在大多数情况下，当数据验证分割的性能达到最小值时，我们停止训练模型。

就像 training_step, 我们可以定义一个 validation_step 来检查我们关心的任何指标，生成样本，或者向我们的日志中添加更多的内容。

def validation_step(self, batch, batch_idx):
    loss = MSE_loss(...)
    self.log('val_loss', loss)

现在我们也可以带验证循环训练。

from pytorch_lightning import Trainer

model = LitMNIST()
trainer = Trainer(tpu_cores=8)
trainer.fit(model, train_loader, val_loader)

你可能已经注意到 Validation sanity check 被记录。这是因为 Lightning 在开始训练之前运行两个batch的 validation。这是一种单元测试，以确保如果在验证循环中出现错误，您不需要等待完整的epoch来发现。

Note: Lightning 禁用梯度，将模型置成eval模式，并且做任何用于validation的事情。

Val loop under the hood

在 hood 下， Lightning 作了如下事情：

model = Model()
model.train()
torch.set_grad_enabled(True)

for epoch in epochs:
    for batch in data:
        # ...
        # train

    # validate
    model.eval()
    torch.set_grad_enabled(False)

    outputs = []
    for batch in val_data:
        x, y = batch                        # validation_step
        y_hat = model(x)                    # validation_step
        loss = loss(y_hat, x)               # validation_step
        outputs.append({'val_loss': loss})  # validation_step

    total_loss = outputs.mean()             # validation_epoch_end

Optional methods

如果您仍然需要精细的控制，那么为循环定义其他的方法。

def validation_step(self, batch, batch_idx):
    preds = ...
    return preds

def validation_epoch_end(self, val_step_outputs):
    for pred in val_step_outputs:
        # do something with all the predictions from each validation_step

Testing

一旦我们的研究完成，我们即将发布或部署一个模型，我们通常想弄清楚它如何在现实世界中推广。为此，我们使用了一个数据分割来进行测试。

就像 validation loop 一样，我们定义了一个 test loop

class LitMNIST(LightningModule):
    def test_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.nll_loss(logits, y)
        self.log('test_loss', loss)

为了确保测试集不会被无意中使用，Lightning有一个单独的API来运行测试。当你在训练你的模型时调用 .test() 。

from pytorch_lightning import Trainer

model = LitMNIST()
trainer = Trainer(tpu_cores=8)
trainer.fit(model)

# run test set
result = trainer.test()
print(result)

输出：

--------------------------------------------------------------
TEST RESULTS
{'test_loss': 1.1703}
--------------------------------------------------------------

您还可以从保存的lightning模型运行测试

model = LitMNIST.load_from_checkpoint(PATH)
trainer = Trainer(tpu_cores=8)
trainer.test(model)

Note: Lightning 禁用梯度，将模型置于eval模式，并且做任何用于训练的事情。

Warning： .test() 在 TPUs 上还不稳定。

Predicting

同样，LightningModule与PyTorch模块完全相同。这意味着您可以加载它并使用它进行预测。

model = LitMNIST.load_from_checkpoint(PATH)
x = torch.randn(1, 1, 28, 28)
out = model(x)

从表面上看，forward 和 training_step 是相似的。一般来说，我们要确保我们想让模型做的是在forward 中发生的事。而 training_step 则从内部调用 forward。

class MNISTClassifier(LightningModule):

    def forward(self, x):
        batch_size, channels, width, height = x.size()
        x = x.view(batch_size, -1)
        x = self.layer_1(x)
        x = F.relu(x)
        x = self.layer_2(x)
        x = F.relu(x)
        x = self.layer_3(x)
        x = F.log_softmax(x, dim=1)
        return x

    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.nll_loss(logits, y)
        return loss

model = MNISTClassifier()
x = mnist_image()
logits = model(x)

在这种情况下，我们设置了这个 LightningModel 来预测logits。但我们也可以用它来预测特征图：

class MNISTRepresentator(LightningModule):

    def forward(self, x):
        batch_size, channels, width, height = x.size()
        x = x.view(batch_size, -1)
        x = self.layer_1(x)
        x1 = F.relu(x)
        x = self.layer_2(x1)
        x2 = F.relu(x)
        x3 = self.layer_3(x2)
        return [x, x1, x2, x3]

    def training_step(self, batch, batch_idx):
        x, y = batch
        out, l1_feats, l2_feats, l3_feats = self(x)
        logits = F.log_softmax(out, dim=1)
        ce_loss = F.nll_loss(logits, y)
        loss = perceptual_loss(l1_feats, l2_feats, l3_feats) + ce_loss
        return loss

model = MNISTRepresentator.load_from_checkpoint(PATH)
x = mnist_image()
feature_maps = model(x)

或者我们有一个模型，我们用它来做生成

class LitMNISTDreamer(LightningModule):

    def forward(self, z):
        imgs = self.decoder(z)
        return imgs

    def training_step(self, batch, batch_idx):
        x, y = batch
        representation = self.encoder(x)
        imgs = self(representation)

        loss = perceptual_loss(imgs, x)
        return loss

model = LitMNISTDreamer.load_from_checkpoint(PATH)
z = sample_noise()
generated_imgs = model(z)

为了进行大规模推理，可以使用 trainer.predict，它使用 LightningModule 中的 predict_step 函数。默认情况下，LightningModule predict_step 调用 forward，但它可以被重写以添加任何处理逻辑。

class LitMNISTDreamer(LightningModule):

    def forward(self, z):
        imgs = self.decoder(z)
        return imgs

    def predict_step(self, batch, batch_idx: int , dataloader_idx: int = None):
        return self(batch)


model = LitMNISTDreamer()
trainer.predict(model, datamodule)

你如何划分 forward ，training_step 和 predict 取决于你想如何使用这个模型进行预测。但是，我们建议 forward 仅包含你的模型的tensor操作， training_step 使用 logging, metrics 和 loss计算封装 forward 逻辑，并且 predict 使用 preprocess， postprocess 函数封装 forward。

The nonessentials

Extensibility

虽然 lightning 让一切都变得超级简单，但它并不牺牲任何灵活性或控制。 lightning 提供多种管理训练状态的方法。

Training overrides

training、validation和testing循环的任何部分都可以修改。例如，如果您想执行自己的反向传播，则可以重写默认实现。

def backward(self, use_amp, loss, optimizer):
    loss.backward()

使用你自己的

class LitMNIST(LightningModule):

    def backward(self, use_amp, loss, optimizer, optimizer_idx):
        # do a custom way of backward
        loss.backward(retain_graph=True)

training的每个部分都可以通过这种方式进行配置。参见 LightningModule 的完整列表。

Callbacks

添加任意功能的另一种方法是为您可能关心的 hooks 添加自定义回调：

from pytorch_lightning.callbacks import Callback

class MyPrintingCallback(Callback):

    def on_init_start(self, trainer):
        print('Starting to init trainer!')

    def on_init_end(self, trainer):
        print('Trainer is init now')

    def on_train_end(self, trainer, pl_module):
        print('do something when training ends')

然后将 callbacks 传给 trainer

trainer = Trainer(callbacks=[MyPrintingCallback()])

参见 callbacks 的 12+ hooks 的完整列表。

Child Modules

研究项目倾向于对同一数据集测试不同的方法。这在Lightning中使用继承是很容易做到的。

例如，假设我们现在想训练一个自动编码器来作为MNIST图像的特征提取器。我们正在从 LitMNIST-module 扩展我们的自动编码器，该模块已经定义了所有的数据加载。在自动编码器模型中唯一改变的是init、forward、training、validation和test步骤。

class Encoder(torch.nn.Module):
    pass

class Decoder(torch.nn.Module):
    pass

class AutoEncoder(LitMNIST):

    def __init__(self):
        super().__init__()
        self.encoder = Encoder()
        self.decoder = Decoder()
        self.metric = MSE()

    def forward(self, x):
        return self.encoder(x)

    def training_step(self, batch, batch_idx):
        x, _ = batch

        representation = self.encoder(x)
        x_hat = self.decoder(representation)

        loss = self.metric(x, x_hat)
        return loss

    def validation_step(self, batch, batch_idx):
        self._shared_eval(batch, batch_idx, 'val')

    def test_step(self, batch, batch_idx):
        self._shared_eval(batch, batch_idx, 'test')

    def _shared_eval(self, batch, batch_idx, prefix):
        x, _ = batch
        representation = self.encoder(x)
        x_hat = self.decoder(representation)

        loss = self.metric(x, x_hat)
        self.log(f'{prefix}_loss', loss)

我们可以用同样的训练器来训练它

autoencoder = AutoEncoder()
trainer = Trainer()
trainer.fit(autoencoder)

请记住，forward方法应该定义LightningModule的实际使用。在这种情况下，我们希望使用自动编码器来提取图像表示

some_images = torch.Tensor(32, 1, 28, 28)
representations = autoencoder(some_images)

Transfer Learning

Using Pretrained Models

有时我们想使用一个LightningModule作为一个预训练的模型。这是很容易实现的，因为LightningModule只是一个torch.nn.Module。

Note：记住，LightningModule就是一个 torch.nn.Module，但有更多的功能。

让我们在一个单独的模型中使用自动编码器作为特征提取器。

class Encoder(torch.nn.Module):
    ...

class AutoEncoder(LightningModule):
    def __init__(self):
        self.encoder = Encoder()
        self.decoder = Decoder()

class CIFAR10Classifier(LightningModule):
    def __init__(self):
        # init the pretrained LightningModule
        self.feature_extractor = AutoEncoder.load_from_checkpoint(PATH)
        self.feature_extractor.freeze()

        # the autoencoder outputs a 100-dim representation and CIFAR-10 has 10 classes
        self.classifier = nn.Linear(100, 10)

    def forward(self, x):
        representations = self.feature_extractor(x)
        x = self.classifier(representations)
        ...

我们使用预训练好的自动编码器(LightningModule)进行迁移学习

Example: Imagenet (computer Vision)

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)
        ...

微调：

model = ImagenetTransferLearning()
trainer = Trainer()
trainer.fit(model)

然后用它来预测你感兴趣的数据

model = ImagenetTransferLearning.load_from_checkpoint(PATH)
model.freeze()

x = some_images_from_cifar10()
predictions = model(x)

我们在imagenet上使用预训练的模型，并在CIFAR-10上进行微调，以在CIFAR-10上进行预测。在非学术领域，我们会根据你的微小数据集进行微调，并根据你的数据集进行预测。

Example: BERT (NLP)

Lightning 是完全不关心用于迁移学习的是什么，只要它是一个 torch.nn.Module 子类。

这里是一个使用 Huggingface transformers 的模型。

class BertMNLIFinetuner(LightningModule):

    def __init__(self):
        super().__init__()

        self.bert = BertModel.from_pretrained('bert-base-cased', output_attentions=True)
        self.W = nn.Linear(bert.config.hidden_size, 3)
        self.num_classes = 3


    def forward(self, input_ids, attention_mask, token_type_ids):

        h, _, attn = self.bert(input_ids=input_ids,
                         attention_mask=attention_mask,
                         token_type_ids=token_type_ids)

        h_cls = h[:, 0]
        logits = self.W(h_cls)
        return logits, attn

Common Use Cases

16-BIT TRAINING

GPU 16-bit

16位精度可以将您的内存占用空间减少一半。如果使用volta架构GPU，它也可以显着提高培训速度。

Note：建议将PyTorch 1.6+用于16位

Native torch

使用PyTorch 1.6+时，Lightning使用native amp实现来支持16位。

# turn on 16-bit
trainer = Trainer(precision=16, gpus=1)

Apex 16-bit

如果您使用的是PyTorch的早期版本，则Lightning将使用Apex支持16位。

请按照以下说明安装Apex。要使用16位精度，请执行以下两项操作：

Install Apex
Set the “precision” trainer flag.

$ git clone https://github.com/NVIDIA/apex
$ cd apex

# ------------------------
# OPTIONAL: on your cluster you might need to load CUDA 10 or 9
# depending on how you installed PyTorch

# see available modules
module avail

# load correct CUDA before install
module load cuda-10.0
# ------------------------

# make sure you've loaded a cuda version > 4.0 and < 7.0
module load gcc-6.1.0

$ pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./

Warning： NVIDIA Apex和DDP存在不稳定问题。我们建议在PyTorch 1.6+中使用native 16位

Enable 16-bit

# turn on 16-bit
trainer = Trainer(amp_level='O2', precision=16)

如果您需要为您的特定用例配置apex init，或者希望使用不同的方式进行16位训练，覆盖 pytorch_lightning.core.LightningModule.configure_apex()。

Early stopping

Stopping an epoch early

您可以通过重写 on_train_batch_start() 以在满足某些条件时返回-1来提前停止一个epoch。

如果您重复执行此操作，则对于您最初请求的每个epoch，都将停止整个运行。

Early stopping based on metric using the EarlyStopping Callback

EarlyStopping callback 可用于监视验证指标，并在未观察到任何改进时停止训练。

要使用它：

Import EarlyStopping callback.
Log the metric you want to monitor using log() method.
Init the callback, and set monitor to the logged metric of your choice.
Pass the EarlyStopping callback to the Trainer callbacks flag.

from pytorch_lightning.callbacks.early_stopping import EarlyStopping

def validation_step(...):
    self.log('val_loss', loss)

trainer = Trainer(callbacks=[EarlyStopping(monitor='val_loss')])

你可以通过改变它的参数来自定义回调行为。

early_stop_callback = EarlyStopping(
   monitor='val_accuracy',
   min_delta=0.00,
   patience=3,
   verbose=False,
   mode='max'
)
trainer = Trainer(callbacks=[early_stop_callback])

如果您需要在训练的不同部分中early stopping，子类EarlyStopping并更改其调用位置：

class MyEarlyStopping(EarlyStopping):

    def on_validation_end(self, trainer, pl_module):
        # override this to disable early stopping at the end of val loop
        pass

    def on_train_end(self, trainer, pl_module):
        # instead, do it at the end of training loop
        self._run_early_stopping_check(trainer, pl_module)

Note： EarlyStopping callback 在每个验证epoch的末尾运行，在默认配置下，它在每个训练epoch之后发生。但是，可以通过在Trainer中设置各种参数（例如check_val_every_n_epoch和val_check_interval）来修改验证频率。必须注意的是，patience参数对没有改进的验证epochs数进行计数，而不是对训练epochs数进行计数。因此，在参数check_val_every_n_epoch = 10且 patience = 3的情况下，Trainer 将在停止之前执行至少40个训练 epochs。

Learning Rate Finder

对于训练深度神经网络，选择良好的学习率对于更好的性能和更快的收敛都是必不可少的。甚至可以自动调整学习速度的优化器（例如Adam）也可以从更多的最佳选择中受益。

为了减少与选择良好的初始学习率有关的猜测量，可以使用学习率查找器。如本文所述，学习速率查找器进行一个小的运行，在每次处理batch后，学习速率都会提高，并记录相应的loss。这样的结果是一个lr与loss的关系图，可用作选择最佳初始lr的指南。

Warning：目前，此功能仅适用于具有单个优化器的模型。 LR Finder对DDP的支持尚未实现，即将推出。

Using Lightning’s built-in LR finder

要启用学习速率查找器，您的 learning 模块需要具有learning_rate或lr属性。然后，在训练器构建过程中设置Trainer(auto_lr_find = True)，然后调用trainer.tune(model)运行LR查找器。建议的learning_rate将被写入控制台，并将自动设置为您的lightning模块，可以通过self.learning_rate或self.lr对其进行访问。

class LitModel(LightningModule):

    def __init__(self, learning_rate):
        self.learning_rate = learning_rate

    def configure_optimizers(self):
        return Adam(self.parameters(), lr=(self.lr or self.learning_rate))

model = LitModel()

# finds learning rate automatically
# sets hparams.lr or hparams.learning_rate to that learning rate
trainer = Trainer(auto_lr_find=True)

trainer.tune(model)

如果您的模型使用的是任意值而不是self.lr或self.learning_rate，则将该值设置为auto_lr_find：

model = LitModel()

# to set to your own hparams.my_value
trainer = Trainer(auto_lr_find='my_value')

trainer.tune(model)

如果要检查学习率查找器的结果，或者只是尝试使用算法的参数，则可以通过调用Trainer的lr_find方法来完成。一个典型的例子：

model = MyModelClass(hparams)
trainer = Trainer()

# Run learning rate finder
lr_finder = trainer.tuner.lr_find(model)

# Results can be found in
lr_finder.results

# Plot with
fig = lr_finder.plot(suggest=True)
fig.show()

# Pick point based on plot, or get suggestion
new_lr = lr_finder.suggestion()

# update hparams of the model
model.hparams.lr = new_lr

# Fit model
trainer.fit(model)

lr_finder.plot() 生成的图应类似于下图。建议不要选择获得最低损失的学习率，而应选择最陡峭的向下斜率(红点)中间的值。这是由 lr_finder.suggestion() 返回的点。

该算法的参数如下所示。

pytorch_lightning.tuner.lr_finder.lr_find(trainer, model, train_dataloader=None, val_dataloaders=None, min_lr=1e-08, max_lr=1, num_training=100, mode='exponential', early_stop_threshold=4.0, datamodule=None, update_attr=False)

lr_find 使用户能够进行良好初始学习率的范围测试，以减少在选择良好初始学习率时的猜测量。

参数：

model (LightningModule) – 测试范围的模型
train_dataloader (Optional[DataLoader]) – 一个有训练样本的 Pytorch DataLoader。如果模型具有预定义的train_dataloader方法，则将跳过此方法。
min_lr (float) – 调查的最小学习率
max_lr (float) – 调查的最大学习率
num_training (int) – 测试的学习率数量
mode (str) – 每个 batch 后更新学习率的搜索策略
- ‘exponential’ (default): 会指数地提高学习速度。
- ‘linear’: 将线性地提高学习速度。
early_stop_threshold (float) – 停止搜索的阈值。如果在任何时候的损失大于 early_stop_threshold * best_loss，则停止搜索。要禁用，请设置为“None”。
datamodule (Optional[LightningDataModule]) – 可选的LightningDataModule，用于保存训练和验证数据加载器。请注意，train_dataloader和val_dataloaders参数不能与此参数同时使用，否则将引发MisconfigurationException。
update_attr (bool) – 是否更新学习速率属性。

Raises：

MisconfigurationException – 如果auto_lr_find = True时没有覆盖model或model.hparams中的 learning rate/ lr，或者您使用了多个优化器。

例子：

# Setup model and trainer
model = MyModelClass(hparams)
trainer = pl.Trainer()

# Run lr finder
lr_finder = trainer.tuner.lr_find(model, ...)

# Inspect results
fig = lr_finder.plot(); fig.show()
suggested_lr = lr_finder.suggestion()

# Overwrite lr and create new model
hparams.lr = suggested_lr
model = MyModelClass(hparams)

# Ready to train with new learning rate
trainer.fit(model)

Saving and loading weights

Lightning 自动保存和加载 checkpoints。Checkpoints 捕获模型使用的所有参数。

Checkpointing 训练允许您在训练过程中断时恢复训练过程，对模型进行微调，或者使用预训练过的模型进行推理，而不必重新训练模型。

Checkpoint saving

Lightning checkpoint 拥有恢复 training session 所需的一切：

16-bit scaling factor (apex)
Current epoch
Global step
Model state_dict
State of all optimizers
State of all learningRate schedulers
State of all callbacks
The hyperparameters used for that model if passed in as hparams (Argparse.Namespace)

Automatic saving

Lightning会自动将 checkpoint 保存在您当前的工作目录中，并带有您的上一个训练epoch的状态。这样可以确保您可以在中断训练的情况下继续训练。

更改 checkpoint 路径：

# saves checkpoints to '/your/path/to/save/checkpoints' at every epoch end
trainer = Trainer(default_root_dir='/your/path/to/save/checkpoints')

您可以自定义 checkpointing 行为，以监视所有 training 或 validation steps 的质量。例如，如果您要基于验证损失来更新 checkpoints，请执行以下操作：

计算您希望监视的任何 metric 或其他 quantity ，例如 validation loss。
使用 log() 方法记录 quantity ，并带有val_loss之类的key。
初始化 ModelCheckpoint callback，并将Monitor设置为quantity的key。
将callback传递给回调Trainer标志。

from pytorch_lightning.callbacks import ModelCheckpoint

class LitAutoEncoder(LightningModule):
    def validation_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self.backbone(x)

        # 1. calculate loss
        loss = F.cross_entropy(y_hat, y)

        # 2. log `val_loss`
        self.log('val_loss', loss)

# 3. Init ModelCheckpoint callback, monitoring 'val_loss'
checkpoint_callback = ModelCheckpoint(monitor='val_loss')

# 4. Add your callback to the callbacks list
trainer = Trainer(callbacks=[checkpoint_callback])

您还可以控制更高级的选项，例如save_top_k，以保存最佳的k个模型，以及受监视 quantity 的模式（最小/最大值），save_weights_only或period来设置 checkpoints 之间的时间间隔，以免降低速度。

from pytorch_lightning.callbacks import ModelCheckpoint

class LitAutoEncoder(LightningModule):
    def validation_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self.backbone(x)
        loss = F.cross_entropy(y_hat, y)
        self.log('val_loss', loss)

# saves a file like: my/path/sample-mnist-epoch=02-val_loss=0.32.ckpt
checkpoint_callback = ModelCheckpoint(
    monitor='val_loss',
    dirpath='my/path/',
    filename='sample-mnist-{epoch:02d}-{val_loss:.2f}',
    save_top_k=3,
    mode='min',
)

trainer = Trainer(callbacks=[checkpoint_callback])

您可以在训练后通过调用来检索 checkpoint

checkpoint_callback = ModelCheckpoint(dirpath='my/path/')
trainer = Trainer(callbacks=[checkpoint_callback])
trainer.fit(model)
checkpoint_callback.best_model_path

Disabling checkpoints

您可以通过传递来禁用 checkpointing

trainer = Trainer(checkpoint_callback=False)

Lightning checkpoint 还保存传入 LightningModule 初始化的参数在 checkpoint 的 hyper_parameters 键下。

class MyLightningModule(LightningModule):

   def __init__(self, learning_rate, *args, **kwargs):
        super().__init__()
        self.save_hyperparameters()

# all init args were saved to the checkpoint
checkpoint = torch.load(CKPT_PATH)
print(checkpoint['hyper_parameters'])
# {'learning_rate': the_value}

Manual saving

您可以手动保存 checkpoints ，并从 checkpoints 状态恢复您的模型。

model = MyLightningModule(hparams)
trainer.fit(model)
trainer.save_checkpoint("example.ckpt")
new_model = MyModel.load_from_checkpoint(checkpoint_path="example.ckpt")

Manual saving with accelerators

Lightning还可以处理多个进程同时运行的加速器，比如DDP。例如，当使用DDP加速器时，我们的训练脚本同时在多个设备上运行。Lightning会自动确保模型只保存在主进程上，而其他进程不会干扰保存 checkpoints。如下所示，这并不需要更改代码。

trainer = Trainer(accelerator="ddp")
model = MyLightningModule(hparams)
trainer.fit(model)
# Saves only on the main process
trainer.save_checkpoint("example.ckpt")

不使用 trainer.save_checkpo 可能会导致意外行为和潜在的死锁。使用其他保存函数将导致所有设备试图保存 checkpoint。因此，我们强烈建议使用 trainer 的保存功能。如果无法避免使用自定义保存函数，我们建议使用rank_zero_only() 来确保只在主进程上进行保存。

Checkpoint loading

使用以下方法加载模型及其权重、偏差和超参数：

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)

但是，如果您不想使用保存在 checkpoint 中的值，则在这里传入您自己的值

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 this
LitModel(in_dim=32, out_dim=10)

# uses in_dim=32, out_dim=10
model = LitModel.load_from_checkpoint(PATH)

# uses in_dim=128, out_dim=10
model = LitModel.load_from_checkpoint(PATH, in_dim=128, out_dim=10)

LightningModule.load_from_checkpoint(checkpoint_path, map_location=None, hparams_file=None, strict=True, **kwargs)

从 checkpoint 加载模型的主要方法。当Lightning保存 checkpoint 时，它将传递给 __init__ 的参数存储在 checkpoint 的 hyper_parameters 下。

任何通过 *args 和 **kwargs指定的参数都将覆盖存储在 hyper_parameters 中的参数。

参数：

checkpoint_path (Union[str, IO]) – checkpoint 的路径。这也可以是一个URL，或类似文件的对象
map_location (Union[Dict[str, str], str, device, int, Callable, None]) – 如果你的 checkpoint 保存了一个GPU模型，你现在加载到cpu或不同数量的GPU，使用这个来映射到新的设置。它的行为和 torch.load() 一样。
hparams_file (Optional[str]) – 一个.yaml文件的可选路径，该文件具有本例中所示的层次结构。你很可能不需要这个，因为 lightning 总是会把 hyperparameters 保存到 checkpoint。但是，如果您的 checkpoint 权重没有保存hyperparameters，可以使用这个方法传入一个带有您想要使用的hparams的.yaml文件。这些将被转换成一个字典，并传递到您的LightningModule中使用。如果您的模型的hparams参数是名称空间，并且.yaml文件具有层次结构，那么您需要重构您的模型，将hparams视为dict。
```
1
2
3
drop_prob: 0.2
dataloader:
  batch_size: 32
```
strict (bool) – 是否严格强制 checkpoint_path 中的 keys 与 module 的 state dict 返回的key匹配。默认为True。
kwargs – 初始化模型所需的任何额外关键字参数。也可以用来覆盖保存的超参数值。

加载了权重和超参数的LightningModule(如果可用)。

例子：

# load weights without mapping ...
MyLightningModule.load_from_checkpoint('path/to/checkpoint.ckpt')

# or load weights mapping all weights from GPU 1 to GPU 0 ...
map_location = {'cuda:1':'cuda:0'}
MyLightningModule.load_from_checkpoint(
    'path/to/checkpoint.ckpt',
    map_location=map_location
)

# or load weights and hyperparameters from separate files.
MyLightningModule.load_from_checkpoint(
    'path/to/checkpoint.ckpt',
    hparams_file='/path/to/hparams_file.yaml'
)

# override some of the params with new values
MyLightningModule.load_from_checkpoint(
    PATH,
    num_layers=128,
    pretrained_ckpt_path: NEW_PATH,
)

# predict
pretrained_model.eval()
pretrained_model.freeze()
y_hat = pretrained_model(x)

Restoring Training State

如果你不只是想进行加载权重训练，而是要恢复全部训练，那么就做下面的步骤

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)

EMA

from copy import deepcopy
from typing import Optional, Union, Dict, Any

import pytorch_lightning as pl
import torch
from overrides import overrides
from pytorch_lightning.utilities import rank_zero_only


class EMA(pl.Callback):
    """Implements EMA (exponential moving average) to any kind of model.
    EMA weights will be used during validation and stored separately from original model weights.

    How to use EMA:
        - Sometimes, last EMA checkpoint isn't the best as EMA weights metrics can show long oscillations in time. See
          https://github.com/rwightman/pytorch-image-models/issues/102
        - Batch Norm layers and likely any other type of norm layers doesn't need to be updated at the end. See
          discussions in: https://github.com/rwightman/pytorch-image-models/issues/106#issuecomment-609461088 and
          https://github.com/rwightman/pytorch-image-models/issues/224
        - For object detection, SWA usually works better. See   https://github.com/timgaripov/swa/issues/16

    Implementation detail:
        - See EMA in Pytorch Lightning: https://github.com/PyTorchLightning/pytorch-lightning/issues/10914
        - When multi gpu, we broadcast ema weights and the original weights in order to only hold 1 copy in memory.
          This is specially relevant when storing EMA weights on CPU + pinned memory as pinned memory is a limited
          resource. In addition, we want to avoid duplicated operations in ranks != 0 to reduce jitter and improve
          performance.
    """
    def __init__(self, decay: float = 0.9999, ema_device: Optional[Union[torch.device, str]] = None, pin_memory=True):
        super().__init__()
        self.decay = decay
        self.ema_device: str = f"{ema_device}" if ema_device else None  # perform ema on different device from the model
        self.ema_pin_memory = pin_memory if torch.cuda.is_available() else False  # Only works if CUDA is available
        self.ema_state_dict: Dict[str, torch.Tensor] = {}
        self.original_state_dict = {}
        self._ema_state_dict_ready = False

    @staticmethod
    def get_state_dict(pl_module: pl.LightningModule):
        """Returns state dictionary from pl_module. Override if you want filter some parameters and/or buffers out.
        For example, in pl_module has metrics, you don't want to return their parameters.
        
        code:
            # Only consider modules that can be seen by optimizers. Lightning modules can have others nn.Module attached
            # like losses, metrics, etc.
            patterns_to_ignore = ("metrics1", "metrics2")
            return dict(filter(lambda i: i[0].startswith(patterns), pl_module.state_dict().items()))
        """
        return pl_module.state_dict()
        
    @overrides
    def on_train_start(self, trainer: "pl.Trainer", pl_module: pl.LightningModule) -> None:
        # Only keep track of EMA weights in rank zero.
        if not self._ema_state_dict_ready and pl_module.global_rank == 0:
            self.ema_state_dict = deepcopy(self.get_state_dict(pl_module))
            if self.ema_device:
                self.ema_state_dict = {k: tensor.to(device=self.ema_device) for k, tensor in self.ema_state_dict.items()}

            if self.ema_device == "cpu" and self.ema_pin_memory:
                self.ema_state_dict = {k: tensor.pin_memory() for k, tensor in self.ema_state_dict.items()}

        self._ema_state_dict_ready = True

    @rank_zero_only
    def on_train_batch_end(self, trainer: "pl.Trainer", pl_module: pl.LightningModule, *args, **kwargs) -> None:
        # Update EMA weights
        with torch.no_grad():
            for key, value in self.get_state_dict(pl_module).items():
                ema_value = self.ema_state_dict[key]
                ema_value.copy_(self.decay * ema_value + (1. - self.decay) * value, non_blocking=True)

    @overrides
    def on_validation_start(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None:
        if not self._ema_state_dict_ready:
            return  # Skip Lightning sanity validation check if no ema weights has been loaded from a checkpoint.

        self.original_state_dict = deepcopy(self.get_state_dict(pl_module))
        ema_state_dict = pl_module.trainer.training_type_plugin.broadcast(self.ema_state_dict, 0)
        self.ema_state_dict = ema_state_dict
        assert self.ema_state_dict.keys() == self.original_state_dict.keys(), \
            f"There are some keys missing in the ema static dictionary broadcasted. " \
            f"They are: {self.original_state_dict.keys() - self.ema_state_dict.keys()}"
        pl_module.load_state_dict(self.ema_state_dict, strict=False)

        if pl_module.global_rank > 0:
            # Remove ema state dict from the memory. In rank 0, it could be in ram pinned memory.
            self.ema_state_dict = {}

    @overrides
    def on_validation_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
        if not self._ema_state_dict_ready:
            return  # Skip Lightning sanity validation check if no ema weights has been loaded from a checkpoint.

        # Replace EMA weights with training weights
        pl_module.load_state_dict(self.original_state_dict, strict=False)

    @overrides
    def on_save_checkpoint(
        self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", checkpoint: Dict[str, Any]
    ) -> dict:
        return {"ema_state_dict": self.ema_state_dict, "_ema_state_dict_ready": self._ema_state_dict_ready}

    @overrides
    def on_load_checkpoint(
        self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", callback_state: Dict[str, Any]
    ) -> None:
        self._ema_state_dict_ready = callback_state["_ema_state_dict_ready"]
        self.ema_state_dict = callback_state["ema_state_dict"]

Manual optimize

先在 ightningModule 的 __init__ 中设置 self.automatic_optimization=False

手动调用以下函数：

self.optimizers() 得到优化器
optimizer.zero_grad() 清楚之前训练的梯度
self.manual_backward(loss) 代替 loss.backward()
optimizer.step() 更新模型参数

下面是一个简单的例子：

from pytorch_lightning import LightningModule


class MyModel(LightningModule):
    def __init__(self):
        super().__init__()
        # Important: This property activates manual optimization.
        self.automatic_optimization = False

    def training_step(self, batch, batch_idx):
        opt = self.optimizers()
        opt.zero_grad()
        loss = self.compute_loss(batch)
        self.manual_backward(loss)
        opt.step()

WeightedRandomSampler

train_label = self.label[self.label['fold']!=self.hparams.fold]
        train_label["category"] = train_label["category_id"].map(lambda x: category_id_to_lv2id(f"{x:0>4}"))
        labels_unique, counts = np.unique(train_label['category'], return_counts=True)
        class_weights = [sum(counts) / c for c in counts]
        example_weights = [class_weights[e] for e in train_label["category"]]
        sampler = WeightedRandomSampler(example_weights, len(train_label['category']))
        sampler = DistributedSamplerWrapper(sampler,
                                            shuffle=True,
                                            num_replicas=get_world_size(),
                                            rank=get_rank())

DistributedSamplerWrapper

from typing import Optional
from torch.utils.data import DistributedSampler, Dataset, Sampler
from operator import itemgetter

class DatasetFromSampler(Dataset):
    """Dataset to create indexes from `Sampler`.
    Args:
        sampler: PyTorch sampler
    """

    def __init__(self, sampler: Sampler):
        """Initialisation for DatasetFromSampler."""
        self.sampler = sampler
        self.sampler_list = None

    def __getitem__(self, index: int):
        """Gets element of the dataset.
        Args:
            index: index of the element in the dataset
        Returns:
            Single element by index
        """
        if self.sampler_list is None:
            self.sampler_list = list(self.sampler)
        return self.sampler_list[index]

    def __len__(self) -> int:
        """
        Returns:
            int: length of the dataset
        """
        return len(self.sampler)

class DistributedSamplerWrapper(DistributedSampler):
    """
    Wrapper over `Sampler` for distributed training.
    Allows you to use any sampler in distributed mode.
    It is especially useful in conjunction with
    `torch.nn.parallel.DistributedDataParallel`. In such case, each
    process can pass a DistributedSamplerWrapper instance as a DataLoader
    sampler, and load a subset of subsampled data of the original dataset
    that is exclusive to it.
    .. note::
        Sampler is assumed to be of constant size.
    """

    def __init__(
        self,
        sampler,
        num_replicas: Optional[int] = None,
        rank: Optional[int] = None,
        shuffle: bool = True,
    ):
        """
        Args:
            sampler: Sampler used for subsampling
            num_replicas (int, optional): Number of processes participating in
              distributed training
            rank (int, optional): Rank of the current process
              within ``num_replicas``
            shuffle (bool, optional): If true (default),
              sampler will shuffle the indices
        """
        super(DistributedSamplerWrapper, self).__init__(
            DatasetFromSampler(sampler),
            num_replicas=num_replicas,
            rank=rank,
            shuffle=shuffle,
        )
        self.sampler = sampler

    def __iter__(self):
        """@TODO: Docs. Contribution is welcome."""
        self.dataset = DatasetFromSampler(self.sampler)
        indexes_of_indexes = super().__iter__()
        subsampler_indexes = self.dataset
        return iter(itemgetter(*indexes_of_indexes)(subsampler_indexes))

communication

# Taken from Detectron 2, licensed under Apache 2.0.
# https://github.com/facebookresearch/detectron2/blob/989f52d67d05445ccd030d8f13d6cc53e297fb91/detectron2/utils/comm.py
# Changes:
# - removed half of the functions
# - init_distributed()
# - `reduce` and `all_reduce` for various datatypes
# - `shared_random_seed` uses `torch.distributed.broadcast` instead of `all_gather` from Detectron2.

import os
import functools

import torch
import numpy as np





def get_rank() -> int:
    """
    Get rank of the process, even when torch.distributed is not initialized.
    Returns
    -------
    int
    """
    if not torch.distributed.is_available():
        return 0
    if not torch.distributed.is_initialized():
        return 0

    return torch.distributed.get_rank()


def get_local_rank() -> int:
    """
    Get rank of the process, even when torch.distributed is not initialized.
    Returns
    -------
    int
    """
    if not torch.distributed.is_available():
        return 0
    if not torch.distributed.is_initialized():
        return 0

    return int(os.environ['LOCAL_RANK'])


def get_world_size() -> int:
    """
    Get number of compute device in the world, returns 1 in case multi device is not initialized.
    Returns
    -------
    int
    """
    if not torch.distributed.is_available():
        return 1
    if not torch.distributed.is_initialized():
        return 1
    return torch.distributed.get_world_size()


def get_backend_compatible_device():
    """Use device that is compatible with the backend (e.g. nccl does not support CPU tensors).
    Used when initializing placeholder tensors before performing communication operations.
    """
    return torch.device('cuda' if torch.distributed.get_backend() == 'nccl' else 'cpu')


def shared_random_seed() -> int:
    """
    All workers must call this function, otherwise it will deadlock.
    Returns
    -------
    A random number that is the same across all workers. If workers need a shared RNG, 
    they can use this shared seed to create one.
    """
    # torch.Generator advises to use a high values as seed, hence 2**31
    # The seed is reproducible when torch seed is set with `torch.manual_seed()`
    seed = torch.randint(2**31, (1,))
    if torch.distributed.is_initialized():
        device = get_backend_compatible_device()
        seed = seed.to(device)
        torch.distributed.broadcast(seed, 0)
    return int(seed)


@functools.lru_cache()
def _get_global_gloo_group():
    """
    Return a process group based on gloo backend, containing all the ranks
    The result is cached.
    """
    if torch.distributed.get_backend() == "nccl":
        return torch.distributed.new_group(backend="gloo")
    return torch.distributed.group.WORLD


# ------------------ Gather -----------------------


def gather(input_data):
    # TODO: Assert for PyTorch version since `gather_object` present from 1.8.0
    if get_world_size() < 2:
        return input_data

    input_data = move_to(input_data, torch.device('cpu'))
    group = _get_global_gloo_group()

    if torch.distributed.get_rank(group=group) == 0:
        gather_list = [None for _ in range(get_world_size())]
        torch.distributed.gather_object(input_data, gather_list, dst=0, group=group)
        return gather_list
    else:
        torch.distributed.gather_object(input_data, dst=0, group=group)
        return input_data


# ------------ Reduce and All Reduce --------------


def reduce(input_data, average=False, all_reduce=False):
    """
    Interface function for performing reduce on any type of 
    data [int, float, tensor, dict, list, tuple] by summing or 
    averaging the value(s) using one of the methods:
    (1) rank 0 reduce (torch.distributed.reduce):
        communicates the sum or average of
        all processes to the process of rank 0 only
    (2) all reduce (torch.distributed.all_reduce)
        communicates the sum or average of
        all processes to each process
    Parameters:
        input_dict (int, float, tensor, dict, list, tuple) -- data for reduction
        average (bool) -- true if the results should be averaged
        all_reduce (bool) -- true if communicating the reduced value to all processes,
                             otherwise process of rank 0 only
    Returns the sum or average of the input data from across processes.
    """
    if get_world_size() < 2:
        return input_data

    device = get_backend_compatible_device()
    with torch.no_grad():
        if isinstance(input_data, torch.Tensor):
            reduced_data = reduce_tensor(input_data, average, all_reduce, device)

        elif isinstance(input_data, dict):
            reduced_data = reduce_dict(input_data, average, all_reduce, device)

        elif isinstance(input_data, (list, tuple)):
            reduced_data = reduce_list_tuple(input_data, average, all_reduce, device)

        elif isinstance(input_data, (float, int)):
            reduced_data = reduce_int_float(input_data, average, all_reduce, device)

        else:
            data_type = str(type(input_data))
            raise NotImplementedError(f"Reduction on data type `{data_type}` is not implemented.")
    return reduced_data


def reduce_tensor(tensor, average, all_reduce, device):
    """Reduce a tensor"""
    tensor = tensor.clone().to(device)
    if all_reduce:
        torch.distributed.all_reduce(tensor)
    else:
        torch.distributed.reduce(tensor, dst=0)

    if average and (get_rank() == 0 or all_reduce):
        tensor /= get_world_size()
    return tensor


def reduce_int_float(input_value, average, all_reduce, device):
    """Reduce an integer or float by converting it to tensor, 
    performing reduction and, finally, casting it back to the initial type. 
    """
    data_type = type(input_value)  # save the original data type
    tensor = torch.Tensor([input_value])  # convert to tensor
    tensor = reduce_tensor(tensor, average, all_reduce, device)
    reduced_value = data_type(tensor.item())  # cast back to int or float
    return reduced_value


def reduce_dict(input_dict, average, all_reduce, device):
    """ Reduce a dict by extracting all of its values into a tensor and communicating it.
    Returns a dict with the same fields as input_dict, after reduction. If its values were 
    int or float, they are converted to tensors.
    """
    names = []
    values = []
    for k in input_dict.keys():
        names.append(k)
        value = input_dict[k]
        # if float or integer and not tensor, convert to tensor
        if is_not_tensor(value):
            if is_float_or_int(value):
                value = torch.Tensor([value])
            elif is_numpy_scalar(value):
                value = torch.Tensor([value.item()])
            else:
                raise NotImplementedError("Dictionary reduction supported only if its values \
                                           are tensors, numpy scalars floats or integers.")
        values.append(value)
    values = torch.stack(values, dim=0).to(device)  # convert the list of tensors to a single tensor

    if all_reduce:
        torch.distributed.all_reduce(values)
    else:
        torch.distributed.reduce(values, dst=0)

    if average and (get_rank() == 0 or all_reduce):
        values /= get_world_size()

    reduced_dict = {k: v for k, v in zip(names, values)}
    return reduced_dict


def reduce_list_tuple(input_data, average, all_reduce, device):
    """ Reduce a list or tuple whose elements are either tensors, floats or integers.
    Returns reduced list/tuple with its elements as tensors.
    """
    data_type = type(input_data)  # save the original data type
    # convert tuple/list values to tensors if they are floats or integers
    for i in range(len(input_data)):
        value = input_data[i]
        if is_not_tensor(value):
            if is_float_or_int(value):
                input_data[i] = torch.Tensor([value])
            elif is_numpy_scalar(value):
                input_data[i] = torch.Tensor([value.item()])
            else:
                raise NotImplementedError("List/tuple reduction supported only if"
                                          " its values are tensors, floats or integers.")
    # Convert list/tuple of tensors to a single tensor
    values = torch.stack(input_data, dim=0).to(device)

    if all_reduce:
        torch.distributed.all_reduce(values)
    else:
        torch.distributed.reduce(values, dst=0)

    if average and (get_rank() == 0 or all_reduce):
        values /= get_world_size()

    # Cast it back to tuple or list
    reduced_sequence = data_type(values)
    return reduced_sequence


# ------------------ Helpers -----------------------


def move_to(obj, device):
    """Move any combination of list or dict containing tensors to the specific device."""
    if not isinstance(obj, (torch.Tensor, dict, list)):
        return obj
    elif torch.is_tensor(obj):
        return obj.to(device)
    elif isinstance(obj, dict):
        res = {}
        for k, v in obj.items():
            res[k] = move_to(v, device)
        return res
    elif isinstance(obj, list):
        res = []
        for v in obj:
            res.append(move_to(v, device))
        return res


def is_not_tensor(x):
    return not isinstance(x, torch.Tensor)


def is_float_or_int(x):
    return isinstance(x, (float, int))


def is_numpy_scalar(x):
    return np.isscalar(x)

TorchMetric

torchmetrics.Accuracy()

torchmetrics.Accuracy(threshold=0.5, top_k=None, subset_accuracy=False, compute_on_step=True, dist_sync_on_step=False, process_group=None, dist_sync_fn=None)

计算准确率：

\[Accuracy = \frac{1}{N} \sum_{i}^N 1(y_i = \hat y_i)\]

其中 $y$ 是目标值的tensor， $\hat y$ 是预测值的tensor。

对于多类和多维多类的概率预测，参数 top_k

mport torch
import torchmetrics

# initialize metric
metric = torchmetrics.Accuracy()

n_batches = 10
for i in range(n_batches):
    # simulate a classification problem
    preds = torch.randn(10, 5).softmax(dim=-1)
    target = torch.randint(5, (10,))
    # metric on current batch
    acc = metric(preds, target)
    print(f"Accuracy on batch {i}: {acc}")

# metric on all batches using custom accumulation
acc = metric.compute()
print(f"Accuracy on all data: {acc}")