在 unet 架构中使用自定义权重图的正确方法

Correct way to use custom weight maps in unet architecture

u-net 架构中有一个著名的技巧,即使用自定义权重图来提高准确性。以下是它的详细信息:

现在,通过在这里和其他多个地方询问,我了解了两种方法。我想知道哪个是正确的,或者有没有其他更正确的方法?

  1. 首先是在训练循环中使用torch.nn.Functional方法:

    loss = torch.nn.functional.cross_entropy(output, target, w) 其中 w 将是计算的自定义权重。

  2. 其次是在训练循环外损失函数的调用中使用reduction='none' criterion = torch.nn.CrossEntropy(reduction='none')

    然后在训练循环中乘以自定义权重:

    gt # Ground truth, format torch.long
    pd # Network output
    W # per-element weighting based on the distance map from UNet
    loss = criterion(pd, gt)
    loss = W*loss # Ensure that weights are scaled appropriately
    loss = torch.sum(loss.flatten(start_dim=1), axis=0) # Sums the loss per image
    loss = torch.mean(loss) # Average across a batch
    

现在,我有点困惑哪一个是正确的,或者有没有其他方法,或者两者都是正确的?

请注意 torch.nn.CrossEntropyLoss() 是调用 torch.nn.functional 的 class。 参见 https://pytorch.org/docs/stable/_modules/torch/nn/modules/loss.html#CrossEntropyLoss

您可以在定义条件时使用权重。从功能上比较它们,两种方法是相同的。

现在,我不明白你在方法 1 中计算训练循环内的损失和方法 2 中训练循环外的损失的想法。如果你在循环外计算损失,那么你将如何反向传播?

加权部分看起来只是简单的加权交叉熵,它对 类 的数量(下例中为 2)如此执行。

weights = torch.FloatTensor([.3, .7])
loss_func = nn.CrossEntropyLoss(weight=weights)

编辑:

你看过 Patrick Black 的 this implementation 吗?

# Set properties
batch_size = 10
out_channels = 2
W = 10
H = 10

# Initialize logits etc. with random
logits = torch.FloatTensor(batch_size, out_channels, H, W).normal_()
target = torch.LongTensor(batch_size, H, W).random_(0, out_channels)
weights = torch.FloatTensor(batch_size, 1, H, W).random_(1, 3)

# Calculate log probabilities
logp = F.log_softmax(logits)

# Gather log probabilities with respect to target
logp = logp.gather(1, target.view(batch_size, 1, H, W))

# Multiply with weights
weighted_logp = (logp * weights).view(batch_size, -1)

# Rescale so that loss is in approx. same interval
weighted_loss = weighted_logp.sum(1) / weights.view(batch_size, -1).sum(1)

# Average over mini-batch
weighted_loss = -1. * weighted_loss.mean()