Tensorflow 将张量组织成一批动态形状的张量

我有以下情况：

• 我想使用Tensorflow服务部署一个人脸检测器模型：
• 在Tensorflow服务中，有一个名为

  ——enable_batching

  的命令行选项。这会导致模型服务器自动批处理请求，以最大限度地提高吞吐量。我希望启用此功能
• 我的模型接收一组图像（称为图像），这是形状的张量

  （批量大小，640480,3）

• 该模型有两个输出：

  （面数，4）

  和

  （面数，）

  。第一个输出将被称为面。最后一个输出，我们可以称之为分区（partitions），是对应面的原始批中的索引。例如，如果我传入一批4张图像，得到7张面，那么这个张量可能是

  [0,0,1,2,2,2,3]

  。前两个面对应于第一个图像，第三个面对应于第二个图像，第三个图像有3个面，等等

我的问题是：

• 为了使

  --enable_batching

  标志起作用，我的模型的输出需要具有与输入相同的第0维。也就是说，我需要一个具有以下形状的张量：

  （批大小，…）

  。我想这是为了让模型服务器知道将批处理中的每个输出发送到哪个grpc连接
• 我想做的是将我从人脸检测器的输出张量从这个形状

  （面数，4）

  转换成这个形状

  （批量大小，无，4）

  。也就是说，一个批次数组，其中每个批次可以具有可变数量的面（例如，批次中的一个图像可能没有面，而另一个图像可能有3个面）

我尝试的是：

• tf.动态分区

  。从表面上看，这个功能看起来很完美。但是，在意识到

  num_partitions

  参数不能是张量，只能是整数后，我遇到了困难：

  tensorflow\u service\u output=tf.动态分区（面、分区、批量大小）

如果

tf.dynamic_partition

函数接受

num_partition

的张量值，那么我的问题似乎就解决了。然而，我回到原点，因为事实并非如此

谢谢大家的帮助！如果有什么不清楚的地方，请告诉我

这里是预期过程的视觉表示：

---

我最终用

TensorArray

和

tf找到了解决方案。while\u loop

：

def batch_reconstructor(tensor, partitions, batch_size):
    """
    Take a tensor of shape (batch_size, 4) and a 1-D partitions tensor as well as the scalar batch_size
    And reconstruct a TensorArray that preserves the original batching

    From the partitions, we can get the maximum amount of tensors within a batch. This will inform the padding we need to use.
    Params:
        - tensor: The tensor to convert to a batch
        - partitions: A list of batch indices. The tensor at position i corresponds to batch # partitions[i]
    """
    tfarr = tf.TensorArray(tf.int32, size=batch_size, infer_shape=False)

    _, _, count = tf.unique_with_counts(partitions)
    maximum_tensor_size = tf.cast(tf.reduce_max(count), tf.int32)

    padding_tensor_index = tf.cast(tf.gather(tf.shape(tensor), 0), tf.int32)

    padding_tensor = tf.expand_dims(tf.cast(tf.fill([4], -1), tf.float32), axis=0) # fill with [-1, -1, -1, -1]
    tensor = tf.concat([tensor, padding_tensor], axis=0)

    def cond(i, acc):
        return tf.less(i, batch_size)

    def body(i, acc):
        partition_indices = tf.reshape(tf.cast(tf.where(tf.equal(partitions, i)), tf.int32), [-1])

        partition_size = tf.gather(tf.shape(partition_indices), 0)

        # concat the partition_indices with padding_size * padding_tensor_index
        padding_size = tf.subtract(maximum_tensor_size, partition_size)
        padding_indices = tf.reshape(tf.fill([padding_size], padding_tensor_index), [-1])

        partition_indices = tf.concat([partition_indices, padding_indices], axis=0)

        return (tf.add(i, 1), acc.write(i, tf.gather(tensor, partition_indices)))

    _, reconstructed = tf.while_loop(
        cond,
        body,
        (tf.constant(0), tfarr),
        name='batch_reconstructor'
    )

    reconstructed = reconstructed.stack()
    return reconstructed