原始笔记

Pytorch

Pytorch

正文

直接整理原始教程,精炼简要的写好笔记就行,不要过多解释,不需要冗余的解释,便于快速查阅开始就行。 注意需要给函数和参数写好简要的解释

pytorch_lightning = 老入口,专指 PyTorch Lightninglightning = 新入口,官方现在更推荐,包含 PyTorch Lightning + Fabric

你现在学习新教程时,优先用:

import lightning as L

除非某个老项目明确要求 pytorch_lightning

A simple example

import pytorch_lightning as pl
import torch
from torch import nn
from torch.nn import functional as F

class LitModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.layer = nn.Linear(28 * 28, 10)

    def forward(self, x):
        return self.layer(x.view(x.size(0), -1))

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

    # Uncomment to add validation step
    # def validation_step(self, batch, batch_idx):
    #     x, y = batch
    #     logits = self(x)
    #     loss = F.cross_entropy(logits, y)
    #     # Add logging
    #     self.log('val_loss', loss)

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.001)

model = LitModel()
# Initialize our model
model = LitModel()
print(model)

# Initialize a trainer
trainer = pl.Trainer(max_epochs=3)

# Train the model
trainer.fit(model, train_loader)

model 负责定义“学什么” train_loader 负责提供“用什么数据学” trainer 负责控制“怎么训练”

模型逻辑:写在 LightningModule 数据逻辑:写在 DataLoader 或 LightningDataModule 训练控制:交给 Trainer

LightningDataModule的作用以及怎么用? 如果想要用自己的数据集要怎么办?

Switch to English