【pytorch扩展】CUDA自定义pytorch算子(简单demo入手)

Pytorch作为一款优秀的AI开发平台,提供了完备的自定义算子的规范。我们用torch开发时,经常会因为现有算子的不足限制我们idea的迸发。于是,CUDA/C++自定义pytorch算子是不得不磕了。

今天通过一个小实验来梳理自定义pytorch算子都需要做哪些准备。比如,我们做一个张量加法。
vim test_add.py

from add import sum_double_op
import torch
import timeclass Timer:def __init__(self, op_name):self.begin_time = 0self.end_time = 0self.op_name = op_namedef __enter__(self):torch.cuda.synchronize()self.begin_time = time.time()def __exit__(self, exc_type, exc_val, exc_tb):torch.cuda.synchronize()self.end_time = time.time()print(f"Average time cost of {self.op_name} is {(self.end_time - self.begin_time) * 1000:.4f} ms")if __name__ == '__main__':n = 1000000device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")tensor1 = torch.ones(n, dtype=torch.float32, device=device, requires_grad=True)tensor2 = torch.ones(n, dtype=torch.float32, device=device, requires_grad=True)with Timer("sum_double"):ans = sum_double_op(tensor1, tensor2)

这里的"sum_double_op"就是我们用CUDA写的算子。那这个可以直接调用,并且可以传递梯度的算子,需要怎么做呢?


众所周知,CUDA/C++都是编译性语言,编译以后再调用会比python这种解释性语言更快。所以,我们需要对CUDA有一个编译过程。这个编译过程用setuptools来实现(可以pip安装)。
先vim setup.py

from setuptools import find_packages, setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtensionsetup(name='myAdd',packages=find_packages(),version='0.1.0',author='muzhan',ext_modules=[CUDAExtension('sum_double',['./add/add.cpp','./add/add_cuda.cu',]),],cmdclass={'build_ext': BuildExtension}
)

直接“python setup.py install”即可完成cuda算子的编译和安装。等等,你的add.cpp和add_cuda.cu还没呢?
vim add_cuda.cu

#include <cstdio>
#define THREADS_PER_BLOCK 256
#define WARP_SIZE 32
#define DIVUP(m, n) ((m + n - 1) / n)__global__ void two_sum_kernel(const float* a, const float* b, float * c, int n){int idx = blockIdx.x * blockDim.x + threadIdx.x;if (idx < n){c[idx] = a[idx] + b[idx];}
}void two_sum_launcher(const float* a, const float* b, float* c, int n){dim3 blockSize(DIVUP(n, THREADS_PER_BLOCK));dim3 threadSize(THREADS_PER_BLOCK);two_sum_kernel<<<blockSize, threadSize>>>(a, b, c, n);
}

vim add.cpp

#include <torch/extension.h>
#include <torch/serialize/tensor.h>#define CHECK_CUDA(x) \TORCH_CHECK(x.type().is_cuda(), #x, " must be a CUDAtensor ")
#define CHECK_CONTIGUOUS(x) \TORCH_CHECK(x.is_contiguous(), #x, " must be contiguous ")
#define CHECK_INPUT(x) \CHECK_CUDA(x);       \CHECK_CONTIGUOUS(x)void two_sum_launcher(const float* a, const float* b, float* c, int n);void two_sum_gpu(at::Tensor a_tensor, at::Tensor b_tensor, at::Tensor c_tensor){CHECK_INPUT(a_tensor);CHECK_INPUT(b_tensor);CHECK_INPUT(c_tensor);const float* a = a_tensor.data_ptr<float>();const float* b = b_tensor.data_ptr<float>();float* c = c_tensor.data_ptr<float>();int n = a_tensor.size(0);two_sum_launcher(a, b, c, n);
}PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {m.def("forward", &two_sum_gpu, "sum two arrays (CUDA)");
}

我们看一下文件结构:

.
├── add
│   ├── add.cpp
│   ├── add_cuda.cu
│   ├── __init__.py
│   └── sum.py
├── README.md
├── setup.py
└── test_add.py

有了add.cpp和add_cuda.cu以后,我们就可以用"python setup.py install"来进行编译和安装了。编译和安装以后,我们需要用python类封装一下:

vim __init__.py
from .sum import *

vim sum.py

