onnxruntime 中的 Gather 算子

2024-03-28 13:28
文章标签 算子 gather onnxruntime

本文主要是介绍onnxruntime 中的 Gather 算子,希望对大家解决编程问题提供一定的参考价值,需要的开发者们随着小编来一起学习吧!

上一篇文章中介绍了 Division by Invariant Integers using Multiplication 的原理,很多框架均才用该算法优化除法运算。onnxruntime 是已知实现中最为简洁的,因此本文结合 onnxruntime 的 Gather 实现进行介绍。 Gather 算子是一个索引类算子,kernel 中每个线程计算偏移时使用 fast_divmod 避免除法运算。

注意:ONNX 中的 Gather 功能与 numpy.take 相同,torch.index_select 是其简化版。而 ONNX 中的 GatherElements 与 torch.gather 和 paddle. take_along_axis 相对应。

Gather

Gather
CudaKernel
OpKernel
GatherBase

会话运行时,ExecuteKernel 函数会调用 OpKernel。
CudaKernel 是 CUDA kernel 的基类,提供了 CudaKernel::Compute 函数。
OpKernelInfo 是一个非常轻量级的类,它作为构建 Kernel 实例所需的所有数据的聚合视图。 注意:它不拥有/持有任何对象。

class Gather : public CudaKernel, public GatherBase {public:Gather(const OpKernelInfo& info) : CudaKernel(info), GatherBase(info) {}Status ComputeInternal(OpKernelContext* context) const override;
};

Gather::ComputeInternal

Gather::ComputeInternal
GatherBase::PrepareForCompute
GatherImpl

创建一个 GatherBase::Prepare 结构体,包含了两个输入和一个输出张量的指针。
GatherBase::PrepareForCompute 准备输入输出。输出张量的秩为input_rank - 1 + indices_rank,即将axis参数指定的轴替换为indices张量的形状。
ORT_RETURN_IF_ERROR 在表达式失败时返回错误。
TensorShape::SizeFromDimension 计算从指定维度开始的乘积大小。
axis参数会将输入张量划分为3部分:batch 维度、索引维度、分块维度。
block_size为每个索引对应的分块大小。
N为索引数量。
input_block_size为在输入上的分块大小。
indices_max即索引上限。

Status Gather::ComputeInternal(OpKernelContext* context) const {Prepare p;ORT_RETURN_IF_ERROR(PrepareForCompute(context, p));const TensorShape& input_shape = p.input_tensor->Shape();const int64_t block_size = input_shape.SizeFromDimension(p.axis + 1);size_t N = p.indices_tensor->Shape().Size();const int64_t input_block_size = input_shape.SizeFromDimension(p.axis);const int64_t output_block_size = N * block_size;const int64_t indices_max = input_shape[p.axis];
  const void* input_data = p.input_tensor->DataRaw();const void* indices_data = p.indices_tensor->DataRaw();void* output_data = p.output_tensor->MutableDataRaw();if (p.output_tensor->Shape().Size() == 0) {return Status::OK();}

gsl::narrow可确保无损失转换,并在无法转换时引发gsl::narrowing_error
fast_divmod 即 DivMod,用于快速计算除法。

  const fast_divmod divmod_output_block_size(gsl::narrow_cast<int>(output_block_size));const fast_divmod divmod_block_size(gsl::narrow_cast<int>(block_size));const size_t element_size = p.input_tensor->DataType()->Size();const size_t index_element_size = p.indices_tensor->DataType()->Size();

GatherImpl 函数索仅支持int32_tint64_t引类型。
传入的p.output_tensor->Shape().Size()即输出元素总数。

  // CUDA Kernel implementation supports element sizes of:// int8_t, int16_t, int32_t and int64_t which covers all supported// types since there is no computations necessary just data movementif (p.indices_tensor->IsDataType<int32_t>() ||p.indices_tensor->IsDataType<int64_t>()) {GatherImpl(Stream(context),input_block_size,indices_max,divmod_output_block_size,divmod_block_size,indices_data,index_element_size,input_data,element_size,output_data,p.output_tensor->Shape().Size());return Status::OK();}

ORT_MAKE_STATUS 创建一个 Status 对象。

  return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, "Type for Tind not supported yet in Gather.");
}

