Speed up collating very large tensors?

username: Kevin_Shen
cooked: <p>I’m working with a dataset where each sample is 120x1024 dimensions. I’m hoping to use a large batch size 4096 (each batch is 4096x120x1024) but am experiencing very slow dataloading even with num_workers=20.</p>
<p>Below is an example of my code. I’m getting 4 iterations/second on my machine. This is a bit too slow to finish training in a reasonable amount of time.</p>
<p>I’ve narrowed the bottleneck down to the <code>torch.stack()</code> call in <code>collate</code>. This seems to be 10x slower than everything else. I believe what is happening is the data is being copied to a contiguous block, hence the slowdown?</p>
<p>Is there any way to speed up dataloading in my case?</p>
<pre><code class="lang-auto">class MyDataset(Dataset):
    def __init__(self):
          ...

    def __getitem__(self, index):

         return torch.ones([120, 1024])        


train_dataset = MyDataset()

def collate(batch):
    
    data = torch.stack([b[0] for b in batch])
    
    return data

train_sampler = RandomSampler(train_dataset)

train_dataloader = DataLoader(
    train_dataset, sampler=train_sampler, batch_size=4096,
    num_workers=20, pin_memory=True, collate_fn=collate
)

</code></pre>

username: ptrblck
cooked: <aside class="quote no-group" data-username="Kevin_Shen" data-post="1" data-topic="84365">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/kevin_shen/48/23263_2.png" class="avatar"> Kevin_Shen:</div>
<blockquote>
<p>even with num_workers=20.</p>
</blockquote>
</aside>
<p>Note that too many workers might slow down your system, so you should test different values for your current setup.</p>
<p>That being said, to avoid the <code>torch.stack</code> call, you could use a <a href="https://pytorch.org/docs/stable/data.html?highlight=batchsampler#torch.utils.data.BatchSampler"><code>BatchSampler</code></a> to pass a batch of indices to your <code>Dataset.__getitem__</code>, preallocate the final tensor via <code>torch.empty(batch_size, 120, 1024)</code> and copy the tensors into it.</p>

username: Kevin_Shen
cooked: <p>Thanks <a class="mention" href="/u/ptrblck">@ptrblck</a>!  The idea of copying the data into the preallocated tensor helps (it’s now 2x faster) but it seems to get slowed down because I need to do a for-loop copy.</p>
<pre><code class="lang-auto">for i, b in eumerate(batch):
    final_tensor[i] = b[0]
</code></pre>
<p>I think if I try to construct a (batch_size, 120, 1024) tensor at any point it will be slow so it seems like I actually need to loop over individual small tensors (1, 120, 1024)? I’m not sure how to use BatchSampler here such that it would help.</p>

username: ptrblck
cooked: <p>I think the approach should be right, since you would need the loop at one point anyway to load and process each image, no?</p>
<p>In the standard approach, the <code>DataLoader</code> will load each image one by one and use the <code>collate_fn</code> to create the batch. Now you could push this loop into the <code>__getitem__</code> to load each sample in the loop and copy the data into the preallocated tensor.</p>
<p>Let me know, if I misunderstood the use case.</p>

username: Kevin_Shen
cooked: <p>If we push the for-loop into the <code>__getitem__</code> function, does that mean it’s 1 worker doing the loop? Would that be faster than just doing the for-loop in collate?</p>
<p>I was wondering if it is possible to assign a row index of the preallocated tensor to each worker and have them copy the data into that index of the tensor in <code>__getitem__</code>?</p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="Kevin_Shen" data-post="5" data-topic="84365">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/kevin_shen/48/23263_2.png" class="avatar"> Kevin_Shen:</div>
<blockquote>
<p>If we push the for-loop into the <code>__getitem__</code> function, does that mean it’s 1 worker doing the loop?</p>
</blockquote>
</aside>
<p>Yes, but also in the default setup, where you are using a single index, a single worker will create the batch, so there shouldn’t be a difference I assume (each worker creates its own batch).</p>
<aside class="quote no-group" data-username="Kevin_Shen" data-post="5" data-topic="84365">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/kevin_shen/48/23263_2.png" class="avatar"> Kevin_Shen:</div>
<blockquote>
<p>I was wondering if it is possible to assign a row index of the preallocated tensor to each worker and have them copy the data into that index of the tensor in <code>__getitem__</code> ?</p>
</blockquote>
</aside>
<p>Multiple workers do not create the same batch, but each worker will build its own batch.<br>
There is a feature request letting multiple workers work on the same batch, but I don’t think it’s ready yet.</p>

