去年年底的时候, 我需要复现一篇文章的工作, 那个网络框架里有变形卷积的内容, 所以需要首先实现变形卷积op. 一开始我觉得不会太费力气, 后来才发现还是有很大的难度的, 尤其是对于我这种刚接触dl不久的新人, 要完成一整个c++部分, cuda部分以及python的代码编写真的挺不容易. 后来年后又将其在pytorch上复现一遍, 又踩了不少的坑, 这篇文章, 我会把pytorch上进行c++ extension的步骤, deformable_conv2d的pytorch以及踩过的坑都详细的说一遍. 希望对以后有需要的人有所帮助.
首先, deformable部分的具体内容我就不给大家细说了, 可以看去年刚出的
Deformable ConvNets v2: More Deformable, Better Resultsarxiv.org以及官方的github
msracver/Deformable-ConvNetsgithub.com
. 我只实现了我所需要的deformable_conv2d部分, 至于deformable roi部分我并没有实现, 但只要会了过程, 也大同小异. 下面是我实现的TensorFlow版本和pytorch版本的变形卷积(刚上传, 过几天我写个详细的readme).
https://github.com/BIGKnight/deformable_conv2d_v2_tensorflowgithub.com BIGKnight/deformable_conv2d_pytorchgithub.com
. 我在第一节讲一下如何在pytorch下进行c++拓展, 大家具体可以看官方教程, 我大致讲一下流程以及遇到的一些坑.我在第一节讲一下如何在pytorch下进行c++拓展, 大家具体可以看官方教程, 我大致讲一下流程以及遇到的一些坑.
PYTORCH CUSTOM C++ AND CUDA EXTENSIONS
不可否认的一点是, 变形卷积是必须用gpu加速的, 所以如果单纯的用框架上层python去实现和封装效率是非常低的. 之前也看到过一些在python层面的实现, 但是效率都非常低. 所以不管在什么框架下实现变形卷积, 首先我们要会的就是如何在该框架中进行op的c++和cuda编写, 这部分说实话pytorch其实非常友好, 相比于tensorflow, pytorch的扩展操作要方便和简单很多.
Custom C++ and CUDA Extensionspytorch.org.
总的来说, 在c++和cuda部分我们需要完成2部分代码, 一个是具体定义了前向forward和后向backward操作的.cpp文件, 我在这里就称其为deformableconv2d.cpp, 还有一个写具体需要gpu加速的函数的cuda文件, 我们称其为deformableconv2dcuda.cu. (当然也会有其他各种.h或者.cuh等辅助文件, 但是这些不是主体, 而且只是为了方便编写时的结构清晰, 都是可以写到一个.cpp和.cu中来的). 若只想实现cpu版本, .cu文件是不需要的, 所有具体实现在.cpp中完成即可. 我再这篇文章里就只讨论gpu版本了.
这小节中我主要讲一下关于pytorch custom op的流程和一些需要注意的点, 一些所需调用的库如Aten, thc这种的用法我就不具体说了, 详细内容可以看官方的源码(应该也有c++的api), 此外, 具体的deformable部分我将放在第二节讲.
1.1 CPP文件
CPP文件中主要可以分为四个部分:
- include头文件以及一些辅助函数以及cuda中定义的函数的声明, 这部分其实可以额外提到.h文件中写, 那样结构会更清晰一些. 不过若是放到.h文件中, 一定要注意添加宏#ifndef以防止重复编译, 因为有时候头文件会被其他多个文件调用, 若不用#ifndef, 则会将头文件中内容多次编译, 这会导致一些class和struct重名冲突, 编译不通过.
#ifndef FUNCTION_DECLARE
#define FUNCTION_DECLARE
void deformable_im2col(cudaStream_t stream,
const float* data_im, const float* data_offset, const float* data_mask,
const TShape& im_shape, const TShape& col_shape, const TShape& kernel_shape,
const TShape& pad, const TShape& stride, const TShape& dilation,
const int deformable_group, float* data_col);
......
void SwapAxis(cudaStream_t stream, float* input_data, const TShape& origin_shape, const int axis_x, const int axis_y);
#endif- forward函数, 就是op的前向部分, 参数就是网络中向这个节点的所有tensor输入(如input, weight, bias等, 注意类型为at::Tensor)以及其他一些像步长, padding, 以及dilation这些mata-parameter. 一般的实现流程都是先进行输入的正确性检查, 然后编写函数体最后返回输出, 输出可以是一个tensor也可以是多个tensor, 若是多个tensor就返回一个vector<at::Tensor>, 否则就返回at::Tensor.
at::Tensor deformable_conv2d_forward(
at::Tensor input,
.....
bool no_bias
){
<check the inputs>
.....
<main part>
.....
<return the output tensor>
}- backward函数, 就是op的后向部分, 这个会在auto_grad的时候调用, 用于计算梯度, 返回的是所有需要更新的输入tensor的梯度变化tensor. 具体实现的流程和forward类似, 也是先进行输入的正确性检查, 然后编写主体, 最后返回计算出的grads
std::vector<at::Tensor> deformable_conv2d_backward(
at::Tensor input,
......
bool no_bias
){
<check the inputs>
.....
<main part>
.....
<return the output tensor>
}- PYBIND, 这个部分在代码上非常简单, 其作用是将c++中写这两个函数在python中注册并生成python中封装后的函数. pybind是用来实现python和c++之间的互操作的, 主要用于创建已有c++代码的python封装版本。其实就是用C++写pyd。编译成功后会生成一个.pyd文件，python可以直接import这个库。
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("forward", &deformable_conv2d_forward, "deformable_conv2d forward (CUDA)");
m.def("backward", &deformable_conv2d_backward, "deformable_conv2d backward (CUDA)");
}1.2 CUDA文件
cuda如何编写不在本文范围内, 如果不清楚的话可以上网去查相应资料, 不管是官方的doc还是博客文章都很多, 而且cuda语法和c几乎一样, 学习也不难. 而且具体的cuda 内容要根据具体的情况而定, 比如卷积操作, 需要将im2col部分进行gpu加速. 此外, 值得注意的一个点是cuda自身就有许多常见的gpu计算的函数, 比如说cublas库就是非常常用的线性代数计算库, 里面有很多像矩阵乘法这种必用gpu的操作的实现, 一般要比我们个人实现的版本效率高些. 此外, pytorch也有很多cuda相关的库, 如THC包下的很多内容(但是貌似很多就只是对于cuda函数的封装).
我的习惯是先写__host__函数, 然后实现具体的kernel核.
cuda我觉得还有一个好处是, 这部分和具体的深度学习框架无关, 所以我的tensorflow版本实现和pytorch版本实现这部分基本上是一样的(自然和msra的版本也基本上一样), 除了一些附带的宏函数(但其实都可以用cuda c自己写).
有一些注意的点(都是我踩过的坑), 如果你是大佬就不用看了, 如果同是新手说不定对你有帮助:
- 需要注意内存和显存不要搞混, 这个我在一开始弄的时候经常疏忽.
- 不要把一些nvcc不能编译的库包进来, 比如torch/extension.h(也就是torch/torch.h ), 一般就需要用到ATEN和THC就可以了.
1.3 编译链接生成动态库
在完成上面两部分后, 一个自定义的操作就完成了, 接下来需要完成编译链接, 目的是得到.so(windows下.dll)动态库.
如果你对g++和nvcc编译比较熟悉的, 完全可以自己做. 如果不熟悉的话, 就老老实实使用setuptools, 如下所示,
from setuptools import setup
from torch.utils.cpp_extension import CppExtension, BuildExtension, CUDAExtension
setup(name='deformable_conv2d',
ext_modules=[CUDAExtension('deformable_conv2d_gpu', ['deformable_conv2d.cc', 'deformable_conv2d_cuda.cu']),],
cmdclass={'build_ext': BuildExtension})目的是将自定义的模块集成到框架中. 当然还可以用torch.utils.cpp_extension, 只不过这样是动态链接.
from torch.utils.cpp_extension import load
deformable_conv2d = load(name='deformable_conv2d', sources=['deformable_conv2d.cpp', 'deformable_conv2d_cuda.cu'])这里有一个我个人觉得比较坑的地方, pytorch在1.0以后(貌似从0.4.1之后就不支持了), 默认不使用老版本c++的ABI了, 而是支持c++11的ABI, 所以在安装pytorch的时候, 默认D_GLIBCXX_USE_CXX11_ABI=1, 但是非常坑的一个地方是setuptools的setup生成的编译语句中不管是nvcc还是gcc都是-D_GLIBCXX_USE_CXX11_ABI=0, 这就导致编译得到的.so和pytorch的ABI不匹配了, 因为不同的ABI对应的应用程序核操作系统之间的协定是不同的. 所以这样setup生成的.so是不支持c++11的ABI的, 这就导致.so和实际pytorch脚本是两个不同的c++版本, 就gg了.
所以, 我的解决的方式:
- 暴力的方式就是直接改回pytorch 0.4.0.
- 或者修改setup生成的.so, 用自己编译链接后的.so替代掉这个. 我将这部分脚本放在下面
#!/usr/bin/env bash
CUSTOM_MODULE_PATH=(/home/zzn/anaconda3/envs/pytorch/lib/python3.6/site-packages/deformable_conv2d-0.0.0-py3.6-linux-x86_64.egg)
DYNAMIC_LIB_NAME=(deformable_conv2d_gpu.cpython-36m-x86_64-linux-gnu.so)
CURRENT_NEW_DYNAMIC_LIB=(build/lib.linux-x86_64-3.6/deformable_conv2d_gpu.cpython-36m-x86_64-linux-gnu.so)
python setup.py install
gcc -pthread -B /home/zzn/anaconda3/envs/pytorch/compiler_compat -Wl,--sysroot=/ -Wsign-compare -DNDEBUG -g -fwrapv -O3 -Wall -Wstrict-prototypes -fPIC -I/home/zzn/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/lib/include -I/home/zzn/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/lib/include/torch/csrc/api/include -I/home/zzn/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/lib/include/TH -I/home/zzn/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/lib/include/THC -I/usr/local/cuda/include -I/home/zzn/anaconda3/envs/pytorch/include/python3.6m -c deformable_conv2d.cc -o build/temp.linux-x86_64-3.6/deformable_conv2d.o -DTORCH_API_INCLUDE_EXTENSION_H -DTORCH_EXTENSION_NAME=deformable_conv2d_gpu -D_GLIBCXX_USE_CXX11_ABI=1 -std=c++11
/usr/local/cuda/bin/nvcc -I/home/zzn/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/lib/include -I/home/zzn/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/lib/include/torch/csrc/api/include -I/home/zzn/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/lib/include/TH -I/home/zzn/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/lib/include/THC -I/usr/local/cuda/include -I/home/zzn/anaconda3/envs/pytorch/include/python3.6m -c deformable_conv2d_cuda.cu -o build/temp.linux-x86_64-3.6/deformable_conv2d_cuda.o -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --compiler-options '-fPIC' -DTORCH_API_INCLUDE_EXTENSION_H -DTORCH_EXTENSION_NAME=deformable_conv2d_gpu -D_GLIBCXX_USE_CXX11_ABI=1 -std=c++11
g++ -pthread -shared -B /home/zzn/anaconda3/envs/pytorch/compiler_compat -L/home/zzn/anaconda3/envs/pytorch/lib -Wl,-rpath=/home/zzn/anaconda3/envs/pytorch/lib -Wl,--no-as-needed -Wl,--sysroot=/ build/temp.linux-x86_64-3.6/deformable_conv2d.o build/temp.linux-x86_64-3.6/deformable_conv2d_cuda.o -L/usr/local/cuda/lib64 -lcudart -o build/lib.linux-x86_64-3.6/deformable_conv2d_gpu.cpython-36m-x86_64-linux-gnu.so
cp -f ${CURRENT_NEW_DYNAMIC_LIB} ${CUSTOM_MODULE_PATH}1.4 python封装
可以将其封装为一个Function类, 将其作为一个functional来使用.
import torch
......
# our module
import deformable_conv2d_gpu
class DeformableConv2DFunction(Function):
@staticmethod
def forward(ctx, *args):
.....
ctx.save_for_backward(input, filter, offset, mask)
return output
# backward 的返回值个数要和forward的输入个数等同
@staticmethod
def backward(ctx, *grad_outputs):
....
grad_input, grad_weight, grad_offset, grad_mask = deformable_conv2d_gpu.backward(...)
return grad_input, grad_weight, grad_offset, grad_mask,
None, None, None, None, None, None, None, None, None, None类似于torch.nn.functional.Conv2d, 如果想进一步封装像torch.nn.Conv2d那样使用, 就再用一个nn.Module去封装.
第一节就先到这里吧.
刚更新了第二节
王火华至秦：deformable变形卷积pytorch实现(第二节deformable_conv2d 实现一)zhuanlan.zhihu.com
668
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
