maskrcnn-benchmark 代码详解之 balanced_positive_negative_sampler.py

前言

　　在二阶段目标检测的卷积神经网络模型中，在经过RPN筛选出相应的候选边框（Proposal)中应既包含内容为背景的边框又包含内容为目标的边框。这就需要一次对这两种边框数目的调整，以使得模型训练效果更好，这两种边框的选择比例一般为１：１，但是随着检测任务的不同而不同。在maskrcnn benchmark中实现这一操作的代码为balanced_positive_negative_sampler.py：

# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
import torch


class BalancedPositiveNegativeSampler(object):
    """
    This class samples batches, ensuring that they contain a fixed proportion of positives
    """
    """
    这个类是用来平衡选出来的候选框中正样本和负样本的比例平衡一些
    """

    def __init__(self, batch_size_per_image, positive_fraction):
        """
        Arguments:
            batch_size_per_image (int): number of elements to be selected per image
            positive_fraction (float): percentace of positive elements per batch
        """
        """
        参数：
            batch_size_per_image：每一张图片上提取的候选框的个数
            positive_fraction：在提取的候选框中正样本应占有的比例
        """
        self.batch_size_per_image = batch_size_per_image
        self.positive_fraction = positive_fraction

    def __call__(self, matched_idxs):
        """
        Arguments:
            matched idxs: list of tensors containing -1, 0 or positive values.
                Each tensor corresponds to a specific image.
                -1 values are ignored, 0 are considered as negatives and > 0 as
                positives.

        Returns:
            pos_idx (list[tensor])
            neg_idx (list[tensor])

        Returns two lists of binary masks for each image.
        The first list contains the positive elements that were selected,
        and the second list the negative example.
        """
        """
        参数：
            matched idxs:与获得的候选框相匹配的目标标签，值为-1,0或者是正的值。
                        　它包含n个张量，n为每次训练的图片个数，每个张量对应一张图片
                        -1为需要忽略的边框，一般为0.3= 1).squeeze(1)
            negative = torch.nonzero(matched_idxs_per_image == 0).squeeze(1)

            # 得到根据默认参数算出的目标边框的个数
            num_pos = int(self.batch_size_per_image * self.positive_fraction)
            # protect against not enough positive examples
            # 当得到的目标边框少于默认参数需要的边框时，从实际出发，将边框数设置为较小的那一个
            num_pos = min(positive.numel(), num_pos)
            # 得到根据默认参数算出的背景边框的个数
            num_neg = self.batch_size_per_image - num_pos
            # protect against not enough negative examples
            # 当得到的背景边框少于默认参数需要的边框时，从实际出发，将边框数设置为较小的那一个
            num_neg = min(negative.numel(), num_neg)

            # randomly select positive and negative examples
            # 随机打乱得到的边框
            perm1 = torch.randperm(positive.numel(), device=positive.device)[:num_pos]
            perm2 = torch.randperm(negative.numel(), device=negative.device)[:num_neg]

            pos_idx_per_image = positive[perm1]
            neg_idx_per_image = negative[perm2]

            # create binary mask from indices
            # 根据索引为确定得到图片边框内是否建立mask
            pos_idx_per_image_mask = torch.zeros_like(
                matched_idxs_per_image, dtype=torch.uint8
            )
            neg_idx_per_image_mask = torch.zeros_like(
                matched_idxs_per_image, dtype=torch.uint8
            )
            pos_idx_per_image_mask[pos_idx_per_image] = 1
            neg_idx_per_image_mask[neg_idx_per_image] = 1

            pos_idx.append(pos_idx_per_image_mask)
            neg_idx.append(neg_idx_per_image_mask)

        return pos_idx, neg_idx