Question

我在C ++中有tensorflow::Tensor batch形状[2,720,1280,3]（#images x height x width x #channels）。

我想得到另一个只有第一张图像的张量，因此我会有一个形状张量[1,720,1280,3]。换句话说，我想：

tensorflow::Tensor first = batch[0]

实现它的最有效方法是什么？

我知道如何在python中执行此操作，但C ++ api和文档不如python的那么好。

Answer 1

在花了一些时间尝试通过副本实现后，我意识到API中支持此操作为Slice：

tensorflow::Tensor first = batch.Slice(0, 1);

请注意，如文档所述，返回的张量与内容缓冲区共享内部缓冲区，如果与您相关，则两个张量的对齐方式可能不同。

编辑：

因为我已经完成了，所以这是我尝试再现相同的功能，基于复制。我认为它应该有用（它与我在其他环境中使用的非常相似）。

#include <cstdlib>
#include <cassert>
#include <tensorflow/core/framework/tensor.h>
#include <tensorflow/core/framework/tensor_shape.h>

tensorflow::Tensor get_element(const tensorflow::Tensor data, unsigned int index, bool keepDim)
{
    using namespace std;
    using namespace tensorflow;

    typedef typename tensorflow::DataTypeToEnum<T> DataType;
    auto dtype = DataType::v();
    assert(dtype == data.dtype());

    auto dtype = data.dtype();
    auto dataShape = data.shape();

    TensorShape elementShape;
    if (keepDim)
    {
        elementShape.addDim(1);
    }
    for (int iDim = 1; iDim < dataShape.dims(); iDim++) {
      elementShape.AddDim(dataShape.dim_size(iDim));
    }
    Tensor element(dtype, elementShape);
    auto elementBytes = elementShape.num_elements() * DataTypeSize(dtype);

    memcpy(element.flat<void>().data(),
           batch.flat<void>().data() + elementBytes * index,
           elementBytes);
    return element;
}

int main()
{
    Tensor batch =  ...;
    Tensor first = get_element(batch, 0);
    return 0;
}

如果您只想将数据提取到例如矢量或其他内容，也可以更改代码。

Answer 2

这很好

#include "tensorflow/core/framework/tensor_slice.h"

Tensor t2 = t1.Slice(0,1);

张量流量张量的子推广（C ++）

2 个答案: