pytorch lightning只保存最新的epoch

最近在做controlnet的实验的时候要保存pytorch lightning的trainer。但是每次保存一个checkpoint就需要20G的存储空间。对于只有50G的数据盘的autodl来说显得捉襟见肘。
为了节省一些存储空间,我决定只保存最新epoch的checkpoint。但是,在百度和Google上搜了一下也没有什么好的解决办法。于是,我查阅了pytorch lightning的官方文档,发现里面有这样一句话
图1
于是,我想:如果我设置every_n_epoch = 0save_last = True,那不就能保存最新的epoch了嘛。
ModelCheckpoint的完整代码:

checkpoint_callback = ModelCheckpoint(
    every_n_epochs = 0,
    verbose=True,
    save_last=True,
    dirpath='./checkpoints',
    filename='{epoch}-checkpoint', 
)

trainer的写法:

trainer = pl.Trainer(devices=1, accelerator="gpu", max_epochs=300,  callbacks=[logger, checkpoint_callback])

参考文献:

pytorch lightning官方文档

  • 1
    点赞
  • 0
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
PyTorch Lightning是一个用于训练PyTorch模型的轻量级开发框架。它提供了一种将PyTorch代码组织到Lightning模块中的方法,以便更简洁、可读性更强、可扩展性更好地编写和管理模型代码。使用PyTorch Lightning可以帮助你在训练深度学习模型时更加专注于核心的研究和开发。 要将PyTorch代码组织到Lightning模块中,首先需要安装PyTorch Lightning。中提供了一个参考链接,其中包含了如何安装PyTorch Lightning的详细说明。安装完成后,可以开始定义Lightning模块。 在PyTorch Lightning中,我们需要定义一个继承自`pl.LightningModule`的类来表示我们的模型。这个类中包含了模型的结构、前向传播方法、损失函数等。通过继承`pl.LightningModule`,我们可以利用Lightning的许多功能,例如自动分布式训练、自动优化器管理、自动日志记录等。 定义Lightning模块后,我们可以使用`pl.Trainer`进行模型的训练。中的参考链接提供了如何使用`pl.Trainer`进行训练的详细说明。`pl.Trainer`提供了许多训练相关的选项,例如训练的epoch数、批次大小、学习率调度等。 总结起来,使用PyTorch Lightning可以通过将PyTorch代码组织到Lightning模块中来更好地管理和开发深度学习模型。通过继承`pl.LightningModule`定义模型结构,并使用`pl.Trainer`进行训练。有关更详细的信息和使用示例,请参考和提供的参考链接。
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值