apply_along_axis张量等效是什么？

Question

我目前有一个数组，其中每一行包含x1，y1，x2，y2，然后是class和confidence的值。我设法通过＆＃34; apply_along_axis＆＃34;删除numpy中的for循环。但是现在我想在Tensorflow中做同样的事情，但我似乎无法弄清楚如何让它发挥作用。

当前正在运行的numpy逻辑：

def calc_boxes_wh(boxes):
    return [np.int32(boxes[0]), np.int32(boxes[1]), np.int32(boxes[2] - boxes[0] + 1), np.int32(boxes[3] - boxes[1] + 1), boxes[4], boxes[5]]

def transform_boxes(all_boxes):
    retval = np.apply_along_axis(calc_box_width_height, 1, all_boxes)
    return retval

我开始关注的路线是tensorflow的map_fn函数，如下所示，但我在尝试中遇到了一系列不同的错误。我使用下面显示的代码得到的当前错误是：

ValueError: The two structures don't have the same number of elements.

First structure (1 elements): dtype: 'float32'

Second structure (6 elements): tf.Tensor 'map/while/ToInt32:0' shape=() dtype=int32>, <tf.Tensor 'map/while/ToInt32_1:0' shape=() dtype=int32>, <tf.Tensor 'map/while/ToInt32_2:0' shape=() dtype=int32>, <tf.Tensor 'map/while/ToInt32_3:0' shape=() dtype=int32>, <tf.Tensor 'map/while/strided_slice_6:0' shape=() dtype=float32>, <tf.Tensor 'map/while/strided_slice_7:0' shape=() dtype=float32>]

Tensorflow尝试的当前代码：

def calc_boxes_wh(boxes):
    return [tf.to_int32(boxes[0]), tf.to_int32(boxes[1]), tf.to_int32(boxes[2] - boxes[0] + 1), tf.to_int32(boxes[3] - boxes[1] + 1), boxes[4], boxes[5]]

def transform_boxes(all_boxes):
    return tf.map_fn(calc_box_width_height, all_boxes)

这是样本数据集的快照：

我是否认为这一切都错了？大多数时候，我已经在TF中发现了numpy等价物，但到目前为止还没有运气。