本文作者:李杰
TF 计算图从逻辑层来讲,由 op 与 tensor 构成。op 是项点代表计算单元,tensor 是边代表 op 之间流动的数据内容,两者配合以数据流图的形式来表达计算图。那么 op 对应的物理层实现是什么?TF 中有哪些 op,以及各自的适用场景是什么?op 到底是如何运行的?接下来让我们一起探索和回答这些问题。
一、初识 op
1.1 op 定义
op 代表计算图中的节点,是 tf.Operation 对象,代表一个计算单元。用户在创建模型和训练代码时,会创建一系列 op 及其依赖关系,并将这些 op 和依赖添加到 tf.Graph 对象中(一般为默认图)。比如:tf.matmul()就是一个 op,它有两个输入 tensor 和一个输出 tensor。
1.2 op 分类
op 的分类一般有多个视角,比如按是否内置划分、按工作类型划分。
按是否内置划分,一般分为:内置 op 和自定义 op(见“二、自定义 op”部分介绍)。
按工作类型划分,一般分为:常见数学 op、数组 op、矩阵 op、有状态 op、神经网络 op、检查点 op、队列与同步 op、控制流 op。TF 白皮书对内置 op 的分类总结如下:
1.3 op 与 kernel
op 一般都有名称且代表一个抽象的计算过程。op 可以设置若干属性,但这些属性必须在编译期提供或推理得到,因为它们用来实例化一个节点对象从而执行真正的计算。属性的经典用法就是拿来支持类型多态,比如两个浮点张量的矩阵乘法与两个整型张量的矩阵乘法。
kernel 是 op 在指定设备类型(CPU/GPU)上的具体实现。TF 二进制库通过注册机制定义了一系列 op 及对应的 kernel 实现,用户可以提供额外的 op 定义与 kernel 实现进行扩充。一般来说,一个 op 对应多个 kernel 实现。
接下来让我们一起用矩阵乘法 MatMul 算子的相关代码来理解 op 与 kernel 的关系(此处不必纠结代码细节,只需体会 op 与 kernel 关系即可):
// 首先给出op注册的定义。其中输入输出支持泛型,其合法类型在Attr中进行枚举。
// 代码位置 tensorflow1.15.5\tensorflow\core\ops\math_ops.cc
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的实现,采用类模板机制
// 代码位置 tensorflow1.15.5\tensorflow\core\kernels\matmul_op.cc
template <typename Device, typename T, bool USE_CUBLAS>
class MatMulOp : public OpKernel {
public:
explicit MatMulOp(OpKernelConstruction* ctx)
: OpKernel(ctx), algorithms_set_already_(false) {
OP_REQUIRES_OK(ctx, ctx->GetAttr("transpose_a", &transpose_a_));
OP_REQUIRES_OK(ctx, ctx->GetAttr("transpose_b", &transpose_b_));
LaunchMatMul<Device, T, USE_CUBLAS>::GetBlasGemmAlgorithm(
ctx, &algorithms_, &algorithms_set_already_);
use_autotune_ = MatmulAutotuneEnable();
}
// 省略了很多代码...
private:
std::vector<int64> algorithms_;
bool algorithms_set_already_;
bool use_autotune_;
bool transpose_a_;
bool transpose_b_;
};
// MatMul的op定义与kernel实现绑定处理
// 代码位置 tensorflow1.15.5\tensorflow\core\kernels\matmul_op.cc
#define REGISTER_CPU_EIGEN(T) /*cpu与eigen组合对应实现*/ \
REGISTER_KERNEL_BUILDER( \
Name("MatMul").Device(DEVICE_CPU).TypeConstraint<T>("T").Label("eigen"), \
MatMulOp<CPUDevice, T, false /* cublas, ignored for CPU */>);
#define REGISTER_CPU(T) /*cpu对应实现(eigen与非eigen)*/ \
REGISTER_KERNEL_BUILDER( \
Name("MatMul").Device(DEVICE_CPU).TypeConstraint<T>("T"), \
MatMulOp<CPUDevice, T, false /* cublas, ignored for CPU */>); \
REGISTER_CPU_EIGEN(T);
#define REGISTER_GPU(T) /*gpu对应实现(cublas与非cublas)*/ \
REGISTER_KERNEL_BUILDER( \
Name("MatMul").Device(DEVICE_GPU).TypeConstraint<T>("T"), \
MatMulOp<GPUDevice, T, true /* cublas, true by default */>); \
REGISTER_KERNEL_BUILDER(Name("MatMul") \
.Device(DEVICE_GPU) \
.TypeConstraint<T>("T") \
.Label("cublas"), \
MatMulOp<GPUDevice, T, true /* cublas */>)
复制代码
二、自定义 op
用户编写的模型训练代码一般由 TF 原生的 op 算子及其依赖关系组成,但有时候我们定义的计算逻辑在 TF 中没有相应的 op 实现。根据 TensorFlow 官网的建议,我们应当先组合 python op 算子或 python 函数进行尝试。完成尝试之后再决定要不要自定义 op。
2.1 自定义 op 场景
一般来说,需要自定义 op 的场景有如下 3 个:
•用 TF 原生 op 组合来表达新计算逻辑的过程比较复杂或不可能
•用 TF 原生 op 组合来表达新计算逻辑,其计算性能较低
•在新版编译器中也较难实现 op 融合的计算逻辑需要我们手动实现融合
在此举个例子方便大家理解。假如我们要实现一个新计算实逻:中位数池化(median pooling),过程中要在滑动窗口不断求得中位数。检索 TF 文档没有发现对应 op,因此我们先考虑用 TF python op 组合来实现它,果然通过 ExtractImagePatches and TopK 就可以实现这个功能。经测试前述组合方案并不是计算和存储高效的,因此我们就有必要将 median pooling 在一个 op 中进行高效实现。
2.2 自定义 op 流程
自定义 op 一般遵循 5 个基本步骤:
1.注册 op,具体包括:指定名称、输入/输出声明、形状函数。
2.定义 kernel(即 op 的实现)并与 op 绑定。一个 op 有多个 kernel 实现,具体由输入输出类型、硬件(CPU、GPU)决定。
3.创建 python 包装器,一般由 op 注册机制自动完成。
4.编写 op 的梯度计算函数(可选项)。
5.测试 op,通过 python 测试较为方便,当然也可通过 C++进行测试。
接下来我们就以官网最简单的 ZeroOut 同步式自定义 op(继承 OpKernel)为例,结合代码来讲述上述 5 个步骤。下面先给出步骤 1 和步骤 2 用 C++实现的代码(官方推荐用 bazel 编译 so 文件):
// 步骤1:注册op
REGISTER_OP("ZeroOut")
.Input("to_zero: int32")
.Output("zeroed: int32")
.SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
c->set_output(0, c->input(0)); //c's input and output type is std::vector<ShapeHandle>
return Status::OK();
});
// 步骤2:定义kernel(常规CPU设备),并把kernel与op绑定
class ZeroOutOp : public OpKernel {
public:
explicit ZeroOutOp(OpKernelConstruction* context) : OpKernel(context) {}
void Compute(OpKernelContext* context) override {
// Grab the input tensor from OpKernelContext instance
const Tensor& input_tensor = context->input(0);
auto input = input_tensor.flat<int32>();
// Create an output tensor
Tensor* output_tensor = NULL;
OP_REQUIRES_OK(context, context->allocate_output(0, input_tensor.shape(),
&output_tensor)); // OP_REQUIRES_OK第二个参数一般为方法调用,此处为输出张量分配内存空间
auto output_flat = output_tensor->flat<int32>();
// Set all but the first element of the output tensor to 0.
const int N = input.size();
for (int i = 1; i < N; i++) {
output_flat(i) = 0;
}
// Preserve the first input value if possible.
if (N > 0) output_flat(0) = input(0);
}
};
REGISTER_KERNEL_BUILDER(Name("ZeroOut").Device(DEVICE_CPU), ZeroOutOp);
复制代码
步骤 3 加载上述 so 文件(自动完成前后端 op 映射);步骤 4 是可选项,此处不需要;步骤 5 基于 python api 测试 op 功能。相应代码如下:
import tensorflow as tf
zero_out_module = tf.load_op_library('./zero_out.so') # 加载so文件生成python module
with tf.Session(''):
zero_out_module.zero_out([[1, 2], [3, 4]]).eval()
# Prints
array([[1, 0], [0, 0]], dtype=int32)
复制代码
2.3 高级话题
关于 op 的技术话题还有很多,我们在此简述一些要点:
1.如果实现了一个多线程 CPU kernel,则可以利用 work_sharder.h 中的 Shard 函数。
2.大多数 op 以同步方式工作,只需继承 OpKernel 改写 Compute()方法,且此方法必须线程安全。
3.如果一个 op 因为其它 op 的运行而阻塞,则这个 op 可以采用异步方式工作,继承 AsyncOpKernel 改写 ComputeAsync()方法,且此方法必须线程安全。异步 op 最经典的例子就是跨设备通信 send/recv pair 中的 RecvOp。
4.如果要为 op 配置一些静态属性,可使用 Attr,它有一套特有的支持类型。典型应用是支持泛型。
5.实现 GPU kernel 有两部分内容:OpKernel 和 CUDA kernel,相应的加载代码。
6.编译自定义 op,首先要配置头文件搜索路径与库文件搜索路径,接着指定编译和链接选项,最后还要确保 ABI 兼容性。
7.Resource(资源)代表相同设备上 op 共享的内容,比如:张量值、kv 存储表、队列、读取器、网络连接等。代表资源的类必须继承 ResourceBase,然后注册 ResourceHandleOp 生成资源句柄,普通 op 以 resouce 类型的 Input 进行引入。
三、op 工作原理
3.1 op 运行框架
整体来看,op 与 kernel 都有其结构描述与统一的注册管理中心。而 OpDefBuilder 有两个包装类 OpDefBuilderWrapper 和 OpDefBuilderReceiver,前者支持 op 构建的链式语法,后者接受 op 构建结果并进行注册。众所周知,op 是编译期概念,而 kernel 是运行期概念,在 AI 编译器的后端处理流程中会进行 op 的算子选择,此过程会基于一系列策略为 op 匹配最合适的 kernel 实现。
3.2 若干技术细节
首先,我们来看一下大家在使用 TensorFlow 过程中经常碰到的 libtensorflow_framework.so。按照 tf1.15.5/tensorflow/BUILD 中的描述,libtensorflow_framework.so 定义了 op 和 kernel 的注册机制而不涉及具体实现。
// rootdir=tensorflow1.15.5
// ${rootdir}/tensorflow/BUILD
/*
# A shared object which includes registration mechanisms for ops and
# kernels. Does not include the implementations of any ops or kernels. Instead,
# the library which loads libtensorflow_framework.so
# (e.g. _pywrap_tensorflow_internal.so for Python, libtensorflow.so for the C
# API) is responsible for registering ops with libtensorflow_framework.so. In
# addition to this core set of ops, user libraries which are loaded (via
# TF_LoadLibrary/tf.load_op_library) register their ops and kernels with this
# shared object directly.
*/
tf_cc_shared_object(
name = "tensorflow_framework",
framework_so = [],
linkopts = select({
"//tensorflow:macos": [],
"//tensorflow:windows": [],
"//tensorflow:freebsd": [
"-Wl,--version-script,$(location //tensorflow:tf_framework_version_script.lds)",
"-lexecinfo",
],
"//conditions:default": [
"-Wl,--version-script,$(location //tensorflow:tf_framework_version_script.lds)",
],
}),
linkstatic = 1,
per_os_targets = True,
soversion = VERSION,
visibility = ["//visibility:public"],
deps = [
"//tensorflow/cc/saved_model:loader_lite_impl",
"//tensorflow/core:core_cpu_impl",
"//tensorflow/core:framework_internal_impl", /* 展开此target进行查看 */
"//tensorflow/core:gpu_runtime_impl",
"//tensorflow/core/grappler/optimizers:custom_graph_optimizer_registry_impl",
"//tensorflow/core:lib_internal_impl",
"//tensorflow/stream_executor:stream_executor_impl",
"//tensorflow:tf_framework_version_script.lds",
] + tf_additional_binary_deps(),
)
// ${rootdir}/tensorflow/core/BUILD
tf_cuda_library(
name = "framework_internal_impl",
srcs = FRAMEWORK_INTERNAL_PRIVATE_HEADERS + glob( // 可以查看FRAMEWORK_INTERNAL_PRIVATE_HEADERS内容
[
"example/**/*.cc",
"framework/**/*.cc",
"util/**/*.cc",
"graph/edgeset.cc",
"graph/graph.cc",
"graph/graph_def_builder.cc",
"graph/node_builder.cc",
"graph/tensor_id.cc",
"graph/while_context.h",
"graph/while_context.cc",
],
// 省略了诸多代码
)
// FRAMEWORK_INTERNAL_PRIVATE_HEADERS的内容
FRAMEWORK_INTERNAL_PRIVATE_HEADERS = [
"graph/edgeset.h",
"graph/graph.h",
"graph/graph_def_builder.h",
"graph/node_builder.h",
"graph/tensor_id.h",
] + glob(
[
"example/**/*.h",
"framework/**/*.h", // 这里就是重点,查看${rootdir}/tensorflow/core/framework/op.h和opkernel.h
"util/**/*.h",
]
)
// 先来看op.h
#define REGISTER_OP(name) REGISTER_OP_UNIQ_HELPER(__COUNTER__, name)
#define REGISTER_OP_UNIQ_HELPER(ctr, name) REGISTER_OP_UNIQ(ctr, name)
#define REGISTER_OP_UNIQ(ctr, name) \
static ::tensorflow::register_op::OpDefBuilderReceiver register_op##ctr \
TF_ATTRIBUTE_UNUSED = \
::tensorflow::register_op::OpDefBuilderWrapper<SHOULD_REGISTER_OP( \
name)>(name)
// 再来看看opkernel.h
#define REGISTER_KERNEL_BUILDER(kernel_builder, ...) \
REGISTER_KERNEL_BUILDER_UNIQ_HELPER(__COUNTER__, kernel_builder, __VA_ARGS__)
#define REGISTER_KERNEL_BUILDER_UNIQ_HELPER(ctr, kernel_builder, ...) \
REGISTER_KERNEL_BUILDER_UNIQ(ctr, kernel_builder, __VA_ARGS__)
#define REGISTER_KERNEL_BUILDER_UNIQ(ctr, kernel_builder, ...) \
constexpr bool should_register_##ctr##__flag = \
SHOULD_REGISTER_OP_KERNEL(#__VA_ARGS__); \
static ::tensorflow::kernel_factory::OpKernelRegistrar \
registrar__body__##ctr##__object( \
should_register_##ctr##__flag \
? ::tensorflow::register_kernel::kernel_builder.Build() \
: nullptr, \
#__VA_ARGS__, \
[](::tensorflow::OpKernelConstruction* context) \
-> ::tensorflow::OpKernel* { \
return new __VA_ARGS__(context); \
});
复制代码
参照上述同样的流程,我们可以发现 libtensorflow.so 中涉及 op 与 kernel 的具体实现,同时也包括 Session 的具体实现。
最后,我们再来讲讲 REGISTER_OP 宏背后的具体原理。我们在上面已经给出了此宏的定义,此处针对它的实现展开谈谈:
// 先来看op.h
#define REGISTER_OP(name) REGISTER_OP_UNIQ_HELPER(__COUNTER__, name)
#define REGISTER_OP_UNIQ_HELPER(ctr, name) REGISTER_OP_UNIQ(ctr, name)
#define REGISTER_OP_UNIQ(ctr, name) \
static ::tensorflow::register_op::OpDefBuilderReceiver register_op##ctr \
TF_ATTRIBUTE_UNUSED = \
::tensorflow::register_op::OpDefBuilderWrapper<SHOULD_REGISTER_OP( \
name)>(name)
// REGISTER_OP的一般用法如下
REGISTER_OP("ZeroOut")
.Input("to_zero: int32")
.Output("zeroed: int32")
.SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
c->set_output(0, c->input(0));
return Status::OK();
});
// op定义的链式规则是通过OpDefBuilderWrapper类实现的
class OpDefBuilderWrapper<true> {
public:
explicit OpDefBuilderWrapper(const char name[]) : builder_(name) {}
OpDefBuilderWrapper<true>& Input(string spec) {
builder_.Input(std::move(spec));
return *this; // 显而易见,调用Input仍然返回OpDefBuilderWrapper<true>本身
}
OpDefBuilderWrapper<true>& Output(string spec) {
builder_.Output(std::move(spec));
return *this;
}
OpDefBuilderWrapper<true>& SetShapeFn(
Status (*fn)(shape_inference::InferenceContext*)) {
builder_.SetShapeFn(fn);
return *this;
}
const ::tensorflow::OpDefBuilder& builder() const { return builder_; }
private:
mutable ::tensorflow::OpDefBuilder builder_;
};
// 当通过链式规划构建好op后,再通过OpDefBuilderReceiver完成op的注册
// op.h
struct OpDefBuilderReceiver {
// To call OpRegistry::Global()->Register(...), used by the
// REGISTER_OP macro below.
// Note: These are implicitly converting constructors.
OpDefBuilderReceiver(
const OpDefBuilderWrapper<true>& wrapper); // NOLINT(runtime/explicit)
constexpr OpDefBuilderReceiver(const OpDefBuilderWrapper<false>&) {
} // NOLINT(runtime/explicit)
};
// op.cc,然后在OpDefBuilderReceiver构造函数内部完成OpDefBuilderWrapper的全局注册
OpDefBuilderReceiver::OpDefBuilderReceiver(
const OpDefBuilderWrapper<true>& wrapper) {
OpRegistry::Global()->Register(
[wrapper](OpRegistrationData* op_reg_data) -> Status {
return wrapper.builder().Finalize(op_reg_data);
});
}
复制代码
四、总结
本文为大家系统讲解了 TensorFlow 的核心抽象 op 及其 kernel 实现。需要自定义 op 的具体场景,以及 op 的运行框架及若干技术细节。读罢此文,读者应该有如下几点收获:
•TensorFlow 中 op 是编译期概念,kernel 是运行期概念,两者各自的定义与注册方式,以及相应的映射逻辑。
•掌握 TensorFlow 的高阶玩法:自定义 op。这将使你之前工作的不可能变为可能,由低效转化为高效。
•掌握 op 与 kernel 注册的宏定义来自何方,以及宏定义背后具体的运行框架。
参考资料
1.《TensorFlow: Large-Scale Machine Learning on Heterogeneous Distributed Systems》: https://arxiv.org/abs/1603.04467
2.Graphs and Sessions: https://github.com/tensorflow/docs/blob/master/site/en/r1/guide/graphs.md
3.Adding a New Op: https://github.com/tensorflow/docs/blob/master/site/en/r1/guide/extend/op.md
4.跨设备通信 send/recv: https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/kernels/sendrecv_ops.h
5.OpKernel definition: https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/op_kernel.h
6.tensorflow 源码解析之 framework-resource: https://www.cnblogs.com/jicanghai/p/9535504.html
7.tensorflow 源码解析之 framework-op: https://www.cnblogs.com/jicanghai/p/9539513.html
评论