Mismatching dims in GRU for classification

Question

I'm trying to complete a task and write simple RNN. Here's the class:

class RNNBaseline(nn.Module):
def __init__(self, vocab_size, embedding_dim, hidden_dim, output_dim, n_layers, 
             bidirectional, dropout, pad_idx):
    
    super().__init__()
    
    self.embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx = pad_idx)
    
    self.rnn = nn.GRU(input_size=embedding_dim, hidden_size=hidden_dim) #RNN(embedding_dim, hidden_dim)
    
    self.fc = nn.Linear(hidden_dim, output_dim)  # YOUR CODE GOES HERE

    self.dropout = nn.Dropout(dropout)
    
    
def forward(self, text, text_lengths, hidden = None):
    
    #text = [sent len, batch size]
    
    embedded = self.embedding(text)
    
    #embedded = [sent len, batch size, emb dim]
    
    #pack sequence
    packed_embedded = nn.utils.rnn.pack_padded_sequence(embedded, text_lengths)
    
    # cell arg for LSTM, remove for GRU
    # packed_output, (hidden, cell) = self.rnn(packed_embedded)
    # unpack sequence
    # output, output_lengths = nn.utils.rnn.pad_packed_sequence(packed_output)  

    #output = [sent len, batch size, hid dim * num directions]
    #output over padding tokens are zero tensors
    
    #hidden = [num layers * num directions, batch size, hid dim]
    #cell = [num layers * num directions, batch size, hid dim]
    
    #concat the final forward (hidden[-2,:,:]) and backward (hidden[-1,:,:]) hidden layers
    #and apply dropout
    output, hidden = self.rnn(packed_embedded, hidden)
    #hidden = None  # concatenate
            
    #hidden = [batch size, hid dim * num directions] or [batch_size, hid dim * num directions]
        
    return self.fc(hidden)

For now I'm not using LSTM or trying to do bidirectional RNN, I just want simple GRU to train without errors. This is the training function:

import numpy as np

min_loss = np.inf

cur_patience = 0

for epoch in range(1, max_epochs + 1):
train_loss = 0.0
model.train()
pbar = tqdm(enumerate(train_iter), total=len(train_iter), leave=False)
pbar.set_description(f"Epoch {epoch}")
for it, ((text, txt_len), label) in pbar: 
    #YOUR CODE GOES HERE
    opt.zero_grad()
    input = text.to(device)
    labels = label.to(device)
    output = model(input, txt_len.type(torch.int64).cpu())
    train_loss = loss_func(output, labels)
    train_loss.backward()
    opt.step()


train_loss /= len(train_iter)
val_loss = 0.0
model.eval()
pbar = tqdm(enumerate(valid_iter), total=len(valid_iter), leave=False)
pbar.set_description(f"Epoch {epoch}")
for it, ((text, txt_len), label) in pbar:
    # YOUR CODE GOES HERE
    input = text.to(device)
    labels = label.to(device)
    output = model(input, txt_len.type(torch.int64).cpu())
    val_loss = loss_func(output, labels)

val_loss /= len(valid_iter)
if val_loss < min_loss:
    min_loss = val_loss
    best_model = model.state_dict()
else:
    cur_patience += 1
    if cur_patience == patience:
        cur_patience = 0
        break

print('Epoch: {}, Training Loss: {}, Validation Loss: {}'.format(epoch, train_loss, val_loss))
model.load_state_dict(best_model)

And some variables:

vocab_size = len(TEXT.vocab)
emb_dim = 100
hidden_dim = 256
output_dim = 1
n_layers = 2
bidirectional = False
dropout = 0.2
PAD_IDX = TEXT.vocab.stoi[TEXT.pad_token]
patience=3

opt = torch.optim.Adam(model.parameters())
loss_func = nn.BCEWithLogitsLoss()

max_epochs = 1

But I get this error:

ValueError: Target size (torch.Size([64])) must be the same as input size (torch.Size([1, 64, 1]))

... in this line:

---> 18 train_loss = loss_func(output, labels)

What am I doing wrong?

Answer 1

nn.BCEWithLogitsLoss expects both outputs and targets (or in your case labels) to be of size [b,d] where b is the batch size and d is the number of classes (or dimension of whatever you are trying to predict). Currently, your outputs are of size [b,d,1] and your targets are of size [d]. Two fixes are necessary, and both are very simple:

1. Add a batch dimension to your targets (labels). This is a common error when using a dataset that returns data elements because it generally does not add a batch dimension. Encapsulating your dataset class within a pytorch dataloader, but if you don't want to do this simply add an unsqueeze() operation. Note that the unsqueeze operation only works with a batch size of 1, otherwise using dataloader is probably a better bet.

2. Your output has an empty 3rd dimension, which can easily be flattened with a squeeze() operation. Both unsqueeze and squeeze are differentiable so shouldn't present problems for backpropagation.

... code before here

for it, ((text, txt_len), label) in pbar:
    # YOUR CODE GOES HERE
    input = text.to(device)
    labels = label.to(device).unsqueeze(0)                 # added unsqueeze operation
    output = model(input, txt_len.type(torch.int64).cpu())
    output = output.squeeze(-1)                            # added squeeze on last dim
    val_loss = loss_func(output, labels)

... code after here