Pytorch Lightning 框架运行顺序

最新推荐文章于 2025-12-10 18:58:51 发布

原创最新推荐文章于 2025-12-10 18:58:51 发布 · 171 阅读

2 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #人工智能 #python

之前学习Lightning在打断点debug的时候一直不懂为什么trainer.fit()之后就开始运行了其中的hooks 什么的概念也不理解在看了若干文章后大致理解其中顺序。如有错误请斧正。
Hooks: 实际就是lighting框架的运行顺序, train阶段的运行顺序大致如下（其中的函数都可以重写）:

def fit(self):
    configure_callbacks()

    if local_rank == 0:
        prepare_data()

    setup("fit")
    configure_model()
    configure_optimizers()

    on_fit_start()

    # the sanity check runs here

    on_train_start()
    for epoch in epochs:
        fit_loop()
    on_train_end()

    on_fit_end()
    teardown("fit")

需要注意的是:

the sanity check runs here 这行注释 lighting会先执行validation_step() 跑两个batch 进行检查。（之前一直不懂为什么train的时候先执行validation）
prepare_data() 和 setup() 这两个function 在Data的class里重写其余的在model的class里重写。