Pytorch-Lightning中模型保存与加载

Pytorch-Lightning中模型保存与加载

保存

自动保存

from lightning.pytorch.callbacks import ModelCheckpoint

class LitAutoEncoder(LightningModule):
    def validation_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self.backbone(x)

        # 1. 计算loss
        loss = F.cross_entropy(y_hat, y)

        # 2. log `val_loss`
        self.log("val_loss", loss)

# 3. 初始化回调，监视 'val_loss'
checkpoint_callback = ModelCheckpoint(monitor="val_loss")

# 4. 向Trainer添加回调
trainer = Trainer(callbacks=[checkpoint_callback])

• ModelCheckpoint更多用法

手动保存

model = MyLightningModule(hparams)
trainer.fit(model)
trainer.save_checkpoint("example.ckpt")

• 不推荐，在分布性训练环境中有产生死锁的风险

加载

使用load_from_checkpoint()

model = MyLightingModule.load_from_checkpoint(PATH)

使用Trainer中恢复

model = LitModel()
trainer = Trainer()
# 自动恢复模型
trainer.fit(model, ckpt_path="some/path/to/my_checkpoint.ckpt")

• 更多Trainer用法

超参数保存与覆盖

class LitModel(LightningModule):
    def __init__(self, in_dim, out_dim):
        super().__init__()
        # 保存超参数到hparams属性中
        self.save_hyperparameters()
        self.l1 = nn.Linear(self.hparams.in_dim, self.hparams.out_dim)
# 假设训练并保存了这个模型到PATH路径
LitModel(in_dim=32, out_dim=10)

# 使用in_dim=32, out_dim=10
model = LitModel.load_from_checkpoint(PATH)

# 覆盖ckpt文件中的超参数，in_dim=128, out_dim=10
model = LitModel.load_from_checkpoint(PATH, in_dim=128, out_dim=10)