Caffe2 Custom Operators

你已经看过Caffe2中提供的种类广泛的算子了吗？仍想推出自己的算子？请继续阅读，但别忘了将您优秀的新算子贡献到项目中！

构建一个基本算子

尽管operator之间会有不同，但几乎所有的operator会使用一个.cc文件来注册operator，并在.h文件中具体实现。例如，在一些情况下，实现代码会在.cc文件中。除此之外，一些operator在.cu文件中存储其GPU/CUDA下的实现。

一个CUDA实现如果利用到了实际的CUDA kernel，需要以.cu命名以便它符合NVCC。如果仅仅用到已有的CUDA库，我们将其命名为_gpu.cc以节约编译时间。

下面我们以fully_connected_op.cc为例，介绍如何编写.cc文件。

#include "caffe2/operators/fully_connected_op.h"

namespace caffe2 {
namespace {

REGISTER_CPU_OPERATOR(FC, FullyConnectedOp<float, CPUContext>);
REGISTER_CPU_OPERATOR(FCGradient, FullyConnectedGradientOp<float, CPUContext>);

首先，通过下面的宏来注册operator和相应gradient operator的名字；在Python中使用FC函数时,其总是被绑定到FullyConnectedOp算子，其中float指明输入类型，CPUContext 指示context 类型；context可以是CPUContext或CUDAContext，具体取决于是在CPU还是GPU设备上使用。

Fully Connected同样拥有一个位于fully_connected_op_gpu.cc中的GPU实现。

#include "caffe2/core/context_gpu.h"
#include "caffe2/operators/fully_connected_op.h"

namespace caffe2 {
namespace {
REGISTER_CUDA_OPERATOR(FC, FullyConnectedOp<float, CUDAContext>);
REGISTER_CUDA_OPERATOR(FCGradient,
                       FullyConnectedGradientOp<float, CUDAContext>);
}  // namespace
}  // namespace caffe2

请注意，GPU实现与CPU实现之间的主要区别是使用REGISTER_CUDA_OPERATOR和 CUDAContext而非REGISTER_CPU_OPERATOR和CPUContext。还要注意包含额外的头文件context_gpu.h，其中承载GPU的任何实现。

回顾fully_connected_op.cc，我们将查看剩余的代码并讨论算子模式。此处告知operator有多少输入和输出产生。本节也用于生成operator目录中的文档，因此请在描述参数和功能时进行详细说明。下面还要注意.Arg、.Input和.Output的最后一个参数是描述，也用于生成文档。

fully_connected_op.cc

OPERATOR_SCHEMA(FC)
  .NumInputs(3)
  .NumOutputs(1)
  .SetDoc(R"DOC(
Computes the result of passing an input vector X into a fully connected layer with 2D weight matrix W and 1D bias vector b.

The layer computes Y = X * W + b, where X has size (M x K), W has size (K x N), b has size (N), and Y has size (M x N), where M is the batch size. Even though b  is 1D, it is resized to size (M x N) implicitly and added to each vector in the batch. These dimensions must be matched correctly, or else the operator will throw errors.
)DOC")
  .Arg("axis", "(int32_t) default to 1; describes the axis of the inputs; "
  "defaults to one because the 0th axis most likely describes the batch_size")
  .Input(0, "X", "2D input of size (MxK) data")
  .Input(1, "W", "2D blob of size (KxN) containing fully connected weight "
  "matrix")
  .Input(2, "b", "1D blob containing bias vector")
  .Output(0, "Y", "1D output tensor");

正如你在上面的架构代码中看到的那样，该算子有3个输入和1个输出，分别由.NumInputs和.NumOutputs指定。文档用.SetDoc指定并且非常全面。它还有一个用.Arg指定的额外可选参数，默认为1。

.SetDocR"DOC(docs go here)DOC"是你提供算子说明文档的地方。

.Input设置operator中使用的主要数据，例如全连接层的权重矩阵。上面的例子显示了.Input的三个条目。注意，第一个参数是输入的索引，从0开始为第一个输入。第二个参数是变量的名称，如X，W或b。最后，第三个参数是描述。

.Arg通常是不涉及原始数据操作的辅助输入。

.Output指定输出。类型参数与 .Input相同：（索引，名称，描述）。

接下来，该模式描述第二个运算符FCGradient。
fully_connected_op.cc

OPERATOR_SCHEMA(FCGradient).NumInputs(3).NumOutputs(2, 3);
class GetFCGradient : public GradientMakerBase {
  using GradientMakerBase::GradientMakerBase;
  vector<OperatorDef> GetGradientDefs() override {
    CHECK_EQ(def_.input_size(), 3);
    return SingleGradientDef(
        "FCGradient", "",
        vector<string>{I(0), I(1), GO(0)},
        vector<string>{GI(1), GI(2), GI(0)});
  }
};
REGISTER_GRADIENT(FC, GetFCGradient);
}  // namespace
}  // namespace caffe2

GradientOp的输入和输出必须使用GradientMakerBase::GetGradientDefs()进行标记。如此，我们告知Caffe2梯度算子的输入和输出如何与对应的算子相关联。特别地，第一个向量标记梯度算子的输入，第二个向量标记输出。请注意，通常，梯度运算符不需要doc方案，除非您认为合适。

实现的细节

如前所述，一般情况下大多数实现细节都在头文件中。实现细节可能直接放在.cc文件中。对于所有CUDA实现，其逻辑和代码主要都在.cu文件中。

Caffe2算子单元测试

编写一些单元测试来验证算子实现是否正确是个好主意。Caffe2中提供了一些辅助程序库，以确保您的算子测试具有良好的覆盖率。

Hypothesis是一个非常有用的基于属性的测试库。其主要思想是表达被测试代码的属性（例如它通过了一个梯度检查，它实现了一个引用函数等），然后生成随机实例并验证它们是否满足这些属性。

我们感兴趣的主要功能暴露于caffe2/python/hypothesis_test_util.py中定义的HypothesisTestCase。

您应该将单元测试添加到caffe2/caffe2/python/operator_tests/文件夹。在这个目录中你可以找到很多已有的例子。

关键函数为：

assertDeviceChecks(devices，op，inputs，outputs)：这表明不管算子在哪个设备上执行，其计算的输出相同。
assertGradientChecks(device，op，inputs，output_，outputs_with_grads)：这为所讨论的算子实现了一个标准的数值梯度检查器。
assertReferenceChecks(device，op，inputs，reference)：它运行引用函数（有效地调用reference(*inputs)，并将其与output的输出进行比较。hypothesis_test_util.py公开了一些有非常用的预先构建的采样器。
hu.gcs - 梯度检查设备（gc）和设备检查设备（dc）
hu.gcs_cpu_only - CPU-only算子的梯度检查设备（gc）和设备检查设备（dc）

举个简单的例子：

@given(X=hu.tensor(), **hu.gcs)
def test_averaged_loss(self, X, gc, dc):
    op = core.CreateOperator("AveragedLoss", ["X"], ["loss"])
    self.assertDeviceChecks(dc, op, [X], [0])
    self.assertGradientChecks(gc, op, [X], 0, [0])

另一个演示assertReferenceChecks用法的例子：

@given(inputs=hu.tensors(n=3),
       in_place=st.booleans(),
       beta1=st.floats(min_value=0.1, max_value=0.9),
       beta2=st.floats(min_value=0.1, max_value=0.9),
       lr=st.floats(min_value=0.1, max_value=0.9),
       iters=st.integers(min_value=1, max_value=10000),
       epsilon=st.floats(min_value=1e-5, max_value=1e-2),
       **hu.gcs)
def test_adam(self, inputs, in_place, beta1, beta2, lr, iters, epsilon,
              gc, dc):
    grad, m1, m2 = inputs
    m2 += np.abs(m2) + 0.01
    lr = np.asarray([lr], dtype=np.float32)
    iters = np.asarray([iters], dtype=np.int32)
    op = core.CreateOperator(
        "Adam",
        ["grad", "m1", "m2", "lr", "iters"],
        ["grad" if in_place else "grad_o",
         "m1" if in_place else "m1_o",
         "m2" if in_place else "m2_o"],
        beta1=beta1, beta2=beta2, epsilon=epsilon,
        device_option=gc)
    input_device_options = {"lr": hu.cpu_do, "iters": hu.cpu_do}
    self.assertDeviceChecks(
        dc, op, [grad, m1, m2, lr, iters], [0], input_device_options)

    # Reference
    def adam(grad, m1, m2, lr, iters):
        lr = lr[0]
        iters = iters[0]
        t = iters + 1
        corrected_local_rate = lr * np.sqrt(1. - np.power(beta2, t)) / \
            (1. - np.power(beta1, t))

        m1_o = (beta1 * m1) + (1. - beta1) * grad
        m2_o = (beta2 * m2) + (1. - beta2) * np.square(grad)
        grad_o = corrected_local_rate * m1_o / \
            (np.sqrt(m2_o) + epsilon)
        return (grad_o, m1_o, m2_o)

    self.assertReferenceChecks(gc, op, [grad, m1, m2, lr, iters],
                               adam, input_device_options)

一个演示绘制更复杂元素的示例：

@given(prediction=hu.arrays(dims=[10, 3],
                            elements=st.floats(allow_nan=False,
                                               allow_infinity=False,
                                               min_value=0,
                                               max_value=1)),
       labels=hu.arrays(dims=[10],
                        dtype=np.int32,
                        elements=st.integers(min_value=0,
                                             max_value=3 - 1)),
        **hu.gcs)
def test_accuracy(self, prediction, labels, gc, dc):
    op = core.CreateOperator(
        "Accuracy",
        ["prediction", "labels"],
        ["accuracy"]
    )

    def op_ref(prediction, labels):
        N = prediction.shape[0]
        correct = 0
        max_ids = np.argmax(prediction, axis=1)
        for i in range(0, N):
            if max_ids[i] == labels[i]:
                correct += 1
        accuracy = correct / N
        return (accuracy,)

    self.assertReferenceChecks(
        device_option=gc,
        op=op,
        inputs=[prediction, labels],
        reference=op_ref)

不要忘记通过创建一个Issue来描述你的算子并链接到你的项目。

Caffe2自定义Operator

Caffe2 Custom Operators

构建一个基本算子

实现的细节

Caffe2算子单元测试

猜你喜欢