TensorFlow结构分析及自定义Op

本文内容仅供参考，如有错误请予以指正

tensorflow 架构

kernel

tensorflow中最底层的实现应该是kernel，在tensorflow/tensorflow/core/kernels/目录下，包括常用操作conv， pooling的cpu， gpu以及梯度计算的实现，继承OpKernel类，实现Computer函数，最后调用REGISTER_KERNEL_BUILDER注册到tensorflow。

在tensorflow/tensorflow/core/kernels/BUILD中，为每个kernel定义了rule：tf_kernel_libraries:

tf_kernel_library(
    name = "control_flow_ops",
    prefix = "control_flow_ops",
    deps = [
        "//tensorflow/core:control_flow_ops_op_lib",
        "//tensorflow/core:framework",
        "//tensorflow/core:lib",
    ],
)

tensorflow/tensorflow/core/ops是在kernel的基础上，调用REGISTER_OP添加属性限制，shape函数等：

REGISTER_OP("Conv2D")
    .Input("input: T")
    .Input("filter: T")
    .Output("output: T")
    .Attr("T: {half, float, double}")
    .Attr("strides: list(int)")
    .Attr("use_cudnn_on_gpu: bool = true")
    .Attr(GetPaddingAttrString())
    .Attr(GetConvnetDataFormatAttrString())
    .SetShapeFn(shape_inference::Conv2DShape)
    .Doc(R"doc(
Computes a 2-D convolution given 4-D `input` and `filter` tensors.
Given an input tensor of shape `[batch, in_height, in_width, in_channels]`
and a filter / kernel tensor of shape
`[filter_height, filter_width, in_channels, out_channels]`, this op
performs the following:
1. Flattens the filter to a 2-D matrix with shape
   `[filter_height * filter_width * in_channels, output_channels]`.
2. Extracts image patches from the input tensor to form a *virtual*
   tensor of shape `[batch, out_height, out_width,
   filter_height * filter_width * in_channels]`.
3. For each patch, right-multiplies the filter matrix and the image patch
   vector.
In detail, with the default NHWC format,
    output[b, i, j, k] =
        sum_{di, dj, q} input[b, strides[1] * i + di, strides[2] * j + dj, q] *
                        filter[di, dj, q, k]
Must have `strides[0] = strides[3] = 1`.  For the most common case of the same
horizontal and vertices strides, `strides = [1, stride, stride, 1]`.
strides: 1-D of length 4.  The stride of the sliding window for each dimension
  of `input`. Must be in the same order as the dimension specified with format.
padding: The type of padding algorithm to use.
data_format: Specify the data format of the input and output data. With the
    default format "NHWC", the data is stored in the order of:
        [batch, in_height, in_width, in_channels].
    Alternatively, the format could be "NCHW", the data storage order of:
        [batch, in_channels, in_height, in_width].
)doc");

ops

在$(tensorflow)/tensorflow/core/BUILD中，The ops are linked into the core TensorFlow library here