Pytorch-Lightning在模型训练时记录中间值报错RuntimeError：CUDA out of memory

最新推荐文章于 2025-05-22 02:04:09 发布

原创

最新推荐文章于 2025-05-22 02:04:09 发布 · 2.0k 阅读

0 ·

CC 4.0 BY-SA版权

文章标签：

#python

本文介绍如何使用Pytorch-Lightning进行模型训练，并详细解释了如何在训练过程中记录模型输出，避免内存溢出的问题。

在训练model的时候，Pytorch-Lightning通过定义System（以model作为输入）和Trainer的方式实现模型训练。

一、定义

（1）首先定义System：

res_from_recursion = []
system = System(
        model=model,#在system之前定义的
        loss_func=loss_func,
        optimizer=optimizer,
        train_loader=train_loader,
        res_from_recursion=res_from_recursion,#我想让模型在训练时输出的值，用一个list保存
        val_loader=val_loader,
        scheduler=scheduler,
        config=conf,