kNN-CUDA 源码分析
3140105324
文家豪
kNN-CUDA 是一个用 CUDA 实现的计算 k 近邻搜索的项目。这个项目用 GPGPU 实现了有两种暴力的算法——它们通过 NVIDA 的 API 在显卡上并行的求解 kNN 搜索。
1. 问题说明
kNN 搜索是指通过计算 query 的向量和 reference 集合内的向量的相似度,并找出最相似的 k 个。
它是许多研究和工业领域的一个问题,如三维物体渲染,基于内容的图像检索,统计(熵和分歧估计),生物学(基因分类)等。
一般的解决方法是串行的算出相似度,然后用容量为 k 的最小堆(或最大堆)来找找出最相似的 k 个。但这种传统的计算方法往往耗时很长。让人们难以接受。
而每两个向量之间的相似度计算是独立的,这就让这个问题可以很好地被并行的实现。
2. GPU 并行计算 kNN
2.1 并行计算思路
kNN 搜索主要分为两个部分:
1. 计算相似度。
2. 找到最相似的 k 个。
因为相似度的计算是独立的,所以这一步可以很好的被并行计算。可以让一个线程计算一个相似度。
而找到最相似的 k 个,则涉及到了排序,一个 query 的排序计算不能被并行,只能让一个线程去计算一个 query 的排序。
而 kNN-CUDA 的作者在他们的论文 Fast k Nearest Neighbor Search using GPU 中研究发现,在 k 很小的时候应该用插入排序,在 k 很大的时候应该使用梳排序(Comb Sort)。他们实验的结果如下图:
而 kNN-CUDA 中则只实现了插入排序。
2.2 源码分析说明
2.2.1 总体说明
kNN-CUDA 一共实现了四种算法,带序号和不带序号的,其中又有为 C 实现的和为 MATLAB 实现的。
我分析的是带序号的为 C 实现的,即 knn_cuda_with_indexes.cu。
这个程序会计算相似度并返回最近的 k 个点的距离和它们的序号。该程序中的函数有:
(并行函数)
- cuComputeDistanceTexture:计算矩阵 A(reference points)和 B(query points)中向量的相似度,其中矩阵 A 是一个 texture。
- cuComputeDistanceGlobal:功能同上,但是 A 是一个矩阵而不是 texture。
- cuInsertionSort:插入排序的实现。
- cuParallelSqrt:计算前 k 个距离矩阵中的平方根。
(一般函数)
- printErrorMessage:输出错误信息。
- knn:计算 knn。
- mexFunction:给 MATLAB 用的接口,只在 MATLAB_CODE == 1
时编译。
- main:kNN 搜索的例子。
2.2.2 具体函数分析
* cuComputeDistanceTexture *
__global__ void cuComputeDistanceTexture(int wA, float * B, int wB, int pB, int dim, float* AB){
unsigned int xIndex = blockIdx.x * blockDim.x + threadIdx.x;
unsigned int yIndex = blockIdx.y * blockDim.y + threadIdx.y;
if ( xIndex<wB && yIndex<wA ){
float ssd = 0;
for (int i=0; i<dim; i++){
float tmp = tex2D(texA, (float)yIndex, (float)i) - B[ i * pB + xIndex ];
ssd += tmp * tmp;
}
AB[yIndex * pB + xIndex] = ssd;
}
}
功能:
该函数计算矩阵 A