MindSpore导入CUDA算子
技术背景
当今众多的基于Python的AI框架(如MindSpore、PyTorch等)给了开发者非常便利的编程的条件,我们可以用Python的简单的语法写代码,然后由框架在后端自动编译成可以在GPU上高效计算的程序。而对于一些定制化比较高的算法,MindSpore也支持了相关的接口,允许开发者自己开发相应的CUDA算子(需要统一接口),然后编译成.so
动态链接库,再用MindSpore内置的函数加载为本地算子。本文针对这种方案写一个简单的示例。
程序结构
本地自己手写一个CUDA算子,一般至少需要两个文件和一个nvcc的环境,最好是在安装完成MindSpore的GPU版本之后,再尝试CUDA算子的引入。具体MindSpore的安装方法,可以参考MindSpore官网,这里不做赘述。我这里使用的环境是10.1版本的nvcc:
$ nvcc -V
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2019 NVIDIA Corporation
Built on Sun_Jul_28_19:07:16_PDT_2019
Cuda compilation tools, release 10.1, V10.1.243
$ python3 -m pip show mindspore
Name: mindspore
Version: 2.1.0
Summary: MindSpore is a new open source deep learning training/inference framework that could be used for mobile, edge and cloud scenarios.
Home-page: https://www.mindspore.cn
Author: The MindSpore Authors
Author-email: contact@mindspore.cn
License: Apache 2.0
Location: /home/dechin/anaconda3/envs/mindspore-latest/lib/python3.7/site-packages
Requires: numpy, protobuf, pillow, astunparse, scipy, psutil, asttokens, packaging
Required-by:
需要准备的两个文件,一个是CUDA算子本身的.cu
文件,另一个是用来调用CUDA算子的.py
文件。操作流程是:先按照自己的需求写好CUDA算子,然后用nvcc进行编译,编译输出为.so
的动态链接库,然后在python脚本中使用mindspore.ops.Custom
生成相应的算子。在MindSpore2.1之后的版本中,对于本地CUDA算子的调用定义了统一的接口,其格式为:
extern "C" int CustomFunc(int nparam, void **params, int *ndims, int64_t **shapes, const char **dtypes, void *stream, void *extra);
具体方法可以参考官网的这一篇文档说明。这样的话,在.cu
文件中至少有两个函数,一个是原本用于计算的Kernel函数
,另一个是用于统一标准接口的CustomFunc函数
。需要说明的是,旧版本的MindSpore是没有这样的规范的,所以旧版本的算子没有CustomFunc函数
也能够用nvcc编译,但是无法在新版本的MindSpore中调用。
一维张量求和
我们用一个一维的张量求和的示例来演示一下如何在本地写一个可以用MindSpore来调用的CUDA算子,一维张量求和的算法是比较简单的:
\]
那么对应的CUDA算子的代码如下所示:
// custom_add.cu
// nvcc --shared -Xcompiler -fPIC -o custom_add.so custom_add.cu
// 常量,一般可以放在.cuh头文件中
constexpr int THREADS = 1024;
// 用于CUDA计算的Kernel函数
__global__ void CustomAddKernel(float *input1, float *input2, float *output, size_t size) {
auto idx = blockIdx.x * THREADS + threadIdx.x;
if (idx < size) {
// 逐元素操作,CUDA算子的基本写法
output[idx] = input1[idx] + input2[idx];
}
}
// 标准算子接口
extern "C" int CustomAdd(int nparam, void **params, int *ndims, int64_t **shapes, const char **dtypes, void *stream,
void *extra) {
cudaStream_t custream = static_cast<cudaStream_t>(stream);
// 输出变量的位置
constexpr int OUTPUT_INDEX = 2;
// 传入的参数都是指针形式
float *input1 = static_cast<float *>(params[0]);
float *input2 = static_cast<float *>(params[1]);
float *output = static_cast<float *>(params[2]);
// 获取输入张量的大小
int size = shapes[OUTPUT_INDEX][0];
// GPU运算中的block和thread,一般要求block*thread大于或者等于size即可
int blocks = ceil(size / THREADS) + 1;
// 调用Kernel函数
CustomAddKernel<<<blocks, THREADS, 0, custream>>>(input1, input2, output, size);
return 0;
}
值得注意的是,上述CustomAdd函数中的params有3个输入,但是实际上其中一个是返回值,这也是MindSpore对于标准接口的设定,不需要我们额外传入一个变量。保存好上述的CUDA算子代码之后,可以用如下指令直接编译成python可以调用的动态链接库:
$ nvcc --shared -Xcompiler -fPIC -o custom_add.so custom_add.cu
编译完成后,会在当前目录下生成一个新的.so
文件,然后就可以在python代码中进行调用:
# test_custom_ops.py
# python3 test_custom_ops.py
import mindspore as ms
from mindspore import ops, Tensor, context
ms.set_context(device_target="GPU", mode=context.GRAPH_MODE)
t1 = Tensor([1., 2., 3.], ms.float32)
t2 = Tensor([3., 2., 1.], ms.float32)
CustomAdd = ops.Custom("./custom_add.so:CustomAdd",
out_shape=[t1.shape[0]],
out_dtype=ms.float32,
func_type="aot"
)
res = CustomAdd(t1, t2)
ops.print_(res)
上述的CustomAdd就是我们导入的基于CUDA算子来写的本地MindSpore算子,并且这种算子还可以使用MindSpore进行注册,这样就不需要每次使用都去加载这个动态链接库,感兴趣的童鞋可以自己研究一下算子注册的方法。上述Python代码的运行结果如下:
$ python3 test_custom_ops.py
Tensor(shape=[3], dtype=Float32, value= [ 4.00000000e+00, 4.00000000e+00, 4.00000000e+00])
可见跟我们预期的结果是一致的,那么就完成了一个本地CUDA算子的实现和调用。
自定义算子反向传播
在前面的章节里面我们已经实现了一个本地的一维张量求和算子,但是这还不是一个完整的算子实现,因为我们在AI框架中一般要求算子可自动微分,光实现一个算子是不完整的,因此这里我们再通过bprop函数,来实现一下自定义算子的反向传播:
# test_custom_ops.py
# python3 test_custom_ops.py
import mindspore as ms
from mindspore import ops, Tensor, context
from mindspore import nn, grad
ms.set_context(device_target="GPU", mode=context.GRAPH_MODE)
t1 = Tensor([1., 2., 3.], ms.float32)
t2 = Tensor([3., 2., 1.], ms.float32)
CustomAdd = ops.Custom("./custom_add.so:CustomAdd",
out_shape=[t1.shape[0]],
out_dtype=ms.float32,
func_type="aot"
)
为了自动微分,我们需要定义一个Cell类来封装我们的自定义算子:
class Add(nn.Cell):
# 反向传播函数
def bprop(self, x, y, out, dout):
return (y, )
# 计算函数
def construct(self, x, y):
return CustomAdd(x, y)
# 把Cell类加载为custom_add函数
custom_add = Add()
# 计算求和结果
res = custom_add(t1, t2)
# 计算自动微分结果
res_g = grad(custom_add, grad_position=(0, ))(t1, t2)
print(res)
print(res_g)
在这个代码中,主要就是增加了一个Cell类和两个新的函数bprop
和construct
,分别用于计算函数的反向传播和正向值,代码运行的结果如下:
$ python3 test_custom_ops.py
[4. 4. 4.]
[3. 2. 1.]
当然,这里我们没有再额外写一个用于返回反向传播值的CUDA算子,但是原则上对于较为复杂的函数,是需要自己手动写一个用于求微分数值的CUDA算子的。
总结概要
本文介绍了在MindSpore标准格式下进行CUDA算子开发的方法和流程,可以让开发者在现有的AI框架下仍然可以调用基于CUDA实现的高性能的算子。并且,除了常规的数值计算之外,在MindSpore框架下,我们还可以通过实现一个bprop函数,使得我们手写的这个CUDA算子也可以使用MindSpore框架本身自带的自动微分-端到端微分技术。
版权声明
本文首发链接为:https://www.cnblogs.com/dechinphy/p/mindspore-cuda.html
作者ID:DechinPhy
更多原著文章:https://www.cnblogs.com/dechinphy/
请博主喝咖啡:https://www.cnblogs.com/dechinphy/gallery/image/379634.html
MindSpore导入CUDA算子的更多相关文章
- MindSpore技术理解(上)
MindSpore技术理解(上) 引言 深度学习研究和应用在近几十年得到了爆炸式的发展,掀起了人工智能的第三次浪潮,并且在图像识别.语音识别与合成.无人驾驶.机器视觉等方面取得了巨大的成功.这也对算法 ...
- MindSpore:自动微分
MindSpore:自动微分 作为一款「全场景 AI 框架」,MindSpore 是人工智能解决方案的重要组成部分,与 TensorFlow.PyTorch.PaddlePaddle 等流行深度学习框 ...
- 使用MindSpore计算旋转矩阵
技术背景 坐标变换.旋转矩阵,是在线性空间常用的操作,在分子动力学模拟领域有非常广泛的应用.比如在一个体系中切换坐标,或者对整体分子进行旋转平移等.如果直接使用Numpy,是很容易可以实现的,只要把相 ...
- git clone开启云上AI开发
摘要:相比于传统的软件开发,AI开发存在以下4个痛点:算法繁多:训练时间长:算力需求大:模型需手动管理,我们可以使用云上AI开发的方式来缓解以上4个痛点. 本文分享自华为云社区<git clon ...
- DeepSpeed Chat: 一键式RLHF训练,让你的类ChatGPT千亿大模型提速省钱15倍
DeepSpeed Chat: 一键式RLHF训练,让你的类ChatGPT千亿大模型提速省钱15倍 1. 概述 近日来,ChatGPT及类似模型引发了人工智能(AI)领域的一场风潮. 这场风潮对数字世 ...
- ubuntu16.04服务器上无root权限,配置个人tensorflow环境--cuda9.0+cuDNN7+tensorflow-gpu-1.8
本人在服务器上已经用Anconda创建好python3.5的环境,这个网上有一大堆教程.接下来是重点. 1. cuda的安装 https://developer.nvidia.com/cuda-dow ...
- 技术干货丨卷积神经网络之LeNet-5迁移实践案例
摘要:LeNet-5是Yann LeCun在1998年设计的用于手写数字识别的卷积神经网络,当年美国大多数银行就是用它来识别支票上面的手写数字的,它是早期卷积神经网络中最有代表性的实验系统之一.可以说 ...
- 用 Numba 加速 Python 代码
原文出自微信公众号:Python那些事 一.介绍 pip install numba Numba 是 python 的即时(Just-in-time)编译器,即当你调用 python 函数时,你的全部 ...
- MindInsight张量可视设计介绍
MindInsight张量可视设计介绍 特性背景 张量可视,能够帮助用户直观查看训练过程中的Tensor值,既支持以直方图的形式呈现Tensor的变化趋势,也支持查看某次step的具体Tensor值. ...
- 开放式神经网络交换-ONNX(上)
目的 本文档包含ONNX语义的规范性规范. "onnx"文件夹下的.proto和.proto3文件构成了用协议缓冲区定义语言编写的语法规范..proto和.proto3文件中的注释 ...
随机推荐
- 五、mongo备份篇 mongoexport、mongoimport 以及mongodump、mongorestore
系列导航 一.linux单机版mongo安装(带密码验证) 二.mongo集群搭建 三.java连接mongo数据库 四.java对mongo数据库增删改查操作 五.mongo备份篇 mongoexp ...
- vue学习笔记 十八、父子组件相互传递参数
系列导航 vue学习笔记 一.环境搭建 vue学习笔记 二.环境搭建+项目创建 vue学习笔记 三.文件和目录结构 vue学习笔记 四.定义组件(组件基本结构) vue学习笔记 五.创建子组件实例 v ...
- vue学习笔记 三、文件和目录结构
系列导航 vue学习笔记 一.环境搭建 vue学习笔记 二.环境搭建+项目创建 vue学习笔记 三.文件和目录结构 vue学习笔记 四.定义组件(组件基本结构) vue学习笔记 五.创建子组件实例 v ...
- asio 使用 openssl 示例
#include <boost/asio.hpp> #include <boost/asio/ssl.hpp> #include <boost/algorithm/str ...
- wireshark 显示过滤表达式
转载请注明出处: 1.根据协议过滤: 在显示过滤表达式的输入框中直接输入对应的协议类型即可:http tcp udp 2.根据 IP 过滤: 根据源IP地址过滤:如源地址IP为:127.0.0. ...
- RSA趣题篇(简单型)
1.n与p的关系 题目 ('n=', 288990088827100766680640490138486855101396196362885475612662192799072729620922966 ...
- web - 解决 formdata 打印空对象
获取单个值可以使用formData对象.get();而直接打印是看不到的.因为外界访问不到,你使用append方法后,对应的键值对就已经添加到表单里面了,你在控制台看到的是FormData原型,存储的 ...
- linux-介绍
- [转帖]centos7离线安装postgresql13
https://www.cnblogs.com/summer-88/p/15341918.html 在一台可以联网的centos上安装postgresql源 yum install -y https: ...
- [转帖]基于Fuse的用户态文件系统性能优化几点建议
https://zhuanlan.zhihu.com/p/68085075 目前很多文件系统基于Fuse( http://fuse.sourceforge.net/ )开发,在较为深入钻研Fuse实现 ...