python 矩阵运算 gpu_使用向量化、矩阵化、cuda等不同方法加速python程序

使用向量化、矩阵化、cuda等不同方法加速python程序

第一次写知乎文章，本文中如有错误请不吝赐教，各位大大多多包涵。

本文是一次实验室组会的分享内容，对于机器视觉各类问题中经常需要处理大量计算任务的情况，以一个计算点云距离的例子来观察不同处理方法在python编程中对速度的影响。运行程序的电脑CPU型号：i9-9900X，GPU型号：RTX 2080 Ti ，python版本3.8.1 ，使用Jupyter lab编写。

问题描述：

在python中，使用普通方法、向量化、矩阵化、cuda等方法加速计算两组点云间各点的距离

# 初始化数据

import numpy as np

import time

n, m = 10000, 15000

a = np.random.random((n,3)) # 点云A中有n组点

b = np.random.random((m,3)) # 点云B中有m组点普通方法

通过嵌套的循环，计算所有点的距离，距离计算公式为：

# 普通计算距离

def cal_distance(a,b):

'''输入a是第一组点，b是第二组点 a=[x,y,z] b=[x,y,z]输出距离'''

d = np.sqrt(np.square(a[0] - b[0]) + np.square(a[1] - b[1]) + np.square(a[2] - b[2]))

return d

# 开始计算

t_1_start = time.time()

common_result = []

for j in range(m):

for i in range(n):

common_result.append(cal_distance(a[i],b[j]))

# 结束

t_1_end = time.time()

t_1 = t_1_end - t_1_start

print('普通计算时间:', t_1, 's')普通计算时间: 950.3101623058319 s

common_result = np.array(common_result).reshape(m,n).T # 元素按照dij表示点云A中第i个点到点云B中第j个点的距离值的格式排列向量化方法

将数据A以重复m次的方式扩展，数据B以复制的方式扩展，实现A中任一点对B中所有点的对应关系。即：

然后将距离计算公式

改写为向量计算式子

# 向量化计算

t_2_start = time.time()

# 数据预处理

a_2 = a.repeat(m, axis=0) # 重复

b_2 = np.tile(b,(n,1)) # 复制

# 开始计算

sub = a_2 - b_2 # 求差

square = np.square(sub) # 求平方

sum_ = np.sum(square, axis=1) # 求和

vector_result = np.sqrt(sum_) # 求平方根

# 结束

t_2_end = time.time()

t_2 = t_2_end - t_2_start

print('向量化计算时间:', t_2, 's')向量化计算时间: 7.28899621963501 s矩阵化方法

将距离计算公式中差的平方和展开，可得：

重新排列各元素：

可以发现，重新排列后的各元素具有“积的和”的性质，则距离计算公式可以改写为矩阵计算。

首先计算第一个括号，发现其为元素自己的平方，则令矩阵A为

再观察第二个括号，令矩阵B为

由于第三个括号的元素是两组点的乘积和，故令矩阵C为

将矩阵A扩展m列，矩阵B扩展n行，则距离矩阵D为

可以发现，矩阵D中第i行第j列元素值表示点云A中第i个点到点云B中第j个点的距离值。

# 矩阵化计算

t_3_start = time.time()

# 开始计算

a_3 = np.sum(np.square(a), axis=1).reshape(n,1) # n*1

b_3 = np.sum(np.square(b), axis=1).reshape(1,m) # 1*m

c_3 = -2 * np.dot(a, b.T) # n*m

matrix_result = np.sqrt(a_3 + b_3 + c_3) # n*m

# 结束

t_3_end = time.time()

t_3 = t_3_end - t_3_start

print('矩阵化计算时间:',t_3, 's')矩阵化计算时间: 2.1971452236175537 s使用cuda计算

使用cuda计算的难点是设计多线程函数，由向量化和矩阵化方法可知，各线程需要计算点云A中第i个点到点云B中第j个点的距离值，每个线程的编号是累加1的，由此2个特点，使用取整方法计算i实现一个周期内重复m次相同的数，使用取余的方法计算j实现一个周期内循环0~m, 总共重复n个周期。设线程编号为idx，点云A有n个点，点云B有m个点，则

from numba import cuda

import math

@cuda.jit

def gpu_add(a, b, result, m, num):

idx = cuda.threadIdx.x + cuda.blockDim.x * cuda.blockIdx.x

if idx < num : # numba对python原生数学方法支持的更好

sum_ = (a[idx//m][0] - b[idx%m][0])**2 + (a[idx//m][1] - b[idx%m][1])**2 + (a[idx//m][2] - b[idx%m][2])**2

result[idx] = math.sqrt(sum_)

def cal():

# 拷贝数据到设备端

a_device = cuda.to_device(a)

b_device = cuda.to_device(b)

# 在显卡设备上初始化一块用于存放GPU计算结果的空间

num = n*m

gpu_result = cuda.device_array(num)

# 计算cuda所需threads和block的大小

threads_per_block = 1024

blocks_per_grid = math.ceil(num / threads_per_block)

# 开始计算

t_4_start = time.time()

gpu_add[blocks_per_grid, threads_per_block](a_device, b_device, gpu_result, m, num)

cuda.synchronize()

# 结束

t_4_end = time.time()

print('GPU计算时间:', t_4_end-t_4_start, 's')

return gpu_result.copy_to_host(), t_4_end-t_4_start

if __name__ == "__main__":

cuda_result, t_4 = cal()GPU计算时间: 0.2234182357788086 s

不同方法的结果对比

# 输出点云A中第2个点到点云B中第14995至15000个点的距离

print('使用一般方法计算的时间与结果:\n', t_1, 's')

print(common_result[1][15000-6:15000-1], '\n')

print('向量化计算的时间与结果:\n', t_2, 's')

print(vector_result[30000-6:30000-1], '\n')

print('矩阵化计算的时间与结果:\n', t_3, 's')

print(matrix_result[1][15000-6:15000-1], '\n')

print('使用cuda计算的时间与结果:\n', t_4, 's')

print(cuda_result[30000-6:30000-1], '\n')

# 也可直接判断不同组数据间是否相同

# if (np.array_equal(common_result, matrix_result)):

# print('结果相同')使用一般方法计算的时间与结果:

950.3101623058319 s

[0.24945065 0.78931863 0.46593733 0.29889308 0.5607863 ]

向量化计算的时间与结果:

7.28899621963501 s

[0.24945065 0.78931863 0.46593733 0.29889308 0.5607863 ]

矩阵化计算的时间与结果:

2.1971452236175537 s

[0.24945065 0.78931863 0.46593733 0.29889308 0.5607863 ]

使用cuda计算的时间与结果:

0.2234182357788086 s

[0.24945065 0.78931863 0.46593733 0.29889308 0.5607863 ]