pytorch - 如何在 PyTorch Lightning 中编写多个训练设置

Question

我想迭代地量化我的模型。这意味着在通常实现的正常训练循环之后，training_step我想迭代地量化一些参数并重新训练模型几个步骤。

我真的不明白这将如何在 Lightning 中完成。我可以向班级添加一个新实例，但是我必须自己再次编写训练循环？

Answer 1

也许，您可以将这样的内容添加到training_step：

step = self.global_step
if step % quantization_period == 0 and step > 0:
    self.submodules[step // quantization_period].half()

您还可以LightningModule.on_epoch_end(self)在每个训练周期后重写以执行某些操作。