python - FocalLoss TypeError：预期的 CPU（得到 CUDA）

Question

我在训练我的网络时遇到了这个错误。从其他帖子中，我知道我应该使用 .cuda() 或 .to("cuda:0") 将张量放在 cuda 设备上。但是，当我这样做时，同样的错误不断出现。

---> 23         pt = Tensor(logpt.data.exp())
     24 
     25         if self.alpha is not None:

TypeError: expected CPU (got CUDA)

我尝试在第 23 行之后添加 .cuda()，但也在 self.alpha = alpha 和其他一些地方之后添加。我错过了什么吗？或者是我应该检查的模型代码。很确定整个模型都在 cuda 上。感谢您的任何见解。

import torch
from torch import nn, Tensor
import torch.nn.functional as F
class FocalLoss(nn.Module):
    def __init__(self, gamma=0, alpha=None, size_average=True):
        super(FocalLoss, self).__init__()
        self.gamma = gamma
        self.alpha = alpha.cuda()
        if isinstance(alpha,(float,int)): self.alpha = torch.Tensor([alpha,1-alpha])
        if isinstance(alpha,list): self.alpha = torch.Tensor(alpha)
        self.size_average = size_average

    def forward(self, input, target):
        if input.dim()>2:
            input = input.view(input.size(0),input.size(1),-1)  # N,C,H,W => N,C,H*W
            input = input.transpose(1,2)    # N,C,H*W => N,H*W,C
            input = input.contiguous().view(-1,input.size(2))   # N,H*W,C => N*H*W,C
        target = target.view(-1,1)

        logpt = F.log_softmax(input)
        logpt = logpt.gather(1,target)
        logpt = logpt.view(-1)
        pt = Tensor(logpt.data.exp())

        if self.alpha is not None:
            if self.alpha.type()!=input.data.type():
                self.alpha = self.alpha.type_as(input.data)
            at = self.alpha.gather(0,target.data.view(-1))
            logpt = logpt * Tensor(at)

        loss = -1 * (1-pt)**self.gamma * logpt
        if self.size_average: return loss.mean()
        else: return loss.sum()

标签： pythonpytorch