Misc fixes to contrib.epidemiology

examples/contrib/epidemiology/sir.py

@@ -79,6 +79,7 @@ def generate_data(args):
 
 
 def infer(args, model):
+    parallel = args.num_chains > 1
     energies = []
 
     def hook_fn(kernel, *unused):
@@ -92,16 +93,17 @@ def hook_fn(kernel, *unused):
                      warmup_steps=args.warmup_steps,
                      num_samples=args.num_samples,
                      num_chains=args.num_chains,
+                     mp_context="spawn" if parallel else None,
                      max_tree_depth=args.max_tree_depth,
                      arrowhead_mass=args.arrowhead_mass,
                      num_quant_bins=args.num_bins,
                      haar=args.haar,
                      haar_full_mass=args.haar_full_mass,
                      jit_compile=args.jit,
-                     hook_fn=hook_fn)
+                     hook_fn=None if parallel else hook_fn)
 
     mcmc.summary()
-    if args.plot:
+    if args.plot and energies:
         import matplotlib.pyplot as plt
         plt.figure(figsize=(6, 3))
         plt.plot(energies)

pyro/infer/mcmc/api.py

@@ -408,8 +408,12 @@ def model(data):
         # If transforms is not explicitly provided, infer automatically using
         # model args, kwargs.
         if self.transforms is None:
+            # Try to initialize kernel.transforms using kernel.setup().
+            if getattr(self.kernel, "transforms", None) is None:
+                warmup_steps = 0
+                self.kernel.setup(warmup_steps, *args, **kwargs)
             # Use `kernel.transforms` when available
-            if hasattr(self.kernel, 'transforms') and self.kernel.transforms is not None:
+            if getattr(self.kernel, "transforms", None) is not None:
                 self.transforms = self.kernel.transforms
             # Else, get transforms from model (e.g. in multiprocessing).
             elif self.kernel.model:

tests/contrib/epidemiology/test_models.py

@@ -33,9 +33,9 @@
     {"jit_compile": True},
     {"jit_compile": True, "haar_full_mass": 2},
     {"jit_compile": True, "num_quant_bins": 2},
-    {"num_chains": 2},
-    {"num_chains": 2, "num_quant_bins": 2},
-    {"num_chains": 2, "jit_compile": True},
+    {"num_chains": 2, "mp_context": "spawn"},
+    {"num_chains": 2, "mp_context": "spawn", "num_quant_bins": 2},
+    {"num_chains": 2, "mp_context": "spawn", "jit_compile": True},
 ], ids=str)
 def test_simple_sir_smoke(duration, forecast, options):
     population = 100