Probability at each timestep

I’m using the MNIST dataset for image classification as done here https://medium.com/dair-ai/building-rnns-is-fun-with-pytorch-and-google-colab-3903ea9a3a79 (in the “RNN for image classification” bit).
The data has shape 28x28, which means that i have 28 timesteps (rows in the image) and 28 inputs (number of columns) . The data goes into the rnn row by row.
I would like my code to print the outputs after each timestep i.e. I would like the code to give me the probabilities after each row goes into the rnn.

Can this be done?

You could pass the time steps one by one instead of the complete batch in these lines of code:

lstm_out, self.hidden = self.basic_rnn(X, self.hidden) # apply loop here
out = self.FC(self.hidden)

Note that you could pass each row to the linear layer to get the logits, but this would change your model architecture, as currently only the hidden state of the last time step is used.

1 Like