PyTorch自定义张量操作开发指南【CFFI+CUDA】

PyTorch 与 TensorFlow 一起成为深度学习研究人员和从业者的标准。虽然 PyTorch 在张量运算或深度学习层方面提供了多种选择,但一些专门的操作仍然需要手动实现。在运行时至关重要的情况下,应使用 C 或 CUDA 来完成此操作,以支持 CPU 和 GPU 计算。

在本文中,我想提供一个简单的示例和框架,用于使用 CFFI for Python 和 CuPy 通过自定义 C 和 CUDA 操作扩展 PyTorch。

NSDT工具推荐: Three.js AI纹理开发包 - YOLO合成数据生成器 - GLTF/GLB在线编辑 - 3D模型格式在线转换 - 可编程3D场景编辑器 - REVIT导出3D模型插件 - 3D模型语义搜索引擎 - Three.js虚拟轴心开发包 - 3D模型在线减面 - STL模型在线切割 

1、简介

PyTorch 已成为深度学习研究和开发的标准工具。即使是不常见的张量运算或神经网络层也可以使用 PyTorch 提供的各种操作轻松实现。但是,对于某些特殊操作,求助于高效的 C 和 CUDA 实现可能是有意义的。在本文中,我想展示如何使用 CFFI 轻松扩展 PyTorch 进行 CPU 操作和使用 CuPy 进行 GPU 操作。作为示例,我将展示如何实现张量运算,计算两个张量之间的元素级、位级汉明距离。

本文的代码可在 GitHub 上找到,并且可以轻松扩展和调整。

2、环境设置

我将使用 PyThon 的 C 外部函数接口 (CFFI) 在 CPU 上实现数据类型为 torch.int32 的张量之间的汉明距离。为了支持 GPU 计算,汉明距离也可以实现为 CUDA 内核。可以使用 CuPy 导入 CUDA 内核。CFFI 和 CuPy 都可以轻松安装,例如,使用 pip install。但是,对于 CuPy,安装需要适合使用的 CUDA 版本(对于 PyTorch 也是如此)。可以在此处找到详细的安装说明。

3、C 中的位汉明距离

下面提供了 C 中两个 32 位整数之间的位级汉明距离的简单实现:

清单 1:两个整数之间的汉明距离。

int a = 15;
int b = 19;
int dist = 0;
int x = a ^ b;while(x != 0) {x = x & (x-1);dist++;
}

为了计算两个张量之间的逐元素汉明距离,可以将清单 1 包装在一个简单的循环中。生成的函数(如清单 2 所示)需要三个 int 数组作为输入:第一个输入张量、第二个输入张量和将用汉明距离填充的输出张量。所有数组都假定为相同的长度 n:

清单 2:两个整数数组之间的逐元素汉明距离。可选地,可以使用 OpenMP 来加快计算速度。

    #pragma omp parallel forfor (int elem_idx = 0; elem_idx < n; elem_idx++) {dist[elem_idx] = 0;int x = a[elem_idx] ^ b[elem_idx];while(x != 0) {x = x & (x-1);dist[elem_idx]++;}}
}

使用 CFFI,此函数可直接用于操作 NumPy 数组或 PyTorch 张量。为此,可以将实现保存在 cffi.c 中,并将相应的头文件保存在 cffi.h 中。

4、位汉明距离 CUDA 内核

清单 1 中概述的相同算法可轻松放入 CUDA 内核中:

清单 3:用于计算两个整数数组之间汉明距离的 CUDA 内核。这本质上是清单 1 中由内核块/dim/id 确定的数组元素。

extern "C" __global__ void cupy_int32hammingdistance(const int n,const int* a,const int* b,int* dist
) {int elem_idx = blockIdx.x * blockDim.x + threadIdx.x;if (elem_idx >= n) {return;}int x = a[elem_idx] ^ b[elem_idx];while(x != 0) {x = x & (x-1);dist[elem_idx]++;}
}

CuPy 只需要内核;内核不需要存储在单独的代码文件中。相反,它可以作为 Python 中的字符串提供给 CuPy。

5、将所有内容放在一起

为了组装所有部分,我将使用一个简单的模块化结构,将实际实现(使用 cffi/ 中的 CFFI 或 cupy.py 中的 CuPy)与 torch.py​​ 中的高级方法分开:

common/
- __init__.py
- cffi/|- cffi.h|- cffi.c|- __init__.py
- cupy.py
- torch.py

5.1 CPU 实现

首先,我将整理 CPU 实现,即 cffi.c 和 cffi.h。为简单起见,可以将它们放入自己的目录中,Python 接口将在相应的 __init__.py 中定义。

清单 4:使用 CFFI,可以即时编译 C 代码并直接在 Python 中访问。详情请参阅注释。

    # https://stackoverflow.com/questions/22931147/stdisinf-does-not-work-with-ffast-math-how-to-check-for-infinityelse:if use_openmp:ffi.set_source('_cffi',my_source.read(),extra_compile_args=['-fopenmp', '-D use_openmp', '-O3','-march=native'],extra_link_args=['-fopenmp'],)else:ffi.set_source('_cffi',my_source.read(),extra_compile_args=['-O3','-march=native'],
)# 4.
# Compile using the parameters above.
ffi.compile()
#ffi.compile(verbose=True)
# 5.
# This simply imports all compiled functions and makes them available in this module.
from _cffi import *

代码负责编译清单 2 中的函数并将其与 Python 接口。之后,可以通过 cffi.lib.cffi_int32hammingdistance 访问该函数,其中清单 4 对应于 cffi/__init__.py。我将在下面详细介绍基本步骤:

  • 获取 __init__.py 文件目录的绝对路径。这对于定位要编译的头文件和实现文件是必要的。根据设置,也可以通过不同的方式解决这个问题,例如,通过硬编码绝对路径。
  • 读取头文件,以便 CFFI 知道函数定义。
  • 读取源文件并设置编译选项。在这里,代码允许几种不同的设置,包括不带优化的调试设置和支持 OpenMP 的设置。
  • 了解函数定义(通过头文件)、确定编译选项并阅读源代码后,CFFI 可以编译所有内容。
  • 最后,导入所有编译函数,以便以后可以更轻松地访问它们。

5.2 GPU 实现

对于 CuPy 部分,我将在 cupy.py 中创建一个单独的模块:

清单 5:与 CFFI 类似,CuPy 也允许即时编译 CUDA 内核。详情请参阅注释。

import torchtry:import cupy# 1. This will be used to call a kernel with source code provided as Python string.@cupy.util.memoize(for_each_device=True)def cunnex(strFunction):return cupy.cuda.compile_with_cache(globals()[strFunction]).get_function(strFunction)
except ImportError:print("CUPY cannot initialize, not using CUDA kernels")class Stream:ptr = torch.cuda.current_stream().cuda_stream# 2. The kernel as Python string from Listing 3
cupy_int32hammingdistance = '''extern "C" __global__ void cupy_int32hammingdistance(const int n,const int* a,const int* b,int* dist) {int elem_idx = blockIdx.x * blockDim.x + threadIdx.x;if (elem_idx >= n) {return;}int x = a[elem_idx] ^ b[elem_idx];while(x != 0) {x = x & (x-1);dist[elem_idx]++;}}
'''

CuPy 接口甚至更简单:

  • 此实用函数将负责编译和接口函数。作为参数,该函数需要一个变量的名称,该变量包含实际的 CUDA 内核源代码。
  • 源代码保存在此变量中,而不是单独的源文件中。

最后,在 torch.py​​ 中,将合并两个实现。结果是一个高级函数 int32_hamming_distance,需要两个 torch.int32 张量作为输入。该函数将自动为输出分配内存,并根据输入是否在 GPU 内存上调用适当的接口。为了确定张量是否在 GPU 内存上,提供了一个简单的 is_cuda 函数(此处未显示)。

清单 6:将 CFFI 和 CuPy 实现放在一个高级方法 int32_hamming_distance 中,该方法根据输入张量自动使用 CPU 或 GPU 实现。

def int32_hamming_distance(a, b):"""Bit-wise hamming distance.:param a: first tensor:type a: torch.Tensor:param b: first tensor:type b: torch.Tensor:return: hamming distance:rtype: torch.Tensor"""if not a.is_contiguous():a.contiguous()assert (a.dtype == torch.int32)cuda = is_cuda(a)if not b.is_contiguous():b.contiguous()assert (b.dtype == torch.int32)assert is_cuda(b) is cudaassert len(a.shape) == len(a.shape)for d in range(len(a.shape)):assert a.shape[d] == b.shape[d]# 1. Initialize output tensor to hold the element-wise hamming distances.dist = a.new_zeros(a.shape).int()n = dist.nelement()if cuda:# 2. Call the cupy implementation using the helper function cupy.cunnex.# The function returned by cupy.cunnex expects, among others, the grid/block division to use# and the kernel arguments; here a.data_ptr() will return the point to the tensor a# and is assumed to be on GPU memory.cupy.cunnex('cupy_int32hammingdistance')(grid=tuple([int((n + 512 - 1) / 512), 1, 1]),block=tuple([512, 1, 1]),args=[n,a.data_ptr(),b.data_ptr(),dist.data_ptr()],stream=cupy.Stream)else:# 3. For CFFI, the inputs have to be cast to the target C equivalents using cffi.ffi.cast.# Afterwards, the C function can be called like a regular Python function using the converted arguments._n = cffi.ffi.cast('int', n)_a = cffi.ffi.cast('int*', a.data_ptr())_b = cffi.ffi.cast('int*', b.data_ptr())_dist = cffi.ffi.cast('int*', dist.data_ptr())cffi.lib.cffi_int32hammingdistance(_n, _a, _b, _dist)return dist

本质上,该函数创建输出张量来保存元素汉明距离,然后根据输入是否存储在 GPU 内存中调用 CuPy 或 CFFI 接口:

  • 创建输出张量,如果需要,也可以在 GPU 上创建。它也将是相同大小的 torch.int32 张量。
  • CuPy 实现通过辅助函数 cupy.cunnex 调用,该函数获取相应的源代码,对其进行编译(如果未缓存)并返回相应的函数。返回的函数需要输入——这里, a.data_ptr() 用于访问给定张量底层的指针——以及内核的块/dim/id 细分。

调用 CFFI 实现需要将输入显式转换为等效的 C 类型。然后,可以像常规 Python 函数一样调用该函数。

6、结束语

总体而言,本文表明,使用 CFFI 和 CuPy 可以非常简单地在 C 和 CUDA 中实现支持 CPU 和 GPU 计算的自定义 PyTorch 操作。此外,我提供了一个简单的框架,可以轻松扩展到自定义操作。


原文链接:PyTorch自定义张量操作 - BimAnt

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

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

相关文章

linux系统——htop命令检测

在之前提到top命令可以检测进程情况&#xff0c;但需要额外一些参数才能更清晰得到一些数据&#xff0c;htop也是进程监测命令&#xff0c;但更为准确&#xff0c;给出信息更为详实

H2RSVLM:引领遥感视觉语言模型的革命

随着人工智能技术的飞速发展&#xff0c;遥感图像理解在环境监测、气候变化、粮食安全和灾害预警等多个领域扮演着越来越重要的角色。然而&#xff0c;现有的通用视觉语言模型&#xff08;VLMs&#xff09;在处理遥感图像时仍面临挑战&#xff0c;主要因为遥感图像的独特性和当…

stream-实践应用-统计分析

背景 业务部门提供了一个数据&#xff0c;数据甚至不是excel类型的&#xff0c;是data.txt&#xff0c;每一行都是一个数据&#xff0c;需要对此数据进行统计分析 统计各个月份的销量 因为直接获取resources下的data.txt&#xff0c;所以要借助输入流进行获取数据&#xff0c;再…

【手把手带你微调 Llama3】 改变大模型的自我认知,单卡就能训

微调Llama3的自我认知后 当你问Llama3中文问题&#xff1a; “你叫什么名字&#xff1f;”、 “做个自我介绍”、 “你好” Llama3 会用中文回答 &#xff1a; “我是AI在手” &#xff08;如下图&#xff09; 1、环境安装 # nvidia 显卡 显存16G# pytorch 2.2.2 …

自然语言处理学习中英文翻译语料库

在自然语言处理&#xff08;NLP&#xff09;领域&#xff0c;学习中英文翻译需要高质量的双语语料库。以下是一些常用的中英文翻译语料库资源&#xff1a; OpenSubtitles&#xff1a; 网站&#xff1a; OpenSubtitles 描述&#xff1a;OpenSubtitles 提供了大量的电影和电视剧…

JavaScript四种输出

window.alert()&#xff1a; 使用警告框来显示数据 如&#xff1a; <!DOCTYPE html> <html> <body> <script>window.alert("我来了"); </script></body> </html> 浏览器打开就会显示一个弹窗“我来了” innerHTML&…

VUE项目引入微信JSSDK 使用微信api

大佬地址&#xff1a;VUE项目引入微信JSSDK 实现微信自定义分享_vue中如何引入wx sdk-CSDN博客 支付功能&#xff1a; 第一个接口 是初始化 ssdk的 第二个接口 是调用微信支付的 jsApiPayBefore().then(res > {let getMsg res.dataswx.config({debug: false, //生产环境需…

【Java】【python】leetcode刷题记录--双指针

双指针也一般称为快慢指针&#xff0c;主要用于处理链表和数组等线性数据结构。这种技巧主要涉及到两个指针&#xff0c;一个快指针&#xff08;通常每次移动两步&#xff09;和一个慢指针&#xff08;通常每次移动一步&#xff09;。快指针可以起到’探路‘的作用&#xff0c;…

S32K --- FLS MCAL配置

一、前言 二、MCAL配置 添加一个Mem_43_infls的模块, infls是访问内部flash, exfls是访问外部flash 2.1 General 这边暂时保持的默认,还没详细的去研究,等有空研究了,我再来更新 2.2 MemInstance 双金“index”的下标“0”可以进里面详细配置,这个是基本操作了。 2.2.1 G…

关于Word目录的更新

左侧标题顺序如有调整&#xff0c;自动目录并不会同步更新&#xff0c;每次都要记得在正文目录左上角点击更新目录

2024-05-29 服务器开发-c++线程池与task-思考

摘要: 无论是什么系统&#xff0c;线程池和task都是给上层所提供的基础的功能单元。本文记录一些核心的设计思想。 线程池要面对的场景: 调用下层接口时&#xff0c;被IO阻塞&#xff0c;导致整个服务无法对外提供服务更上层调用本模块接口时&#xff0c;是需要做到同步&#…

【nnUNetv2进阶】十、nnUNetv2 魔改网络-小试牛刀-引入注意力机制EPSA

nnUNet是一个自适应的深度学习框架,专为医学图像分割任务设计。以下是关于nnUNet的详细解释和特点: 自适应框架:nnUNet能够根据具体的医学图像分割任务自动调整模型结构、训练参数等,从而避免了繁琐的手工调参过程。 自动化流程:nnUNet包含了从数据预处理到模型训练、验证…

Vue 3 中的 emit(‘update_modelValue‘):更灵活的双向绑定

你好&#xff0c;我是小白Coding日志&#xff0c;一个热爱技术的程序员。在这里&#xff0c;我分享自己在编程和技术世界中的学习心得和体会。希望我的文章能够给你带来一些灵感和帮助。欢迎来到我的博客&#xff0c;一起在技术的世界里探索前行吧&#xff01; 随着 Vue 3 的发…

面向链接预测的知识图谱表示学习方法综述

源自&#xff1a;软件学报 作者&#xff1a;杜雪盈, 刘名威, 沈立炜, 彭鑫 注&#xff1a;若出现无法显示完全的情况&#xff0c;可搜索“人工智能技术与咨询”查看完整文章 摘 要 作为人工智能的重要基石, 知识图谱能够从互联网海量数据中抽取并表达先验知识, 极大程度解决…

开源基于Node编写的批量HTML转PDF

LTPP批量HTML转PDF工具 Github 地址 LTPP-GIT 地址 官方文档 功能 LTPP 批量 HTML 转 PDF 工具支持将当前目录下所有 HTML 文件转成 PDF 文件&#xff0c;并且在新目录中保存文件结构与原目录结构一致 说明 一共两个独立版本&#xff0c;html-pdf 目录下是基于 html-pdf 模…

代码随想录算法训练营day38 | 435. 无重叠区间、763.划分字母区间、56. 合并区间

435. 无重叠区间 按照左边界排序&#xff0c;直接求重叠区间 class Solution:def eraseOverlapIntervals(self, intervals: List[List[int]]) -> int:if len(intervals) < 1:return 0intervals.sort(keylambda x: x[0])end intervals[0][1]result 0for i in range(1,…

【CALayer-时钟练习-旋转 Objective-C语言】

一、好,接下来呢,我们要让它旋转出来, 1.让它先旋转起来啊,这根秒针,让它先转着, 把之前的代码复制粘贴一份,改个名字,叫:07-时钟练习(旋转) 旋转的话,我现在应该让它,一秒钟,旋转一次,一秒钟,旋转一次, 那么,这个时候,我们应该怎么样去做, 我现在这个是…

便携式应急气象站:应急气象监测装备

TH-BQX5便携式应急气象站&#xff0c;作为现代气象监测的重要装备&#xff0c;以其独特的便携性、高效性和灵活性&#xff0c;在应急气象监测领域发挥着至关重要的作用。这类气象站不仅为灾害预警、环境保护、农业生产等多个领域提供了实时、准确的气象数据&#xff0c;还在突发…

PHP模块pdo_sqlite.so: undefined symbol: sqlite3_column_table_name

安装 php-sqlite3 之后&#xff0c;执行php -m 命令有警告&#xff0c;如下 PHP Warning: PHP Startup: Unable to load dynamic library pdo_sqlite (tried: /usr/lib64/php/modules/pdo_sqlite (/usr/lib64/php/modules/pdo_sqlite: cannot open shared object file: No su…

理论知识.质数打表

啊&#xff0c;哈喽&#xff0c;小伙伴们大家好。我是#张亿&#xff0c;今天呐&#xff0c;学的是理论知识.质数打表 为什么需要质数打表 我们已经学习了如何判断一个数是不是质数了&#xff0c;但是还不够。假设要判断很多很多个数是不是质数的时候&#xff0c;之前的学习的…