最近在做controlnet的实验的时候要保存pytorch lightning的trainer。但是每次保存一个checkpoint就需要20G的存储空间。对于只有50G的数据盘的autodl来说显得捉襟见肘。
为了节省一些存储空间,我决定只保存最新epoch的checkpoint。但是,在百度和Google上搜了一下也没有什么好的解决办法。于是,我查阅了pytorch lightning的官方文档,发现里面有这样一句话
于是,我想:如果我设置every_n_epoch = 0
且save_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])