PyroSample和cuda GPU 概率编程语言Pyro如何使用cuda,使用显卡，贝叶斯神经网络

PyroSample和cuda GPU

我正在实验在GPU上训练一个简单的贝叶斯前馈网络。该模型在CPU上训练得很好，但当我尝试在GPU上训练时，它给出了一个关于张量在多个设备上的错误。

我已经能够追踪到错误的原因是取代nn的PyroSample语句。带Pyro参数的线性层参数。当我注释掉这些语句时，模型在CPU和GPU上都运行良好。但是当我包含这些语句时，模型遇到了问题，因为我相信我使用的self.cuda()没有将Pyro参数放在GPU上。

所以，我的问题是，在GPU上放置热解样品重量/偏差参数(覆盖线性层的PyTorch参数)的最佳方式是什么？我在下面粘贴了我的简单代码。似乎有一个简单的解决方案，但我的self.cuda()似乎不起作用。

# Put data on GPU.
X, y = X.cuda(), y.cuda()

# Specify model.
class TestNN(PyroModule):
    def __init__(self, in_features, out_features=1):
        super().__init__()
        self.fc1 = PyroModule[nn.Linear](in_features, out_features)
        # Replace network layer parameters with Pyro priors.
        self.fc1.weight = PyroSample(dist.Normal(0., 1.).expand([out_features, in_features]).to_event(2))  # -->>> how to put these parameters on the GPU??
        self.fc1.bias = PyroSample(dist.Normal(0., 10.).expand([out_features]).to_event(1))  # -->>> how to put these parameters on the GPU??
        self.cuda()
        
    def forward(self, x, y=None):
        p = torch.sigmoid(self.fc1(x))
        # Likelihood.
        with pyro.plate('data', x.shape[0]):
            obs = pyro.sample('obs', dist.Bernoulli(p), obs=y)
        return p

# The below also doesn't seem to put the parameters on the GPU.
test_nn = TestNN(in_features=X.shape[1])
test_nn.cuda()  # doesn't work??

解决办法

PyroSample(dist.Normal(0., torch.tensor(1.0, device=X.device)).expand([out_features, in_features]).to_event(2))

啊，看起来确实是通过将一个参数设置为torch.tensor并明确说明设备来修复的。谢谢！

我尝试了均值参数和比例参数，两种情况下都有效。不确定是否有更好的方式来用torch.tensor(device)指定平均值或标度，但无论哪种方式，似乎只有一个参数需要它，而不是两个都需要。