如何在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)，以便在每个训练周期之后执行某些操作。