Batch Sampler does not seem to work

username: PhysicsIsFun
cooked: <p>Greetings,</p>
<p>I would like to do experiments with varying batch sizes during model training. I am working with simple feed forward networks at the moment.</p>
<p>I came across the sampler class that one can pass to the Dataloader. I tried to use one of the examples I found online (code below), but when trying to compute the loss (BCE loss), I get the error</p>
<p><code>ValueError: Using a target size (torch.Size([1, 50])) that is different to the input size (torch.Size([50])) is deprecated. Please ensure they have the same size.</code></p>
<p>The same code, however, worked for the usual Dataloaders with constant batch size. Something about the dimensions is going wrong with my sampler. Would appreciate any help!</p>
<pre data-code-wrap="python"><code class="lang-python">## Sampler Class
class VariableBatchSampler(torch.utils.data.Sampler):
    def __init__(self, dataset_len: int, batch_sizes: list):
        self.dataset_len = dataset_len
        self.batch_sizes = batch_sizes
        self.batch_idx = 0
        self.start_idx = 0
        self.end_idx = self.batch_sizes[self.batch_idx]
        
    def __iter__(self):
        return self
       
    def __next__(self):
        if self.start_idx &gt;= self.dataset_len:
            raise StopIteration()
 
        batch_indices = torch.arange(self.start_idx, self.end_idx, dtype=torch.int)
        self.start_idx += (self.end_idx - self.start_idx)
        self.batch_idx += 1

        try:
            self.end_idx += self.batch_sizes[self.batch_idx]
        except IndexError:
            self.end_idx = self.dataset_len
             
        return batch_indices


## Initizalization in main
train_set = torch.utils.data.TensorDataset(Xtrain, ytrain)

batch_sampler = VariableBatchSampler(len(Xtrain), [50,500])
train_loader = torch.utils.data.DataLoader(train_set, sampler=batch_sampler) # allows iterating
</code></pre>
<p>The error is raised in the training loop:</p>
<pre data-code-wrap="python"><code class="lang-python"> for data, target in data_loader:
                data, target = data.to(device), target.to(device)
                output = self.forward(data).squeeze()
                loss += criterion(output, target).data.item()   
                ## and so on
</code></pre>

username: ptrblck
cooked: <p>Based on the error message it seems the shapes of your model output and targets do not match and thus the loss calculation is failing. I don’t think the issue is necessarily related to the <code>BatchSampler</code> and you can <code>unsqueeze()</code> the output or <code>permute</code> the target depending what the actual batch size and feature dimension is.<br>
I would assume you are using <code>50</code> samples and are working on a binary classification.<br>
If so, both the model output and target should have the shape <code>[50, 1]</code>.</p>

username: PhysicsIsFun
cooked: <p>Thanks ptrblck!</p>
<p>It is indeed binary classification and the batch size is 50 in the example.</p>
<p>Changing the line<br>
`python</p>
<pre><code class="lang-auto">output = self.forward(data).squeeze()
</code></pre>
<p>to</p>
<pre><code class="lang-auto">output = self.forward(data).reshape(target.shape) 
</code></pre>
<p>solves the issue and also works in the case with constant batch size.</p>
<p>I do get another error with the sampler, though:</p>
<pre><code class="lang-auto">
  File ~/path_to_file/code.py:582 in train_func
    (data, target) = next(iter(self.train_loader))          # compute initial gradients

  File ~/miniconda3/envs/PYTORCH/lib/python3.11/site-packages/torch/utils/data/dataloader.py:631 in __next__
    data = self._next_data()

  File ~/miniconda3/envs/PYTORCH/lib/python3.11/site-packages/torch/utils/data/dataloader.py:674 in _next_data
    index = self._next_index()  # may raise StopIteration

  File ~/miniconda3/envs/PYTORCH/lib/python3.11/site-packages/torch/utils/data/dataloader.py:621 in _next_index
    return next(self._sampler_iter)  # may raise StopIteration
</code></pre>
<p>What I try to do in <code>train_func</code> is to obtain the first gradient manually. I need it to initialize a help variable. Once again, it works for normal dataloaders, but not when passing the sampler argument.</p>

username: ptrblck
cooked: <p>Could you post the error message you are seeing using your <code>BatchSampler</code>?</p>

username: PhysicsIsFun
cooked: <p>I created a small example. I think I understand something wrong about what the sampler does.</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
import torch.nn as nn


# for variable batch sizes
class VariableBatchSampler(torch.utils.data.Sampler):
    def __init__(self, dataset_len: int, batch_sizes: list):
        self.dataset_len = dataset_len
        self.batch_sizes = batch_sizes
        self.batch_idx = 0
        self.start_idx = 0
        self.end_idx = self.batch_sizes[self.batch_idx]
        
    def __iter__(self):
        return self
       
    def __next__(self):
        if self.start_idx &gt;= self.dataset_len:
            raise StopIteration()
 
        batch_indices = torch.arange(self.start_idx, self.end_idx, dtype=torch.int)
        self.start_idx += (self.end_idx - self.start_idx)
        self.batch_idx += 1

        try:
            self.end_idx += self.batch_sizes[self.batch_idx]
        except IndexError:
            self.end_idx = self.dataset_len
             
        return batch_indices


## define some random data set
Npoints = 500
Xtrain = torch.normal(mean=torch.zeros((Npoints,2)))
Ytrain = torch.zeros(Npoints)
Ytrain[0:Npoints//2] = 1

train_set = torch.utils.data.TensorDataset(Xtrain, Ytrain) # turns it to TensorDataset


batch_sampler = VariableBatchSampler(len(Xtrain), [50,100])
train_loader = torch.utils.data.DataLoader(train_set, sampler=batch_sampler) 

(data, target) = next(iter(train_loader))  
print(data.size())                              # prints torch.Size([1, 50, 2])


(data, target) = next(iter(train_loader))
print(data.size())                              # prints torch.Size([1, 100, 2])


(data, target) = next(iter(train_loader)) 
print(data.size())                              # prints torch.Size([1, 350, 2]), why?


(data, target) = next(iter(train_loader)) 
print(data.size())                              # raises stop iteration error
</code></pre>
<p>I expected it to give me batches of size 50, 100, 50, 100 and so on. Yet one of the batches is of size 350. Also, if the end of a the dataset is reached, asking for a new batch leads to the error. I would like it to simply keep giving batches as with the dataloader without the sampler.</p>