PyroModelGuideWarmup fails on GPU - probably need to be manually run before `trainer.fit()` about scvi-tools HOT 4 OPEN

This can probably be addressed using the following modification of the TrainRunner or to the model.train() method:

class TrainRunner:

    def __call__(self):
        # other code .....
        
        from copy import copy
        dl = copy(self.data_splitter)
        dl.setup()
        dl = dl.train_dataloader()
        PyroModelGuideWarmup(dl).setup(
            self.trainer, self.training_plan, stage="fit"
        )

        self.trainer.fit(
            self.training_plan, self.data_splitter, ckpt_path=self.ckpt_path
        )
        # other code .....

At this stage self.data_splitter.setup() has not been called yet and PyTorch Lightning expects to call self.data_splitter.setup() later. So we need to copy self.data_splitter, call self.data_splitter.setup() on a copy and create the dataloader needed for this callback.

Using model.train() is probably better because the model would have all parameters created before Lightning sees it.

from scvi-tools.

Would be great to hear what you think @martinkim0 and I can add the proposed changes

from scvi-tools.

Hey sorry, took a look at this and forgot to respond. I think it makes sense to add the fixes to train instead of the TrainRunner since this will be specific to Pyro models. Happy to take a PR if you'd like to take a stab at it!

from scvi-tools.

Sounds good! Later this week, I will make a PR about this issue - as well as another issue with the second GuideWarmup callback (pyro doesn't track deterministic variables initialised after setup).

I think we need to get rid of both pyro GuideWarmup callbacks and just run guide once in model.train(). This would break how people use them now but IMO a better solution.

from scvi-tools.