基于pytorch分割模型使用OHEM
OHEM在分割模型的应用,让模型更专注于难分类样本点(预测的概率值低)
def sample(self, seg_logit, seg_label):
"""Sample pixels that have high loss or with low prediction confidence.
Args:
seg_logit (torch.Tensor): segmentation logits, shape (N, C, H, W)
seg_label (torch.Tensor): segmentation label, shape (N, 1, H, W)
Returns:
torch.Tensor: segmentation weight, shape (N, H, W)
"""
with torch.no_grad():
assert seg_logit.shape[2:] == seg_label.shape[2:]
assert seg_label.shape[1] == 1
seg_label = seg_label.squeeze(1).long()
batch_kept = self.min_kept * seg_label.size(0)
valid_mask = seg_label != self.context.ignore_index
# 生成和seg_logit相同类型、相同device的0 tensor,形状和seg_label相同:(N, 1, H, W)
seg_weight = seg_logit.new_zeros(size=seg_label.size())
# 将seg_weight拉直为Nx1xHxW
valid_seg_weight = seg_weight[valid_mask]
# 使用softmax对预测值计算概率值
seg_prob = F.softmax(seg_logit, dim=1)
tmp_seg_label = seg_label.clone().unsqueeze(1)
tmp_seg_label[tmp_seg_label == self.context.ignore_index] = 0
# 通过gather在第1维按照索引tmp_seg_label取值
seg_prob = seg_prob.gather(1, tmp_seg_label).squeeze(1)
sort_prob, sort_indices = seg_prob[valid_mask].sort()
# 按照最低数量和排序重新计算阈值
if sort_prob.numel() > 0:
min_threshold = sort_prob[min(batch_kept,
sort_prob.numel() - 1)]
else:
min_threshold = 0.0
threshold = max(min_threshold, self.thresh)
valid_seg_weight[seg_prob[valid_mask] < threshold] = 1.
# 按照索引重新赋值回去
seg_weight[valid_mask] = valid_seg_weight
return seg_weight