PyTorch U-net使用Focal Loss出现CUDA张量转NumPy错误如何解决

你引入的第三方focal_loss包内部实现依赖numpy数组运算，没有适配PyTorch的CUDA张量与自动梯度机制：

• GPU运行时，显存中的CUDA张量无法直接转换为内存中的numpy数组，触发can't convert cuda:0 device type tensor to numpy报错
• CPU运行时，模型输出的张量带有梯度追踪属性，PyTorch禁止直接将带梯度的张量转为numpy数组，触发Can't call numpy() on Tensor that requires grad报错

直接放弃第三方focal_loss包，改用纯PyTorch张量操作实现的Focal Loss，同时修改代码中不合理的归一化逻辑，具体操作如下：

• 先卸载已安装的第三方focal_loss包，避免导入冲突：
  pip uninstall focal-loss

1. 实现PyTorch版二分类Focal Loss

import torch
import torch.nn as nn
import torch.nn.functional as F

class BinaryFocalLoss(nn.Module):
    def __init__(self, alpha=0.25, gamma=2, reduction='mean'):
        super().__init__()
        self.alpha = alpha
        self.gamma = gamma
        self.reduction = reduction

    def forward(self, inputs, targets):
        #  inputs为模型输出经sigmoid激活后的结果，shape与targets一致
        BCE_loss = F.binary_cross_entropy(inputs, targets, reduction='none')
        pt = torch.exp(-BCE_loss)
        F_loss = self.alpha * (1-pt)**self.gamma * BCE_loss

        if self.reduction == 'mean':
            return torch.mean(F_loss)
        elif self.reduction == 'sum':
            return torch.sum(F_loss)
        return F_loss

2. 修改训练逻辑中的输出处理代码
   你当前对模型输出手动做min-max归一化的操作会破坏模型的学习分布，完全没有必要，替换为sigmoid激活即可：
   把fit函数里两处y_hat = (y_hat - torch.min(y_hat))/(torch.max(y_hat) - torch.min(y_hat))
   修改为：

y_hat = torch.sigmoid(y_hat)

3. 检查IOU计算函数
   如果你自行实现的iou函数中有将张量转为numpy数组的操作，要么改为纯张量运算，要么转换前先调用.detach().cpu()剥离梯度并转到CPU，参考实现如下：

def iou(y_hat, masks, threshold=0.5):
    y_hat = (y_hat > threshold).float()
    intersection = (y_hat * masks).sum()
    union = y_hat.sum() + masks.sum() - intersection
    return (intersection / (union + 1e-6)).item()

修改完成后即可正常在GPU上训练，不会再触发对应报错。

内容的提问来源于stack exchange，提问作者Lars Bosma