时间戳

首页 > 解决方案 > 即使使用显式 cuda() 调用，pytorch 实例张量也不会移动到 gpu

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

问题描述

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

model = Model()
model.cuda()

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

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

我明白了

RuntimeError：预期为 torch.cuda.FloatTensor 类型的对象，但为参数 #2 'mat2' 找到了类型 torch.FloatTensor

表明 matmul 的第二个参数（我声明的实例张量）在 CPU 上，并且预计在 GPU 上（作为模型和数据的其余部分）。

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

我什至尝试通过以下方式覆盖 cuda() 方法：

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

数据已经在 GPU 中，这意味着以下代码行已经执行：

X = X.cuda()

错误还明确说明了 matmul 的参数 2，在这种情况下是张量（称为矩阵）而不是 X。

标签： pythonmachine-learningartificial-intelligencepytorch

解决方案

让我们假设以下内容：

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

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

nn.Module.cuda()由于方法和方法之间的细微差别，这也不起作用torch.Tensor.cuda()。

虽然nn.Module.cuda()将所有的Parameters和移动Buffers到ModuleGPU 并返回自身，torch.Tensor.cuda()但只返回 GPU 上张量的副本。

原始张量不受影响。

总之，要么：

  1. 将您的属性包装matrix为Parameteror
  2. 通过以下方式将 GPU 副本分配回矩阵：
self.matrix = self.matrix.cuda()

在您的覆盖中。

我建议第一个。

推荐阅读