如何方便的获取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)

不幸的是，没有属性device的torch.nn.Module（提高AttributeError）。

我应该怎么做才能获得设备类型base？或者任何其他方法可以自动制作base并extra在同一设备上，即使其结构base不特定？

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.

推荐的工作流程（如 PyTorch 博客中所述）是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.

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'

Answer 3

@Duane 的答案在模型中创建一个参数（尽管是一个小张量）。

我认为这个答案稍微更加Pythonic和优雅：

class Model(nn.Module):
    def __init__(self, *args, **kwargs):
        super().__init__()
        self.device = torch.device('cpu') # device parameter not defined by default for modules

    def _apply(self, fn):
        # /sf/ask/3829430251/
        # override apply by moving the attribute device of the class object as well.
        # This allows to directly know where the class is when creating new attribute for the class object.
        super()._apply(fn)
        self.device = fn(self.device)
        return self

net.cuda()、net.float()等都可以正常工作，因为它们都调用_apply而不是（如源代码to中所示）。

@Kani（已接受的答案）评论中的替代解决方案也非常优雅：

class Model(nn.Module):
    def __init__(self, *args, **kwargs):
        """
        Constructor for Neural Network.
        """
        super().__init__()

    @property
    def device(self):
        return next(self.parameters()).device

model.device您可以通过参数访问设备。当模型内没有参数时，此解决方案不起作用。