tf.py_func在定义多输出函数时,输出变量类型需要用[ ]框起来;
tf.py_func在定义单输出函数时,输出变量
类型不能再用[ ]框起来;
这个需要各外注意!
例如:
多变量时
def _proposal_layer(self, rpn_cls_prob, rpn_bbox_pred, name):
with tf.variable_scope(name) as scope:
rois, rpn_scores, inds= tf.py_func(proposal_layer,
[rpn_cls_prob, rpn_bbox_pred, self._im_info, self._mode,
self._feat_stride, self._anchors, self._num_anchors],
[tf.float32, tf.float32,tf.int64])
# rois.set_shape([None, 5])
# rpn_scores.set_shape([None, 1])
rois.set_shape([1,None,None,self._num_anchors*5])
rpn_scores.set_shape([1,None,None,self._num_anchors*1])
return rois, rpn_scores,inds
单变量时
def _draw_proposals_to_image(self,rois,scores,inds,keep_inds,stride,name):
with tf.variable_scope(name) as scope:
mask = tf.py_func(
proposals_to_image,
[rois, scores, inds, keep_inds,stride],
tf.float32)
mask = tf.stop_gradient(mask)
mask.set_shape([1, None, None, cfg.TRAIN.BATCH_SIZE])
return mask