翻译自官方文档
前置知识:推荐pytorch
目标:通过PL中7个关键步骤了解PL工作流程
PL是基于pytorch的高层API,自带丰富的工具为AI学者和工程师快速创建高性能模型,去除繁琐的重复流程同时保持灵活性。
使用组织好的pytorch代码,PL可以:
pip 安装
pip install lightning
conda 安装
conda install lightning -c conda-forge
子类中可以使用原生pytorch nn.Module创建的模块搭建PL的模型,然后在training_step()方法中实现损失计算的过程。将LightningModule的子类传入Trainer后,即会自动调用计算损失并反向传播。
# 导入包
import os
from torch import optim, nn, utils, Tensor
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor
import lightning.pytorch as pl
# 定义一个线性层和激活函数组成的编码器-解码器,你可以使用任何nn.Module创建的模块
encoder = nn.Sequential(nn.Linear(28 * 28, 64), nn.ReLU(), nn.Linear(64, 3))
decoder = nn.Sequential(nn.Linear(3, 64), nn.ReLU(), nn.Linear(64, 28 * 28))
# 定义一个继承LightningModule的子类
class LitAutoEncoder(pl.LightningModule):
def __init__(self, encoder, decoder):
super().__init__()
# 设置属性,模型包含两个子模块
self.encoder = encoder
self.decoder = decoder
# 单步训练过程,里面包含forward流程,可以单独写出
def training_step(self, batch, batch_idx):
# training_step defines 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 = nn.functional.mse_loss(x_hat, x)
# 记录训练过程中的损失。如果要使用tensorboard,wandb等工具,需要在trainer创建时指定,并如下设置logger=Ture
self.log("train_loss", loss, logger=Ture)
return loss
# 配置优化器
def configure_optimizers(self):
optimizer = optim.Adam(self.parameters(), lr=1e-3)
return optimizer
# 创建一个模型
autoencoder = LitAutoEncoder(encoder, decoder)
任何可迭代iterable 的对象(list, dataloader, dict, numpy等)均可作为加载器。
# 这里使用的torchvision中包含的MNIST数据集,可以自定义,方法和pytorch一样实现get_item()方法。
dataset = MNIST(os.getcwd(), download=True, transform=ToTensor())
train_loader = utils.data.DataLoader(dataset)
创建一个训练管理器trainer,调用fit()方法时传入模型和数据集加载器,开始训练。
# 创建Trainer时有不同参数,详见API说明
trainer = pl.Trainer(limit_train_batches=100, max_epochs=1)
# 传入模型和数据集加载器并开始训练
trainer.fit(model=autoencoder, train_dataloaders=train_loader)
Trainer 包含train/validation/test loop,日志记录、gpu设置、模型保存策略
等逻辑,避免花费过多精力在重复的流程上。
训练完成后可以将模型转换格式部署至生产环境中,也只可加载权重预测。
# 加载checkpoint
checkpoint = "./lightning_logs/version_0/checkpoints/epoch=0-step=100.ckpt"
autoencoder = LitAutoEncoder.load_from_checkpoint(checkpoint, encoder=encoder, decoder=decoder)
# choose your trained nn.Module
encoder = autoencoder.encoder
encoder.eval()
# 预测
fake_image_batch = Tensor(4, 28 * 28)
embeddings = encoder(fake_image_batch)
print("⚡" * 20, "\nPredictions (4 image embeddings):\n", embeddings, "\n", "⚡" * 20)
如果创建Trainer时指定了logger,如tensorboard, wandb并给定相关设置,即可打开日志文件查看训练状态。比如使用了tensorboard作为logger,执行以下命令后,浏览器打开 http://localhost:6006/
cd path/to/your/log/file
tensorboard --logdir .
在创建Trainer时你可以使用不同的设备(cpu、gpu、tpu等)、不同的策略、精度,以及callback。
# 4块GPU上训练
trainer = Trainer(
devices=4,
accelerator="gpu",
)
# 通过Deepspeed/fsdp策略,在16精度下由4块GPU训练
trainer = Trainer(
devices=4,
accelerator="gpu",
strategy="deepspeed_stage_2",
precision=16
)
# Trainer 中20多个有用的配置项,比如
trainer = Trainer(
max_epochs=10,# 最大训练10个epochs
min_epochs=5,# 最小训练5个epoch
overfit_batches=0.01 # 使用0.01的数据训练,以快速测试代码,默认为1
)
# 将实现了callback接口的模块整合进流程中,callback中不同流程的处理逻辑
trainer = Trainer(callbacks=[StochasticWeightAveraging(...)])
参考文档