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

Question

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

例如A 是cuda 模型，B 是cpu 模型(但在我得到之前我不知道设备类型)。那么新模型分别是C和D，其中

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 和 extra 移动到相同 设备，即 base 和 extra C 的 是 cuda 模型，D 的是 cpu 模型。所以我尝试了这个__inin__:

def __init__(self, base):
    super(NewModule, self).__init__()
    self.base = base
    self.extra = my_layer().to(base.device)

不幸的是，torch.nn.Module 中没有属性device(引发AttributeError)。

如何获取base的设备类型？或者任何其他方法使 base 和 extra 自动位于同一设备上，即使 base 的结构不明确？

Answer 1

这个问题已经被问过很多次了(1，2)。引用 PyTorch 开发者的回复:

那是不可能的。模块可以在不同的设备上保存不同类型的参数，因此并不总是能够明确地确定设备。

推荐的工作流程 (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 属性。

正如 Kani 所提到的(在评论中)，如果模型中的所有参数都在同一台设备上，则可以使用 next(model.parameters()).device。