当前位置: 首页 > news >正文

python调用底层c++算子示例

test.cpp

#include <torch/extension.h>// 定义一个简单的加法函数
at::Tensor add(at::Tensor a, at::Tensor b) {return a + b;
}PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {m.def("add", &add, "Add two tensors");
}

test.py

import torch
from torch.utils.cpp_extension import loadtest_load = load(name='test_load', sources=['test.cpp'],extra_cflags=['-O2'],verbose=True,
)
a = torch.tensor([1, 2, 3])
b = torch.tensor([4, 5, 6])
result = test_load.add(a, b)
print(result)  # Should print tensor([5, 7, 9])

注意:load中name实际上传给test.cpp中TORCH_EXTENSION_NAME, 会编译生成一个test_load.so 动态库

test2.py:调用test.py中test_load模块

from test import test_load
import torch
a = torch.tensor([1, 2, 3])
b = torch.tensor([4, 5, 6])
result = test_load.add(a, b)
print(result)  # Should print tensor([5, 7, 9])
http://www.dtcms.com/a/206327.html

相关文章:

  • Docker常用命令介绍
  • [欠拟合过拟合]机器学习-part10
  • Phantom 视频生成的流程
  • HarmonyOS学习——UIAbility组件(下)
  • 鸿蒙App开发学习路径
  • (第95天)OGG 微服务搭建 Oracle 19C 到 MySQL 8 双向同步
  • 【Jaspersoft studio 生成的模板,无法分页的问题】
  • 卓力达靶标:精密制造赋能材料沉积技术革新
  • docker-volume-backup 备份 ragflow volumes
  • 五元组+协议分层:拆解网络通信的底层密码
  • Metal入门,使用Metal实现纹理效果
  • [C++面试] 基础题
  • const修饰指针
  • 【网络篇】TCP协议的三次握手和四次挥手
  • 如何让Wi-Fi设备传输距离达到1100米?涂鸦新方案让通信距离远超传统5倍
  • Go 语言中的 Struct Tag 的用法详解
  • 从零开始:用Python语言基础构建宠物养成游戏:从核心知识到完整实战
  • MySQL 数据库表结构修改与字段添加
  • 常见的游戏服务器架构有哪些?
  • 【MySQL】06.MySQL表的增删查改
  • (1)深度学习基础知识(八股)——常用名词解释
  • gd32e230c8t6 驱动ws2812
  • vue2实现元素拖拽
  • 自由开发者计划 002:创建一个贷款计算器的微信小程序
  • Elasticsearch 写入性能优化有哪些常见手段?
  • 2025版 JavaScript性能优化实战指南从入门到精通
  • 【机器学习基础】机器学习入门核心算法:线性回归(Linear Regression)
  • 用vue canvas画一个能源电表和设备的监测图
  • 《STL--string的使用及其底层实现》
  • (第94天)OGG 微服务搭建 Oracle 19C CDB 架构同步