即使调用 model.to(device)，torch 模块的某些成员也不会移动到 GPU

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)为什么whis不起作用？我怎样才能解决这个问题？请注意以下事项：

1. 尽管这个特定的例子可以通过使用self.mat = nn.Linear(2,2)来修复x = self.mat(x)，但在我的原始程序中，我需要一个临时张量来存储一些数据，forward()这些数据也用于一些算术中。如何构造这样的张量并在调用时将其发送到 GPUm.to(device)

2. 事先并不知道计算机是否有 GPU。因此，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))