How to use a different test batch size for RNN in PyTorch?

Question

I want to train an RNN over 5 training points where each sequence also has a size of 5. At test time, I want to send in a single data point and compute the output.

The task is to predict the next character in a sequence of five characters (all encoded as 1-hot vectors). I have tried duplicating the test data point five times. However, I am sure that this is not the right way to solve this problem.

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

# Define the parameters
H = [ 1, 0, 0, 0 ]
E = [ 0, 1, 0, 0 ]
L = [ 0, 0, 1, 0 ]
O = [ 0, 0, 0, 1 ]

# Define the model
net = nn.RNN(input_size=4, hidden_size=4, batch_first=True)

# Generate data
data = [[H,E,L,L,O],
        [E,L,L,O,H],
        [L,L,O,H,E],
        [L,O,H,E,L],
        [O,H,E,L,L]]
inputs = torch.tensor(data).float()
hidden = torch.randn(1,5,4) # Random initialization
correct_outputs = torch.tensor(np.array(data[1:]+[data[0]]).astype(float).tolist(), requires_grad=True)

# Set the loss function
criterion = torch.nn.MSELoss()

# Set the optimizer
optimizer = torch.optim.SGD(net.parameters(), lr=0.1)

# Perform gradient descent until convergence
for epoch in range(1000):
    # Forward Propagation
    outputs, hidden = net(inputs, hidden)
    # Compute and print loss
    loss = criterion(nn.functional.softmax(outputs,2), correct_outputs)
    print('epoch: ', epoch,' loss: ', loss.item())
    # Zero the gradients
    optimizer.zero_grad()
    # Backpropagation
    loss.backward(retain_graph=True)
    # Parameter update
    optimizer.step()

# Predict
net(torch.tensor([[H,E,L,L,O]]).float(),hidden)

I get the following error:

RuntimeError: Expected hidden size (1, 1, 4), got (1, 5, 4)

I understand that torch wants a tensor of size (1,1,4) but I am not sure how I can convert the initial hidden state from (1, 5, 4) to (1, 1, 4). Any help would be highly appreciated!

Answer 1

You are getting the error because you are using:

hidden = torch.randn(1,5,4) # Random initialization

Instead, you should use:

hidden = torch.randn(1,inputs.size(0),4) # Random initialization

to cope up with the batch size of the inputs. So, do the following:

# Predict
inputs = torch.tensor([[H,E,L,L,O]]).float()
hidden = torch.randn(1,inputs.size(0),4)
net(inputs, hidden)

Suggestion: improve your coding style by following some good examples in PyTorch.

Answer 2

Another option would be to just remove the keyword argument, batch_first=True when you define the model.

# Define the model
net = nn.RNN(input_size=4, hidden_size=4)