用自写的loss计算很慢的问题

Question

背景：一个图像预测任务，图像并不是指常见的jpg等格式的图片，其实就是网格内每个像素点有数值的那种类似于热点图的图像，如下图

问题：如上图所示，一张图里其实有很大一片区域是0值。但搭的网络输入输出是整图，维度为（32，224，224）。最近意识到计算loss不可以将整图放入计算，只想要关注body内有值的区域，所以就用上自写的loss function，出现了计算极其缓慢的现象。
——
train预测过程：x为输入，y为GT，y_pred为预测结果，b是与y维度相同的、只有0或1的mask，b==1就是body内区域。我的做法是把b展平成一维，再用np.where读取b==1的坐标到BODY，同时y，y_pred也展平成一维。放入自写的smooth_L1计算loss

        x = np.array(inputList)
        y = np.array(labelList)
        b = np.array(bodyList)
        b = b.flatten()
        BODY = np.where(b == 1)
        BODY = np.array(BODY)
        BODY = BODY.flatten()
        x = np.float32(x)
        y = np.float32(y)
        x = torch.tensor(x)
        y = torch.tensor(y)
        BODY = torch.tensor(BODY)
        x = x.to(device)
        y = y.to(device)
        BODY = BODY.to(device)
        optimizer.zero_grad()
        y_pred = model(x)
        y_pred = y_pred.view(-1)
        y = y.view(-1)
        train_loss = smooth_l1(y_pred, y, BODY)
        print(f"第 {t + 1}/{epoch} epoch, 第 {i + 1}/{batch_train} batch, train_loss: {train_loss.item()}")
        train_loss.backward()
        optimizer.step()

自写的loss：就是一个smoothL1 Loss，但是我只想计算body内区域。所以用了一个for循环读取mask里（也就是BODY里）存好的坐标。

# 定义损失函数
def smooth_item(x, beta):
    if x < beta:
        loss_item = (0.5*torch.square(x))/beta
    else:
        loss_item = x-0.5*beta
    return loss_item
def smooth_l1(y_pred, y, mask, beta=1):
    n = len(mask)
    loss = 0
    for i in mask:
        loss += smooth_item(torch.abs(y_pred[i]-y[i]), beta)
    return loss/n

debug发现这样子计算loss好慢好慢，有什么方法优化吗？要说用GPU加速，但这样子一个个坐标读取再累加也不是一个并行过程呀

Answer 1

参考GPT和自己的思路，您可以尝试使用torch.where()函数，而不是用np.where()函数，这可以使您在GPU上并行化处理，从而提高您的计算速度。具体来说，您可以使用torch.where()函数来创建一个bool类型的掩码张量，该张量的形状与输入张量相同，但是只有掩码中的值为True时，才会计算损失。这样就不必将掩码扁平化并使用for循环逐个读取坐标，而是可以使用张量的广播机制，使每个掩码元素与相应的预测值和目标值相匹配，从而并行计算损失。

下面是使用torch.where()函数修改您的smooth_l1函数的示例：

def smooth_l1(y_pred, y, mask, beta=1):
    loss_item = torch.abs(y_pred - y)
    loss_item = torch.where(loss_item < beta, 0.5 * torch.pow(loss_item, 2) / beta, loss_item - 0.5 * beta)
    loss = torch.sum(loss_item * mask) / torch.sum(mask)
    return loss

在这个版本的smooth_l1函数中，我们首先计算了预测值与目标值之间的差异，然后使用torch.where()函数来计算smooth L1损失。最后，我们将掩码应用于损失，并在返回之前求出平均值。