Pyro integration

docs/api/core.rst

@@ -14,6 +14,7 @@ Compose
    compose.Encoder
    compose.LossRecorder
    compose.BaseModuleClass
+   compose.PyroBaseModuleClass
 
 .. autosummary::
    :toctree: reference/
@@ -68,6 +69,7 @@ steps for modules like `TOTALVAE`, `SCANVAE`, etc.
    lightning.TrainingPlan
    lightning.SemiSupervisedTrainingPlan
    lightning.AdversarialTrainingPlan
+   lightning.PyroTrainingPlan
    lightning.Trainer
 
 Utilities

pyproject.toml

@@ -45,6 +45,7 @@ pandas = ">=1.0"
 pre-commit = {version = ">=2.7.1", optional = true}
 prospector = {version = "*", optional = true}
 pydata-sphinx-theme = {version = ">=0.4.0", optional = true}
+pyro-ppl = ">=1.1.0"
 pytest = {version = ">=4.4", optional = true}
 python = ">=3.6.1,<4.0"
 python-igraph = {version = "*", optional = true}

scvi/compose/__init__.py

@@ -9,7 +9,7 @@
     MultiDecoder,
     MultiEncoder,
 )
-from ._base_module import BaseModuleClass, LossRecorder
+from ._base_module import BaseModuleClass, LossRecorder, PyroBaseModuleClass
 from ._decorators import auto_move_data
 from ._utils import one_hot  # Do we want one_hot here?
 
@@ -27,4 +27,5 @@
     "BaseModuleClass",
     "one_hot",
     "auto_move_data",
+    "PyroBaseModuleClass",
 ]

scvi/compose/_base_module.py

@@ -1,5 +1,5 @@
 from abc import abstractmethod
-from typing import Dict, Optional, Tuple, Union
+from typing import Dict, Optional, List, Tuple, Union
 
 import torch
 import torch.nn as nn
@@ -194,3 +194,68 @@ def _get_dict_if_none(param):
     param = {} if not isinstance(param, dict) else param
 
     return param
+
+
+class PyroBaseModuleClass(nn.Module):
+    """
+    Base module class for Pyro models.
+
+    We use the forward method of this class as the "model" in Pyro terms,
+    and also implement the guide here, through the guide method.
+
+    Parameters
+    ----------
+    loss
+        Tensor with loss for minibatch. Should be one dimensional with one value.
+        Note that loss should be a :class:`~torch.Tensor` and not the result of `.item()`.
+    reconstruction_loss
+        Reconstruction loss for each observation in the minibatch.
+    kl_local
+        KL divergence associated with each observation in the minibatch.
+    kl_global
+        Global kl divergence term. Should be one dimensional with one value.
+    """
+
+    def __init__(
+        self,
+    ):
+        super().__init__()
+
+    @abstractmethod
+    def _get_guide_tensors(
+        self, tensors: Dict[str, torch.Tensor]
+    ) -> List[torch.Tensor]:
+        pass
+
+    @abstractmethod
+    def _guide(self, *args, **kwargs):
+        pass
+
+    @abstractmethod
+    def guide(self, tensors: Dict[str, torch.Tensor]) -> dict:
+        """
+        Pyro Guide method.
+
+        This is a wrapper for the `_guide` method. This function parses
+        the tensors dictionary and passes to `_guide` as positional arguments.
+        """
+        self._guide(*self._get_guide_tensors(tensors))
+
+    @abstractmethod
+    def _get_forward_tensors(
+        self, tensors: Dict[str, torch.Tensor]
+    ) -> List[torch.Tensor]:
+        pass
+
+    @abstractmethod
+    def _forward(self, *args, **kwargs):
+        pass
+
+    def forward(self, tensors: Dict[str, torch.Tensor]) -> dict:
+        """
+        Pyro model method.
+
+        This is a wrapper for the `_model` method. This function parses
+        the tensors dictionary and passes to `_model` as positional arguments.
+        """
+        self._forward(*self._get_forward_tensors(tensors))

scvi/lightning/__init__.py

@@ -1,13 +1,15 @@
 from ._trainer import Trainer
 from ._trainingplans import (
     AdversarialTrainingPlan,
+    PyroTrainingPlan,
     SemiSupervisedTrainingPlan,
     TrainingPlan,
 )
 
 __all__ = [
     "TrainingPlan",
     "Trainer",
+    "PyroTrainingPlan",
     "SemiSupervisedTrainingPlan",
     "AdversarialTrainingPlan",
 ]

scvi/lightning/_trainingplans.py

@@ -1,13 +1,14 @@
 from inspect import getfullargspec
-from typing import Union
+from typing import Callable, Union
 
+import pyro
 import pytorch_lightning as pl
 import torch
 from torch.optim.lr_scheduler import ReduceLROnPlateau
 
 from scvi import _CONSTANTS
 from scvi._compat import Literal
