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

Question

我想弄清楚模型是否PyTorch处于打开状态cpu，或者cuda是为了根据模型初始化某些其他Torch.Tensor变量Torch.cuda.Tensor。

然而，查看dir()函数的输出，我只看到.cpu(), .cuda(),to()将模型放置在 to 中指定的设备、GPU 或其他设备上的方法。对于PyTorch张量来说，有is_cuda属性，但整个模型没有类似物。

是否有某种方法可以为模型推断出这一点，或者需要参考特定的重量？

Answer 1

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

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

next(model.parameters()).is_cuda

正如这里所描述的。