python - 如何在 GPU 上制作 PyTorch 分发版

Question

是否可以让 PyTorch 发行版直接在 GPU 上创建样本。

如果我这样做了

from torch.distributions import Uniform, Normal
normal = Normal(3, 1)
sample = normal.sample()

然后样本将在 CPU 上。当然，可以在 GPU 上执行 sample = sample.to(torch.device("cuda")) 。但是有没有办法让样本直接进入 GPU 而无需先在 CPU 上创建？

PyTorch 分发继承自 Object，而不是 nn.Module，因此它没有将分发实例放在 GPU 上的 to 方法。

有什么想法吗？

Answer 1

分布使用重新参数化技巧。因此，将 GPU 上的大小为 0 的张量提供给分布构造函数是可行的。如下:

normal = Normal(torch.tensor(0).to(device=torch.device("cuda")), torch.tensor(1).to(device=torch.device("cuda")))