Simplest FitLoop customization doesn't work

[whaowhao]

🐛 Bug

Setting trainer.fit_loop to default FitLoop causes error

To Reproduce

Take any trainer and model, and add one line that sets trainer.fit_loop to default FitLoop causes error

... code before
from pytorch_lightning.loops import FitLoop
trainer.fit_loop = FitLoop()
... code after

Exception has occurred: AttributeError       (note: full exception trace is shown but execution is paused at: <module>)
'NoneType' object has no attribute '_results'
  File "/home/whao/.virtualenv/mixedclassify/lib/python3.6/site-packages/pytorch_lightning/loops/fit_loop.py", line 145, in _results
    return self.epoch_loop._results
  File "/home/whao/.virtualenv/mixedclassify/lib/python3.6/site-packages/pytorch_lightning/trainer/trainer.py", line 2143, in _results
    return active_loop._results
  File "/home/whao/.virtualenv/mixedclassify/lib/python3.6/site-packages/pytorch_lightning/trainer/connectors/logger_connector/logger_connector.py", line 308, in reset_results
    if self.trainer._results is not None:
  File "/home/whao/.virtualenv/mixedclassify/lib/python3.6/site-packages/pytorch_lightning/trainer/trainer.py", line 1174, in _run
    self.logger_connector.reset_results()
  File "/home/whao/.virtualenv/mixedclassify/lib/python3.6/site-packages/pytorch_lightning/trainer/trainer.py", line 770, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File "/home/whao/.virtualenv/mixedclassify/lib/python3.6/site-packages/pytorch_lightning/trainer/trainer.py", line 682, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
  File "/home/whao/.virtualenv/mixedclassify/lib/python3.6/site-packages/pytorch_lightning/trainer/trainer.py", line 736, in fit
    self._fit_impl, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path
  File "/home/whao/whao/MultiHead/multihead/cli.py", line 153, in run_active_learn
    trainer.fit(model)

Expected behavior

Nothing should change since we set fitloop to default FitLoop

Environment

pl==1.5.0

cc @tchaton @carmocca @justusschock

[awaelchli]

Hi @whaowhao

Thanks for checking out the loop features.
The reason you see the error is because the FitLoop requires some subloops to function properly. This is the default configuration that the trainer uses:

fit_loop = FitLoop()  # optionally pass max_epochs
epoch_loop = TrainingEpochLoop()  # optionally pass max_steps
batch_loop = TrainingBatchLoop()
optimizer_loop = OptimizerLoop()
val_loop = EvaluationLoop()

batch_loop.connect(optimizer_loop=optimizer_loop)
epoch_loop.connect(batch_loop=batch_loop, val_loop=val_loop)
fit_loop.connect(epoch_loop=epoch_loop)

trainer.fit_loop = fit_loop

Any specific customization you have in mind?

[whaowhao]

Thanks @awaelchli ! So I am trying to custom the Loop such that after each fit loop, trainer.predict runs on an arbitrary file based on model.hparams

[rohitgr7]

you mean after each epoch? @whaowhao

[whaowhao]

Yes, after train and validation epoch are done

[tchaton]

Hey @whaowhao,

You could do the following:

class CustomTrainingEpochLoop(TrainingEpochLoop):

    def on_run_end(self):
        super().on_run_end()
        predictions = self.trainer.predict_loop.run()

        # do something with the predictions

fit_loop = FitLoop()  # optionally pass max_epochs
epoch_loop = CustomTrainingEpochLoop()  # optionally pass max_steps
batch_loop = TrainingBatchLoop()
optimizer_loop = OptimizerLoop()
val_loop = EvaluationLoop()

batch_loop.connect(optimizer_loop=optimizer_loop)
epoch_loop.connect(batch_loop=batch_loop, val_loop=val_loop)
fit_loop.connect(epoch_loop=epoch_loop)

trainer.fit_loop = fit_loop

[carmocca]

#10309 should greatly improve the usability here

[btseytlin]

Still doesn't work with cryptic errors:

fit_loop = FitLoop()
training_epoch_loop = TrainingEpochLoop(min_steps=0, max_steps=-1)
fit_loop.connect(epoch_loop=training_epoch_loop)
trainer.fit_loop = fit_loop

[carmocca]

@btseytlin I just tried it and it's working well for me. What errors are you seeing?

Code

import os

import torch
from torch.utils.data import DataLoader, Dataset

from pytorch_lightning import LightningModule, Trainer
from pytorch_lightning.loops import TrainingEpochLoop, FitLoop


class RandomDataset(Dataset):
    def __init__(self, size, length):
        self.len = length
        self.data = torch.randn(length, size)

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return self.len


class BoringModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(32, 2)

    def forward(self, x):
        return self.layer(x)

    def training_step(self, batch, batch_idx):
        return self(batch).sum()

    def configure_optimizers(self):
        return torch.optim.SGD(self.layer.parameters(), lr=0.1)


def run():
    train_data = DataLoader(RandomDataset(32, 64), batch_size=2)

    model = BoringModel()
    trainer = Trainer(
        default_root_dir=os.getcwd(),
        limit_train_batches=1,
        limit_val_batches=1,
        limit_test_batches=1,
        num_sanity_val_steps=0,
        max_epochs=1,
        enable_model_summary=False,
    )

    fit_loop = FitLoop()
    training_epoch_loop = TrainingEpochLoop(min_steps=0, max_steps=-1)
    fit_loop.connect(epoch_loop=training_epoch_loop)
    trainer.fit_loop = fit_loop

    trainer.fit(model, train_dataloaders=train_data)


if __name__ == "__main__":
    run()

You could also use:

training_epoch_loop = TrainingEpochLoop(...)
trainer.fit_loop.connect(epoch_loop=training_epoch_loop)

or

trainer.fit_loop.replace(epoch_loop=TrainingEpochLoop(...))