How to combine guides?

if I had two neural network models f1: x → y and f2: y → z and have generated autoguide for each, how should I combine these guides to form a guide for a sequential model f1 f2: x → z?

More concretely:

import torch
import pyro
f1 = torch.nn.Linear(4, 8)
f2 = torch.nn.Linear(8, 16)
pyro.nn.module.to_pyro_module_(f1)
pyro.nn.module.to_pyro_module_(f2)
f1_guide = pyro.infer.autoguide.AutoNormal(f1)
f2_guide = pyro.infer.autoguide.AutoDelta(f2)

class F1F2(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.f1, self.f2 = f1, f2

    def forward(self, x):
        return self.f2(self.f1(x))

f1f2 = F1F2()
pyro.nn.module.to_pyro_module_(f1f2)

how should I express the guide for this f1f2 model if I wanted to keep its original guides?

Hmm I believe you can use AutoGuideList bug I think you’ll need to use poutine.block to get this working with Pyro’s name-based autoguide logic, something like this:

import torch
import pyro

class F1F2(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.f1 = torch.nn.Linear(4, 8)
        self.f2 = torch.nn.Linear(8, 16)

    def forward(self, x):
        return self.f2(self.f1(x))

f1f2 = F1F2()
pyro.nn.module.to_pyro_module_(f1f2)  # you need this only on the top level module
guide = AutoGuideList(f1f2)
guide.append(
    pyro.infer.autoguide.AutoNormal(
        pyro.poutine.block(f1f2, expose_fn=lambda site: "f1." in site["name"])
    )
)
guide.append(
    pyro.infer.autoguide.AutoDelta(
        pyro.poutine.block(f1f2, expose_fn=lambda site: "f2." in site["name"])
    )
) 

Thanks!

What if f1 and f2 don’t have autoguides but some carefully crafted guides that I wanted to preserve. Would it be possible to copy those guides into f1f2?

Yes if you have guides already constructed, you should be able to use AutoGuideList without the poutine.block (the poutine.block stuff is just how the AutoNormal and other autoguides determine which latent variables to model). It may a little tricky getting the pyro sample site names to match when you train a partial guide with a partial model then use that guide in a larger model, due to the autonaming behavior of pyro.nn.to_pyro_module_(), so just be aware you may need to tweak some names.

Thank you!

What’s the best way to change names from PyroModule children parameters? Say if I wanted to add a prefix to all names?

PyroModule naming is automatic based on the module hierarchy. You can change prefixes but only in a very restrictive way, e.g. if you have a module

class Model(PyroModule):
    def __init__(self, ...):
        super().__init__()
        self.encoder = MyEncoder()

then the parameters of MyEncoder will have the prefix encoder.. You can add another prefix by changing the hierarchy

class Model(PyroModule):
    def __init__(self, ...):
        super().__init__()
        self.foo = PyroModule()
        self.foo.encoder = MyEncoder()

which will change the prefix to foo.encoder.

Probably the best way to understand this is to play around with different names and use poutine.trace() or .named_parameters() to inspect how things work. If you do find a clear explanation, we’d love help clarifying docs, say as a new section in the modules tutorial.

Thanks! Does AutoGuideList only work with AutoGuides? What if I have hand-crafted guides that I would like to combine? For example if I have a hierarchical VAE where each guide depend upon the sample in the previous layer in an amortized fashion?

Yes, AutoGuidList works with both AutoGuides and user-defined callable guides by wrapping callable guides in AutoCallable.