from torch.autograd import Function
import sum_doubleclass SumDouble(Function):@staticmethoddef forward(ctx, array1, array2):"""sum_double function forward.Args:array1 (torch.Tensor): [n,]array2 (torch.Tensor): [n,]Returns:ans (torch.Tensor): [n,]"""array1 = array1.float()array2 = array2.float()ans = array1.new_zeros(array1.shape)sum_double.forward(array1.contiguous(), array2.contiguous(), ans)# ctx.mark_non_differentiable(ans) # if the function is no need for backpropogationreturn ans@staticmethoddef backward(ctx, g_out):# return None, None   # if the function is no need for backpropogationg_in1 = g_out.clone()g_in2 = g_out.clone()return g_in1, g_in2sum_double_op = SumDouble.apply

最后,直接

python test_add.py

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处:http://www.mzph.cn/diannao/40774.shtml

如若内容造成侵权/违法违规/事实不符,请联系多彩编程网进行投诉反馈email:809451989@qq.com,一经查实,立即删除!

相关文章

软设之类的继承与泛化,多重继承

在类中&#xff0c;假如父类已经写好属性或方法&#xff0c;子类想要实现相同的功能&#xff0c;不用专门写代码&#xff0c;直接用专门的继承语言继承就可以了。 比如说有一个动物类&#xff0c;有毛色和叫这两个属性和方法&#xff0c;又写了一个子类是猫类&#xff0c;猫类…

腾讯云COS分布式对象存储

腾讯云COS分布式对象存储 腾讯云对象存储&#xff08;Cloud Object Storage&#xff0c;COS&#xff09;是腾讯云提供的一种用于存储海量文件的分布式存储服务。 腾讯云 COS 适用于多种场景&#xff0c;如静态网站托管、大规模数据备份和归档、多媒体存储和处理、移动应用数据存…

Kafka搭建(单机版)

部署前提 VMware环境 : 两台centos系统 Jdk包:jdk-8u202-linux-x64.tar.gz Kafka包:kafka_2.12-3.5.0.tgz Zookeeper包:apache-zookeeper-3.7.2-bin.tar.gz 百度网盘自取: 链接: https://pan.baidu.com/s/11EWuhBoSmH3musd_3Rgodw?pwde32t 提取码: e32t Kafka搭建&#xff08;…

Camtasia 2024新功能 Camtasia2024更新介绍:AI剪辑助力微课制作 Camtasia2024密钥 Camtasia2023免费升级更新

Camtasia 是一款功能强大的屏幕录制和视频编辑软件&#xff0c;广泛应用于教育、商业和娱乐领域。无论是创建教学视频、产品演示、教程还是营销内容&#xff0c;Camtasia都能提供专业的工具和功能&#xff0c;帮助用户制作高质量的视频内容。 Camtasia 2024 中文免费安装包百度…

暑假学习DevEco Studio第2天

学习目标&#xff1a; 掌握页面跳转 学习内容&#xff1a; 跳转页面 创建页面&#xff1a; 在“project”窗口。打开“entry>src>main>ets”,右击“pages”&#xff0c;选择“New>ArkTS File”,命名“Second”&#xff0c;点击回车键。 在页面的路由&#xff0…

昇思25天学习打卡营第16天|文本解码原理——以MindNLP为例

在大模型中&#xff0c;文本解码通常是指在自然语言处理&#xff08;NLP&#xff09;任务中使用的大型神经网络模型&#xff08;如Transformer架构的模型&#xff09;将编码后的文本数据转换回可读的原始文本的过程。这些模型在处理自然语言时&#xff0c;首先将输入文本&#…

【Unix/Linux】Unix/Linux如何查看系统版本

Unix和Linux查看系统版本的指令有些区别&#xff0c;下面分别介绍: 一.Unix查看系统版本 在Unix系统中&#xff0c;查看系统版本的方法可能会根据具体的Unix操作系统而有所不同。以下是一些通用的方法&#xff0c;适用于多种Unix系统&#xff0c;包括但不限于Solaris、AIX、H…

vienna整流器过零畸变原因分析

Vienna整流器是一种常见的三电平功率因数校正&#xff08;PFC&#xff09;整流器&#xff0c;广泛应用于电源和电能质量控制领域。由于其高效率、高功率密度和低谐波失真的特点&#xff0c;Vienna整流器在工业和电力电子应用中具有重要地位。然而&#xff0c;在实际应用中&…

ssh:(xshell)远程连接失败

项目场景&#xff1a; 提示&#xff1a;这里简述项目相关背景&#xff1a; 云服务器远程连接失败 xshell 远程连接失败 xshell (ssh客户端&#xff09; ---------------------------------------------安全组----------防火墙-------黑白名单-----SSH服务 问题排查 1. 安全…

