时间戳

首页 > 解决方案 > 如何直接在 GPU 上或在另一个张量的设备上创建张量？

python - 如何直接在 GPU 上或在另一个张量的设备上创建张量？

问题描述

我发现了关于这个的讨论，其中的代码

if std.is_cuda:
    eps = torch.FloatTensor(std.size()).cuda().normal_()
else:
    eps = torch.FloatTensor(std.size()).normal_()

变成美好的

eps = std.new().normal_()

但据说那里

不推荐使用 new() 方法。

  1. 如何直接在特定设备上创建新张量？
  2. 如何在同一个设备上创建一个新的张量，就像另一个没有丑陋的张量一样if？

标签： pythonpytorchgpu

解决方案

我认为文档现在非常清楚。这里描述了创建新张量的 4 种主要方法，您只需指定要在 gpu 上创建它的设备：

t1 = torch.zeros((3,3), device=torch.device('cuda'))
t2 = torch.ones_like(t1, device=torch.device('cuda'))
t3 = torch.randn((3,5), device=torch.device('cuda'))

此链接添加了有关torch.tensor()构造函数的更多信息。同样，设备是要指定的参数。

如果你想使用另一个张量的设备，你可以通过以下方式访问它tensor.device：

t4 = torch.empty((2,2), device=t3.device)

推荐阅读