Custom Dataloader/ dataset to load several samples at once

JuanFMontesinos: <p>Hi,<br>
I’m new using PyTorch. I am implementing and testing a new paper called <a href="https://arxiv.org/pdf/1804.03160.pdf" rel="nofollow noopener">Sound of Pixels</a>.</p>
<p>In short it’s a net which works with a 2-tower stream. One tower is fed with a stack of images and the other one is fed with audio spectrograms.<br>
I’m using a private dataset, in which each sample is a numpy binary file which contains a python dictionary with both, audio and images.</p>
<p>I’ve created my own dataset subclass in which I’m enlisting all the files inside a folder and using this list to call<br>
<code>__getitem__</code> function.<br>
The matter is I have to combine features of N samples (since the net is trained by summing audio samples). As far as I understand, I have to  combine these N samples inside the <code>__getitem__</code> function. However, this function is designed to deal with one sample and pass the info to the dataloader.</p>
<p>I considered the option of doing a post-processing of the batch doing what I need or making my own dataloader using one of the implemented samplers. Before doing that I would like to know if there is a more efficient way of dealing with this in pytorch because I have to compute fourier transforms of audio and I don’t want to bottleneck the training with a bad data preprocessing.<br>
Thank you very much</p>
<pre><code class="lang-auto">Dataset class designed for just one sample

class BinaryData(torch.utils.data.Dataset):

    def __init__(self, root_dir,transform):
        self.input_list = []
        self.transform = transform
        for path, subdirs, files in os.walk(root_dir):
            for name in files:
                self.input_list.append(os.path.join(path, name))

    def __len__(self):
        return len(self.input_list)

    def __getitem__(self, idx):
        dic = data2dic(np.load(self.input_list[idx]))
        audio = dic['audio']
        audio = Sound_standarize(audio)
        frames = dic['frames']
        size = np.shape(frames)
        images = []
        for i in range(size[3]):
            images.append(self.transform(Image.fromarray(frames[:,:,:,i])))
        frames = torch.stack(images)
        return audio,frames 
</code></pre>

Nicolo_Savioli: <p>Hi, this was not a pytorch framework problem.<br>
If I understood correctly how to speed up your post-processing during the download?<br>
I suggest you use a multi-thread approach.</p>
<p><a href="http://chriskiehl.com/article/parallelism-in-one-line/" class="onebox" target="_blank" rel="nofollow noopener">http://chriskiehl.com/article/parallelism-in-one-line/</a></p>

JuanFMontesinos: <p>Well, It’s not a matter of post-processing or speed up. In sample words:</p>
<p>I have ie 1000 files (dataset).<br>
I want that my dataloader opens (ie) 2 of these files (randomly) and sum them. This sum is what i will feed the net with, so in the end my net will percive there are 1000/2 = 500 samples.  A batch have to be done of sums (but not raw files)</p>
<p>The matter is I don’t know what is the proper way to do this in pytorch. I can roughly emulate this behavior but I would like to know if pytorch framework is prepared to do this is a easy way.</p>

ptrblck: <p>I’m not sure, if this covers your whole use case, but you could just get the desired number of samples from your dataset and process them in <code>__getitem__</code>.<br>
Here is a small example:</p>
<pre><code class="lang-python">class MyDataset(Dataset):
    def __init__(self, nb_samples):
        self.images = torch.randn(100, 3, 24, 24)
        self.audio = torch.randn(100, 1024, 12)
        self.nb_samples = nb_samples
        
    def __getitem__(self, index):
        # Load all nb_samples
        images = self.images[index*self.nb_samples:index*self.nb_samples+self.nb_samples]
        audio_specs = self.audio[index*self.nb_samples:index*self.nb_samples+self.nb_samples]
        
        # Transform
        
        # Sum samples
        x1 = torch.sum(images)
        x2 = torch.sum(audio_specs)
        
        return x1, x2
        
        
    def __len__(self):
        return len(self.images) / self.nb_samples - self.nb_samples


dataset = MyDataset(2)
img_sum, audio_sum = dataset[0]
</code></pre>
<p>Currently adjacent samples are used. Let me know, if you need another strategy.</p>

JuanFMontesinos: <p>Oh in fact it seems perfect, Thank you very much</p>

Max1: <p><a class="mention" href="/u/ptrblck">@ptrblck</a>, thank you for an example. My dataset is “endless” and stored externally, thus, unfeasible to store in the class instance. So currently <code>__getitem__()</code> retrieves just one sample from the external database.</p>
<p>I’m concerned with the high overhead for this sample-by-sample querying; for batch size of e.g. 64, <code>__getitem__()</code> gets called 64 times consequently, doesn’t it?</p>
<p>I thought I could build a cache in <code>MyDataset</code>, fetching and storing bunches of samples in the class instance before returning them individually via <code>__getitem__()</code>. Is there a better way / better place to implement it?</p>

ptrblck: <p>Yes, you are right.<br>
If you are concerned about the overhead of the (DB) query, you could create the whole batch directly in <code>__getitem__</code> (with a single query) and set <code>batch_size=1</code> in your <code>DataLoader</code>.<br>
Would this work for you?</p>

