pytorch中调用C进行扩展

/* src/my_lib.h */
int my_lib_add_forward(THFloatTensor *input1, THFloatTensor *input2, THFloatTensor *output);
int my_lib_add_backward(THFloatTensor *grad_output, THFloatTensor *grad_input);

/* src/my_lib.c */
#include <TH/TH.h>

int my_lib_add_forward(THFloatTensor *input1, THFloatTensor *input2,
THFloatTensor *output)
{
    if (!THFloatTensor_isSameSizeAs(input1, input2))
        return 0;
    THFloatTensor_resizeAs(output, input1);
    THFloatTensor_cadd(output, input1, 1.0, input2);
    return 1;
}

int my_lib_add_backward(THFloatTensor *grad_output, THFloatTensor *grad_input)
{
    THFloatTensor_resizeAs(grad_input, grad_output);
    THFloatTensor_fill(grad_input, 1);
    return 1;
}

# build.py
from torch.utils.ffi import create_extension
ffi = create_extension(
name=‘_ext.my_lib‘,        # 输出文件地址及名称
headers=‘src/my_lib.h‘,    # 编译.h文件地址及名称
sources=[‘src/my_lib.c‘],  # 编译.c文件地址及名称
with_cuda=False            # 不使用cuda
)
ffi.build()

import torch
from torch.autograd import Function
from _ext import my_lib
import torch.nn as nn

class MyAddFunction(Function):
    def forward(self, input1, input2):
        output = torch.FloatTensor()
        my_lib.my_lib_add_forward(input1, input2, output)
        return output

    def backward(self, grad_output):
        grad_input = torch.FloatTensor()
        my_lib.my_lib_add_backward(grad_input, grad_output)
        return grad_input

class MyAddModule(nn.Module):
    def forward(self, input1, input2):
        return MyAddFunction()(input1, input2)

class MyNetWork(nn.Module):
    def __init__(self):
        super(MyNetWork, self).__init__()
        self.add = MyAddModule()

    def forward(self, input1, input2):
        return self.add(input1, input2)

model = MyNetWork()
input1, input2 = torch.randn(5, 5), torch.randn(5, 5)
print(model(input1, input2))
print(input1 + input2)