RuntimeError_ Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor)

两个解决方案

方案一:检查网络和数据是否都在GPU上

RuntimeError:输入类型(torch.FloatTensor)和权重类型(torch.cuda.FloatTensor)应该相同,或者输入应该是一个MKLDNN张量,而权重是一个密集张量。

我们的权重是cuda类型(GPU训练得到),而输入(要测试的数据)不是cuda类型。 采用GPU训练的模型,不能直接在CPU上使用,也要放到GPU中预测。

如果用的模型参数是在cuda上(gpu)训练的,在使用其进行测试时,需要将要测试的数据也放到GPU上,即:data.cuda()

仔细检查一下神经网络和相关的数据是否都放在了GPU上!

方案二:模型定义全放在类的初始化里

此外,根据大佬的分析,也有可能,是因为你的模型在定义的时候,没有定义好,导致模型的一部分在加载的时候没有办法转移到 cuda上。 具体测试就略过了。

总的来说,如果上面的方法不奏效,那么试一试把所有的网络层(只要有参数需要训练的网络层)都放到 __init__()函数里面去定义,只在 forward()中写运行时的逻辑,即:

class A(nn.Module):
    def __init__(self):
        super(A,self).__init__()
        self.conv = nn.Conv2d(in_channels=3, out_channels=8, kernel_size=3)
        self.relu = nn.ReLU(inplace=True)
        self.b_module = B()

    def forward(self,x):
        out = self.conv(x)
        out = self.relu(out)
        out = self.b_module(out)
        return out

class B(nn.Module):
    def __init__(self):
        super(B,self).__init__()
        self.conv = nn.Conv2d(in_channels=8, out_channels=16, kernel_size=3)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        out = self.conv(x)
        out = self.relu(out)
        return out

参考:

  • Pytorch避坑之:RuntimeError: Input type(torch.cuda.FloatTensor) and weight type(torch.FloatTensor) shoul
  • 来源:lucky-wz

    物联沃分享整理
    物联沃-IOTWORD物联网 » RuntimeError_ Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor)

    发表评论