Playwright之录制脚本转Page Object类

Playwright之录制脚本转Page Object类 设计思路 &#xff1a; 我们今天UI自动化设计的时候&#xff0c;通常会遵循一些设计模式&#xff0c;例如Page Object模式。但是自己找元素再去填写有一些麻烦&#xff0c;所以我们可以通过拆解录制的脚本&#xff0c;将其中的元素提取出来…

DALL-E、Stable Diffusion 等 20+ 图像生成模型综述

二、任务场景 2.1. 无条件生成 无条件生成是指生成模型在生成图像时不受任何额外条件或约束的影响。模型从学习的数据分布中生成图像&#xff0c;而不需要关注输入条件。 2.2. 有条件生成 有条件生成是指生成模型在生成图像时受到额外条件或上下文的影响。这些条件可以是类别…

Vscode 保存代码,代码自动格式化

我这里使用的插件是Prettier-Code formatter&#xff1a;自动缩进整理代码的格式&#xff0c;使用方法如下&#xff1a; 先在vscode商店找到插件并安装&#xff1a;安装插件之后&#xff0c;随便找到一个项目文件&#xff0c;右键选择格式化文档&#xff1a;选中我们安装的插件…

掌握Vim的会话之道:深度解析会话管理功能

掌握Vim的会话之道&#xff1a;深度解析会话管理功能 在高效的文本编辑工作流中&#xff0c;能够保存和恢复编辑会话是极其重要的。Vim&#xff0c;作为一个功能强大的文本编辑器&#xff0c;提供了会话管理功能&#xff0c;允许用户保存当前的工作状态&#xff0c;并在之后重…

spring6框架解析(by尚硅谷)

文章目录 spring61. 一些基本的概念、优势2. 入门案例实现maven聚合工程创建步骤分析实现过程 3. IoC&#xff08;Inversion of Control&#xff09;基于xml的bean环境搭建获取bean获取接口创建实现类依赖注入 setter注入 和 构造器注入原生方式的setter注入原生方式的构造器注…

Java 多线程stream流按行读取文件

stream并行流快&#xff08;文件11g&#xff09; try (Stream<String> lines Files.lines(filePath)) {lines.parallel().forEach(str -> operatePartData(str, allDataList)); } catch (IOException e) {throw new RuntimeException(e); }线程池慢&#xff08;文件…

PyPDF2合并PDF文件的高级应用:指定合并方式

本文目录 前言一、合并PDF的高级应用1、逻辑讲解2、合并效果图3、完整代码二、异常校验1、合并过程中的错误校验前言 本文我们主要来讲解一下PyPDF2合并PDF文件的高级应用,就是指定合并方式进行合并,构建函数支持模式选择,主要不管咋折腾,其实就是不想去付费买那个PDF编辑…

PDF怎么分割成一页一页的?原来可以这么轻松

PDF怎么分割成一页一页的&#xff1f;PDF文档因其跨平台兼容性和可打印性而被广泛使用&#xff0c;但有时为了便于发送电子邮件、管理文档或保护敏感信息&#xff0c;我们需要将一个大型的PDF文件分割成多个小文件。幸运的是&#xff0c;分割PDF文件并不复杂。下文中就介绍了三…

webp2jpg网页在线图片格式转换源码

源码介绍 webp2jpg-免费在线图片格式转化器, 可将jpeg、jpg、png、gif、 webp、svg、ico、bmp文件转化为jpeg、png、webp、webp动画、gif文件。 无需上传文件&#xff0c;本地即可完成转换! 源码特点&#xff1a; 无需上传&#xff0c;使用浏览器自身进行转换批量转换输出we…

easyexcel使用小结-未完待续

官网&#xff1a;https://easyexcel.opensource.alibaba.com/docs/current/ <dependency><groupId>com.alibaba</groupId><artifactId>easyexcel</artifactId><version>4.0.1</version></dependency>一、读 1.1简单读 Getter…

系统安全体系架构规划框架

安全技术体系架构是对组织机构信息技术系统的安全体系结构的整体描述。安全技术体系架构框架是拥有信息技术系统的组织机构根据其策略的要求和风险评估的结果&#xff0c;参考相关技术体系构架的标准和最佳实践&#xff0c;结合组织机构信息技术系统的具体现状和需求&#xff0…