PyTorch Lightning 使用教程

戚展焰Beatrix

于 2024-08-12 08:41:58 发布

阅读量403

点赞数 3

本文链接：https://blog.csdn.net/gitblog_00590/article/details/141119046

版权

PyTorch Lightning 使用教程

pytorch-lightningLightning-AI/pytorch-lightning: PyTorch Lightning 是一个轻量级的高级接口，用于简化 PyTorch 中深度学习模型的训练流程。它抽象出了繁杂的工程细节，使研究者能够专注于模型本身的逻辑和实验设计，同时仍能充分利用PyTorch底层的灵活性。项目地址:https://gitcode.com/gh_mirrors/py/pytorch-lightning

项目介绍

PyTorch Lightning 是一个轻量级的 PyTorch 框架扩展，旨在简化深度学习模型的训练和部署过程。它通过提供高级抽象和自动化功能，帮助开发者减少样板代码，专注于模型的科学部分。PyTorch Lightning 的核心包包括：

PyTorch Lightning: 用于大规模训练和部署 PyTorch 模型。
Lightning Fabric: 提供专家级别的控制，适用于复杂模型的训练和扩展策略。

项目快速启动

安装

首先，通过 pip 安装 PyTorch Lightning：

pip install pytorch-lightning

示例代码

以下是一个简单的 PyTorch Lightning 示例，展示如何训练一个基本的分类模型：

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

# 定义一个简单的神经网络
class SimpleNN(pl.LightningModule):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.layer = torch.nn.Linear(10, 1)

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

    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = torch.nn.functional.mse_loss(y_hat, y)
        return loss

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

# 生成一些随机数据
x_train = torch.randn(100, 10)
y_train = torch.randn(100, 1)
train_dataset = TensorDataset(x_train, y_train)
train_loader = DataLoader(train_dataset, batch_size=10)

# 初始化模型和训练器
model = SimpleNN()
trainer = pl.Trainer(max_epochs=10)

# 开始训练
trainer.fit(model, train_loader)