Почему PyTorch nn.Module.cuda () не перемещает тензор модуля, а только параметры и буферы в графический процессор? - PullRequest
Новая
       27

Почему PyTorch nn.Module.cuda () не перемещает тензор модуля, а только параметры и буферы в графический процессор?

3 голосов
/ 29 марта 2020

nn.Module.cuda() перемещает все параметры модели и буферы в графический процессор.

Но почему не тензор члена модели? 

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')

1 Ответ

3 голосов
/ 29 марта 2020

Если вы задаете тензор внутри модуля, его необходимо зарегистрировать как параметр или буфер, чтобы модуль знал об этом.

Параметры являются тензорами которые должны быть обучены и будут возвращены 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)
...