『MXNet』第十二弹_再谈新建计算节点

上一节我们已经谈到了计算节点，但是即使是官方文档介绍里面相关内容也过于简略，我们使用Faster-RCNN代码中的新建节点为例，重新介绍一下新建节点的调用栈。

1、调用新建节点

参数分为三部分，op_type是节点名称，对应于辅助class的装饰器的输入；其他参数一部分传递给辅助class的初始化函数（这部分参数的虚参名和初始化函数的需参名要对应上），一部分直接作为一个list传给节点定义class的forward函数的in_data参数。

group = mx.symbol.Custom(rois=rois,                     # 2000*5的roi信息，CustomOpProp无此参数
                         gt_boxes=gt_boxes,             # n*5的ground truth信息，n表示object数量，CustomOpProp无此参数
                         op_type='proposal_target',     # <-----对应辅助节点装饰器参数
                         # CustomOp，CustomOpProp初始化参数
                         num_classes=num_classes,       # num_classes是实际要分类的类别数加上背景类
                         batch_images=rcnn_batch_size,  # 1
                         batch_rois=rcnn_batch_rois,    # 128
                         fg_fraction=rcnn_fg_fraction,  #s 0.25，正样本所占的比例
                         fg_overlap=rcnn_fg_overlap,    # 0.5
                         box_stds=rcnn_bbox_stds        # (0.1, 0.1, 0.2, 0.2)
                         )

2、建立辅助class：CustomOpProp

本部分的class方法参数固定，不可以随意修改

@mx.operator.register('proposal_target')  # <-----对应调用的op_type
class ProposalTargetProp(mx.operator.CustomOpProp):
    def __init__(self, num_classes='21', batch_images='1', batch_rois='128', fg_fraction='0.25',  # 接受上面调用的后5个参数
                 fg_overlap='0.5', box_stds='(0.1, 0.1, 0.2, 0.2)'):
        super(ProposalTargetProp, self).__init__(need_top_grad=False)                 # <-----本节点是否需要后面的梯度
        self._num_classes = int(num_classes)                                          # num_classes是实际要分类的类别数加上背景类
        self._batch_images = int(batch_images)                                        # 1
        self._batch_rois = int(batch_rois)                                            # 128
        self._fg_fraction = float(fg_fraction)                                        # 0.25，正样本所占的比例
        self._fg_overlap = float(fg_overlap)                                          # 0.5
        self._box_stds = tuple(np.fromstring(box_stds[1:-1], dtype=float, sep=','))   # (0.1, 0.1, 0.2, 0.2)

    def list_arguments(self):
        return ['rois', 'gt_boxes']  # 向前传播需要的参数

    def list_outputs(self):
        return ['rois_output', 'label', 'bbox_target', 'bbox_weight']  # 向前传播输出参数名

    def infer_shape(self, in_shape):
        assert self._batch_rois % self._batch_images == 0, \
            'BATCHIMAGES {} must devide BATCH_ROIS {}'.format(self._batch_images, self._batch_rois)

        rpn_rois_shape = in_shape[0]
        gt_boxes_shape = in_shape[1]

        output_rois_shape = (self._batch_rois, 5)
        label_shape = (self._batch_rois, )
        bbox_target_shape = (self._batch_rois, self._num_classes * 4)
        bbox_weight_shape = (self._batch_rois, self._num_classes * 4)

        return [rpn_rois_shape, gt_boxes_shape], \
               [output_rois_shape, label_shape, bbox_target_shape, bbox_weight_shape]

    def create_operator(self, ctx, shapes, dtypes):  # 返回初始化了的自定义节点类
        return ProposalTargetOperator(self._num_classes, self._batch_images, self._batch_rois, self._fg_fraction,
                                      self._fg_overlap, self._box_stds)

    def declare_backward_dependency(self, out_grad, in_data, out_data):
        return []

__init__

初始化方法会首先从调用位置接收参数，调用位置的op_type参数用于指定选用哪个辅助class，然后其他参数优先传入本初的__init__方法，剩下的没有和本方法参数对应上的参数会做为节点class的forward方法的in_data参数。

下面一行表示本节点不需要接收后面层的梯度，对应到节点定义class的backward方法，out_grad参数就不应在函数体内调用了，in_grad（向前传播回去的梯度）计算完全依赖本层的参数。

super(ProposalTargetProp, self).__init__(need_top_grad=False)

list_arguments

使用list_arguments()方法返回时，输出并不直接就是上述代码list_arguments方法的return列表，实际上会将整个net结构截至的本节点为止，全部的variable变量名称输入。

对于本节点group，定义输入有两个：

rois=rois  # 2000*5的roi信息，CustomOpProp无此参数
gt_boxes=gt_boxes

其中gt_boxes变量本身是个variable（定义为gt_boxes = mx.symbol.Variable(name="gt_boxes")），但是rois为symbol，是之前网络的输出symbol，所以实际输出的arguments为gt_boxes本身，以及rois所依赖的全部variables，包含认为定义的占位符variable和网络层自带的参数variable。

['data', 'conv1_1_weight', 'conv1_1_bias', 'conv1_2_weight', 'conv1_2_bias', 'conv2_1_weight', 'conv2_1_bias', 'conv2_2_weight', 'conv2_2_bias', 'conv3_1_weight', 'conv3_1_bias', 'conv3_2_weight', 'conv3_2_bias', 'conv3_3_weight', 'conv3_3_bias', 'conv4_1_weight', 'conv4_1_bias', 'conv4_2_weight', 'conv4_2_bias', 'conv4_3_weight', 'conv4_3_bias', 'conv5_1_weight', 'conv5_1_bias', 'conv5_2_weight', 'conv5_2_bias', 'conv5_3_weight', 'conv5_3_bias', 'rpn_conv_3x3_weight', 'rpn_conv_3x3_bias', 'rpn_cls_score_weight', 'rpn_cls_score_bias', 'rpn_bbox_pred_weight', 'rpn_bbox_pred_bias', 'im_info',

'gt_boxes']

