目录
背景
在我之前的文章Tensorflow自定义算子实现原理 中说明了自定义OP的注册过程,其中开始的第一步通过REGISTER_OP注册算子,有一个类型为OpShapeInferenceFn的参数通过SetShapeFn注册到OpRegistrationData对象中,本文解释OpShapeInferenceFn的作用和实例
OpShapeInference分析
OP定义了Input和Output,但是没有定义它们的形状(shape),以矩阵乘法MatMul OP为例,REGISTER_OP的参数中并没有指定它是个[2,3] x [3,4] = [2, 4]的矩阵,还是[5,6] x [6,7] = [5, 7]的矩阵。所以对于OP需要有一个形状推断的函数,即OpShapeInferenceFn的作用。
typedef std::function<Status(shape_inference::InferenceContext* c)>
OpShapeInferenceFn;
OpShapeInferenceFn定义在tensorflow/core/framework/op_def_builder.h,是一个std::function类型
![](https://i-blog.csdnimg.cn/blog_migrate/78cae63e0c54f5b23d70f7dfaf50863f.png)
在tensorflow/core/framework/shape_inference.h中定义了做shape 推断的数据结构,如上图所示,它们关系很简单,以三维矩阵 A[3, 4, 5] 为例
代表A的Shape中rank_ 为3,dims_中有3个DimensionHandle,每个DimenionHandle指向的Dimension中value_分别为3,4,5
ShapeManager统一管理所有的Shape和Dimension的数据,而InferenceContext作为OpShapeInferenceFn的关键参数提供了各种做形状处理的API 来使用Sape和Dimension数据(Google的工程师写代码分层解耦能力刚刚的)
举个栗子
抽象的文字描述总是苍白的,用说人话的方式举个2维矩阵乘法MatMul的栗子🌰
REGISTER_OP("MatMul")
.Input("a: T")
.Input("b: T")
.Output("product: T")
.Attr("transpose_a: bool = false")
.Attr("transpose_b: bool = false")
.Attr(
"T: {bfloat16, half, float, double, int32, int64, complex64, "
"complex128}")
.SetShapeFn(shape_inference::MatMulShape);
MatMul 算子定义在tensorflow/core/ops/math_ops.cc,其中SetShapeFn指定了形状推断函数为shape_inference::MatMulShape
Status MatMulShape(shape_inference::InferenceContext* c) {
ShapeHandle a;
TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 2, &a));
ShapeHandle b;
TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 2, &b));
bool transpose_a, transpose_b;
TF_RETURN_IF_ERROR(c->GetAttr("transpose_a", &transpose_a));
TF_RETURN_IF_ERROR(c->GetAttr("transpose_b", &transpose_b));
DimensionHandle output_rows = transpose_a ? c->Dim(a, 1) : c->Dim(a, 0);
DimensionHandle output_cols = transpose_b ? c->Dim(b, 0) : c->Dim(b, 1);
// Validate that the inner shapes are compatible.
DimensionHandle inner_a = transpose_a ? c->Dim(a, 0) : c->Dim(a, 1);
DimensionHandle inner_b = transpose_b ? c->Dim(b, 1) : c->Dim(b, 0);
DimensionHandle merged;
TF_RETURN_IF_ERROR(c->Merge(inner_a, inner_b, &merged));
c->set_output(0, c->Matrix(output_rows, output_cols));
return Status::OK();
}
shape_inference::MatMulShape定义在tensorflow/core/framework/common_shape_fns.cc
以2维矩阵 A[3, 4] x B[4, 5] = D[3, 5] 为例分析代码如下:
- c->input(0)获得A的ShapeHandle,c->WithRank比较第一个参数A的rank( = 2)与第二参数2相等,从而创建了一个rank=2的ShapeHandle a输出
- c->input(1)获得B的ShapeHandle,c->WithRank比较第一个参数B的rank( = 2)与第二参数2相等,从而创建了一个rank=2的ShapeHandle b输出
- 不考虑矩阵转置的情况,即transpose_a 和 transpose_b 都是false
- c->Dim(a, 0)取出a的第0维(行)作为output_rows ( = 3 )
- c->Dim(b, 1)取出b的第1维(列)作为output_cols ( = 5 )
- 验证矩阵乘的合法性,即A的列 = B的行
- c->Dim(a, 1)取出a的第1维(列)作为inner_a ( = 4 )
- c->Dim(b, 0)取出b的第0维(行)作为inner_b ( = 4 )
- 通过c->Merge比较inner_a = inner_b 即A的列 = B的行
- c->Matrix创建一个新的ShapeHandle (rank = 2,Dimension分别为3和5)返回
- c->set_output将Matrix创建的ShapeHandle作为output形状,至此完成了MatMul算子的输入与输出形状推断