Merge pull request #1059 from vitkl/pyro-mixin

Adding mixin classes for pyro training and posterior sampling
scverse · Jun 19, 2021 · bc415d8 · bc415d8
2 parents 3a6ca06 + 52759bb
commit bc415d8
Show file tree

Hide file tree

Showing 9 changed files with 770 additions and 50 deletions.
diff --git a/docs/_templates/class_no_inherited.rst b/docs/_templates/class_no_inherited.rst
@@ -7,6 +7,21 @@
 .. autoclass:: {{ objname }}
    :show-inheritance:
 
+   {% block attributes %}
+   {% if attributes %}
+   .. rubric:: Attributes
+
+   .. autosummary::
+      :toctree: .
+    {% for item in attributes %}
+        {%- if item not in inherited_members%}
+            ~{{ fullname }}.{{ item }}
+        {%- endif -%}
+    {%- endfor %}
+    {% endif %}
+    {% endblock %}
+
+
    {% block methods %}
    {% if methods %}
    .. rubric:: Methods

diff --git a/docs/api/developer.rst b/docs/api/developer.rst
@@ -64,6 +64,9 @@ These classes should be used to construct user-facing model classes.
     model.base.RNASeqMixin
     model.base.ArchesMixin
     model.base.UnsupervisedTrainingMixin
+    model.base.PyroSviTrainMixin
+    model.base.PyroSampleMixin
+    model.base.PyroJitGuideWarmup
 
 Module
 ------

diff --git a/pyproject.toml b/pyproject.toml
@@ -66,7 +66,7 @@ sphinx = {version = ">=3.4", optional = true}
 sphinx-autodoc-typehints = {version = "*", optional = true}
 sphinx-gallery = {version = ">0.6", optional = true}
 sphinx-tabs = {version = "*", optional = true}
-sphinx_copybutton = {version = "*", optional = true}
+sphinx_copybutton = {version = "<=0.3.1", optional = true}
 torch = ">=1.8.0"
 tqdm = ">=4.56.0"
 typing_extensions = {version = "*", python = "<3.8"}

diff --git a/scvi/model/base/__init__.py b/scvi/model/base/__init__.py
@@ -1,5 +1,6 @@
 from ._archesmixin import ArchesMixin
 from ._base_model import BaseModelClass
+from ._pyromixin import PyroJitGuideWarmup, PyroSampleMixin, PyroSviTrainMixin
 from ._rnamixin import RNASeqMixin
 from ._training_mixin import UnsupervisedTrainingMixin
 from ._vaemixin import VAEMixin
@@ -10,4 +11,7 @@
     "RNASeqMixin",
     "VAEMixin",
     "UnsupervisedTrainingMixin",
+    "PyroSviTrainMixin",
+    "PyroSampleMixin",
+    "PyroJitGuideWarmup",
 ]
diff --git a/scvi/model/base/_base_model.py b/scvi/model/base/_base_model.py
@@ -99,6 +99,8 @@ def _make_data_loader(
             Minibatch size for data loading into model. Defaults to `scvi.settings.batch_size`.
         shuffle
             Whether observations are shuffled each iteration though
+        data_loader_class
+            Class to use for data loader
         data_loader_kwargs
             Kwargs to the class-specific data loader class
         """
@@ -344,8 +346,10 @@ def load(
             model.module.load_state_dict(model_state_dict)
         except RuntimeError as err:
             if isinstance(model.module, PyroBaseModuleClass):
+                old_history = model.history_
                 logger.info("Preparing underlying module for load")
                 model.train(max_steps=1)
+                model.history_ = old_history
                 pyro.clear_param_store()
                 model.module.load_state_dict(model_state_dict)
             else: