https://pytorch-lightning.readthedocs.io/en/latest/starter/introduction_guide.html

模型部分

定义模型

设计一个三层的神经网络，定义类时继承自LightningModule
对比：Pytorch中直接定义网络时，继承自nn.Module

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, height, width)
        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, height, width = 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

LightningModule 和 Pytorch里的模块相同的，但增加了一些功能
实例化的用法相同

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

训练

增加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

优化器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)

如果有学习率的变化

from torch.optim.lr_scheduler import CosineAnnealingLR
class LitMNIST(LightningModule):
    def configure_optimizers(self):
        opt = Adam(self.parameters(), lr=1e-3)
        scheduler = CosineAnnealingLR(opt, T_max=10)
        return [opt], [scheduler]

数据

Pytorch里加载数据

主要是使用Datasets和DataLoader

from torch.utils.data import DataLoader, random_split
from torchvision.datasets import MNIST
import os
from torchvision import datasets, transforms
from pytorch_lightning import Trainer
# 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)

Lightning里