OHEM代码梳理

àì夳堔傛蜴生んèń 2022-02-22 04:58 255阅读 0赞

*  **传送门**：

1.  相关OHEM的介绍：[检测模型改进—OHEM与Focal-Loss算法总结][OHEM_Focal-Loss]
2.  代码地址：[OHEM][]

# 1. 前言 #

有关OHEM的介绍请参考上面给出的链接，这里主要就OHEM是怎么运行的做一些简单的分析，整个OHEM的代码也不是很多，这里将算法的步骤归纳为：  
1）计算检测器的损失，这部分是使用和最后fc6、fc7预测头一样的共享参数，预测分类与边界框回归的结果，将预测的结果与GT进行比较得到分类和边界框回归的loss，这里的损失是将两种损失相加得到的；  
2）使用阈值为0.7的NMS预先处理一遍检测框，去除一些无效的检测框；  
3）NMS之后的检测框按照loss由大到小排列，选取一定数目（由两个数取最小决定）的边界框返回。  
下面是OHEM在网络定义文件中的定义，方便后面查看相关代码的时候查找对应条目。

layer {
      name: "hard_roi_mining"
      type: "Python"
      bottom: "cls_prob_readonly"
      bottom: "bbox_pred_readonly"
      bottom: "rois"
      bottom: "labels"
      bottom: "bbox_targets"
      bottom: "bbox_inside_weights"
      bottom: "bbox_outside_weights"
      top: "rois_hard"
      top: "labels_hard"
      top: "bbox_targets_hard"
      top: "bbox_inside_weights_hard"
      top: "bbox_outside_weights_hard"
      propagate_down: false
      propagate_down: false
      propagate_down: false
      propagate_down: false
      propagate_down: false
      propagate_down: false
      propagate_down: false
      python_param {
        module: "roi_data_layer.layer"
        layer: "OHEMDataLayer"
        param_str: "'num_classes': 6" #6
      }
    }

# 2. OHEM代码简单梳理 #

## 2.1 OHEMDataLayer ##

class OHEMDataLayer(caffe.Layer):
        """Online Hard-example Mining Layer."""
        def setup(self, bottom, top):
            """Setup the OHEMDataLayer."""
    
            # parse the layer parameter string, which must be valid YAML
            layer_params = yaml.load(self.param_str_)
    
            self._num_classes = layer_params['num_classes']	 # 获取分类数目
    
            self._name_to_bottom_map = {   # 将bottom的blob名称与index使用dict关联
                'cls_prob_readonly': 0,
                'bbox_pred_readonly': 1,
                'rois': 2,
                'labels': 3}
    
            if cfg.TRAIN.BBOX_REG:  # 有边界框回归
                self._name_to_bottom_map['bbox_targets'] = 4
                self._name_to_bottom_map['bbox_loss_weights'] = 5
    
            self._name_to_top_map = { }  # 同理top的blob名称也要与index关联起来
            ……
    	
    	# 前向传播函数
        def forward(self, bottom, top):
            """Compute loss, select RoIs using OHEM. Use RoIs to get blobs and copy them into this layer's top blob vector."""
    
            cls_prob = bottom[0].data  # 获取对应bottom的数据
            bbox_pred = bottom[1].data
            rois = bottom[2].data
            labels = bottom[3].data
            if cfg.TRAIN.BBOX_REG:
                bbox_target = bottom[4].data
                bbox_inside_weights = bottom[5].data
                bbox_outside_weights = bottom[6].data
            else:
                bbox_target = None
                bbox_inside_weights = None
                bbox_outside_weights = None
    
            flt_min = np.finfo(float).eps
            # 计算分类的损失
            loss = [ -1 * np.log(max(x, flt_min)) \
                for x in [cls_prob[i,label] for i, label in enumerate(labels)]]
    
            # 计算边界框回归的损失，并且将两个损失加起来
            if cfg.TRAIN.BBOX_REG:
                # bounding-box regression loss
                # d := w * (b0 - b1)
                # smoothL1(x) = 0.5 * x^2 if |x| < 1
                # |x| - 0.5 otherwise
                def smoothL1(x):
                    if abs(x) < 1:
                        return 0.5 * x * x
                    else:
                        return abs(x) - 0.5
    
                bbox_loss = np.zeros(labels.shape[0])  # 边界框损失
                for i in np.where(labels > 0 )[0]:
                    indices = np.where(bbox_inside_weights[i,:] != 0)[0]
                    bbox_loss[i] = sum(bbox_outside_weights[i,indices] * [smoothL1(x) \
                        for x in bbox_inside_weights[i,indices] * (bbox_pred[i,indices] - bbox_target[i,indices])])
                loss += bbox_loss  # 两者损失相加
    
            # 筛选出损失比较大的返回
            blobs = get_ohem_minibatch(loss, rois, labels, bbox_target, \
                bbox_inside_weights, bbox_outside_weights)
    
    		# 给top blob赋值
            for blob_name, blob in blobs.iteritems():
                top_ind = self._name_to_top_map[blob_name]
                # Reshape net's input blobs
                top[top_ind].reshape(*(blob.shape))
                # Copy data into net's input blobs
                top[top_ind].data[...] = blob.astype(np.float32, copy=False)

