caffe2 添加自定义op

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

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

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

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

c之后我们仿照caffe2/operators/accumultate_op.h和affe2/operators/accumultate_op.cc,仿写一个我们自己的运算atest_op.h和atest_op.cc

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

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

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

然后运行

看是否成功

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

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

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

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

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

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

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

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

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

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

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

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

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

 

 

作者: CrazyKK

ex-ACMer@hust,stackoverflow-engineer@sensetime

说点什么

您将是第一位评论人!

提醒
wpDiscuz