PyTorch RNN is more efficient with `batch_first=False`?

user_9793316 0
Viewed 2783

In machine translation, we always need to slice out the first timestep (the SOS token) in the annotation and prediction.

When using batch_first=False, slicing out the first timestep still keeps the tensor contiguous.

import torch
batch_size = 128
seq_len = 12
embedding = 50

# Making a dummy output that is `batch_first=False`
batch_not_first = torch.randn((seq_len,batch_size,embedding))
batch_not_first = batch_first[1:].view(-1, embedding) # slicing out the first time step

However, if we use batch_first=True, after slicing, the tensor is no longer contiguous. We need to make it contiguous before we can do different operations such as view.

batch_first = torch.randn((batch_size,seq_len,embedding))
batch_first[:,1:].view(-1, embedding) # slicing out the first time step

output>>>
"""
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-8-a9bd590a1679> in <module>
----> 1 batch_first[:,1:].view(-1, embedding) # slicing out the first time step

RuntimeError: view size is not compatible with input tensor's size and stride (at least one dimension spans across two contiguous subspaces). Use .reshape(...) instead.
"""

Does that mean batch_first=False is better, at least, in the context of machine translation? Since it saves us from doing the contiguous() step. Is there any cases that batch_first=True works better?

pytorch nlp python
1 Answers

Performance

There doesn't seem to be a considerable difference between batch_first=True and batch_first=False. Please see the script below:

import time

import torch


def time_measure(batch_first: bool):
    torch.cuda.synchronize()
    layer = torch.nn.RNN(10, 20, batch_first=batch_first).cuda()
    if batch_first:
        inputs = torch.randn(100000, 7, 10).cuda()
    else:
        inputs = torch.randn(7, 100000, 10).cuda()

    start = time.perf_counter()

    for chunk in torch.chunk(inputs, 100000 // 64, dim=0 if batch_first else 1):
        _, last = layer(chunk)

    return time.perf_counter() - start


print(f"Time taken for batch_first=False: {time_measure(False)}")
print(f"Time taken for batch_first=True: {time_measure(True)}")

On my device (GTX 1050 Ti), PyTorch 1.6.0 and CUDA 11.0 here are the results:

Time taken for batch_first=False: 0.3275816479999776
Time taken for batch_first=True: 0.3159054920001836

(and it varies either way so nothing conclusive).

Code readability

batch_first=True is simpler when you want to use other PyTorch layers which require batch as 0th dimension (which is the case for almost all torch.nn layers like torch.nn.Linear).

In this case you would have to permute returned tensor anyway if batch_first=False was specified.

Machine translation

It should be better as the tensor is contiguous all the time and no copy of data has to be done. It also looks cleaner to slice using [1:] instead of [:,1:].

user_108864200
Related