pytorch lightning使用（简要介绍）

最新推荐文章于 2025-09-07 22:50:36 发布

原创

最新推荐文章于 2025-09-07 22:50:36 发布 · 6.9k 阅读

25 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #lightning #深度学习 #python #神经网络

0. 简介

pytorch lightning通过提供LightningModule和LightningDataModule，使得在用pytorch编写网络模型时，加载数据、分割数据集、训练、验证、测试、计算指标的代码全部都能很好的组织起来，显得主程序调用时，代码简洁可读性大幅度提升。

1. pytorch lightning的安装

pip install pytorch-lightning
conda install pytorch-lightning -c conda-forge

2. 定义一个网络模型模型：LightningModule

通过继承LightningModule，并实现几个关键的函数，使得模型在训练、验证和测试过程中能进行模块化调用，具体细节完全被自定义的函数封装，整体十分简洁。定义一个LightningModule的基类，可以实现的函数如下：

from pytorch_lightning import LightningModule

class MyModel(LightningModule):
    """
    The only required methods in the LightningModule are:
    init
    training_step
    configure_optimizers
    """
    def __init__(self, *args, **kwargs): pass
    def forward(self, *args, **kwargs): pass
    def training_step(self, batch, batch_idx, optimizer_idx, hiddens): pass
    def training_step_end(self, *args, **kwargs): pass # 接受train_step的返回值
    def training_epoch_end(self, outpu