How to write a guide function for neural networks with embedding layers?

Hi!

I’m a beginner to Pyro and I want to build a neural network with embedding layers to simulate a collaborative filter system. My model has two embedding layers, which take two categorical variables as inputs and convert them to embedding vectors. Then the two vectors will be concatenated to be fed to four fully connected layers. The output is a single number.

It’s kind of like a Bayesian regression. I’ve gone through the tutorial of Bayesian regression and finished the model. But I got stuck at the step of writing my guide. So how to write a guide function for a neural network with two embedding layers and four FC layers?

Thank you!

Here’s my model:

class EmbeddingNet(PyroModule):
    def __init__(self, n_users, n_movies, n_factors, hidden, dropouts, embedding_dropout=0.02):
        super(EmbeddingNet, self).__init__()
        # embedding layers
        self.u = PyroModule[nn.Embedding](n_users, n_factors)

        self.m = PyroModule[nn.Embedding](n_movies, n_factors)

        self.drop = PyroModule[nn.Dropout](embedding_dropout)
        # activation
        self.act = nn.ELU()

        # FC layers
        lb = -0.05
        ub = 0.05

        self.fc1 = PyroModule[nn.Linear](n_factors * 2, hidden[0])
        self.fc1.weight = PyroSample(dist.Uniform(lb, ub).expand([hidden[0], n_factors * 2]).to_event(2))
        self.fc1.bias = PyroSample(dist.Uniform(lb, ub).expand([hidden[0]]).to_event(1))
        self.dp1 = PyroModule[nn.Dropout](dropouts[0])

        self.fc2 = PyroModule[nn.Linear](hidden[0], hidden[1])
        self.fc2.weight = PyroSample(dist.Uniform(lb, ub).expand([hidden[1], hidden[0]]).to_event(2))
        self.fc2.bias = PyroSample(dist.Uniform(lb, ub).expand([hidden[1]]).to_event(0))
        self.dp2 = PyroModule[nn.Dropout](dropouts[1])

        self.fc3 = PyroModule[nn.Linear](hidden[1], hidden[2])
        self.fc3.weight = PyroSample(dist.Uniform(lb, ub).expand([hidden[2], hidden[1]]).to_event(2))
        self.fc3.bias = PyroSample(dist.Uniform(lb, ub).expand([hidden[2]]).to_event(1))
        self.dp3 = PyroModule[nn.Dropout](dropouts[2])

        self.fc4 = PyroModule[nn.Linear](hidden[2], 1)
        self.fc4.weight = PyroSample(dist.Uniform(lb, ub).expand([1, hidden[2]]).to_event(2))
        self.fc4.bias = PyroSample(dist.Uniform(lb, ub).expand([1]).to_event(1))

    def forward(self, user, movie, ratings, minmax=None):
        # get the embedded features
        features = torch.cat([self.u(user), self.m(movie)], dim=1)
        x = self.drop(features)
        # pass features to fc layers
        x = self.dp1(self.act(self.fc1(x)))
        x = self.dp2(self.act(self.fc2(x)))
        x = self.dp3(self.act(self.fc3(x)))
        x_mean = torch.sigmoid(self.fc4(x)).squeeze(-1)

        x_sigma = pyro.sample('sigma', dist.Uniform(0., 1.0))
        if minmax is not None:
            min_rating, max_rating = minmax
            x_mean = x_mean*(max_rating - min_rating + 1) + min_rating - 0.5

        with pyro.plate('data', x.shape[0]):
            obs = pyro.sample('obs', dist.Normal(x_mean, x_sigma), obs=ratings)
        return x_mean

Hi @Yujian, I usually start with autoguides or easyguides before I write a custom guide. You might first ensure that an AutoDelta works (this is just MAP estimation), then you could try a completely mean field AutoDiagonalNormal guide, and then an AutoLowRankMultivariateNormal guide.