python - 有没有办法确定 PyTorch 模型是在 cpu 上还是在设备上？

Question

我想弄清楚PyTorch模型是打开cpu还是cuda为了根据模型初始化一些其他Torch.Tensor变量Torch.cuda.Tensor。

但是，查看dir()函数的输出时，我只看到将模型放在设备、GPU 或其他设备上的方法，这些方法在 to 中.cpu()指定.cuda()。张量有属性，但整个模型没有类似物to()。PyTorchis_cuda

有什么方法可以为模型推断出这一点，还是需要参考特定的权重？

标签： pythonpytorchgpu

Answer 1

不，没有这样的功能nn.Module，我相信这是因为参数可以同时在多个设备上。

如果您使用的是单个设备，解决方法是检查第一个参数：

next(model.parameters()).is_cuda

如此处所述。