pytorch如何自定义operation

官方教程
参考代码
参考代码

最近的任务会使用到光流，目前有一些深度网络可以用来预测视频帧的光流，包括deepflow、flownet、flownet2.0、PWC-Net等，其中flownet2.0提供了完整的train/inference代码
由于flownet2.0使用到了correlation操作用来计算前后两帧图像特征的volume cost，pytorch没有实现该层，因此需要自己实现。

pytorch.Autograd

想要自定义operation，首先应该了解pytorch的计算图原理。pytorch通过autograd记录计算图中的function，然后再backward的过程中，通过记录进行回溯，通过链式求导法则更新梯度。

Function

Function是具体的函数定义，并且提供了forward和backward的定义，在forward过程中保存一些梯度计算需要的中间结果，在autograd回溯过程中，会调用backward函数，计算出对应输入的梯度。
每个non-leaf节点（required_grad==True的tensor）都会有一个grad_fn，方便自动求导

代码分析

在参考代码中的networks/correlation_package文件夹中，是自定义的correlation包，其中cuda相关文件用来定义forward和backward的具体实现，然后通过pybind导出python接口，供correlation.py文件调用。setup.py文件参考自官方提供的安装方式，将自动进行cu文件的编译以及python包的安装

setup.py

在进行安装时需要编译并自动安装，通常使用的是python的setup工具，参考setup.py

CUDA编程

CUDA原理
threadID计算方式