Converting pytorch LSTM to Pyro using to_pyro_module_

1

Hi! I am experimenting with Pyro LSTM model. I have used MNIST dataset to construct a LSTM classifier in pytorch. Then used to_pyro_module_(model) from documentation to convert to pyro model and used pyro autoguide AutoNormal for guide. But on training I am getting error “forward() takes 2 positional arguments but 3 were given”. So how do we give the input sequence?

#Convert to pyro
to_pyro_module_(model)
for name, module in model.named_parameters():
    module = PyroSample(Normal(
        loc=module.detach(), scale=1
    ).to_event(module.dim()))

#guide
optim = Adam({"lr": 0.01})
guide = AutoNormal(model)
svi = SVI(model, guide, optim, loss=Trace_ELBO())

# training
for j in range(num_iterations):
    loss = 0
    for i, (x, y) in enumerate(train_loader):
        x = x.view(-1, 28, 28).requires_grad_()
        # calculate the loss and take a gradient step
        loss += svi.step(x,y)
        normalizer_train = len(train_loader.dataset)
        total_epoch_loss_train = loss / normalizer_train

Error at: loss += svi.step(x,y)
forward() takes 2 positional arguments but 3 were given

2 Likes
2

Hi @Sree, can you ensure that your model indeed takes two arguments?

# Before converting to a PyroModule
for x, y in train_loader:
    model(x, y)  # just a smoke test
    break

#Convert to pyro
to_pyro_module_(model)
...
3

@Sree you might also be interested in TyXe, a higher-level library for Bayesian neural networks in Pyro that automates things like initializing guide and prior scales, loss construction, and local reparametrization.

1 Like
4

Thank you for your response. My model takes in only x but for inference observed data is also required right. I am confused in the conversion.

class LSTMModel(nn.Module):
    def __init__(self, input_dim, hidden_dim, layer_dim, output_dim):
        super(LSTMModel, self).__init__()
        self.hidden_dim = hidden_dim
        self.layer_dim = layer_dim
        self.lstm = nn.LSTM(input_dim, hidden_dim, layer_dim, batch_first=True)

        # Readout layer
        self.fc = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        # Initialize hidden state with zeros
        h0 = torch.zeros(self.layer_dim, x.size(0), self.hidden_dim).requires_grad_()

        # Initialize cell state
        c0 = torch.zeros(self.layer_dim, x.size(0), self.hidden_dim).requires_grad_()
     
        out, (hn, cn) = self.lstm(x, (h0.detach(), c0.detach()))

        # out[:, -1, :]  last time step hidden states
        out = self.fc(out[:, -1, :]) 
        # out.size() --> 100, 10
        return out

After I convert using to_pyro_module_ ,along with priors for model parameters ,should I also define distribution for observed data (y, categorical in this case ) using pyro.sample?