本文作者:李傑
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