GatherImpl

GatherImpl
_GatherKernel

GridDim 结构体中定义了美剧值。
N为输出元素数量。直接求出所需 threadblock 的数量,没有太多策略。

void GatherImpl(cudaStream_t stream,const int64_t input_block_size,const int64_t indices_max,const fast_divmod& output_block_size,const fast_divmod& block_size,const void* indices_data,size_t index_element_size,const void* input_data,size_t element_size,void* output_data,const size_t N) {int blocksPerGrid = (int)(ceil(static_cast<float>(N) / GridDim::maxThreadsPerBlock));

ToCudaType 模板类将类型枚举转换为数据类型。
根据元素大小调用 _GatherKernel 模板函数,这样减少了实例化类型。

  switch (element_size) {case sizeof(int8_t): {using CudaType = typename ToCudaType<int8_t>::MappedType;_GatherKernel<<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0, stream>>>(input_block_size, indices_max, output_block_size, block_size, indices_data, index_element_size,reinterpret_cast<const CudaType*>(input_data), reinterpret_cast<CudaType*>(output_data), (CUDA_LONG)N);} break;case sizeof(int16_t): {using CudaType = typename ToCudaType<int16_t>::MappedType;_GatherKernel<<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0, stream>>>(input_block_size, indices_max, output_block_size, block_size, indices_data, index_element_size,reinterpret_cast<const CudaType*>(input_data), reinterpret_cast<CudaType*>(output_data), (CUDA_LONG)N);} break;case sizeof(int32_t): {using CudaType = typename ToCudaType<int32_t>::MappedType;_GatherKernel<<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0, stream>>>(input_block_size, indices_max, output_block_size, block_size, indices_data, index_element_size,reinterpret_cast<const CudaType*>(input_data), reinterpret_cast<CudaType*>(output_data), (CUDA_LONG)N);} break;case sizeof(int64_t): {using CudaType = typename ToCudaType<int64_t>::MappedType;_GatherKernel<<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0, stream>>>(input_block_size, indices_max, output_block_size, block_size, indices_data, index_element_size,reinterpret_cast<const CudaType*>(input_data), reinterpret_cast<CudaType*>(output_data), (CUDA_LONG)N);} break;default:ORT_THROW("Unsupported element size by the Gather CUDA kernel");}
}

_GatherKernel

_GatherKernel
GetIndexValue

CALCULATE_ELEMENTWISE_INDEX_OR_EXIT 计算元素索引,并在超出范围时返回。

template <typename T>
__global__ void _GatherKernel(const int64_t input_block_size,const int64_t indices_max,const fast_divmod output_block_size,const fast_divmod block_size,const void* indices_data,const size_t index_element_size,const T* input_data,T* output_data,const CUDA_LONG N) {CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(id, N);CUDA_LONG input_index = 0;

线程号id除以output_block_size得到输出元素所对应的输入分块索引input_block_index和输入分块内的偏移block_offset
根据block_offset计算对应的indices张量的索引以及分块内元素偏移。
GetIndexValue 取出indices张量的值。相比 TensorFlow 的 gather_functor_gpu.cu.h 没有进行向量化访存优化。
idx支持负数。索引值超出范围时赋零。

  int input_block_index, block_offset;output_block_size.divmod(id, input_block_index, block_offset);int indices_index, offset;block_size.divmod(block_offset, indices_index, offset);int64_t idx = GetIndexValue(indices_data, index_element_size, indices_index);idx = idx < 0 ? idx + indices_max : idx;if (idx < 0 || idx >= indices_max) {output_data[id] = 0;return;}

三部分相加得到输入张量索引。

  input_index = input_block_index * input_block_size + idx * block_size.d_ + offset;output_data[id] = input_data[input_index];
}

