Converting pytorch LSTM to Pyro using to_pyro_module_

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
for name, module in model.named_parameters():
    module = PyroSample(Normal(
        loc=module.detach(), scale=1

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

1 Like

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

#Convert to pyro

@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.

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?