Question

上下文

我正在构建超分辨率成像模型（U-Net）。我的数据包括显微血液涂片。我正在使用快速AI。我想尝试一个新的损失函数，该函数将“专注于”细胞的成分。

我做了什么

我创建了一个 pseudo_mask 函数，该函数带有两个参数，即图像和范围在0到255之间的整数β。此函数首先计算单元格分割掩码，然后将其转换为将mask的0值设置为β，并将结果应用于原始图像。

Pseudo_mask output of the same image, with β values of respectively 0, 122 and 255 (e.g. the original image)

我想创建一个自定义损失函数，该函数将在pseudo_mask函数的输出上计算MSE并将其反向传播到原始图像。

我的代码相当简单，因为我只是简单地“完成”了Pytorch MSE损失功能。

class MSE_Mask_Loss(_Loss):

   def __init__(self, size_average=None, reduce=None, reduction: str = 'mean') -> None:
       super(MSE_Mask_Loss, self).__init__(size_average, reduce, reduction)

   def forward(self, input: Tensor, target: Tensor) -> Tensor:

        input = tensor_pseudo_mask(input,255)
        target = tensor_pseudo_mask(target,255)
             
       return F.mse_loss(input, target, reduction=self.reduction)

代码运行良好，但我认为反向传播存在一些问题，因为这种损失和β= 255所获得的结果是可怕的，而它们应该等同于原始图像所获得的（良好）结果。

我做错了什么？

修改

tensor_pseudo_mask定义如下：

def tensor_pseudo_mask(tns,val):

  for i in range(len(tns)):

    tns[i,:,:,:] = pseudo_mask(tns[i,:,:,:],val).data

  return tns.cuda()

其中

def convert(x, val):

  if x == 0:
    return val
  else:
    return x

vecfunc = np.vectorize(convert)

def apply_mask(img, mask):

  tensor_mask = torch.unsqueeze(torch.Tensor(mask) / 255, 2) #For broadcasting
  img_masked = tensor_mask * torch.Tensor(convert_Image_RGB(img))

  return Image(convert_RGB_Image(img_masked)) 

def pseudo_mask(img, val):

  mask_value = torch.Tensor(np.array(mask(img)))
  pseudo_mask = vecfunc(mask_value, val)
  applied_pseudo_mask = apply_mask(img,pseudo_mask)

  return applied_pseudo_mask

自定义损失函数的反向传播问题

0 个答案: