RuntimeError: The size of tensor a (3) must match the size of tensor b (32) at non-singleton dimension 1

Question

How to fix the given issue?

.local/lib/python3.8/site-packages/torch/nn/modules/loss.py:535: UserWarning: Using a target size (torch.Size([32, 2])) that is different to the input size (torch.Size([1, 3, 2])). This will likely lead to incorrect results due to broadcasting. Please ensure they have the same size. return F.mse_loss(input, target, reduction=self.reduction) Traceback (most recent call last): File "/home/etri/Documents/new_pred_data/model_LSTM.py", line 151, in trainer.fit(model, train_loader, val_loader) File "/home/etri/.local/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 543, in fit call._call_and_handle_interrupt( File "/home/etri/.local/lib/python3.8/site-packages/pytorch_lightning/trainer/call.py", line 44, in _call_and_handle_interrupt return trainer_fn(*args, **kwargs) File "/home/etri/.local/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 579, in _fit_impl self._run(model, ckpt_path=ckpt_path) File "/home/etri/.local/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 986, in _run results = self._run_stage() File "/home/etri/.local/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1030, in _run_stage self._run_sanity_check() File "/home/etri/.local/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1059, in _run_sanity_check val_loop.run() File "/home/etri/.local/lib/python3.8/site-packages/pytorch_lightning/loops/utilities.py", line 182, in _decorator return loop_run(self, *args, **kwargs) File "/home/etri/.local/lib/python3.8/site-packages/pytorch_lightning/loops/evaluation_loop.py", line 135, in run self._evaluation_step(batch, batch_idx, dataloader_idx, dataloader_iter) File "/home/etri/.local/lib/python3.8/site-packages/pytorch_lightning/loops/evaluation_loop.py", line 396, in _evaluation_step output = call._call_strategy_hook(trainer, hook_name, *step_args) File "/home/etri/.local/lib/python3.8/site-packages/pytorch_lightning/trainer/call.py", line 309, in _call_strategy_hook output = fn(*args, **kwargs) File "/home/etri/.local/lib/python3.8/site-packages/pytorch_lightning/strategies/strategy.py", line 412, in validation_step return self.lightning_module.validation_step(*args, **kwargs) File "/home/etri/Documents/new_pred_data/model_LSTM.py", line 79, in validation_step val_loss = nn.MSELoss()(y_hat, y) File "/home/etri/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl return self._call_impl(*args, **kwargs) File "/home/etri/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl return forward_call(*args, **kwargs) File "/home/etri/.local/lib/python3.8/site-packages/torch/nn/modules/loss.py", line 535, in forward return F.mse_loss(input, target, reduction=self.reduction) File "/home/etri/.local/lib/python3.8/site-packages/torch/nn/functional.py", line 3338, in mse_loss expanded_input, expanded_target = torch.broadcast_tensors(input, target) File "/home/etri/.local/lib/python3.8/site-packages/torch/functional.py", line 76, in broadcast_tensors return _VF.broadcast_tensors(tensors) # type: ignore[attr-defined] RuntimeError: The size of tensor a (3) must match the size of tensor b (32) at non-singleton dimension 1

import torch
import torch.nn as nn
import pytorch_lightning as pl
from torch.utils.data import Dataset, DataLoader
import pandas as pd
import numpy as np
import os
import pickle
from tqdm import tqdm

class CustomDataset(Dataset):
    def __init__(self, data):
        self.data = data

    def __len__(self):
        return len(self.data) - 1  # Adjust length to handle future position

    def __getitem__(self, idx):
        x_position = torch.tensor(self.data['position_x'][idx]).float()
        y_position = torch.tensor(self.data['position_y'][idx]).float()
        agent_id = torch.tensor(int(self.data['ids'][idx])).float()

        x = torch.cat(
            (x_position.unsqueeze(0), y_position.unsqueeze(0), agent_id.unsqueeze(0)),
            dim=0)

        if idx + 1 < len(self.data):
            future_x_position = torch.tensor(self.data['position_x'][idx + 1]).float()
            future_y_position = torch.tensor(self.data['position_y'][idx + 1]).float()
            future_position = torch.cat(
                (future_x_position.unsqueeze(0), future_y_position.unsqueeze(0)), dim=0)
        else:
            # Handle the case when future position is not available
            future_position = torch.zeros_like(x)

        return x, future_position


