自定义实现C++拓展pytorch功能

ncrelu.cpp

cpp 复制代码
#include <torch/extension.h>					// 头文件引用部分

namespace py = pybind11;

torch::Tensor ncrelu_forward(torch::Tensor input) {
    auto pos = input.clamp_min(0);				       // 具体实现部分
    auto neg = input.clamp_max(0);
    return torch::cat({pos, neg}, 1);
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {	// 绑定部分
    m.def("forward", &ncrelu_forward, py::arg("input"), "NCReLU forward");
}

setup.py

python 复制代码
from setuptools import setup
from torch.utils import cpp_extension


setup(
    name='ncrelu_cpp',
    version='1.0',# 编译后的链接库名称
    py_modules=['ncrelu_cpp'],
    ext_modules=[
        cpp_extension.CppExtension(
            'ncrelu_cpp', ['ncrelu.cpp'],
            extra_compile_args={'cxx': ['-O2']}
            # 待编译文件,及编译函数
        )
    ],
    cmdclass={						       # 执行编译命令设置
        'build_ext': cpp_extension.BuildExtension
    }
)

test.py

python 复制代码
import torch
import ncrelu_cpp
import sys
print(sys.path)
a = torch.randn(4,3)
print(a)
b = ncrelu_cpp.forward(a)

python setup.py install

或pip install .

但是在Windows平台下不知道为什么会报错找不到包,或者找不到函数,很奇怪,但是正常运行没有任何问题

相关推荐
灵智工坊LingzhiAI13 分钟前
人体坐姿检测系统项目教程(YOLO11+PyTorch+可视化)
人工智能·pytorch·python
小汉堡编程1 小时前
数据结构——vector数组c++(超详细)
数据结构·c++
tan180°6 小时前
MySQL表的操作(3)
linux·数据库·c++·vscode·后端·mysql
彭祥.8 小时前
Jetson边缘计算主板:Ubuntu 环境配置 CUDA 与 cudNN 推理环境 + OpenCV 与 C++ 进行目标分类
c++·opencv·分类
烛阴8 小时前
简单入门Python装饰器
前端·python
lzb_kkk8 小时前
【C++】C++四种类型转换操作符详解
开发语言·c++·windows·1024程序员节
好开心啊没烦恼8 小时前
Python 数据分析:numpy,说人话,说说数组维度。听故事学知识点怎么这么容易?
开发语言·人工智能·python·数据挖掘·数据分析·numpy
面朝大海,春不暖,花不开8 小时前
使用 Python 实现 ETL 流程:从文本文件提取到数据处理的全面指南
python·etl·原型模式
2301_8050545610 小时前
Python训练营打卡Day59(2025.7.3)
开发语言·python
胖大和尚10 小时前
clang 编译器怎么查看在编译过程中做了哪些优化
c++·clang