In the latest version of PyTorch, what is best practice to get all tensors to use a particular device by default?

Question:

In pytorch, if I do something like

import torch
x = torch.randn(3)
y = x + 5

all tensors correspond to the "cpu" device by default. Is there some way to make to make it so, by default, all tensors are on another device (e.g. "cuda:0")?

I know I can always be careful to add .cuda() or specify cuda whenever creating a tensor, but it would be great if I could just change the default device directly at the beginning of the program and be done with it, so that torch.randn(3) comes from the desired device without having to specify it every time.

Or would that be a bad thing to do for some reason? E.g. is there any reason I wouldn’t want every tensor/operation to be done on cuda by default?

Asked By: chausies

||

Answers:

Pytorch has an optional function to change the default type of tensor set_default_tensor_type. Applying the default type on the main script:

>>> import torch
>>>
>>> if __name__ == '__main__':
...     cuda = torch.cuda.is_available()
...     if cuda:
...         torch.set_default_tensor_type('torch.cuda.FloatTensor')
...     a = torch.randn(3,3)
...     print(a.device)
...
cuda:0

Or would that be a bad thing to do for some reason? E.g. is there any reason I wouldn’t want every tensor/operation to be done on cuda by default?

I couldn’t find any reference or any document to answer this question. However, in my opinion, it’s to avoid the memory fragmentation in GPU memory.

I’m not an expert, but the data in memory should be arranged in an efficient way, if not, the redundant space will cause OOM. That’s why, in default, Tensorflow will take all of your GPU’s memory no matter how many parameters your model has. You can improve the space and speed just by setting the tensor shape multiples of 8 amp documents.

In practice, higher performance is achieved when A and B dimensions are multiples of 8.

In conclusion, I think it’s better to control the device of tensor manually instead of setting it gpu as default.

Answered By: CuCaRot

Latest Torch has a set_default_device function.

But I wouldn’t recommend setting default device to gpu, gpu doesn’t have that much VRAM, so you may want to keep most data on cpu and only push to gpu the stuff that you are using.

A good option though is to use with torch.device(device) instead to set it locally

Answered By: Nephanth