Некоторые члены модуля факела не перемещаются в графический процессор, даже если вызывается model.to (устройство)

Question

Mwe выглядит следующим образом:

import torch
import torch.nn as nn

class model(nn.Module):
    def __init__(self):
        super(model,self).__init__()
        self.mat = torch.randn(2,2)

    def forward(self,x):
        print('self.mat.device is',self.mat.device)
        x = torch.mv(self.mat,x)
        return x

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
m = model()
m.to(device)

x = torch.tensor([2.,1.])
x = x.to(device)

m(x)

Выход

self.mat.device is cpu

и сразу после этого приходит

Traceback (most recent call last):
  File "Z:\cudatest.py", line 21, in <module>
    print(m(x))
  File "E:\Python37\lib\site-packages\torch\nn\modules\module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "Z:\cudatest.py", line 11, in forward
    x = torch.mv(self.mat,x)
RuntimeError: Expected object of device type cuda but got device type cpu for argument #1 'self' in call to _th_mv

Код работает нормально, если я установить device = torch.device('cpu'). Кажется, проблема в том, что model.mat не перемещен в GPU даже после вызова m.to(device). Почему не работает? Как я могу это исправить? Обратите внимание на следующее:

1. Хотя этот конкретный пример можно исправить, используя вместо этого self.mat = nn.Linear(2,2) и x = self.mat(x), в моей исходной программе мне нужен временный тензор для хранения некоторых данных в forward(), который также используется в некоторой арифметике. Как создать такой тензор и отправить его в графический процессор при вызове m.to(device)

2. Заранее неизвестно, есть ли у компьютера графический процессор или нет. Поэтому написание self.mat = self.mat.cuda() не является хорошим решением для моего случая.

Answer 1

pytorch применяет методы модуля, такие как .cpu(), .cuda() и .to() только к подмодулям, параметры и буферы , но НЕ обычным ученикам. Pytorch не может знать, что self.mat, в вашем случае, является фактическим тензором, который должен перемещаться.

Как только вы решите, должен ли ваш mat быть параметром или буфером, просто зарегистрируйте его соответственно, например

class model(nn.Module):
    def __init__(self):
        super(model,self).__init__()
        self.register_buffer(name='mat', tensor=torch.randn(2,2))