## 2.2 get\_ohem\_minibatch ##

# 获取OHEM训练的batch
    def get_ohem_minibatch(loss, rois, labels, bbox_targets=None,
                           bbox_inside_weights=None, bbox_outside_weights=None):
        """Given rois and their loss, construct a minibatch using OHEM."""
        loss = np.array(loss)
    	
    	# 使用NMS过滤检测框
        if cfg.TRAIN.OHEM_USE_NMS:	# NMS thresh=0.7
            # Do NMS using loss for de-dup and diversity
            keep_inds = []
            nms_thresh = cfg.TRAIN.OHEM_NMS_THRESH  # 0.7
            source_img_ids = [roi[0] for roi in rois] # 0或1（背景与前景）
            for img_id in np.unique(source_img_ids):
                for label in np.unique(labels):
                    sel_indx = np.where(np.logical_and(labels == label, \
                                        source_img_ids == img_id))[0]
                    if not len(sel_indx):
                        continue
                    boxes = np.concatenate((rois[sel_indx, 1:],
                            loss[sel_indx][:,np.newaxis]), axis=1).astype(np.float32)
                    keep_inds.extend(sel_indx[nms(boxes, nms_thresh)])
    
            hard_keep_inds = select_hard_examples(loss[keep_inds])  # 按照损失排序选择样本
            hard_inds = np.array(keep_inds)[hard_keep_inds]  # 最后保留下来的困难样本索引
        else:
            hard_inds = select_hard_examples(loss)
    
        blobs = { 'rois_hard': rois[hard_inds, :].copy(),
                 'labels_hard': labels[hard_inds].copy()}
        if bbox_targets is not None:
            assert cfg.TRAIN.BBOX_REG
            blobs['bbox_targets_hard'] = bbox_targets[hard_inds, :].copy()
            blobs['bbox_inside_weights_hard'] = bbox_inside_weights[hard_inds, :].copy()
            blobs['bbox_outside_weights_hard'] = bbox_outside_weights[hard_inds, :].copy()
    
        return blobs
    
    def select_hard_examples(loss):
        """Select hard rois."""
        # Sort and select top hard examples.
        sorted_indices = np.argsort(loss)[::-1]
        hard_keep_inds = sorted_indices[0:np.minimum(len(loss), cfg.TRAIN.BATCH_SIZE)]
        # (explore more ways of selecting examples in this function; e.g., sampling)
    
        return hard_keep_inds

[OHEM_Focal-Loss]: https://blog.csdn.net/m_buddy/article/details/89042338
[OHEM]: https://github.com/abhi2610/ohem