c++ - Tensorflow 从 C++ 中的 Graph Def 获取形状

Question

在进行任何推理调用之前，我试图在 C++ 中获取输入张量的形状。所以我只能通过 Graph Def 来完成。

我正在尝试:

 auto att = graph_def.node(i).attr();
 att["shape"].PrintDebugString();
 Tensor tensor;
 std::cerr << tensor.FromProto(att["shape"].tensor()) << std::endl;

它返回错误。当 PrintDebugString 打印时:

shape { dim { size: -1 } dim { size: 1024 } dim { size: 1024 } dim { size: 3 } }

所以我只需要获得 1024x1024x3。我对 Protobuf 协议(protocol)不是很熟悉，对我来说完全不清楚该怎么做。

Answer 1

我设法做到了这一点:

auto shape = graph_def.node().Get(0).attr().at("shape").shape();
for (int i = 0; i < shape.dim_size(); i++) {
    std::cout << shape.dim(i).size()<<std::endl;
}

在我的例子中是:1 128 128 3