pytorch\u5df2\u7ecf\u57fa\u672c\u5b9e\u73b0\u4e86\u5e38\u89c1\u7684\u5404\u79cdop\uff0c\u7136\u800c\uff0c\u5f53\u60f3\u5b9e\u73b0\u4e00\u4e2apytorch\u4e2d\u6ca1\u6709\u7684op\u65f6\uff0c\u6709\u4e24\u79cd\u65b9\u5f0f\u3002\u4e00\u79cd\u65b9\u5f0f\u662f\u8fd9\u4e2aop\u53ef\u4ee5\u7531pytorch\u4e2d\u5df2\u6709\u7684op\u8fdb\u884c\u7ec4\u5408\u800c\u6210\uff0c\u56e0\u6b64\u53ea\u9700\u8981\u4f7f\u7528python\u63a5\u53e3\u8fdb\u884c\u7ec4\u5408\u5c31\u53ef\u4ee5\u4e86\u3002\u53cd\u4e4b\uff0c\u5c31\u5fc5\u987b\u4f7f\u7528c++\u6216\u8005cuda\u5b9e\u73b0\u8be5op\uff0c\u7136\u540e\u6dfb\u52a0\u5230pytorch\u4e2d\u3002\u672c\u6587\u5c06\u4ecb\u7ecd\u6dfb\u52a0c++\u5b9e\u73b0\u7684\u81ea\u5b9a\u4e49op\uff0c\u56e0\u4e3a\u6211\u8fd8\u4e0d\u4f1acuda : (<\/p>\n\n\n\n
\u672c\u6587\u4ecb\u7ecd\u4f7f\u7528python\u7684setuptools\u5c06c++\u5b9e\u73b0\u7684op\u6dfb\u52a0\u5230pytorch\u4e2d\u3002\u9996\u5148\u8981\u7528c++\u5b9e\u73b0\u5b9a\u4e49\u7684op\u3002\u6bd4\u5982\u60f3\u5b9e\u73b0\u4e00\u4e2aop\u4e3a z=3x-y \u3002\u5934\u6587\u4ef6\u4e3amy_op.h<\/p>\n\n\n\n
#include <torch\/extension.h> \/\/\u8fd9\u4e00\u53e5\u662f\u65e0\u8bba\u8981\u5b9e\u73b0\u4efb\u4f55op\u90fd\u5fc5\u987b\u6dfb\u52a0\u7684\n#include <vector>\n\n\/\/\u524d\u5411\u4f20\u64ad\ntorch::Tensor my_op_forward(const torch::Tensor& x, const torch::Tensor& y);\n\/\/\u53cd\u5411\u4f20\u64ad\nstd::vector<torch::Tensor> my_op_backward(const torch::Tensor& gradOutput);<\/pre>\n\n\n\n\u6e90\u6587\u4ef6\u4e3amy_op.cpp<\/p>\n\n\n\n
#include \"my_op.h\"\n\ntorch::Tensor my_op_forward(const torch::Tensor& x, const torch::Tensor& y) { \n AT_ASSERTM(x.sizes() == y.sizes(), \"x must be the same size as y\");\n torch::Tensor z = torch::zeros(x.sizes());\n z = 3 * x - y;\n return z; } \n\nstd::vector<torch::Tensor> my_op_backward(const torch::Tensor& gradOutput) {\n torch::Tensor gradOutputX = 3 * gradOutput * torch::ones(gradOutput.sizes());\n torch::Tensor gradOutputY = -1 * gradOutput * torch::ones(gradOutput.sizes());\n return {gradOutputX, gradOutputY}; } \n\n\/\/ pybind11 \u7ed1\u5b9a \nPYBIND11_MODULE(my_op_api, m) {\n m.def(\"forward\", &my_op_forward, \"MY_OP forward\");\n m.def(\"backward\", &my_op_backward, \"MY_OP backward\"); \n} <\/pre>\n\n\n\n\u5176\u4e2d\u6700\u540e\u7684PYBIND11_MODULE\u662f\u7528\u6765\u5c06C++\u51fd\u6570\u7ed1\u5b9a\u5230python\u4e0a\u7684\u3002\u5176\u4e2d\u7b2c\u4e00\u4e2a\u53c2\u6570my_op_api\u4e3a\u8981\u751f\u6210\u7684python\u6a21\u5757\u540d\uff0c\u4ee5\u540eimport my_op_api\u5c31\u53ef\u4ee5\u8c03\u7528\u8be5op\u4e86\u3002\u7b2c\u4e8c\u4e2a\u53c2\u6570\u56fa\u5b9a\u4e3am
\u51fd\u6570\u4f53\u4e2d\u7684\u4e24\u4e2a\u8bed\u53e5\u5206\u522b\u662f\u7ed1\u5b9a\u524d\u5411\u4f20\u64ad\u4e0e\u53cd\u5411\u4f20\u64ad\u5230\u5b9e\u73b0\u7684\u4e24\u4e2a\u51fd\u6570\u4e0a\u3002<\/p>\n\n\n\n\u7136\u540e\u7f16\u5199setup.py\uff0c\u7528\u6765\u6784\u5efapytorch\u7684c++\u6269\u5c55\u3002<\/p>\n\n\n\n
from setuptools import setup\nfrom torch.utils.cpp_extension import BuildExtension, CppExtension\n\nsetup(name='my_op_api',\n version='0.l',\n ext_modules=[CppExtension('my_op_api', sources=['my_op.cpp'], extra_compile_args=['-std=c++11'])],\n cmdclass={'build_ext':BuildExtension})<\/pre>\n\n\n\n\u5176\u4e2d\uff0csetup\u4e2d\u7684name\u4ee5\u53caCppExtension\u4e2d\u7b2c\u4e00\u4e2a\u53c2\u6570\uff08\u4e5f\u662fname\uff09\u8981\u548cPYBIND11_MODULE\u91cc\u8bbe\u7684\u6a21\u5757\u540d\u4fdd\u6301\u4e00\u81f4\uff0c\u8fd9\u91cc\u90fd\u662fmy_op_api\u3002CppExtension\u4e2d\u7684extra_compile_args=[‘-std=c++11’]\u662f\u8d9f\u5751\u53d1\u73b0\u7684\uff0c\u4e0d\u52a0\u7684\u8bddgcc\u53ef\u80fd\u4f1a\u62a5n\u591a\u9519\uff08pytorch\u662f\u7528c++11\u7f16\u8bd1\u7684\uff0c\u56e0\u6b64\u8fd9\u91cc\u7528gcc\u7f16\u8bd1\u7684\u65f6\u5019\u4e5f\u8981\u4f7f\u7528c++11\uff09<\/p>\n\n\n\n
\u7136\u540e\u8fd0\u884cpython setup.py install\uff0c\u5982\u679c\u6ca1\u6709\u95ee\u9898\u7684\u8bdd\uff0c\u5c31\u751f\u6210\u4e86\u6240\u9700\u7684python\u6a21\u5757\u3002\u53ef\u4ee5\u4ece\u8f93\u51fa\u7684\u4fe1\u606f\u770b\u5230\u8be5\u6a21\u5757\u5728\u6240\u5728python\u73af\u5883\u4e0b\u7684site-packages\u6587\u4ef6\u5939\u4e0b\uff0c\u4ee5.egg\u7ed3\u5c3e\u3002\u53e6\u5916\uff0c\u5728\u5f53\u524d\u76ee\u5f55\u4e0b\u4f1a\u67093\u4e2a\u6587\u4ef6\u5939\u751f\u6210\uff0cbuild\u3001dist\u3001my_op_api.egg-info\uff0c\u5176\u4e2ddist\u4e0b\u4e5f\u6709.egg\u6587\u4ef6\uff0c\u53ef\u4ee5\u53d1\u5e03\u5230\u5176\u5b83python\u73af\u5883\u3002<\/p>\n\n\n\n
\u7136\u540e\u5c31\u53ef\u4ee5\u5728python\u4e2dimport my_op_api\u8fdb\u884c\u8c03\u7528\u6269\u5c55\u7684op\u4e86\u3002\u8fd9\u91cc\u9700\u8981\u6ce8\u610f\u4e00\u70b9\u7684\u662f\uff0c\u5728import \u81ea\u5b9a\u4e49\u7684op\u4e4b\u524d\uff0c\u5fc5\u987b\u5148import torch\u3002 \u4f46\u662f\uff0c\u8fd9\u6837\u7684op\u548c\u6211\u4eec\u65e5\u5e38\u4f7f\u7528\u7684\u8fd8\u662f\u4e0d\u592a\u4e00\u6837\uff0c\u8fd9\u65f6\u9700\u8981\u5c06\u5b83\u5305\u88c5\u4e3apytorch\u4e2d\u7684\u51fd\u6570\u548c\u6a21\u5757\uff0c\u4ee5\u4fbf\u6211\u4eec\u50cf\u4f7f\u7528\u5176\u5b83\u6a21\u5757\u4e00\u6837\u4f7f\u7528\u81ea\u5b9a\u4e49\u7684op\u3002\u8981\u5305\u88c5\u4e3a\u6a21\u5757\uff0c\u9996\u5148\u5305\u88c5\u6210\u51fd\u6570\u3002 <\/p>\n\n\n\n
\u5305\u88c5\u6210\u51fd\u6570\uff0c\u9700\u8981\u7ee7\u627ftorch.autograd.Function\u3002\u7136\u540e\u5305\u88c5\u6210\u6a21\u5757\uff0c\u9700\u8981\u7ee7\u627ftorch.nn.Module<\/p>\n\n\n\n
import torch\nfrom torch.autograd import Function\nfrom torch.nn import Module\nimport my_op_api\n\nclass MyOpFunction(Function):\n @staticmethod\n def forward(ctx, x, y):\n #\u5982\u679c\u6709\u4e00\u4e9b\u4fe1\u606f\uff0c\u9700\u8981\u5728\u68af\u5ea6\u53cd\u5411\u4f20\u64ad\u65f6\u7528\u5230\uff0c\u53ef\u4ee5\u4f7f\u7528ctx.save_for_backward()\u8fdb\u884c\u4fdd\u5b58\n return my_op_api.forward(x, y)\n @staticmethod\n def backward(ctx, gradOutput):\n #\u5982\u679c\u5728forward\u4e2d\u4fdd\u5b58\u4e86\u4fe1\u606f\uff0c\u53ef\u4ee5\u4f7f\u7528ctx.saved_tensors\u53d6\u56de\n grad_x, grad_y = my_op_api.backward(gradOutput)\n return grad_x, grad_y\n\nclass MyOpModule(Module):\n def __init__(self):\n super(MyOpModule, self).__init__()\n def forward(self, input_x, input_y):#\u53ea\u9700\u8981\u5b9a\u4e49forward\u7684\u51fd\u6570\u5c31\u53ef\u4ee5\u4e86\n return MyOpFunction.apply(input_x, input_y)<\/pre>\n\n\n\n<\/p>\n","protected":false},"excerpt":{"rendered":"
pytorch\u5df2\u7ecf\u57fa\u672c\u5b9e\u73b0\u4e86\u5e38\u89c1\u7684\u5404\u79cdop\uff0c\u7136\u800c\uff0c\u5f53\u60f3\u5b9e\u73b0\u4e00\u4e2apytorch\u4e2d\u6ca1\u6709\u7684op\u65f6\uff0c\u6709\u4e24\u79cd\u65b9\u5f0f\u3002\u4e00\u79cd\u65b9 … <\/p>\n