-from scvi.compose import BaseModuleClass, one_hot
+from scvi.compose import PyroBaseModuleClass, BaseModuleClass, one_hot
 from scvi.modules import Classifier
 
 
@@ -522,3 +523,55 @@ def validation_step(self, batch, batch_idx, optimizer_idx=0):
             "kl_global": scvi_losses.kl_global,
             "n_obs": reconstruction_loss.shape[0],
         }
+
+
+class PyroTrainingPlan(pl.LightningModule):
+    """
+    Lightning module task to train Pyro scvi-tools modules.
+
+    Parameters
+    ----------
+    pyro_module
+        A model instance from class ``PyroBaseModuleClass``.
+    lr
+        Learning rate used for optimization.
+    """
+
+    def __init__(
+        self,
+        pyro_module: PyroBaseModuleClass,
+        lr: float = 1e-3,
+        loss_fn: Callable = pyro.infer.Trace_ELBO(),
+    ):
+        super().__init__()
+        self.module = pyro_module
+        self.loss_fn = loss_fn
+        self.lr = lr
+
+        self.automatic_optimization = False
+        self.guide = self.module.guide
+
+        self.svi = pyro.infer.SVI(
+            model=self.module,
+            guide=self.guide,
+            optim=pyro.optim.Adam({"lr": self.lr}),
+            loss=self.loss_fn,
+        )
+
+    def forward(self, *args, **kwargs):
+        """Passthrough to `model.forward()`."""
+        return self.module(*args, **kwargs)
+
+    def training_step(self, batch, batch_idx, optimizer_idx=0):
+
+        loss = self.svi.step(batch)
+        self.log("train_loss", loss, prog_bar=True, on_epoch=True)
+
+    def configure_optimizers(self):
+        return None
+
+    def optimizer_step(self, *args, **kwargs):
+        pass
+
+    def backward(self, *args, **kwargs):
+        pass

tests/models/test_pyro.py

@@ -0,0 +1,77 @@
+import pyro
+import pyro.distributions as dist
+import torch
+import torch.nn as nn
+from pyro.infer.autoguide import AutoDiagonalNormal
+from pyro.nn import PyroModule
+
+from scvi import _CONSTANTS
+from scvi.compose import PyroBaseModuleClass
+from scvi.data import synthetic_iid
+from scvi.dataloaders import AnnDataLoader
+from scvi.lightning import PyroTrainingPlan, Trainer
+
+
+class BayesianRegression(PyroModule, PyroBaseModuleClass):
+    def __init__(self, in_features, out_features):
+        super().__init__()
+
+        self._auto_guide = AutoDiagonalNormal(self)
+
+        self.register_buffer("zero", torch.tensor(0.0, requires_grad=False))
+        self.register_buffer("one", torch.tensor(1.0, requires_grad=False))
+        self.register_buffer("ten", torch.tensor(10.0, requires_grad=False))
+
+        self.linear = nn.Linear(in_features, out_features)
+
+    def _get_forward_tensors(self, tensors):
+        x = tensors[_CONSTANTS.X_KEY]
+        y = tensors[_CONSTANTS.LABELS_KEY]
+
+        return x, y
+
+    def _forward(self, x, y):
+        sigma = pyro.sample("sigma", dist.Uniform(self.zero, self.ten))
+        mean = self.linear(x).squeeze(-1)
+        with pyro.plate("data", x.shape[0]):
+            pyro.sample("obs", dist.Normal(mean, sigma), obs=y)
+        return mean
+
+    def _get_guide_tensors(self, tensors):
+        return [tensors]
+
+    def _guide(self, tensors):
+        return self._auto_guide(tensors)
+
+
+def test_pyro_bayesian_regression():
+    use_gpu = int(torch.cuda.is_available())
+    adata = synthetic_iid()
+    train_dl = AnnDataLoader(adata, shuffle=True, batch_size=128)
+    pyro.clear_param_store()
+    model = BayesianRegression(adata.shape[1], 1)
+    plan = PyroTrainingPlan(model)
+    trainer = Trainer(
+        gpus=use_gpu,
+        max_epochs=2,
+    )
+    trainer.fit(plan, train_dl)
+
+
+def test_pyro_bayesian_regression_jit():
+    use_gpu = int(torch.cuda.is_available())
+    adata = synthetic_iid()
+    train_dl = AnnDataLoader(adata, shuffle=True, batch_size=128)
+    pyro.clear_param_store()
+    model = BayesianRegression(adata.shape[1], 1)
+    # warmup guide for JIT
+    for tensors in train_dl:
+        model.guide(tensors)
+        break
+    train_dl = AnnDataLoader(adata, shuffle=True, batch_size=128)
+    plan = PyroTrainingPlan(model, loss_fn=pyro.infer.JitTrace_ELBO())
+    trainer = Trainer(
+        gpus=use_gpu,
+        max_epochs=2,
+    )
+    trainer.fit(plan, train_dl)