How to get parameters to be passed in torch optimizers?

Hi all.

I am new to pyro. I want to use PyTorch optimizers for custom SVI. I am following Custom SVI Objectives tutorial.

Now, I need to pass parameters to the Adam optimizer but I am not sure how to get parameters out of the guide.

An example of how I have set up my model and guide, is similar to LinearRegression tutorial -

class BayesianRegression(PyroModule):
    def __init__(self, in_features, out_features):
        super().__init__()
        self.linear = PyroModule[nn.Linear](in_features, out_features)
        self.linear.weight = PyroSample(dist.Normal(0., 1.).expand([out_features, in_features]).to_event(2))
        self.linear.bias = PyroSample(dist.Normal(0., 10.).expand([out_features]).to_event(1))

    def forward(self, x, y=None):
        sigma = pyro.sample("sigma", dist.Uniform(0., 10.))
        mean = self.linear(x).squeeze(-1)
        with pyro.plate("data", x.shape[0]):
            obs = pyro.sample("obs", dist.Normal(mean, sigma), obs=y)
        return mean

model = BayesianRegression(3, 1)
guide = AutoDiagonalNormal(model)
# How to get guide parameters?

Thanks

Hi @sayam049, For autoguides or other PyroModule guides you should be able to simply use the .parameters() method:

model = ...  # a PyroModule
guide = AutoDiagonalNormal(model)
parameters = list(guide.parameters())
1 Like

Thanks @fritzo
I tried this approach earlier as well, and that returned me an empty list. Here is a reproducible snippet

In [1]: import pyro

In [2]: pyro.__version__
Out[2]: '1.6.0'

In [3]: from pyro.nn import PyroModule, PyroSample

In [4]: import pyro.distributions as dist

In [5]: import torch.nn as nn

In [6]: class BayesianRegression(PyroModule):
   ...:     def __init__(self, in_features, out_features):
   ...:         super().__init__()
   ...:         self.linear = PyroModule[nn.Linear](in_features, out_features)
   ...:         self.linear.weight = PyroSample(dist.Normal(0., 1.).expand([out_features, in_features]).to_event(2))
   ...:         self.linear.bias = PyroSample(dist.Normal(0., 10.).expand([out_features]).to_event(1))
   ...: 
   ...:     def forward(self, x, y=None):
   ...:         sigma = pyro.sample("sigma", dist.Uniform(0., 10.))
   ...:         mean = self.linear(x).squeeze(-1)
   ...:         with pyro.plate("data", x.shape[0]):
   ...:             obs = pyro.sample("obs", dist.Normal(mean, sigma), obs=y)
   ...:         return mean
   ...: 

In [7]: from pyro.infer.autoguide import AutoDiagonalNormal
   ...: 
   ...: model = BayesianRegression(3, 1)
   ...: guide = AutoDiagonalNormal(model)

In [8]: 

In [8]: list(guide.parameters())
Out[8]: []

Do I need to manually call a svi.step() to populate the guide parameters? grab those parameters out for torch Adam, and restart the training loop?

Hi @sayam049, yes, you’ll need to initialize the guide. You can do this by calling

guide(*args, *kwargs)

with the same args,kwargs that you’d pass into svi.step(-), e.g. guide(x, y).

1 Like

Thanks for a reply @fritzo
The guide parameters’ shape is just the flattened model parameters’ shape concatenated together. And then why does it need to see (*args, **kwargs) to get its parameters initialized?

And, Is the guide parameters’ initialization dependent on the data we pass in?

why does [an autoguide instance] need to see (*args, **kwargs) to get its parameters initialized?

Autoguides determine the latent shapes of their models by running and tracing the model once. To trace the model, the guide needs to pass the model proper *args,**kwargs.

Is the guide parameters’ initialization dependent on the data we pass in?

For many models, autoguide parameter initialization does not depend on the data passed into the initial guide call. However parameters can depend on initial data if you use a parameter-dependent initialization strategy such as init_to_median or init_to_sample, and if the parameters of your model depend on input data (e.g. if input data are covariates or features describing distribution parameters).

Note autoguides require full data to be passed in, and do not support amortization, i.e. passing in different minibatches each learning step. To support subsampling / amortization you should either write a custom guide or use an EasyGuide.

2 Likes

Thanks for your reply @fritzo.

I experimented with EasyGuide and AutoNormal for a small Linear Regression problem. But I find that AutoNormal seems to work on minibatched data.
Here is the experiments notebook. I am not sure if AutoNormal follows the behaviour you said above. Can you have a look?

Thanks

Hi @sayam049, I cannot read your linked notebook (error 503), but if you’d like to use AutoNormal with subsampling, I believe you’ll need to pass in a create_plates argument that returns a plate with the subsample argument set just as you do in the model; that way the guide and model plates are properly aligned. We should definitely have better docs on this (feel free to submit a PR)… here’s an example

def model(full_data, subsample_indices):
    with pyro.plate("my_plate", len(full_data), subsample=subsample):
        batch = full_data[subsample]
        ...

def create_plates(full_data, subsample_indices):  # same args as model
    # this line is copied from the model
    return pyro.plate("my_plate", len(full_data), subsample=subsample)

guide = AutoNormal(model, create_plates=create_plates)

Here is the updated link autoguide_experiments.ipynb · GitHub.

I’ll check it out and get back on this. Many thanks @fritzo

1 Like

Hi @sayam049, thanks for linking your notebook. It looks like in your case you don’t need create_plates because there are no latent variables inside the plate. A vanilla AutoNormal(model) should work fine.

Many thanks @fritzo for your reply. :slight_smile: