caffe2 添加自定义operater

Note: This article is available in Chinese only. 本文暂无英文版本。 View original

记录一下，一个之前没有接触过 caffe/caffe2 的人为了添加自定义 op 到 caffe2 需要做的工作。

首先参考 caffe2 tutorial，随便跑个 op 来试试，不妨以比较简单的 Accumulate_op 为例子。

这个 op 的作用就是计算 Y=X+gamma*Y，其中 X 为输入，Y 为输出，gamma 是参数。

跑起来这个运算所需要的代码如下：

 1from caffe2.python import workspace, model_helper
 2import numpy as np
 3# Create the input data
 4data = np.arange(6).reshape(2,3).astype(np.float32)
 5print ("data=",data)
 6
 7# Create labels for the data as integers [0, 9].
 8
 9workspace.FeedBlob("data", data)
10# Create model using a model helper
11m = model_helper.ModelHelper(name="my first net")
12output = m.net.Accumulate(["data"], "output")
13print(m.net.Proto())
14workspace.RunNetOnce(m.param_init_net)
15workspace.CreateNet(m.net)
16workspace.RunNet(m.name,2)  # run 2 times
17print("output=",workspace.FetchBlob('output'))

之后我们仿照 caffe2/operators/accumulate_op.h 和 caffe2/operators/accumulate_op.cc，仿写一个自己的运算 atest_op.h 和 atest_op.cc

实现的功能为 Y=5*X+gamma*Y

详细代码

 1#include "caffe2/operators/atest_op.h"
 2
 3namespace caffe2 {
 4REGISTER_CPU_OPERATOR(Atest, AtestOp<float, CPUContext>);
 5
 6OPERATOR_SCHEMA(Atest)
 7 .NumInputs(1)
 8 .NumOutputs(1)
 9 .IdenticalTypeAndShape()
10 .SetDoc(R"DOC(
11Accumulate operator accumulates the input tensor to the output tensor. If the
12output tensor already has the right size, we add to it; otherwise, we first
13initialize the output tensor to all zeros, and then do accumulation. Any
14further calls to the operator, given that no one else fiddles with the output
15in the interim, will do simple accumulations.
16Accumulation is done using Axpby operation as shown:
17 Y = 1*X + gamma*Y
18where X is the input tensor, Y is the output tensor and gamma is the multiplier
19argument.
20)DOC")
21 .Arg("gamma", "(float, default 1.0) Accumulation multiplier")
22 .Input(0, "input", "The input tensor that has to be accumulated to the "
23 "output tensor. If the output size is not the same as input size, the "
24 "output tensor is first reshaped and initialized to zero, and only "
25 "then, accumulation is done.")
26 .Output(0, "output", "Accumulated output tensor");
27
28SHOULD_NOT_DO_GRADIENT(Atest);
29} // namespace caffe2
30
31
32
33
34
35#ifndef CAFFE2_OPERATORS_ATEST_OP_H_
36#define CAFFE2_OPERATORS_ATEST_OP_H_
37
38#include "caffe2/core/context.h"
39#include "caffe2/core/operator.h"
40#include "caffe2/utils/math.h"
41
42namespace caffe2 {
43
44template <typename T, class Context>
45class AtestOp final : public Operator<Context> {
46 public:
47 AtestOp(const OperatorDef& operator_def, Workspace* ws)
48 : Operator<Context>(operator_def, ws),
49 gamma_(static_cast<T>(
50 OperatorBase::template GetSingleArgument<float>("gamma", 1.0))) {}
51 USE_OPERATOR_CONTEXT_FUNCTIONS;
52
53 bool RunOnDevice() override {
54 auto& input = Input(0);
55 auto* output = Output(0);
56 if (output->dims() != input.dims()) {
57 LOG(INFO) << "Reshaping and initializing output.";
58 output->ResizeLike(input);
59 math::Set<T, Context>(
60 output->size(), 0, output->template mutable_data<T>(), &context_);
61 }
62 math::Axpby<T, Context>(
63 input.size(),
64 static_cast<T>(5),
65 input.template data<T>(),
66 gamma_,
67 output->template mutable_data<T>(),
68 &context_);
69 return true;
70 }
71
72 protected:
73 T gamma_;
74};
75
76} // namespace caffe2
77
78#endif // CAFFE2_OPERATORS_Atest_OP_H_

之后我们编译整个 caffe2，编译方式是运行 pytorch/scripts/build_local.sh

编译成功后，需要将 pytorch 目录添加到 PYTHONPATH 中

1export PYTHONPATH=$PYTHONPATH:/mnt/lustre/renkuanze/workspace/rjm_pytorch

然后运行

1cd ~ && python -c 'from caffe2.python import core' 2>/dev/null && echo "Success" || echo "Failure"

看是否成功

编译的时候可能出现 mpi_test.cc.o: undefined reference to symbol ‘_ZN3MPI8Datatype4FreeEv 的报错，解决办法是把 CMakeList 中的 MPI 关掉就好了。

以及……operators 中的文件都不要删……本想删一些不相关的 op 来减少编译时间，想法是对的，但是似乎只删 op 是行不通的，不如不删，不然会编译出现奇怪的错误！

不然会编译出现奇怪的错误！

以及从 github download 的速度太慢了，干脆开了个 40$/m 的 vps 来搞。

编译成功后修改测试的 python 代码，来测试一下我们定义的 op

 1from caffe2.python import workspace, model_helper
 2import numpy as np
 3# Create the input data
 4data = np.arange(6).reshape(2,3).astype(np.float32)
 5print ("data=",data)
 6
 7# Create labels for the data as integers [0, 9].
 8
 9workspace.FeedBlob("data", data)
10# Create model using a model helper
11m = model_helper.ModelHelper(name="my first net")
12output = m.net.Atest(["data"], "output")
13print(m.net.Proto())
14workspace.RunNetOnce(m.param_init_net)
15workspace.CreateNet(m.net)
16workspace.RunNet(m.name,1)  # run 2 times
17print("output=",workspace.FetchBlob('output'))

发现确实是得到了 Y=5*X+gamma*Y 的结果。

撒花！（然而这只是最近要做的任务中最容易的一条线 orz……）

需要注意的是，运行 bash build_local.sh 脚本之后会在 pytorch 目录下生成 build 文件夹，下次编译的时候直接在 build 目录下执行 make -j20，这样才是增量编译

编译的时候直接在 build 目录下执行 make -j20，这样才是增量编译

不然每次执行 build_local.sh，不知 caffe2 用了什么机制，每次要把所有文件编译一遍，简直没有人性啊，16 个 cpu 一起编也要 5 分钟 orz。

以及要在自己定义的函数最后返回 true，这样在运行的时候才不会报错，否则会报 net error 的错误 orz……我好傻啊。

Related