首先在tensorflow/core/kernels/http://matmul_op.cc中,多次调用了REGISTER_KERNEL_BUILDER()宏来对MatMul这个op,针对于不同的设备类型和数据类型,注册了多个kernel。例如针对于GPU&float,其宏调用代码(只考虑最简单的情况)是 REGISTER_KERNEL_BUILDER(Name("MatMul").Device(DEVICE_GPU).TypeConstraint<float>("floa...
REGISTER_OP 宏来定义 Op 的接口,定义op得输入输出和属性: Attr:该op接口内得属性值,例如:preserver_index 就是该op得属性,以后Zeroout会自带该属性 Input:该op得输入 Output:该op得输出 SetShapeFn:控制输入输出前后得shape变化 通过REGISTER_KERNEL_BUILDER向tensorflow系统注册op,该实现对应得是CPU: 2.上面三项...
基于REGISTER_KERNEL_BUILDER这样的宏,将上面所定义的接口和内部的实现给绑定起来。这是因为 TF 支持基于不同的输入、输出类型和所运行的底层设备架构来定义同一个算子不同的内部实现,所以用户可以定义多种kernel实现,告知给系统什么场景下运行具体哪一个kernel,在实际运行时,TF 就可以根据不同的设备、数据流上下文调用...
REGISTER_OP实用宏注册OP。前端系统,也存在类似的OP注册机制。REGISTER_OP注册OP过程,实际上是一个REGISTER_OP描述到OpDef表示的翻译过程。OpDefBuilder通过链式调用Input,Output,Attr方法分别构造OP的输入、输出列表,及其属性列表。最后,通过调用Finalize成员函数,经过解析字符串表示,将其翻译为OpDef的内在表示,最后注册到...
REGISTER_KERNEL_BUILDER(Name("ZeroOut").Device(DEVICE_CPU), ZeroOutOp); 这里注册的操作名是ZeroOut,通过上面的语句和ZeroOutOp对应吧 GPU版本:重写的compute里一般都只是数据的处理:大小推断、flat成一维、分配内存、设初值等等,具体的计算用cuda的核函数实现 ...
REGISTER_KERNEL_BUILDER 和 REGISTER_OP 后面跟的接口名称是"若干个首字母大写的单词"组成的名称,它对应到python之后,接口名称就变为"每个单词全变为小写,单词间以下划线分割"的名称。这应该是swig处理的。 ### python加载带有custom op的saved_model ```python op = tf.load_op_library('./xxxx.so') tenso...
explicit Name(const char* op) : KernelDefBuilder(op) {} }; REGISTER_KERNEL_BUILDER宏里面就是一些trick,实质是创建一个名称唯一的类型为OpKernelRegistrar的全局静态变量,如果你有兴趣可以看一下: #define REGISTER_KERNEL_BUILDER(kernel_builder, ...) \ ...
而每个算子在Tensorflow上都会提供GPU的算法:关于OP的具体实现...,如果要支持OpenCL,可以参考开源(点击打开链接) 对CUDA的支持使用了基于CUDA平台的第三方开发库,没有直接使用CUDA编程 2.2 直接调用CUDA Tensorflow 同时本身也可以直接调用...REGISTER_KERNEL_BUILDER( Name("arithmetic").Device(DEVICE_GPU)....
还需通过REGISTER_KERNEL_BUILDER定义方法名,这个是在Python里使用的时候的名字。将方法名和上面的类进行绑定。 打包和使用 all: tfop TF_INC=/home/recommend/.local/lib/python2.7/site-packages/tensorflow/include TF_LIB=/home/recommend/.local/lib/python2.7/site-packages/tensorflow ...
1#defineREGISTER_GPU_KERNEL(type) \2REGISTER_KERNEL_BUILDER(Name("Reshape") \3.Device(DEVICE_GPU) \4.HostMemory("shape") \5.TypeConstraint<type>("T") \6.TypeConstraint<int32>("Tshape"), \7ReshapeOp); \8REGISTER_KERNEL_BUILDER(Name("Reshape") \9.Device(DEVICE_GPU) \10.HostMemory...