标签:不同 ptr 对象 kernel 索引 数据支持 shape frame 作用
前面我们讲到op的时候,提到了一个操作的注册器OpRegistry,并且提到,其中注册的数据是一个结构OpRegistrationData,这个结构中除了OpDef之外,还包含了一个OpShapeInferenceFn,这个数据是做什么用的呢?
我们知道,op只是定义了操作的输入输出和参数,但并没有定义操作具体的输入形状,举个例子,MatMul操作,代表矩阵乘法,这只是一个抽象的表示,没有具体说,这个矩阵乘法代表的是[2,3]x[3,4]=[2,4],还是[100,200]x[200,300]=[100,300]。所以在实际应用中,输入的真实形状我们是不知道的,但是为了产生输出,我们必须知道输出的形状,好给它申请对应大小的内存空间。所以,我们需要为每一个操作,配备一个形状推断的函数,这就是ShapeInference的由来。
上面提到了,操作注册器中用到的是OpRegistrationData,而不是ShapeInference,这两者有什么关系呢?回想一下前面讲过的OpKernelContext,其实它们的功能很像。OpKernelContext是作为OpKernel的核心API Compute函数的参数,所有计算相关的参数都会包含在这个对象中。ShapeInference也是一样,我们把所有跟形状推断相关的数据和功能函数封装在一个ShapeInference对象中,然后把这个对象传递给OpShapeInferenceFn,就可以实现形状推断。这种设计实现了数据部分和实现逻辑的解耦。
在具体看ShapeInference类之前,我们先要看一些辅助类:
class Dimension {
private:
//...
const int64 value_;
};
class DimensionHandle {
private:
//...
const Dimension* ptr_ = nullptr;
};
class Shape {
//...
private:
const int32 rank_;
const std::vector<DimensionHandle> dims_;
};
class ShapeHandle {
//...
private:
const Shape* ptr = nullptr;
};
class DimensionOrConstant {
public:
//...
DimensionHandle dim;
int64 val;
};
class ShapeAndType {
ShapeHandle shape;
DataType dtype = DT_INVALID;
};
这几个类都比较简单。在下面用到时能够认得就好了。
下面我们看下InferenceContext这个类:
class InferenceContext {
public:
InferenceContext(int graph_def_version, const NodeDef* node_def, const OpDef& op_def, const std::vector<ShapeHandle>& input_shapes, const std::vector<const Tensor*>& input_tensors, const std::vector<ShapeHandle>& input_tensors_as_shapes, std::vector<std::unique_ptr<std::vector<ShapeAndType>>> input_handle_shapes_and_types);//构造函数
Status Run(const std::function<Status(shape_inference::InferenceContext* c)>& fn);//运行一个以this为参数的函数,没错,这里运行的就是OpShapeInferenceFn
bool MergeInput(int idx, ShapeHandle shape);
bool RelaxInput(int idx, ShapeHandle shape);
private:
ShapeManager shape_manager_;
std::vector<ShapeHandle> inputs_;
std::vector<const Tensor*> input_tensors_;
std::vector<bool> requested_input_tensor_;
std::vector<ShapeHandle> outputs_;
std::vector<ShapeHandle> input_tensors_as_shapes_;
std::vector<bool> requested_input_tensor_as_partial_shape_;
std::vector<std::unique_ptr<std::vector<ShapeAndType>>> input_handle_shapes_and_types_;
std::vector<std::unique_ptr<std::vector<ShapeAndType>>> output_handle_shapes_and_types_;
const int graph_def_version_;
const NodeDef& node_def_;
NameRangeMap input_name_map_;
NameRangeMap output_name_map_;
Status construction_status_;
};
前面已经介绍过了这个类的作用,是作为真正的形状推断函数的参数,为形状推断提供足够的数据和功能函数支持,那么这个类的成员就比较清晰了,首先私有的一大堆成员,为形状推断提供数据支持,而大量的共有API函数,为形状推断提供公用的功能函数,比如上面提到的MergeInput和RelaxOutput,下面我们重点介绍下这两个函数的功能:
MergeInput函数是将输入索引idx处的输入与shape合并,具体的合并规则是:
tensorflow源码解析之framework-shapeinference
标签:不同 ptr 对象 kernel 索引 数据支持 shape frame 作用
原文地址:https://www.cnblogs.com/jicanghai/p/9552490.html