带有形状函数的 TensorFlow 输出形状答案

【问题标题】：TensorFlow output shape with Shape Function带有形状函数的 TensorFlow 输出形状
【发布时间】：2017-09-27 13:22:36
【问题描述】：

我正在按照guide 用 C++ 和 CUDA 编写自定义 TensorFlow Op，我对 ops 输出形状的设置感到困惑。我希望能够提供所需的输出尺寸作为属性，并使用此属性来设置形状与SetShapeFn

我的操作代码很大程度上是对上面链接的示例的直接改编，但是，我不确定如何以这种方式使用提供的属性。

REGISTER_OP("GaussianProcess")
.Attr("output_dim: int") //This is the desired output dimension.
.Input("data_points: float32")
.Input("query_point: float32")
.Input("alpha: float32")
.Input("hyp: float32")
.Output("shape_descriptor: float32")
.SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
      /*
       Use attribute 'output_dim' here.
      */
});

对于上述操作注册，如何访问传递给SetShapeFn的lambda中的output_dim属性？

我需要能够将它传递给这样的调用：

c->set_output(0, output_dim);

【问题讨论】：

标签： c++ tensorflow deep-learning

【解决方案1】：

InferenceContext 允许通过GetAttr 方法获取属性。所以你可以像这样得到output_dim：

int output_dim;
c->GetAttr("output_dim", &output_dim);

希望有帮助！

【讨论】：