Max1: <p>Sure. Thanks! I just didn’t dare to, wondering instead why I need to <code>torch.unsqueeze()</code> single samples into tensors with singleton outer dimension prior to returning them from <code>__getitem__()</code>.</p>
<p>Without <code>unsqueeze()</code>, I was getting the famous <code>RuntimeError: Expected 4-dimensional input for 4-dimensional weight [4, 1, 3, 3], but got 3-dimensional input of size [64, 32, 32] instead</code> from the very first layer, <code>conv2d</code> (3x3 to four features). The input is a 32x32 gray-scale image, batch_size is 64.</p>

ptrblck: <p>This should usually not be necessary since the <code>DataLoader</code> should take care of this.<br>
Did you get an error if you didn’t unsqueeze the data sample?</p>

dinarkino: <p>Thank you for this nice answer! Is there a way to do the same when you do not have self.images, but instead loading random seq of images from the folder. Before that, I used Image.open() from PIL to open an image in <strong>getitem</strong>. Now I want to do the same but for the seq of images (let’s say 5 images) each time. So the shape of the batch would be (8, 5). It seems like a not optimal solution to load images in a loop with Image.open()</p>

ptrblck: <p>I’m not aware of a “batched” <code>open</code> method for multiple image files, so I think you would need to use the loop. However, note that you could speed up the image loading by replacing <code>PIL</code> with <code>PIL-SIMD</code>. You would still need to execute the loading in a loop, but might lower the duration.</p>

dinarkino: <p>I will check the performance of PIL-SIMD. For now, I solved the problem by using imread_collection from <a href="http://skimage.io" rel="noopener nofollow ugc">skimage.io</a>. In principle, everything works properly but I face the problem of slow DataLoader. I use such a code:</p>
<pre><code class="lang-auto">    def __getitem__(self, idx):
        idx_q = int(torch.randint(0 + self.boundary, self.length - self.boundary, (1,))) 
        idx_n = int(torch.randint(0 + self.boundary, self.length - self.boundary, (1,)))

        q = imread_collection([self.image_paths[idx_q-1], self.image_paths[idx_q], self.image_paths[idx_q+1]], conserve_memory=True)
        p = imread_collection([self.image_paths_winter[idx_q-1], self.image_paths_winter[idx_q], self.image_paths_winter[idx_q+1]], conserve_memory=True)
        n = imread_collection([self.image_paths_winter[idx_n-1], self.image_paths_winter[idx_n], self.image_paths_winter[idx_n+1]], conserve_memory=True)
        
        if self.transform:
            q = torch.stack([self.transform(img) for img in q])
            p = torch.stack([self.transform(img) for img in p])
            n = torch.stack([self.transform(img) for img in n])

        return q, p, n
</code></pre>
<p>So, for batch_size=4 I load 4x3x3 images which are 36 images. When I try to run all this code from <strong>getitem</strong> out of torch’s DataLoader it takes around 0.3 sec to run. But with DataLoader it takes 13 seconds which is a huge difference. Of course, num_workers help here when there are more iterations. But I don’t know where is the bottleneck. If you have any thoughts, that could be very helpful.</p>

ptrblck: <p>How large is your batch size when you are using the <code>DataLoader</code>?<br>
Each worker of the <code>DataLoader</code> will create a full batch before the <code>DataLoader</code> loop is executed, so in particular the first iteration might be slow, since all workers would start collecting all samples. Once the workers are done, the ready batches will be added to a queue and the workers will start creating the next batch while the training loop is executed, which might reduce the data loading time assuming that the actual model training is not tiny compared to the data loading.</p>

dinarkino: <p>I checked for batch size = 4. If I run such a code:</p>
<pre><code class="lang-auto">batch_size = 4
trainloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, num_workers=4)
dataiter = iter(trainloader)
for i in range(8):
    start = time.time()
    q, p, n = dataiter.next()
    q, p, n  = q.to(device), p.to(device), n.to(device)
    end = time.time()
    print(end - start)
</code></pre>
<p>I will get such time results:</p>
<pre><code class="lang-auto">10.58901309967041
1.8810162544250488
0.009713411331176758
0.7070198059082031
10.15380048751831
2.085867404937744
0.009247779846191406
0.009662628173828125
</code></pre>
<p>I mean this is ok compared to one worker, but still, I don’t know what is happening inside the loader so that it needs 10 seconds instead of 0.3. Maybe it is because workers collecting all samples as you explained, but for one worker problem stays</p>

dinarkino: <p>Btw for batch size 8:</p>
<pre><code class="lang-auto">23.229380130767822
0.024968624114990234
0.023113250732421875
1.4580109119415283
24.699864387512207
0.018816709518432617
0.01724529266357422
0.017206192016601562
</code></pre>
<p>So it grows linearly</p>

ptrblck: <p>The periodical slowdown every <code>num_workers</code> iterations points towards the aforementioned loading of a full batch in each worker.<br>
In your current script you don’t have any workload, so it’s expected that you would see slower iterations once the queue is empty.</p>
<p>The timings would refer to:</p>
<pre><code class="lang-python">10.58901309967041 # no batches ready in the queue, all 4 workers are preloading
1.8810162544250488 # Grab next batch from queue. Since this is slower than the next step the queue might not have been filled yet.
0.009713411331176758 # grab next batch from queue
0.7070198059082031 # grab next batch from queue
10.15380048751831 # Queue is empty, as no worker was able to load a full batch yet
</code></pre>
<p>If your model training is sufficiently large, the <code>DataLoader</code> might be able to preload the next batch(es) while your model is training.</p>

dinarkino: <p>I see. Thank you for the help!</p>