Question

我必须在具有不同设备的不同类型的pytorch模型上堆叠自己的图层。

例如 A 是 cuda 模型，而 B 是 cpu 模型（但我不了解设备类型）。然后，新模型分别为 C 和 D ，其中

#if DEBUG
    // your log w/ explicit public modifier
#endif

但是我必须将class NewModule(torch.nn.Module): def __init__(self, base): super(NewModule, self).__init__() self.base = base self.extra = my_layer() # e.g. torch.nn.Linear() def forward(self,x): y = self.base(x) z = self.extra(y) return z ... C = NewModule(A) # cuda D = NewModule(B) # cpu和base移动到相同设备，即C的extra和base是 cuda < / strong>模型和D是 cpu 模型。所以我尝试了extra：

__inin__

很遗憾，def __init__(self, base): super(NewModule, self).__init__() self.base = base self.extra = my_layer().to(base.device)中没有属性device（提高torch.nn.Module）。

我该怎么做才能获得AttributeError的设备类型？还是有其他方法可以使base和base自动位于同一设备上，即使extra的结构是不确定的？

Answer 1

这个问题已经问了很多遍了（1，2）。引用来自PyTorch开发人员的回复：

That’s not possible. Modules can hold parameters of different types on different devices, and so it’s not always possible to unambiguously determine the device.

推荐的工作流程（as described on PyTorch blog）是分别创建device对象，并在各处使用它。从博客中复制粘贴示例：

# at beginning of the script
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

...

# then whenever you get a new Tensor or Module
# this won't copy if they are already on the desired device
input = data.to(device)
model = MyModule(...).to(device)

请注意，没有什么可以阻止您向模型添加.device属性。

正如卡尼（在评论中）所述，如果模型中的所有参数都在同一设备上，则可以使用next(model.parameters()).device。

Answer 2

我的解决方案在99％的情况下都有效。

class Net(nn.Module):
  def __init__()
    super().__init__()
    self.dummy_param = nn.Parameter(torch.empty(0))

  def forward(x):
    device = self.dummy_param.device
    ... etc

此后，dummy_param将始终与模块Net具有相同的设备，因此您可以随时获取它。例如：

net = Net()
net.dummy_param.device

'cpu'

net = net.to('cuda')
net.dummy_param.device

'cuda:0'

如何方便地获取pytorch模块的设备类型？

2 个答案: