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

user_1114253 0
Viewed 88

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?

pytorch deep-learning python-3.x python
1 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.

user_107323510
Related