在pytorch中，如果我做类似的事情
import torch
x = torch.randn(3)
y = x + 5
默认情况下，所有张量都对应于“cpu”设备。有什么办法可以做到，默认情况下，所有张量都在另一个设备上（例如“cuda：0”）？
我知道在创建张量时我总是可以小心地添加.cuda()或指定 cuda，但是如果我可以直接在程序开始时更改默认设备并完成它，那就太好了，所以它torch.randn(3)来自所需的设备无需每次都指定。
或者由于某种原因这会是一件坏事吗？例如，有什么理由我不希望默认情况下每个张量/操作都在 cuda 上完成？
Pytorch 有一个可选的函数来改变张量的默认类型set_default_tensor_type。在主脚本上应用默认类型：
>>> 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
或者由于某种原因这会是一件坏事吗？例如，有什么理由我不希望默认情况下每个张量/操作都在 cuda 上完成？
我找不到任何参考资料或任何文件来回答这个问题。但是，在我看来，这是为了避免 GPU 内存中的内存碎片。
我不是专家，但是内存中的数据应该以一种有效的方式排列，否则，多余的空间会导致OOM。这就是为什么默认情况下，无论您的模型有多少参数，Tensorflow 都会占用您所有 GPU 的内存。只需设置 8 amp 文档的张量形状倍数即可提高空间和速度。
在实践中，当 A 和 B 维度是 8 的倍数时，可以获得更高的性能。
总之，我认为最好手动控制张量的设备，而不是将其设置为 gpu 为默认值。
本文收集自互联网，转载请注明来源。
如有侵权，请联系 [email protected] 删除。
我来说两句