Understanding tensor sizes in Dataloader

badger: <p>I have workarounds but I suspect there is something fundamental I am missing.</p>
<p>If I build a simple dataloader using a pandas array as input, I can never get the dimensions quite right–I always have to squeeze and unsqueeze for loss functions.  Example:</p>
<pre><code class="lang-auto">Class CustomSimpleDataset(Dataset):
    def __init__(self, featureDataFrame, targetDataFrame):
        self.features = featureDataFrame
        self.targets = targetDataFrame
    def __len__(self):
        return len(self.features)
    def __getitem__(self, idx):
        feature = self.features.iloc[idx]   #select a row from the dataframe
        feature = torch.tensor(feature, dtype=torch.float32) #turn that row into a tensor
        target = self.targets.iloc[idx]
        target = torch.tensor(target, dtype=torch.long) 
        return feature, target  

test_dataset = CustomSimpleDataset(X_test, y_test)
test_dataloader = DataLoader(test_dataset, batch_size=1000, shuffle=False)

loss=nn.BCEWithLogitsLoss()
loss(my_lightning_model.forward(i),j.unsqueeze(dim=1))  #requires either a squeeze or unsqueeze!!!
</code></pre>

ptrblck: <p>Could you post a minimal and executable code snippet reproducing the issue using random data?</p>

badger: <pre><code class="lang-auto">import torch
from torch.utils.data import Dataset
from torchvision import datasets
from torchvision.transforms import ToTensor
import matplotlib.pyplot as plt
from sklearn.datasets import fetch_california_housing
from torch.utils.data import DataLoader
from torch import nn
import os
import pandas as pd

class CustomDataset(Dataset):
    def __init__(self, featureDataFrame, targetDataFrame):
        self.features = featureDataFrame
        self.targets = targetDataFrame
    def __len__(self):
        return len(self.targets)
    def __getitem__(self, idx):
        feature = self.features.iloc[idx]   #select a row from the dataframe
        feature = torch.tensor(feature, dtype=torch.float32) #turn that row into a tensor
        target = self.targets.iloc[idx]
        target = torch.tensor(target,dtype=torch.long)
        return feature, target  #you can return more, but generally return your features and targets.

data = fetch_california_housing(as_frame=True)
X = data['data']
y= data['target']
test_dataset = CustomDataset(X, y)
test_dataloader = DataLoader(test_dataset, batch_size=10, num_workers=10, shuffle=False)
for x,y in test_dataloader:
    print(x.shape,x.dtype)
    print(y.shape,y.dtype)
    break

class SimpleNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(8,10),
            nn.ReLU(),
            nn.Linear(10, 1),
        )
    def forward(self, x):
        logits = self.linear_relu_stack(x)
        return logits
        
model = SimpleNetwork()

loss = nn.MSELoss()
for x,y in test_dataloader:
    p = model(x)
    print(p.dtype)
    print(y.dtype)
    print(p.shape)
    print(y.shape)
    print("This gives a warning")
    print(loss(p,y)) # gives a warning
    print("This does not!")
    print(loss(p,y.unsqueeze(dim=1))) #no warning
    break
</code></pre>
<p>Classification Example that fails:</p>
<pre><code class="lang-auto">class CustomDataset(Dataset):
    def __init__(self, featureDataFrame, targetDataFrame):
        self.features = featureDataFrame
        self.targets = targetDataFrame
    def __len__(self):
        return len(self.targets)
    def __getitem__(self, idx):
        feature = self.features.iloc[idx]   #select a row from the dataframe
        feature = torch.tensor(feature, dtype=torch.float32) #turn that row into a tensor
        target = self.targets.iloc[idx]
        target = torch.tensor(target,dtype=torch.float)
        return feature, target  #you can return more, but generally return your features and targets.

data = load_breast_cancer(as_frame=True)
X = data['data']
y= data['target']
test_dataset = CustomDataset(X, y)
test_dataloader = DataLoader(test_dataset, batch_size=10, num_workers=10, shuffle=False)
for x,y in test_dataloader:
    print(x.shape,x.dtype)
    print(y.shape,y.dtype)
    break

class SimpleNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(30,10),
            nn.ReLU(),
            nn.Linear(10, 1),
        )
    def forward(self, x):
        logits = self.linear_relu_stack(x)
        return logits
        
model = SimpleNetwork()

loss = nn.BCEWithLogitsLoss()
for x,y in test_dataloader:
    p = model(x)
    print(p.dtype)
    print(y.dtype)
    print(p.shape)
    print(y.shape)
    try:
        print(loss(p,y)) # gives a warning
    except ValueError as e:
        print("This doesn't work, wrong dims!")
    try:
        print("Unsqueeze works fine")
        print(loss(p,y.unsqueeze(dim=1))) #no warning
    except Error as e:
        pass
    break
</code></pre>
<p>I realize pandas throws all those warnings about indicies, but our grading sheets regularly come in csv which lend themselves to using pandas to transform data into tensors.</p>

badger: <p>Also note that one gets different (but closely related) answers in the Regression/MSE example:<br>
This gives a warning<br>
tensor(4854.9785, grad_fn=)<br>
This does not!<br>
tensor(4872.5762, grad_fn=)</p>

ptrblck: <pre data-code-wrap="python"><code class="lang-python">&gt;&gt;&gt; test_dataloader = DataLoader(test_dataset, batch_size=10, num_workers=10, shuffle=False)
&gt;&gt;&gt; for x,y in test_dataloader:
...     print(x.shape,x.dtype)
...     print(y.shape,y.dtype)
...     break
... 
torch.Size([10, 30]) torch.float32
torch.Size([10]) torch.float32
</code></pre>
<p>returns the expected shapes with a <code>batch_size</code> of <code>10</code>.</p>
<p>The warning is raised in the loss function as <code>nn.BCEWithLogitsLoss</code> expects the model output and target to have the same shape as described in the <a href="https://pytorch.org/docs/stable/generated/torch.nn.BCEWithLogitsLoss.html#torch.nn.BCEWithLogitsLoss">docs</a>, so you should either <code>unsqueeze</code> <code>dim1</code> inside the dataset or the training loop as already done.<br>
Keep in mind that <code>BCEWithLogitsLoss</code> is used for binary or multi-label classification use cases, where the latter use case allows you to classify each sample into zero, one, or multiple classes. Since your model output has the shape <code>[batch_size, 1]</code>, it seems you are working on a binary classification so would thus need to <code>unsqueeze</code> the missing “class dimension”.</p>

badger: <p>Fair enough, my question was why the data loader was giving the wrong matrix size ([10] vs [10,1]) for these trivial examples.  Of course they are basically the same (but not quite) and that’s why i thought there was some key detail I was missing.  I realize the answer is to squeeze/unsqueeze, I was under the impression they should match without matrix manipulation (for this case).</p>

ptrblck: <p>The indexing inside the <code>__getitem__</code> creates scalar tensors for the target (0-dim tensors) and the <code>DataLoader</code> will then batch them and add the missing batch dimension. You could try to use a slicing approach in pandas to create actual arrays but calling <code>unsqueeze</code> might actually be easier.</p>