RNN batch_first performance considerations based on backing kernel
Question:
Can I use batch_first=True everywhere without worrying about performance differences on CPU, GPU, TPU?
The defaults PyTorch sets are batch_first=False for all RNNs (RNN, LSTM, GRU). Pytorch Lightning mandates batch_first=True for truncated_bptt_steps, however. Will setting it True this mean relative faster performance on CPU and slower on GPU (for example) as compared to setting it False?
If batch_first affects performance specific to device (CPU, GPU, TPU), should we check for device and set batch_first accordingly? (In this case, seems like a prime target for this framework.)
(Aside: Something I've always wondered: if the CUDA kernel is better (better time/space?) when batch_first=False, why doesn't it internally transpose(0, 1), compute, and transpose(1, 0) back instead of exposing this complexity to the user?
ie., in the general case, might it not be better to use batch_first everywhere, and perform gymnastics internally within the specific implementation.)
Not sure if this project is the best place to ask this question, but I was implementing a project using this wonderful framework and it seems to me this question falls within the "engineering code" purview of Pytorch Lightning.