如何在GPU上创建PyTorch分布

Question

是否有可能使PyTorch分布直接在GPU上创建样本。

如果我这样做：

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

样本将在CPU上运行。当然，可以使用sample = sample.to(torch.device("cuda"))将其转换为GPU上的样本。但是是否有一种方法可以直接将样本发送到GPU，而无需先在CPU上创建它？

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

有什么想法吗？

- patapouf_ai

Answer 1

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

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

- patapouf_ai

---

在正态分布中，出现了错误“normal_cuda_kernel not implemented for long”。正确的代码如下（注意0和1现在是0.0和1.0）：“normal = Normal(torch.tensor(0.0).to(device=device), torch.tensor(1.0).to(device=device))”。 - Boppity Bop

真希望有更简洁的符号来表示这个。 - daknowles

Answer 2

在我的情况下，我在神经网络模型中使用正态分布作为先验。例如，我有一个名为class1的类，在init函数中，我需要初始化我的先验。然而，对于class1实例调用.to('cuda')不能改变分布设备，并导致后续使用时出现错误。因此，我可以使用register buffers来进行管理，如下所示。

class class1(nn.Module):
    def __init__(self):
        super().__init__()
        self.register_buffer("mean", torch.tensor(0.))
        self.register_buffer("var", torch.tensor(1.))
    def get_dist(self):
        return torch.distributions.Normal(self.mean, self.var)

然而，我有几个先前的问题，而且无法注册缓冲列表。因此，一个选择是在get_dist属性中初始化分布，除非您不关心初始化分布的时间复杂度。我决定定义一个函数来初始化分布，并在get_dist中使用try-except来处理不同的状态。如果分布变量未被分配或在CPU上，而我们期望它在GPU上，则跳转到except，在那里我使用torch.zeros(..).to(device)来初始化分布。

总的来说，为了处理CPU/GPU设备的错误，您需要使用适当的设备使用Tensor输入参数来初始化分布。主要原因是torch.Distribution模块不幸没有设备属性。

- Payam Jome Yazdian

Answer 3

我刚遇到了同样的问题，感谢其他答案中提供的指针。如果您想在模块内部进行分布，我想提供另一个选择，即覆盖模块中的 to 方法，并手动调用分布参数张量上的 to 方法。我只测试了 Uniform ，但在这里运行良好。

class MyModule(nn.Module):
    def __init__(self, ...):
        self.rng = Uniform(
            low=torch.zeros(3),
            high=torch.ones(3)
        )
    
    def to(self, *args, **kwargs):
        super().to(*args, **kwargs)
        self.rng.low = self.rng.low.to(*args, **kwargs)
        self.rng.high = self.rng.high.to(*args, **kwargs)

现在您可以像往常一样将模型放在GPU上，self.rng.sample()将返回正确设备上的样本。

- lopsided

---

请注意，根据分布的不同，您必须包括更多参数。 例如，对于torch.distributions.MultivariateNormal，您需要设置：_unbroadcasted_scale_tril，loc， scale_tril， covariance_matrix和 precision_matrix，并且明确检查后三个是否为None。 - Fabricio

Answer 4

你可以通过重写网络的self._apply(self, fn)方法来解决“将非参数/缓冲区属性传输到GPU”的问题。像这样：

def _apply(self, fn):
    # apply fn() to your modules
    for module in self.children():  # like 'ResNet_backbone'
        module._apply(fn)
    # apply fn() to your prior
    self.prior.attr1 = fn(self.prior.attr1)  # like 'MultivariateNormal.loc', need to be Tensor
    self.prior.attr2 = fn(self.prior.attr2) 
    ···
    self.prior.attrN = fn(self.prior.attrN) 
    # if we do not use register_buffer(Tensor)
    #    apply fn() to your non-parameter/buffer attributes
    #    need to be Tensor too
    self.attr1 = fn(self.attr1)
    self.attr2 = fn(self.attr2) 
    ···
    self.attrN = fn(self.attrN) 

- Clive Indigo