PyTorch中的C++扩展实现

#include <torch/extension.h>
#include <vector>

// 前向传播
torch::Tensor Test_forward_cpu(const torch::Tensor& inputA,
              const torch::Tensor& inputB);
// 反向传播
std::vector<torch::Tensor> Test_backward_cpu(const torch::Tensor& gradOutput);

#include "test.h"

// 前向传播，两个 Tensor 相加。这里只关注 C++ 扩展的流程，具体实现不深入探讨。
torch::Tensor Test_forward_cpu(const torch::Tensor& x,
              const torch::Tensor& y) {
  AT_ASSERTM(x.sizes() == y.sizes(), "x must be the same size as y");
  torch::Tensor z = torch::zeros(x.sizes());
  z = 2 * x + y;
  return z;
}

// 反向传播
// 在这个例子中，z对x的导数是2，z对y的导数是1。
// 至于这个backward函数的接口（参数，返回值）为何要这样设计，后面会讲。
std::vector<torch::Tensor> Test_backward_cpu(const torch::Tensor& gradOutput) {
  torch::Tensor gradOutputX = 2 * gradOutput * torch::ones(gradOutput.sizes());
  torch::Tensor gradOutputY = gradOutput * torch::ones(gradOutput.sizes());
  return {gradOutputX, gradOutputY};
}

// pybind11 绑定
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
 m.def("forward", &Test_forward_cpu, "TEST forward");
 m.def("backward", &Test_backward_cpu, "TEST backward");
}

└── csrc
  ├── cpu
  │  ├── test.cpp
  │  └── test.h
  └── setup.py

from setuptools import setup
import os
import glob
from torch.utils.cpp_extension import BuildExtension, CppExtension

# 头文件目录
include_dirs = os.path.dirname(os.path.abspath(__file__))
# 源代码目录
source_cpu = glob.glob(os.path.join(include_dirs, 'cpu', '*.cpp'))

setup(
  name='test_cpp', # 模块名称，需要在python中调用
  version="0.1",
  ext_modules=[
    CppExtension('test_cpp', sources=source_cpu, include_dirs=[include_dirs]),
  ],
  cmdclass={
    'build_ext': BuildExtension
  }
)

python setup.py install

from torch.autograd import Function

import test_cpp

class TestFunction(Function):

  @staticmethod
  def forward(ctx, x, y):
    return test_cpp.forward(x, y)

  @staticmethod
  def backward(ctx, gradOutput):
    gradX, gradY = test_cpp.backward(gradOutput)
    return gradX, gradY

class Function(with_metaclass(FunctionMeta, _C._FunctionBase, _ContextMethodMixin, _HookMixin)):
 
  ...

  @staticmethod
  def forward(ctx, *args, **kwargs):
    r"""Performs the operation.

    This function is to be overridden by all subclasses.

    It must accept a context ctx as the first argument, followed by any
    number of arguments (tensors or other types).

    The context can be used to store tensors that can be then retrieved
    during the backward pass.
    """
    raise NotImplementedError

  @staticmethod
  def backward(ctx, *grad_outputs):
    r"""Defines a formula for differentiating the operation.

    This function is to be overridden by all subclasses.

    It must accept a context :attr:`ctx` as the first argument, followed by
    as many outputs did :func:`forward` return, and it should return as many
    tensors, as there were inputs to :func:`forward`. Each argument is the
    gradient w.r.t the given output, and each returned value should be the
    gradient w.r.t. the corresponding input.

    The context can be used to retrieve tensors saved during the forward
    pass. It also has an attribute :attr:`ctx.needs_input_grad` as a tuple
    of booleans representing whether each input needs gradient. E.g.,
    :func:`backward` will have ``ctx.needs_input_grad[0] = True`` if the
    first input to :func:`forward` needs gradient computated w.r.t. the
    output.
    """
    raise NotImplementedError

import torch

class Test(torch.nn.Module):

  def __init__(self):
    super(Test, self).__init__()

  def forward(self, inputA, inputB):
    return TestFunction.apply(inputA, inputB)

├── csrc
│  ├── cpu
│  │  ├── test.cpp
│  │  └── test.h
│  └── setup.py
└── test.py

import torch
from torch.autograd import Variable

from test import Test

x = Variable(torch.Tensor([1,2,3]), requires_grad=True)
y = Variable(torch.Tensor([4,5,6]), requires_grad=True)
test = Test()
z = test(x, y)
z.sum().backward()
print('x: ', x)
print('y: ', y)
print('z: ', z)
print('x.grad: ', x.grad)
print('y.grad: ', y.grad)