python - 即使使用显式 cuda() 调用，pytorch 实例张量也不会移动到 GPU

Question

我正在开发一个项目，其中模型需要访问我在类的构造函数 init 中声明的张量(我是 torch.nn.Module 类的子类)，然后我需要通过简单的 matmul() 在forward()方法中使用这个张量，模型通过cuda()调用发送到GPU:

model = Model()
model.cuda()

但是，当我通过以下方式对简单输入 X 进行前向传播时:

model(X) # or model.forward(X)

我明白

RuntimeError: Expected object of type torch.cuda.FloatTensor but found type torch.FloatTensor for argument #2 'mat2'

表明 matmul 的第二个参数(我声明的实例张量)在 CPU 上，并且预期在 GPU 上(与模型和数据的其余部分一样)。

在matmul中，张量通过matrix.t()转置

我什至尝试重写 cuda() 方法:

def cuda(self):
    super().cuda()
    self.matrix.cuda()

数据已经在GPU中，意味着下面这行代码已经被执行:

X = X.cuda()

错误还明确指出了 matmul 的参数 2，在本例中是张量(称为矩阵)而不是 X。

Answer 1

让我们假设以下情况:

1. X 已正确移动到 GPU

2. 在 Model 类中声明的张量是一个简单的属性。

   即类似于以下内容:

class Model(nn.Module):
   def __init__(self):
       super().__init__()
       self.matrix = torch.randn(784, 10)
       
   def forward(self, x):
       return torch.matmul(x, self.matrix)

<小时/>

如果是这样，您的第一次尝试将不起作用，因为 nn.Module.cuda() 方法仅移动所有 Parameters 和 Buffers 到 GPU。

您需要将 Model.matrix 设为 Parameter而不是常规属性。将其包装在参数类中。像这样的东西:

self.matrix = nn.Parameter(torch.randn(784, 10))

现在，您尝试在覆盖内手动调用 Model.matrix 上的 .cuda() 方法，而不是像上面那样自动转换到 GPU。

这也不起作用，因为 nn.Module.cuda() 之间存在细微的差异。方法和torch.Tensor.cuda()方法。

当nn.Module.cuda()将Module的所有Parameters和Buffers移动到GPU时并返回自身，torch.Tensor.cuda() 仅返回 GPU 上张量的副本。

原始张量不受影响。

<小时/>

总而言之，要么:

1. 将您的矩阵属性包装为参数或
2. 通过以下方式将 GPU 副本分配回矩阵:

self.matrix = self.matrix.cuda()

在您的覆盖中。

我建议第一个。