GetIndexValue

index_data指针转为相应类型,然后返回偏移位置的值。

__host__ __device__ inline int64_t GetIndexValue(const void* index_data, size_t index_element_size, size_t offset) {switch (index_element_size) {case sizeof(int32_t):return *(reinterpret_cast<const int32_t*>(index_data) + offset);break;case sizeof(int64_t):return *(reinterpret_cast<const int64_t*>(index_data) + offset);break;default:break;}// What is a sensible thing to do here?assert(false);return std::numeric_limits<int64_t>::max();
}

DivMod

除法取余实现基于 Division by Invariant Integers using Multiplication 中的 Figure 4.1。

在这里插入图片描述

// The code below is based on section 4 Unsigned division of paper https://gmplib.org/~tege/divcnst-pldi94.pdf
// In current ORT, fast_divmod is used for calculating the position of a element in tensor,
// so unsigned integer division from the paper is good enough for ORT. The advantage is that div is very simple,
// then GPU compiler can do loop unroll easilly when divmod is called in a loop.
template <>
struct DivMod<int> {DivMod(int d = 1) {d_ = d == 0 ? 1 : d;ORT_ENFORCE(d_ >= 1 && d_ <= static_cast<uint32_t>(std::numeric_limits<int>::max()));

l_ ℓ = ⌈ log ⁡ 2 x ⌉ \ell = \lceil \log_2 x \rceil =log2x

    for (l_ = 0; l_ < 32; l_++)if ((1U << l_) >= d_) break;

m m ′ = ⌊ 2 N ∗ ( 2 ℓ − d ) / d ⌋ + 1 m' = \lfloor 2^N ∗ (2^\ell − d)/d\rfloor + 1 m=2N(2d)/d+1

