如何将Pyro的PyroModule模型正确迁移至GPU（cuda）或CPU？

我正在学习Pyro框架，以及如何用PyTorch构建概率神经网络。通常对于PyTorch的nn.Module模型，我可以使用model.to('cuda')将其迁移至GPU，但该方法对PyroModule似乎无效。请问如何正确将PyroModule模型部署到GPU？

示例模型如下：

import torch
import pyro
import pyro.distributions as dist
from pyro.nn import PyroModule, PyroSample
import torch.nn as nn
from pyro.infer.autoguide import AutoDiagonalNormal
from pyro.infer import SVI, Trace_ELBO, Predictive

class Model(PyroModule):
    def __init__(self, h1=20, h2=20):
        super().__init__()
        self.fc1 = PyroModule[nn.Linear](1, h1)
        self.fc1.weight = PyroSample(dist.Normal(0., 1.).expand([h1, 1]).to_event(2))
        self.fc1.bias = PyroSample(dist.Normal(0., 1.).expand([h1]).to_event(1))
        self.fc2 = PyroModule[nn.Linear](h1, h2)
        self.fc2.weight = PyroSample(dist.Normal(0., 1.).expand([h2, h1]).to_event(2))
        self.fc2.bias = PyroSample(dist.Normal(0., 1.).expand([h2]).to_event(1))
        self.fc3 = PyroModule[nn.Linear](h2, 1)
        self.fc3.weight = PyroSample(dist.Normal(0., 1.).expand([1, h2]).to_event(2))
        self.fc3.bias = PyroSample(dist.Normal(0., 1.).expand([1]).to_event(1))
        self.relu = nn.ReLU()

    def forward(self, x, y=None):
        x = x.reshape(-1, 1)
        x = self.relu(self.fc1(x))
        x = self.relu(self.fc2(x))
        mu = self.fc3(x).squeeze()
        sigma = pyro.sample("sigma", dist.Uniform(0., 1.))
        with pyro.plate("data", x.shape[0]):
            obs = pyro.sample("obs", dist.Normal(mu, sigma), obs=y)
        return mu

执行model = Model()后，使用model.to('cuda')似乎并未真正将模型迁移至GPU。

更新：我不确定这是否是正确的解决方案……我发现如果将PyroSample替换为pyro.nn.PyroParam，参数会被列入ParamDict并可迁移至GPU。

• 根本原因：PyroSample定义的是随机变量（比如权重的先验分布），不属于PyTorch的可训练参数范畴，因此nn.Module的to()方法不会自动处理这些变量。而PyroParam是Pyro对应PyTorchnn.Parameter的可训练参数类型，会被纳入参数管理体系，所以to()方法对其有效。

• 正确迁移步骤：

  1. 移动模型基础结构到GPU：先执行model = Model().to('cuda')，这会处理PyroModule包裹的所有PyTorch层（比如示例中的线性层）。
  2. 统一Pyro设备设置：通过全局设置让Pyro的采样操作默认使用GPU，避免设备不匹配：

     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     if device.type == 'cuda':
         pyro.set_default_tensor_type(torch.cuda.FloatTensor)
     

  3. 同步数据到GPU：训练或推理时，必须将输入数据x、y等移动到GPU，比如x = x.to(device)。

• PyroSample与PyroParam的适用场景：

  • 若需要为参数定义先验分布（贝叶斯模型场景），必须用PyroSample，此时需通过全局设备设置确保采样在GPU执行。
  • 若参数是确定性的可训练变量（非贝叶斯场景），可以用PyroParam，享受和PyTorch参数一致的设备迁移体验。

内容的提问来源于stack exchange，提问作者P-Rod