class My_Baseline(pl.LightningModule):
    def __init__(self, input_dim, hidden_dim, output_dim, seq_len, num_layers=1):
        super(My_Baseline, self).__init__()
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.output_dim = output_dim
        self.seq_len = seq_len
        self.num_layers = num_layers

        self.encoder = nn.LSTM(input_dim, hidden_dim, num_layers=num_layers)
        self.decoder = nn.LSTM(hidden_dim, hidden_dim, num_layers=num_layers)
        self.fc = nn.Linear(hidden_dim, self.output_dim)
        self.activation = nn.ReLU()

    def forward(self, x):
        _, (hidden, cell) = self.encoder(x.to(self.device))
        hidden = hidden.repeat(self.num_layers, x.size(1), 1)
        cell = cell.repeat(self.num_layers, x.size(1), 1)

        decoder_input = torch.zeros(self.seq_len, x.size(1), self.hidden_dim, device=hidden.device)

        outputs = []
        for _ in range(self.seq_len):
            decoder_output, (hidden, cell) = self.decoder(decoder_input, (hidden, cell))
            output = self.fc(decoder_output)
            outputs.append(output)
            decoder_input = output
        outputs = torch.cat(outputs, dim=0)
        return outputs

    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = nn.MSELoss()(y_hat, y)
        self.log('train_loss', loss)
        return loss

    def validation_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        val_loss = nn.MSELoss()(y_hat, y)
        self.log('val_loss', val_loss, prog_bar=True)
        return val_loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=1e-3)

train_data = None
train_data_list= []
directory = './data/train/processed'
for filename in os.listdir(directory):
    if filename.endswith('.pkl'):
        filepath = os.path.join(directory, filename)
        with open(filepath, 'rb') as f:
            scene = pickle.load(f)
            agent_dict = scene['agent']
            agent_data = agent_dict
            positions = agent_data['position']
            ids = agent_data['id']
            positions_np = positions.numpy()
            positions_np_2d = positions_np.reshape(-1, 3)
            ids_repeated = np.repeat(ids, positions_np_2d.shape[0] // len(ids))
            x_values = positions_np_2d[:, 0]
            y_values = positions_np_2d[:, 1]
            min_length = min(len(ids_repeated), len(x_values), len(y_values))
            ids_repeated = ids_repeated[:min_length]
            x_values = x_values[:min_length]
            y_values = y_values[:min_length]
            df = pd.DataFrame({'ids': ids_repeated, 'position_x': x_values, 'position_y': y_values})
            train_data_list.append(df)

train_data = pd.concat(train_data_list, ignore_index=True)
seq_len = 1
input_dim = 3
hidden_dim = 64
output_dim = 2
num_layers = 1

train_dataset = CustomDataset(train_data)
train_loader = DataLoader(train_dataset, batch_size=32, num_workers=4, shuffle=True)

model = My_Baseline(input_dim, hidden_dim, output_dim, seq_len, num_layers)
trainer = pl.Trainer(max_epochs=3, accelerator="auto")

val_data = []
directory = './data/val/processed'

for filename in os.listdir(directory):
    if filename.endswith('.pkl'):
        filepath = os.path.join(directory, filename)
        with open(filepath, 'rb') as f:
            scene = pickle.load(f)
            agent_dict = scene['agent']
            agent_data = agent_dict
            positions = agent_data['position']
            ids = agent_data['id']
            positions_np = positions.numpy()
            positions_np_2d = positions_np.reshape(-1, 3)
            ids_repeated = np.repeat(ids, positions_np_2d.shape[0] // len(ids))
            x_values = positions_np_2d[:, 0]
            y_values = positions_np_2d[:, 1]
            min_length = min(len(ids_repeated), len(x_values), len(y_values))
            ids_repeated = ids_repeated[:min_length]
            x_values = x_values[:min_length]
            y_values = y_values[:min_length]
            df = pd.DataFrame({'ids': ids_repeated, 'position_x': x_values, 'position_y': y_values})
            val_data.append(df)

val_data = pd.concat(val_data, ignore_index=True)
val_dataset = CustomDataset(val_data)
val_loader = DataLoader(val_dataset, batch_size=32, num_workers=4, )

trainer.fit(model, train_loader, val_loader)

model.eval()
predictions = []
ground_truth = []

with torch.no_grad():
    for batch in val_loader:
        x, y = batch
        prediction = model(x)
        predictions.append(prediction)
        ground_truth.append(y)

predictions = torch.cat(predictions).cpu().numpy()
ground_truth = torch.cat(ground_truth).cpu().numpy()

print("Predictions:")
print(predictions)
print("\nGround Truth:")
print(ground_truth)

Answer 1

You have a problem in your model, your predicted tensor is of shape (1,3,2) while your input has a batch size of 32, according to your sample code. You need to set batch_first=True in your nn.LSTM initializations for those layers to process the input batch-first. Having the default batch_first=False is not wrong, but considering the linear layers after it, this makes the model incorrect.