PyTorch 学习笔记 (3) PyTorch 1.0+ C++/CUDA extension

参考教程

PyTorch官方教程
PyTorch官方教程源码

期间参考了这位博主的一个教程。

本机系统

Ubuntu 18.04 LTS
NVIDIA GeForce GTX 1080 with driver 430.50
CUDA V10.1.243
python 3.6.8 virtualenv
torch.version == ‘1.2.0’

源码

所有源码可在这里获取。

Visual Studio Code配置

由于我使用的是python 的virtual environment，PyTorch放置的c++头文件的位置需要手工制定之后vs code的c++插件才会找到相关的头文件。在本机系统上，virtual env的位置是
/home/yaoyu/p3pt，
所需要制定的头文件位置为
/home/yaoyu/p3pt/lib/python3.6/site-packages/torch/include/torch/extension.h

Setup.py

在module的文件内新建setup.py文件，内容为

from setuptools import setup, Extension
from torch.utils import cpp_extension

setup(name="SigmoidCpp",
      ext_modules=[cpp_extension.CppExtension("SigmoidCpp", ['SigmoidCpp.cpp'])],
      cmdclass={
   'build_ext': cpp_extension.BuildExtension})

其中SigmoidCpp 是我们将要创建的c++ PyTorch extension 的python module的名字，SigmoidCpp.cpp是我们的c++源文件。

cpp

创建SigmoidCpp.cpp文件，并保存在setup.py文件边上。

#include <torch/extension.h>

#include <iostream>
#include <vector>

std::vector<torch::Tensor> sigmoid_cpp_forward( torch::Tensor input )
{
   
    return {
    torch::sigmoid( input ) };
}

std::vector<torch::Tensor> sigmoid_cpp_backward( torch::Tensor grad, torch::Tensor s )
{
   
    auto sp  = (1 - s) * s;

    return {
    grad * sp };
}

PYBIND11_MODULE( TORCH_EXTENSION_NAME, m )
{
   
    m.def("forward", &sigmoid_cpp_forward, "SigmoidCpp forward");
    m.def("backward", &sigmoid_cpp_backward, "SigmoidCpp backward");
}

这里sigmoid_cpp_forwad()和sigmoid_cpp_backward()我都使用了std::vector作为返回值，python binding过程会将这些std::vector转换为python list。这个简单实例其实并不需要std::vector，这里仅作一个演示如何输出多个tensor。

此后，在module目录内运行

python setup.py build_ext

若没有任何编译链接错误的话，会显示成功build了build/lib.linux-x86_64-3.6/SigmoidCpp.cpython-36m-x86_64-linux-gnu.so。此处build目录会在当前module目录下自动创建。

执行

python setup.py install --record .installed_files.txt

将刚刚编译完成的package安装到python环境中，由于目前使用的是virtual env，所以将会把文件安装到veritual env的目录内，我们可以显示.installed_files.txt的内容以查看已经被安装的文件。在调试过程中，可以通过如下命令来删除这些安装入python环境的文件

cat .installed_files.txt | xargs rm -rf

测试package

启动python或ipython，尝试如下命令

import torch
import SigmoidCpp

help(SigmoidCpp.forward)

将显示如下结果

Help on built-in function forward in module SigmoidCpp:

forward(...) method of builtins.PyCapsule instance
    forward(arg0: at::Tensor) -> List[at::Tensor]
    
    SigmoidCpp forward

本机显示的python function的signature和PyTorch官方教程的略有出入，主要体现在函数的参数类型上，这里都是ATen的namespace而PyTorch官方教程都是torch的namespace。这里需要注意，首先要import torch否则可能会报一些奇怪的c++动态库加载失败的错误。同样，可以help(SigmoidCpp.backward)

Help on built-in function backward in module SigmoidCpp:

backward(...) method of builtins.PyCapsule instance
    backward(arg0: at::Tensor, arg1: at::Tensor) -> List[at::Tensor]
    
    SigmoidCpp backward

创建实例化autograd.Function

创建一个新的python文件，这里我命名为SigmoidCppAG.py， AG意为AutoGradient，我知道这命名有点二，是我随手想的没有深入考究。文件内容可为

import torch

import SigmoidCpp

class SigmoidCppFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        s = SigmoidCpp.forward(x)
        ctx.save_for_backward( s[0] )
        return s[0]

    @staticmethod
    def backward(ctx, grad):
        sv = ctx.saved_variables

        output = SigmoidCpp.backward( grad, sv[0] )

        return output[0]

class SigmoidCppM(torch.nn.Module):
    def __init__(self):
        super(SigmoidCppM, self).__init__()
    
    def forward(self, x):
        return SigmoidCppFunction.apply( x )

这里，根据PyTorch官方教程的推荐，我们同时实例化torch.autograd.Function和torch.nn.Module并且在实例化torch.nn.Module时使用先实例化好的torch.autograd.Function。有几点需要注意

• 实例化torch.autograd.Function时，forward和backward函数是静态的。
• SigmoidCpp是我们利用c++扩展的package的名称，定义在setup.py文件中。
• SigmoidCpp.forward和SigmoidCpp.backward函数是在c++扩展中实现的sigmoid_cpp_forward和sigmoid_cpp_backward函数，但是利用PYBIND11_MODULE宏映射成了foward和backward。
• 根据我们当前的实例，SigmoidCpp.forward将返回一个list，SigmoidCpp.backward也将返回一个list。
• 无论ctx.save_for_backward打包了几个variable，ctx.saved_variables都将返回一个list。

在本机上设计一个简单的测试脚本，用于测试SigmoidCppAG.py的正确性。这个脚本参考了上一个PyTorch学习笔记。

import torch
from SigmoidCppAG import SigmoidCppM

if __name__ == "__main__":
    dc = SigmoidCppM()

    x = torch.rand((2,2), requires_grad=True)  # The input data.
    Y = torch.rand((2,2), requires_grad=False) # The true data.

    # Forward.
    y = dc(x)

    # Compute the loss.
    L = Y - y

    # Backward.
    L.backward(torch.ones(2,2))

    print("x = {}. ".format(x))
    print("y = {}. ".format(y))
    print("x.grad = {}. ".format( x.grad )