Dataset Multiple Samples per getitem Call

callahman: <p>I have a custom Dataset I’m trying to build out. The actual details of my Dataset are below, but for now I’m going to focus on the following example code.</p>
<p>The goal is to load some data into __getitem__() and segment the array into several samples which I can then stack and output with the batch.</p>
<pre><code class="lang-auto">from torch.utils.data import Dataset, DataLoader
import torch
import numpy as np

class Example_DS(Dataset):

    def __init__(self, data):
        self.data = data
        
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        
        record = self.data[idx]
        
        X1 = record[:5]
        X2 = record[5:]
        
        X1 = torch.from_numpy(X1)
        X2 = torch.from_numpy(X2)
        
        X = torch.stack([X1, X2])
        
        sample = {
            'sample' : X
            }
        
        return sample


########################################
data = [
        np.random.randint(0, 10, 10)
        ,np.random.randint(0, 10, 10)
        ,np.random.randint(0, 10, 10)
        ,np.random.randint(0, 10, 10)
        ,np.random.randint(0, 10, 10)
        ,np.random.randint(0, 10, 10)
        ]


ds = Example_DS(data)

dls = DataLoader(ds, batch_size = 2, shuffle = True, num_workers = 1)

for batch in dls:
    print(batch['sample'].shape)
</code></pre>
<p>As seen when running the example code the output tensors for each batch have a shape of [2, 2, 5]. I would like to have the shape be [4, 5] instead.</p>
<p>What is the best way to make this happen?</p>
<p><strong>More detail as promised:</strong><br>
I’m trying to work with waveform data. I could get multiple samples from 1 audio file, but without this method I would have to open each audio file N times for each of N samples. I would prefer to open the file once, slice the N samples, and concat/stack the resulting tensors along with the rest of the files/samples in the batch.</p>
<p><strong>Current:</strong></p>
<blockquote>
<p>loadtime = time_to_load_file * N_files * N_samples</p>
</blockquote>
<p><strong>Preferred:</strong></p>
<blockquote>
<p>loadtime = time_to_load_file * N_files</p>
</blockquote>
<p>Thanks in advance for any help you can provide!<br>
Will update if I find a solution.</p>

ptrblck: <p>Assuming the returned tensor contains the desired values, you could flatten it in the <code>DataLoader</code> loop via:</p>
<pre><code class="lang-python">
for batch in dls:
    sample = batch['sample']
    sample = sample.view(-1, 5)
    print(sample.shape) # should print [4, 5] now
</code></pre>
<p>Using a custom <code>collate_fn</code> would probably also work, but I think the <code>view</code> operation might be easier.</p>

callahman: <p>Thanks for the response!<br>
Probably not as elegant as the <code>collate_fn</code> idea, but it gets the job done!</p>
<p>Thanks again!</p>

heh: <p>I use Pytorch Lightning, and I have to feed in a dataloader() object into the Lightning Trainer, and thus you won’t get access to the DataLoader loop. That’s where the collate_fn comes in handy.</p>

datalw: <p>I am facing the same issue. I would like to know, how <code>collate_fn</code> works here. Could you give a code example?<br>
And if the batch size is 2, how can it be integrated into <code>collate_fn</code>, so that a batch won’t be length of 4 in the end after cutting each record in the middle?</p>
<p>Thanks!</p>

ptrblck: <p>The <code>collate_fn</code> will receive a <code>list</code> of samples returned by <code>Datasert.__getitem__</code> and creates the final batch. This code shows its input:</p>
<pre data-code-wrap="python"><code class="lang-python">class Example_DS(Dataset):
    def __init__(self):
        self.data = torch.arange(10).float().view(10, 1)
        
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        x = self.data[idx]
        sample = {
            'sample' : x
            }        
        return sample


def my_collate(batch):
    print(batch)
    print(len(batch))
    return batch


ds = Example_DS()
dls = DataLoader(ds, batch_size=2, shuffle=True, num_workers=1, collate_fn=my_collate)

for batch in dls:
    pass

# Output
# [{'sample': tensor([3.])}, {'sample': tensor([6.])}]
# 2
# [{'sample': tensor([8.])}, {'sample': tensor([2.])}]
# 2
# [{'sample': tensor([5.])}, {'sample': tensor([9.])}]
# 2
# [{'sample': tensor([0.])}, {'sample': tensor([7.])}]
# 2
# [{'sample': tensor([1.])}, {'sample': tensor([4.])}]
# 2
</code></pre>
<p>It depends on your goal and use case what exactly you want to manipulate and how.<br>
The original question was easier solved by just calling a <code>view</code> op on the final batch.</p>
<p>Also, take a look at the <a href="https://github.com/pytorch/pytorch/blob/82099ab87bb789670cd95c91fac0681958f3be94/torch/utils/data/_utils/collate.py#L88"><code>default_collate</code></a> as a reference.</p>