username: ardivekar
cooked: <p>Assuming your tensors are uint8 (images with values between 0-255 inclusive), it’s much faster to convert to Numpy, then stack in Numpy, then convert to torch.</p>
<p>Random tensors:</p>
<pre><code class="lang-auto">tensor_list = [torch.randint(size=(3, 120, 1024), low=0, high=255, dtype=torch.uint8) for _ in range(4096)]
</code></pre>
<p><strong>Stacking in torch:</strong></p>
<pre><code class="lang-auto">def stack_via_torch1(tensor_list):
    return torch.stack(tensor_list)

def stack_via_torch2(tensor_list):
    stacked_ten = torch.empty((len(tensor_list), *tensor_list[0].shape), dtype=torch.uint8)
    for i, ten in enumerate(tensor_list):
        stacked_ten[i] = ten
    return stacked_ten

%timeit stack_via_torch1(tensor_list)
946 ms ± 1.28 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit stack_via_torch2(tensor_list)
972 ms ± 1.51 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
</code></pre>
<p><strong>Stacking via NumPy (~1.8x faster for 4096 tensors):</strong></p>
<pre><code class="lang-auto">def stack_via_numpy1(tensor_list):
    return torch.from_numpy(np.array([ten.numpy() for ten in tensor_list]))

def stack_via_numpy2(tensor_list):
    np_data: np.ndarray = np.empty(len(tensor_list), dtype=object)
    for i, ten in enumerate(tensor_list):
        np_data[i] = tensor_list[i]
    return torch.from_numpy(np.stack(np_data))

def stack_via_numpy3(tensor_list):
    np_data: np.ndarray = np.empty((len(tensor_list), *tensor_list[0].shape), dtype=np.uint8)
    for i, ten in enumerate(tensor_list):
        np_data[i] = ten
    return torch.from_numpy(np_data)

def stack_via_numpy4(tensor_list):
    np_data: np.ndarray = np.empty(len(tensor_list), dtype=object)
    for i, ten in enumerate(tensor_list):
        np_data[i] = ten.numpy()
    return torch.from_numpy(np.stack(np_data))

def stack_via_numpy5(tensor_list):
    np_data: np.ndarray = np.empty(len(tensor_list), dtype=object)
    for i, ten in enumerate(tensor_list):
        np_data[i] = tensor_list[i]
    np_empty = np.empty((len(tensor_list), *tensor_list[0].shape), dtype=np.uint8)
    return torch.from_numpy(np.stack(np_data, out=np_empty))


%timeit stack_via_numpy1(tensor_list)
526 ms ± 2.19 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit stack_via_numpy2(tensor_list)
533 ms ± 2.17 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit stack_via_numpy3(tensor_list)
533 ms ± 2.23 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit stack_via_numpy4(tensor_list)
528 ms ± 1.31 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit stack_via_numpy5(tensor_list)
534 ms ± 1.42 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
</code></pre>
<p>Thus stacking via Numpy is ~1.8x faster for 4096 tensors.</p>
<p>Additional:</p>
<ul>
<li>I prefer <code>stack_via_numpy1</code> since it’s the fastest and easy to remember.</li>
<li>I observed this level of speedup even for a small number of tensors (1 to 1024) and upto 65k. An exception was in the range of 64 to 256 tensors, where both torch and numpy took about the same time weirdly.</li>
</ul>