Question

nn.Module.cuda()将所有模型参数和缓冲区移至GPU。

但是为什么不使用模型成员张量？

class ToyModule(torch.nn.Module):
    def __init__(self) -> None:
        super(ToyModule, self).__init__()
        self.layer = torch.nn.Linear(2, 2)
        self.expected_moved_cuda_tensor = torch.tensor([0, 2, 3])

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        return self.layer(input)

toy_module = ToyModule()
toy_module.cuda()

next(toy_module.layer.parameters()).device
>>> device(type='cuda', index=0)

对于模型成员张量，设备保持不变。

>>> toy_module.expected_moved_cuda_tensor.device
device(type='cpu')

Answer 1

如果在模块内部定义了张量，则需要将其注册为参数或缓冲区，以便模块可以识别它。

参数是要训练的张量，将由model.parameters()返回。它们很容易注册，您所需要做的就是将张量包装为nn.Parameter类型，它将自动注册。请注意，只有浮点张量可以作为参数。

class ToyModule(torch.nn.Module):
    def __init__(self) -> None:
        super(ToyModule, self).__init__()
        self.layer = torch.nn.Linear(2, 2)
        # registering expected_moved_cuda_tensor as a trainable parameter
        self.expected_moved_cuda_tensor = torch.nn.Parameter(torch.tensor([0., 2., 3.]))

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        return self.layer(input)

缓冲区是将在模块中注册的张量，因此.cuda()之类的方法会影响它们，但model.parameters()不会返回它们。缓冲区不限于特定的数据类型。

class ToyModule(torch.nn.Module):
    def __init__(self) -> None:
        super(ToyModule, self).__init__()
        self.layer = torch.nn.Linear(2, 2)
        # registering expected_moved_cuda_tensor as a buffer
        # Note: this creates a new member variable named expected_moved_cuda_tensor
        self.register_buffer('expected_moved_cuda_tensor', torch.tensor([0, 2, 3])))

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        return self.layer(input)

在上述两种情况下，以下代码的行为均相同

>>> toy_module = ToyModule()
>>> toy_module.cuda()
>>> next(toy_module.layer.parameters()).device
device(type='cuda', index=0)
>>> toy_module.expected_moved_cuda_tensor.device
device(type='cuda', index=0)

为什么PyTorch nn.Module.cuda（）不将模块张量移动，而仅将参数和缓冲区移动到GPU？

1 个答案: