Skip to main content

10. MUSA AI 框架算子适配及 CV 图像处理开发示例

10.1. muPyTorch-muDNN 对接扩展

muPyTorch-muDNN对接扩展,可以让PyTorch框架使用muDNN算子库进行扩展,从而可以让用户在不对模型进行较大改动的情况下,使用musa-plugin在MTGPU上运行并计算AI模型。本节主要关注:PyTorch的dispatch机制、算子调用以及开发者如何进行muPyTorch与muDNN算子库的对接。对接完成后,PyTorch中的算子就可以通过muDNN算子库,运行在MTGPU设备上,进行算子计算和模型训练及推理。

10.1.1. 算子派发机制

PyTorch中有很多算子,同一种算子又分为很多类型。以add算子为例,按device类型可以分为CPU、GPU、NPU、TPU、FPGA等设备;按layout类型可以分为普通张量和稀疏张量,不同的张量有不同的布局;按dtype类型可以分为fp64、fp32、fp16、int8,甚至bf16、hf32等不同类型。那么当python端调用一个torch.add算子时,最后实际执行的算子该怎么选择呢?此时,就要提到PyTorch的dispatch机制了,接下来便对dispatch机制进行具体介绍。

10.1.1.1. dispatch机制

dispatcher可以理解为分发器,当执行一个operator时,dispatcher会根据tensor输入的一些信息和一些其他信息(参数个数、返回值类型等)计算得到一个dispatch key,然后根据dispatch key从dispatch table中找到对应的kernel函数指针,最后回调执行。dispatch table信息如下图所示。

img101

图1 Dispatch Table表示意图

可以看到dispatch key不仅有硬件后端,还有一些更抽象的概念如autograd、tracing等。dispatch key的计算是通过一个dispatch key set的结构来实现的,dispatch key set可以理解为一个64bit的数组,每一个bit都代表了一个dispatch key,从左到右有优先级关系。同样一算子,有针对不同dispatch key的实现,针对这些散落在PyTorch各处的注册实现,这个数组将所有可用实现都结合到一起,调用其中优先级最高的dispatch key对应的kernel实现。那么这个dispatch table是如何形成的呢?此时,就要提到算子的注册机制了,接下来介绍算子注册的交互方式。

10.1.1.2. dispatch table注册

首选,需要定义一个关于算子的运算符模式schema。此时并没有提供操作符的具体实现,只是提供了一个模式字符串,指定操作符的类型签名。后续其他内核具体实现时,都将遵守这种模式。以add算子为例,具体定义方式如下:

TORCH_LIBRARY(myops, m) {
m.def("add(Tensor self, Tensor other) -> Tensor");
}

如上所示,已经定义了add的schema,接下来需要提供这个操作符的具体后端实现。最简单的注册方法是def(“add”, add_cpu),这会将内核注册为在所有情况下运行,即使这个张量不是CPU张量。为了确保我们注册的特定算子在特定的设备上运行,可以使用TORCH_LIBRARY_IMPL宏。在这个宏中使用m.impl就可以将带有dispatch key信息的算子实现注册到dispatch table中。具体定义方式如下:

TORCH_LIBRARY_IMPL(myops, CPU, m) {
m.impl("add", add_cpu_fun_ptr);
}

如上所示,本例中的CPU便是dispatch key,add_cpu_fun_ptr便是add算子CPU对应的函数指针。同理,也可以注册一个CUDA后端或MTGPU后端的item,这些注册可以跨文件分割或者跨库边界分割。通常来说,注册结构为一个单独的TORCH_LIBRARY文件,里面集中的列出了命名空间中每个自定义的操作符。然后每个dispatch key都有一个TORCH_LIBRARY_IMPL块,当然也可以将TORCH_LIBRARY_IMPL块按operator类型切分为不同的块。特别地,当每个算子都有一个分离的文件时,同时不想在头文件中暴露这些算子时,可以直接将注册文件放在函数定义的cpp文件中。

10.1.2. 算子调用流程

通过上面介绍,可以了解到算子的注册方式和算子的dispatch机制。那么,在python端调用一个torch.add算子是如何衔接到到我们的C++代码执行呢?这就要说到PyTorch中的C++/CUDA扩展了,C++/CUDA扩展一般有预编译和实时编译(JIT)模式,这里主要介绍预编译模式,即假设我们的C++拓展代码已经编译打包成库文件。Python中有一个pybind11的宏,主要用来在C++代码中创建Python的链接库。这里以创建名为ext的扩展库(extension)为例,通过pybind11绑定C++代码到Python示例如下:

Tensor my_add(Tensor a, Tensor b);
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m ){
m.def("my_add", &my_add, "my_add (CPU/CUDA) ", py::arg("a"), py::arg("b"));
}

如上所示,PYBIND11_MODULE的作用是为C++代码接入Python解释器提供入口,TORCH_EXTENSION_NAME是编译器在编译扩展库过程中出现的宏,对应为extension中的name变量,在这里会被解释为ext。m代表TORCH_EXTENSION_NAME所对应的实例模块,{}中的每个m.def都定义了一个ext的成员函数,其一般形式为m.def("函数名",函数指针, "文档", 参数列表)。通过这种形式,my_add也就顺利地成为了ext的成员函数,其具体实现为已经定义好的my_add函数。在Python脚本中通过“from ext import my_add”导入ext模块中的my_add, 就可以调用自定义算子my_add了。

10.1.3. muDNN算子库对接

本节开始主要介绍了muDNN算子库与muPyTorch框架的具体对接流程,涉及设备注册、算子对接、算子正确性测试等内容。

10.1.3.1. 设备注册流程

通过上述介绍,我们更关心如何在PyTorch中注册新的设备后端,即如何让PyTorch支持我们的MTGPU设备。注册的详细步骤如下:

  1. 步骤1:添加backend
  • torch.backends控制PyTorch支持的各种后端的行为,如:torch.backends.cuda,torch.backends.mkl。
  • 在PyTorch/c10/core/Backend.h头文件的"enum class Backend"中添加后端MUSA,并且在backendToDispatchKey()中将backend与DispatchKey(PrivateUse1)绑定,在backendToDeviceType()将backend与DeviceType(MTGPU)绑定。
  1. 步骤2:添加device
  • 在PyTorch/c10/core/DeviceType.h头文件的"enum class DeviceType"中添加MTGPU。
  • 在PyTorch/c10/core/Device.h头文件的"struct Device"中添加is_mtgpu(),获得对应的DeviceType。
  • 在PyTorch/c10/core/DispatchKey.cpp源文件的函数toBackendComponent()中添加MTGPU对应的BackendComponent。
  • 在PyTorch/c10/core/TensorOptions.h头文件的函数computeDispatchKey()中根据MTGPU获得对应的DispatchKey。
  • 在PyTorch/torch/library.h头文件的dispach中将Device_type MTGPU与DispatchKey::PrivateUse1绑定。

10.1.3.2. 算子对接流程

算子对接可以将PyTorch框架与muDNN算子库进行对接,使得在不对用户模型进行较大改动的情况下,将AI模型通过muDNN算子库在mtGPU上进行运行和计算。使用方式也十分简单,仅需将数据和模型搬运到mtgpu上,计算完成后再搬回cpu即可。这里以relu算子为例,介绍算子库对接的基本流程。
首先,需要查看muDNN算子库是否支持relu算子,而relu算子属于Unary类算子,查看muDNN头文件mudnn_math.h可知,Unary.Mode支持relu算子。接下来,需要查看PyTorch框架的C++拓展接口中relu算子对应的dispatch函数,在aten/src/ATen/native/native_functions.yaml文件中,搜索relu可以得到如下代码

- func: relu(Tensor self) -> Tensor
device_check: NoCheck # TensorIterator
variants: function, method
dispatch:
CPU, CUDA: relu
MPS: relu_mps
MkldnnCPU: mkldnn_relu
QuantizedCPU: relu_quantized_cpu
NestedTensorCPU, NestedTensorCUDA: NestedTensor_relu

由上述dispatch可知,在CPU和CUDA平台调用该算子时,会dispatch到名称为relu的函数。因此,下一步需要查看relu函数的函数原型,该算子的函数实现在torch/include/ATen/Functions.h文件中可以找到,可知具体实现在torch/include/ATen/ops/relu.h文件中,由此得到了函数原型Tensor relu(const at::Tensor & self)。
在得到函数原型后,我们就知道了该算子具体调用的函数的运算符模式,在注册自定义算子的时候,需要遵守这种模式。接下来,定义MTGPU上的relu算子时,可以将函数Tensor mtgpu_relu(const at::Tensor & self)作为MTGPU上注册的relu函数名,函数内部实现由我们自定义,最后使用上文所述的TORCH_LIBRARY_IMPL注册MTGPU上的relu算子,如m.impl("relu", &mtgpu),至此算子对接流程基本完成。

10.1.3.3. 自定义函数实现

上文较详细地介绍了算子对接的基本流程,但并未具体介绍自定义函数内部的实现细节,本小节将具体介绍函数内部实现的一些注意事项。
由于函数声明中使用的数据类型是PyTorch框架C++扩展模块ATen中的相关数据类型,而算字库有自己的一套数据类型。因此,自定义函数中首先需要根据框架中的Tensor数据信息来创建muDNN中的Tensor数据类型,这个操作可以借助CreatMTensor函数来完成,该函数主要提取了框架中Tensor的dim、size、type和addr等信息,返回一个muDNN的mTensor数据类型。mTensor数据类型都创建完成后,接着需要创建mHandle句柄和具体op的实例对象,并根据相关op的具体情况,设置相应的如mode、axis、alpha等超参信息,最后将符合muDNN类型要求的参数传入op的Run方法中运行得到输出结果。
另外,由于PyTorch的数据类型的storage具有offset属性,当多个tensor公用一块storage时,tensor的offset可能不为零,此时PyTorch中打印tensor的地址信息可以发现该地址是经过PyTorch计算后得到的地址信息,即storage首地址 + offset * sizeof(type)。另外,算子库目前也不支持不连续的tensor,因此在自定义函数中,往往需要进行tensor的连续性检测和连续性处理,这个操作可以借助MusaContiguous完成。
最后,在编写自定义算子时,有时需要进行一些边界检查,这时我们可以参照native目录下的CPU代码或者CUDA代码进行编写。
注册伪代码如下:

Tensor musa_relu(const Tensor& input) {
// step1: create output tensor first
Tensor result = at::native::empty_mtgpu(input.sizes(), DeviceType::MTGPU, ...);

// step2: create mTensor(in&out), set Attribute
using mTensor = ::musa::dnn::Tensor;
mTensor in_tensor, out_tensor;
in_tensor.SetAddr(input.data_ptr());
out_tensor.SetAddr(result.data_ptr());
in_tensor.SetNdInfo(mTensor::Type::FLOAT, input.size(), ...);
out_tensor.SetNdInfo(mTensor::Type::FLOAT, result.size(), ...);

// step3: create operation descriptor
using mHandle = ::musa::dnn::Handle;
using Unary = ::musa::dnn::Unary;
Unary op;
op.SetMode(Unary::Mode::RELU);

// step4: run op kernel function
op.Run(mHandle, out_tensor, in_tensor);

return result;
}

10.1.3.3. 算子测试和模型测试

算子注册完成后,需要进行正确性测试,因此需要对musa_torch_extension拓展模块进行重新编译和安装。下面介绍两种正确测试,即单算子正确性测试和模型正确性测试。单算子正确性测试及模型正确性测试demo如下:
算子测试demo:

import torch
import musa_torch_extension

data = torch.randn(3,4)
out = torch.relu(data)

with torch.autograd.inference_mode(mode=True):
data_mt = data.to("mtgpu")
out_mt = torch.relu(data_mt)
out_c = out_mt.to("cpu")

模型测试demo:
import torch
import musa_torch_extension

model = alexnet(pretrained = True).eval()
model_mt = model.to("mtgpu")

x = torch.rand((1, 3, 224, 224))
x_mt = x.to("mtgpu")

with torch.autograd.inference_mode(mode=True):
y = model(x)
y_mt = model_mt(x_mt)
y_c = y_mt.to("cpu")

10.2. muTensorFlow-muDNN 对接扩展

TensorFlow 是一个端到端开源机器学习平台。它拥有一个全面而灵活的生态系统,其中包含各种工具、库和社区资源,可助力研究人员推动先进机器学习技术的发展。并使开发者能够轻松地构建和部署由机器学习提供支持的应用。本教程介绍如何在 TensorFlow 中通过自定义算子使其在 MTGPU 上运行从而实现计算的加速。

10.2.1. Kernel注册

10.2.1.1. pluging的注册方式

Op 和 Kernel 是 TensorFlow 框架中重要的两个概念,Op类似于函数声明,Kernel类似函数实现实现。在TensorFlow中,一个 Op 对应的所有 Kernel 都是通过REGISTER_KERNEL_BUILDER宏注册到一个全局的单例注册表中.当需要自定义 Kernel 时,首先创建一个对应的类,该类继承自OpKernel类, 重载其Compute成员函数。构造函数中主要对自定义 Kernel 进行状态初始化(可选!!),Compute执行真正的计算过程。本例中的VecAddOp没有初始状态因此构造函数为空,Compute函数中要调用 muDNN 库实现计算过程主要分为几步

  1. 对 output 分配内存空间
  2. 将 Tensorflow 中的 Tensor 数据类型通过调用CreateMTensor函数转化为 muDNN 的数据类型 mTensor
  3. 调用 muDNN 库函数(传入 mTensor )
    同样也可以通过 musa 编写自定义算子,在 Compute 函数中进行调用
class VecAddOp : public OpKernel {
public:
explicit VecAddOp(OpKernelConstruction* context) : OpKernel(context) {}

void Compute(OpKernelContext* context) override {
const Tensor& input_0 = context->input(0);
const Tensor& input_1 = context->input(1);
OP_REQUIRES(context, input_0.dtype() == input_1.dtype(),
errors::InvalidArgument("input data type should be the same!"));
Tensor* output = nullptr;
//step1: 给output分配内存空间
OP_REQUIRES_OK(context,
context->allocate_output(0, input_0.shape(), &output));
// step2: 将Tensor类型转换为muDNN中的mTensor类型
auto in = CreateMTensor(input_0);
auto other = CreateMTensor(input_1);
auto out = CreateMTensor(*output);
// step3: 调用muDNN中的函数
mHandle& h = GetDeviceByCtx(context)->GetHandle();
mBinary op;
op.SetMode(BINARY_MODE::ADD);
op.Run(h, out, in, other);
}
};

实现了自定义 kernel 类后,下一步就是对 kernel 的注册,下面为对 VecAddOp kernel 的注册,其中 REGISTER_KERNEL_BUILDER 宏内的内容和 Tensorflow 进行原生 kernel 注册的方式一致,需要注意的是 Device 要传入 DEVICE_MTGPU 从而将其注册到 MTGPU 上。MUSA_KERNEL_REGISTER 宏用于在加载 musa-plugin 库的时候辅助 MTGPU 上的 kernel 的注册

MUSA_KERNEL_REGISTER(VecAdd) {
REGISTER_KERNEL_BUILDER(Name("VecAdd").Device(DEVICE_MTGPU), VecAddOp);
}

对于自定义 kernel 的代码,需要将其放在muTensorFlow/musa-plugin/src/kernels目录下,随后在muTensorFlow/musa-plugin目录下调用python setup install完成整个 musa-plugin 的安装。

10.2.2. Op注册

自定义 Op 和设备没有关系,因此参考 TensorFlow 的官方教程中的 Op 的注册方法即可,下面是一个 Op 注册的例子。

#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/shape_inference.h"

using namespace tensorflow;

REGISTER_OP("VecAdd")
.Input("in1: int32") // 输入0,类型为int32
.Input("in2: int32") // 输入1,类型为int32
.Output("out: int32") // 输出0,类型为int32
.SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
c->set_output(0, c->input(0)); // 设置输出0的shape,这里输出0的shape和输入1的shape相同
return Status::OK();
});

首先是将其编译为动态库,TensorFlow 提供了两个接口,用于我们编译自定义算子从而生成动态库。

TF_CFLAGS=( $(python -c 'import tensorflow as tf; print(" ".join(tf.sysconfig.get_compile_flags()))') )
TF_LFLAGS=( $(python -c 'import tensorflow as tf; print(" ".join(tf.sysconfig.get_link_flags()))') )
g++ -std=c++14 -shared vec_add_ops.cc -o _vec_add_ops.so -fPIC ${TF_CFLAGS[@]} ${TF_LFLAGS[@]} -O2

有了动态库后可以直接使用 TensorFlow 提供的函数 tf.load_op_library 来加载动态库,从而用 python 接口使用自定义算子

import tensorflow as tf
vec_add_module = tf.load_op_library('./_vec_add_ops.so')

result = vec_add_module.vec_add([[1, 2], [3, 4]],[[1, 2], [3, 4]])
print(result)

最后,也可以使用 setup 将上述的 python 接口打包为 pip 包进行安装,从而方便使用。

10.3. 图像的预处理

图像预处理的定义:对最低抽象级别图像的操作,输入输出都是亮度图像。

图像预处理的目的:预处理的目的是抑制不想要的失真或者增强某些对于后续处理重要的图像特征。

图像的预处理分为四类:

  • 像素亮度变化;
  • 几何变化;
  • 局部领域预处理;
  • 图像复原;

接下来我们将针对这四类预处理进行讨论。

10.3.1. 像素亮度变化

有两类像素亮度变化:

  • 亮度矫正:修改像素的亮度时需要考虑图像原始像素的亮度和位置。
  • 灰度级变换:修改像素亮度时只需要考虑原始像素的亮度。

10.3.1.1. 亮度矫正

在实际图像的采集和数字化过程中,由于传感器的灵敏度的不同、不均匀的光照等原因会造成图像的退化。如果这种退化是系统性的(有规律可循),就可以通过亮度矫正来加以抑制。这需要有退化图像和原始未退化图像,然后通过逆变化得到误差系数。

10.3.1.2. 灰度级变化

常用的灰度级变化有:

  • 亮度阈值化:像素亮度大于阈值将其亮度赋一个值(可能是255),否则会赋另一个值(可能是0)。
  • 查找表变换:阈值化只有一个阈值,如果需要与多个阈值进行比较,就需要用到显示查找表(LUT)。
  • 伪色彩变化:将单色图像的各个灰度级对应色彩查找表CLUT中某一项的入口地址,根据该地址可查找出包含实际R、G、B的强度值。
  • 灰度化:将彩色图像转化为灰度图像的操作被称为图像灰度化,对RGB图像而言,当三个分量相等时即为灰度图像。常用的转化的方法有:分量法、最值法、平均法、加权平均法。
  • 直方图均衡:对于偏亮、偏暗或亮度集中的图片,可通过直方图均衡提高图像的全局对比度。
  • 对数的灰度级变化:主要用于将图像的低灰度值部分扩展,将高灰度值部分压缩,以达到强调图像低灰度值部分的目的。

10.3.2. 几何变化

几何变化可以消除图像获取时出现的几何失真。几何变化一般由两个基本步骤组成:

  • 像素坐标变换:将输入图像坐标映射到输出图像中的点。
$$
(x,y)=T\left(x^{\prime}, y^{\prime}\right) \rightarrow\left[\begin{array}{lll}x & y & 1\end{array}\right]=
\left[\begin{array}{lll}x^{\prime} & y^{\prime} & 1\end{array}\right] T=\left[\begin{array}{lll}x^{\prime} & y^{\prime} & 1\end{array}\right]\left[\begin{array}{lll}t_{11} & t_{12} & t_{13} \\t_{21} & t_{22} & t_{23} \\t_{31} & t_{32} & t_{33}\end{array}\right]
$$
  • 亮度插值:插值被用来确定输出像素的亮度。当像素坐标变化已完成后,得到新坐标点(x,y)(x^{\prime}, y^{\prime}),变换后的点一般不是整数网格上的点,输出图像网格上的值可通过相邻的非整数样本的亮度插值得到。

10.3.2.1. 像素坐标变化:

下面是几个重要的像素坐标变化及其对应的变换矩阵TT

  • 平移:T=\left[\begin\{matrix}1&0&0\\0&1&0\\t_x&t_y&1\\\end\{matrix}\right]
  • 旋转:T=\left[\begin\{matrix}\cos\{\theta}&\sin\{\theta}&0\\-\sin\{\theta}&\cos\{\theta}&0\\0&0&1\\\end\{matrix}\right]
  • 缩放:T=\left[\begin\{matrix}c_x&0&0\\0&c_y&0\\0&0&1\\\end\{matrix}\right]
  • 垂直偏移变化:T=\left[\begin\{matrix}1&0&0\\s_v&1&0\\0&0&1\\\end\{matrix}\right]
  • 水平偏移变化:T=\left[\begin\{matrix}1&s_k&0\\0&1&0\\0&0&1\\\end\{matrix}\right]

10.3.2.2. 亮度插值:

常用的三种插值方法:最相邻、双线性、双三次插值。

  • 最相邻插值:赋予离现像素点(x,y) (x,y) 最近的原像素点。
  • 双线性插值:根据相邻的四个点,在两个方向通过三次单线性插值得到亮度值。
  • 双三次插值:用双三次多项式表面局部地近似亮度函数来改善模型,用16个相邻的点做插值。

10.3.3. 局部领域预处理

根据处理的目的,可以将局部预处理分为两类:

  • 平滑:平滑目的在于抑制噪声和其他小的波动,但也会模糊所有含有图像重要信息的锐利边缘。
  • 边缘检测(梯度算子):梯度算子是基于图像函数的局部导数。在图像函数中发生快速变化的位置导数较大,而梯度算子的目的是在图像中指示这些位置。但噪声水平也会增加,所以平滑与梯度算子的目标是冲突的。

根据变换的性质,也可以将局部预处理分为两类:

  • 线性
  • 非线性

10.3.3.1. 图像平滑

常见的平滑方式:均值滤波、中值滤波、高斯滤波。

  • 均值滤波:均值滤波采用线性的方法,使用邻域内所有像素的平均值代替邻域中心像素的灰度值。均值滤波不能很好地保护图像细节,会使图像变得模糊,不能很好地去除噪声点。对高斯噪声表现较好,对椒盐噪声表现较差。
  • 中值滤波:中值滤波采用非线性的方法,计算邻域内所有像素中的中值代替邻域中心像素的灰度值。它在抑制平滑脉冲噪声方面非常有效,同时它可以保护图像尖锐的边缘。对椒盐噪声表现较好,对高斯噪声表现较差。
  • 高斯滤波:高斯滤波采用线性的方法,对邻域内不同位置的像素赋予不同的权值,从而计算出邻域中心像素的灰度值。对椒盐噪声和高斯噪声都有一定的效果,能够更多的保留图像的总体灰度分布特征。

10.3.3.2. 梯度算子

图像函数的变化可以用指向图像函数最大增长方向的梯度来描述,边缘是附加到单个像素的属性,是根据像素邻域中的图像函数计算得出的,可用一阶导数或者二阶导数来检测边缘。要得到一幅图像的梯度,可以通过梯度算子来卷积得到图像每个像素点的梯度。常用的梯度算子有:roberts、prewitt、sobel算子等。下面是3*3邻域的8联通梯度算子的卷积掩膜

  • Roberts:g_x\left(x,y\right)=\left[\begin\{matrix}-1&0\\0&1\\\end\{matrix}\right],\ g_y\left(x,y\right)=\left[\begin\{matrix}0&-1\\1&0\\\end\{matrix}\right]
  • Prewitt:g_x\left(x,y\right)=\left[\begin\{matrix}-1&-1&-1\\0&0&0\\1&1&1\\\end\{matrix}\right],\ g_y\left(x,y\right)=\left[\begin\{matrix}-1&0&1\\-1&0&1\\-1&0&1\\\end\{matrix}\right]
  • Sobel:g_x\left(x,y\right)=\left[\begin\{matrix}-1&-2&-1\\0&0&0\\1&2&1\\\end\{matrix}\right],\ g_y\left(x,y\right)=\left[\begin\{matrix}-1&0&1\\-2&0&2\\-1&0&1\\\end\{matrix}\right]
  • 更先进的边缘检测算子:Marr-Hildreth算子、Canny边缘检测算子等

10.3.4. 图像复原

图像复原旨在利用有关退化性质知识来抑制退化。多数图像复原方法是基于整幅图像上的全局性“去卷积”的方法。

10.3.5. 其他预处理

当然上述只是图像预处理的基本操作,在实际运用中还有一些常用图像处理操作便于我们进行计算。
如:

  • 图片数据集的打包和解包
  • 图片转tensor
  • 图片颜色空间的转化