    uint64_t one = 1;uint64_t m = ((one << 32) * ((one << l_) - d_)) / d_ + 1;M_ = static_cast<uint32_t>(m);// according to paper, the value of m' should fit in a unsigned integer.ORT_ENFORCE(M_ > 0 && M_ == m);}

DivMod::div

t t 1 = M U L U H ( m ′ , n ) t_1 = \mathrm{MULUH}(m', n) t1=MULUH(m,n),使用uint64_t计算避免溢出。
对于 q q q

  • 如果 d = 1 d = 1 d=1,那么 ℓ = 0 \ell = 0 =0,所以 m ′ = 1 m' = 1 m=1 s h 1 = s h 2 = 0 sh_1 = sh_2 = 0 sh1=sh2=0。代码计算 t 1 = ⌊ 1 ∗ n / 2 N ⌋ = 0 t_1 = \lfloor 1 ∗ n/2^N \rfloor = 0 t1=1n/2N=0 q = n q = n q=n
  • d > 1 d > 1 d>1,则 ℓ ≥ 1 \ell≥1 1,故 s h 1 = 1 sh_1 = 1 sh1=1 s h 2 = ℓ − 1 sh_2 =\ell −1 sh2=1
    q = S R L ( t 1 + S R L ( n − t 1 , s h 1 ) , s h 2 ) = S R L ( t 1 + S R L ( n − t 1 , 1 ) , ℓ − 1 ) = ⌊ t 1 + ⌊ ( n − t 1 ) 2 ⌋ 2 ℓ − 1 ⌋ = ⌊ ⌊ 2 ∗ t 1 2 + ( n − t 1 ) 2 ⌋ 2 ℓ − 1 ⌋ (4.5) = ⌊ ⌊ ( t 1 + n ) / 2 ⌋ 2 ℓ − 1 ⌋ = ⌊ t 1 + n 2 ℓ ⌋ \begin{aligned} q &= \mathrm{SRL}(t_1 + \mathrm{SRL}(n − t_1, sh_1), sh_2)\\ &= \mathrm{SRL}(t_1 + \mathrm{SRL}(n − t_1, 1), \ell− 1)\\ &=\lfloor \frac{t_1 + \lfloor \frac{(n − t_1)}{2} \rfloor}{2^{\ell− 1}}\rfloor\\ &=\lfloor \frac{\lfloor \frac{2*t_1}{2} + \frac{(n − t_1)}{2} \rfloor}{2^{\ell− 1}}\rfloor \qquad\text{(4.5)}\\ &=\lfloor \frac{\lfloor(t_1 + n)/2\rfloor}{2^{\ell− 1}} \rfloor\\ &=\lfloor \frac{t_1 + n}{2^{\ell}} \rfloor \end{aligned} q=SRL(t1+SRL(nt1,sh1),sh2)=SRL(t1+SRL(nt1,1),1)=21t1+2(nt1)=2122t1+2(nt1)(4.5)=21⌊(t1+n)/2=2t1+n

__umulhi 计算两个 32 位无符号整数的乘积的最高有效 32 位。

  __host__ __device__ inline int div(int n) const {
#if defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__)uint32_t t = __umulhi(M_, n);return (t + n) >> l_;
#else// Using uint64_t for t, then t + n won't overflow.uint64_t t = ((uint64_t)M_ * n) >> 32;return static_cast<int>((t + n) >> l_);
#endif}

DivMod::mod

n m o d d = n − d ∗ ⌊ n / d ⌋ n \enspace \mathrm{mod} \enspace d = n − d ∗ \lfloor n/d \rfloor nmodd=ndn/d

  __host__ __device__ inline int mod(int n) const {return n - div(n) * d_;}

DivMod::divmod

  __host__ __device__ inline void divmod(int n, int& q, int& r) const {q = div(n);r = n - q * d_;}
  uint32_t d_;  // divisoruint32_t M_;  // m' in the paper.uint32_t l_;  // l_ = ceil(log2(d_))
};
  • Gather
  • ONNXRuntime整体概览
  • ONNXRuntime源码之OpKernel注册
  • Ways to specify [[nodiscard]] before C++17
  • microsoft/GSL
  • How to use gsl narrow cast
  • 警告 C26472
  • GSL and C++ Core Guidelines
  • Gather
  • Gather
  • tf.gather
  • torch.gather
  • paddle.gather
  • [菁英计划] 索引取值及gather函数 #36815
  • paddle. take_along_axis
  • torch.gather in pytorch.onnx and onnxruntime #31464
  • Replace torch.gather by other operator?
  • Problem compiling onnx model using GLOW compiler: constant not found
  • pytorch导出onnx的原则-以SwinTransformer和DETR在trt8.0.3.4部署为例
  • GatherElements
  • tf2onnx Gather
  • OpenVINO Gather
  • Pytorch equivalent of numpy.take()
  • torch.index_select
  • wkentaro/pytorch-for-numpy-users
  • Similar operation like numpy.take
  • numpy.take
  • tensorflow/tensorflow/core/kernels/gather_functor.h
  • tensorflow/core/kernels/gather_functor_batched.h
  • abseil中的微操

这篇关于onnxruntime 中的 Gather 算子的文章就介绍到这儿,希望我们推荐的文章对编程师们有所帮助!



http://www.chinasem.cn/article/855736

相关文章

PyInstaller问题解决 onnxruntime-gpu 使用GPU和CUDA加速模型推理

前言 在模型推理时,需要使用GPU加速,相关的CUDA和CUDNN安装好后,通过onnxruntime-gpu实现。 直接运行python程序是正常使用GPU的,如果使用PyInstaller将.py文件打包为.exe,发现只能使用CPU推理了。 本文分析这个问题和提供解决方案,供大家参考。 问题分析——找不到ONNX Runtime GPU 动态库 首先直接运行python程序

yolov8-obb旋转目标检测onnxruntime和tensorrt推理

onnxruntime推理 导出onnx模型: from ultralytics import YOLOmodel = YOLO("yolov8n-obb.pt") model.export(format="onnx") onnx模型结构如下: python推理代码: import cv2import mathimport numpy as npimport onnxr

spark算子集锦

Spark 是大数据领域的一大利器,花时间总结了一下 Spark 常用算子,正所谓温故而知新。 Spark 算子按照功能分,可以分成两大类:transform 和 action。Transform 不进行实际计算,是惰性的,action 操作才进行实际的计算。如何区分两者?看函数返回,如果输入到输出都是RDD类型,则认为是transform操作,反之为action操作。 准备 准备阶段包括s

OrangePi AIpro 香橙派 昇腾 Ascend C 算子开发 与 调用 - Tiling实现 2

OrangePi AIpro 香橙派 昇腾 Ascend C 算子开发 与 调用 - Tiling实现 2 flyfish 前置知识 1 前置知识 2 Host侧CPU和Device侧NPU的主要区别 不同的硬件资源 CPU是为了执行通用计算任务而设计的,但在处理大量的并行计算(如矩阵乘、批数据处理)时效率不高。NPU是为了加速机器学习和深度学习任务而设计的,它擅长执行大量的并行计算。N

fpga图像处理实战-边缘检测 (Roberts算子)

Roberts算子         Roberts算子是一种用于边缘检测的算子,主要用于图像处理中检测图像的边缘。它是最早的边缘检测算法之一,以其计算简单、速度快而著称。Roberts算子通过计算图像像素在对角方向的梯度来检测边缘,从而突出图像中灰度变化最剧烈的部分。 原理             Roberts算子通过对图像应用两个2x2的卷积核(也称为掩模或滤波器)来计算图像在水平和垂直

图像边缘检测技术详解:利用OpenCV实现Sobel算子

图像边缘检测技术详解:利用OpenCV实现Sobel算子 前言Sobel算子的原理代码演示结果展示结语 前言   在数字图像处理的广阔领域中,边缘检测技术扮演着至关重要的角色。无论是在科学研究、工业自动化,还是在日常生活中的智能设备中,我们都需要从图像中提取有用的信息。边缘,作为图像中亮度变化最显著的地方,为我们提供了识别和理解图像内容的关键线索。因此,边缘检测算法成为了计算机视

OrangePi AIpro 香橙派 昇腾 Ascend C 算子开发 与 调用 - Tiling实现

OrangePi AIpro 香橙派 昇腾 Ascend C 算子开发 与 调用 - Tiling实现 flyfish 前置知识 基于Kernel直调工程的算子开发流程图 其中有一个Tiling实现 什么是Tiling、Tiling实现 计算API,包括标量计算API、向量计算API、矩阵计算API,分别实现调用Scalar计算单元、Vector计算单元、Cube计算单元执行计算的功

深度学习常用算子

深度学习常用算子 算子功能ReluReLU(x)=max(0,x)LeakyReluLeakyRelu(x) = (x >= 0 ? x : x*negative_slope)Relu6LeakyRelu(x) = max(max(x, 0), 6)Tantanh(x)=(exp(x)-exp(-x))/(exp(x)+exp(-x))sigmoidsigmoid(x) = 1.

计算机 软件 什么是算子

算法(algorithm)是为了达到某个目标,实施的一系列指令的过程,而指令包含算子(operator)和操作数(operand)。   算子:operator, 简单说来就是进行某种“操作“,动作。算法中的一个函数、几行可以重复使用的代码、一个数学中的平方操作,这些都可以认为是算子 操作数:operand,被操作的对象,称之为操作数。     广义的讲,对任何函数进行某一项操作都可以

Halcon提取边缘线段lines_gauss 算子

Halcon提取边缘线段lines_gauss 算子 edges_color_sub_pix和edges_sub_pix两个算子使用边缘滤波器进行边缘检测。还有一个常用的算子lines_gauss算子,也可以用于提取边缘线段,它的鲁棒性非常好,提取出的线段类型是亚像素精度的XLD轮廓。其原型如下: lines gauss(Image : Lines : Sigma, Low, High, Li