Как написать несколько настроек обучения в 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) что-то делать после каждой тренировочной эпохи.