infer_shape

其输入参数in_shape就是list_arguments中return的那几个变量的shape，对于本例，就是rois和gtboxes的shape，本方法用于推断输出symbol和梯度symbol的shape是否正确。

3、实现节点class

节点class的初始化和调用部分的参数完全无关，是由辅助节点来进行传参调用的。但是其forward方法的in_data参数其值接收是从调用初进行的，in_data中的参数就是上面list_arguments方法的return结果（['rois', 'gt_boxes']），实际传参可以有空缺（例如第一小节可以删掉gt_boxes），缺省参数视为定义了一个Variable占位。

class ProposalTargetOperator(mx.operator.CustomOp):
    def __init__(self, num_classes, batch_images, batch_rois, fg_fraction, fg_overlap, box_stds):
        super(ProposalTargetOperator, self).__init__()
        self._num_classes = num_classes     # num_classes是实际要分类的类别数加上背景类
        self._batch_images = batch_images   # 1
        self._batch_rois = batch_rois       # 128
        self._rois_per_image = int(batch_rois / batch_images)
        self._fg_rois_per_image = int(round(fg_fraction * self._rois_per_image))
        self._fg_overlap = fg_overlap       # 0.5
        self._box_stds = box_stds           # (0.1, 0.1, 0.2, 0.2)

    def forward(self, is_train, req, in_data, out_data, aux):
        """Forward interface. Can override when creating new operators.

        Parameters
        ----------
        is_train : bool
            whether this is for training
        req : list of str
            how to assign to out_data. can be 'null', 'write', or 'add'.
            You can optionally use self.assign(dst, req, src) to handle this.
        in_data, out_data, aux: list of NDArrays
            input, output, and auxiliary states for forward. See document for
            corresponding arguments of Operator::Forward
        """
        assert self._batch_images == in_data[1].shape[0], 'check batch size of gt_boxes'

        all_rois = in_data[0].asnumpy()  # [2000, 5]
        all_gt_boxes = in_data[1].asnumpy()  # [n, 5]

        rois = np.empty((0, 5), dtype=np.float32)
        labels = np.empty((0, ), dtype=np.float32)
        bbox_targets = np.empty((0, 4 * self._num_classes), dtype=np.float32)
        bbox_weights = np.empty((0, 4 * self._num_classes), dtype=np.float32)
        for batch_idx in range(self._batch_images):
            b_rois = all_rois[np.where(all_rois[:, 0] == batch_idx)[0]]
            b_gt_boxes = all_gt_boxes[batch_idx]
            b_gt_boxes = b_gt_boxes[np.where(b_gt_boxes[:, -1] > 0)[0]]

            # Include ground-truth boxes in the set of candidate rois
            batch_pad = batch_idx * np.ones((b_gt_boxes.shape[0], 1), dtype=b_gt_boxes.dtype)
            b_rois = np.vstack((b_rois, np.hstack((batch_pad, b_gt_boxes[:, :-1]))))

            b_rois, b_labels, b_bbox_targets, b_bbox_weights = \
                sample_rois(b_rois, b_gt_boxes, num_classes=self._num_classes, rois_per_image=self._rois_per_image,
                            fg_rois_per_image=self._fg_rois_per_image, fg_overlap=self._fg_overlap, box_stds=self._box_stds)

            rois = np.vstack((rois, b_rois))
            labels = np.hstack((labels, b_labels))
            bbox_targets = np.vstack((bbox_targets, b_bbox_targets))
            bbox_weights = np.vstack((bbox_weights, b_bbox_weights))

        self.assign(out_data[0], req[0], rois)
        self.assign(out_data[1], req[1], labels)
        self.assign(out_data[2], req[2], bbox_targets)
        self.assign(out_data[3], req[3], bbox_weights)

    def backward(self, req, out_grad, in_data, out_data, in_grad, aux):
        """Backward interface. Can override when creating new operators.

        Parameters
        ----------
        req : list of str
            how to assign to in_grad. can be 'null', 'write', or 'add'.
            You can optionally use self.assign(dst, req, src) to handle this.
        out_grad, in_data, out_data, in_grad, aux : list of NDArrays
            input and output for backward. See document for
            corresponding arguments of Operator::Backward
        """
        self.assign(in_grad[0], req[0], 0)
        self.assign(in_grad[1], req[1], 0)

介绍完辅助节点，本部分的介绍就不太多了，注意的就是向前向后两个方法没有返回值，使用assign来给symbol赋值，数量顺序要和辅助class的argument、output对应上，具体实现因为没有研究C++源码，没办法更详细介绍了，不过不影响使用（大概）。

另外，辅助节点class是会在python解释器里直接执行的，也就是说我们添加在函数体中的print什么的能够得到输出，但是在本class中，添加的中间输出不会被print出来，应该是建立符号图时被略去了（有关C++优化计算图的机理，李沐博士有介绍，不过我的C++功底不够，没有看过源码，仍旧觉得符号式编程的运行过程很神奇……），另外，我们在bind等操作做检查时，仅仅会运行辅助节点，不到真实的数据流入，这个class是不会运行乃至报错的，所以辅助节点的设计真的很重要。