如何加速你的代码
🗼

如何加速你的代码

Property
notion image
 
我的这个文章有点点奇怪哈,只要标题中填写了英文单词,文章的html就会用英文名称,导致如果两篇文章标题含有的英文完全相同,最后指向的文章也会相同。当然,标题中的语言肯定是指Python,毕竟其他的语言(例如Java、C++)都挺快的,只有Python在执行动态类型的计算时,会比较慢。
 
之前我的代码计算也很慢,因为涉及了比较多的循环,又没有办法去缩减这些循环(使用数据结构优化的话可能有),因此一个epoch跑了近一天,很不可思议。
 
加速Python的计算的话,主要存在两种方式:
一、使用numpy等使用C扩展写的包,它对矩阵计算有优化,或者使用CPython的方式去运行;
二、使用numba这个JIT加速库对普通的Python代码进行加速;
 

1. 使用Numpy对矩阵运算加速

这个我相信大家都知道,numpy对矩阵的运算加速效果比较显著。
如下,是numpy的矩阵运算代码,对应元素相乘:
import numpy as np

# numpy 随机数矩阵
a = np.random.random((20,30))
b = np.random.random((20,30))

# 对应元素相乘  矩阵相乘是 点乘
start_time = time.time()
for i in range(10000):
		# numpy支持的操作
		c = a * b
end_time = time.time()

print(f"Executed time is {end_time-start_time} s.")
运行时间是0.006s;
notion image
 
对于Python原来的执行逻辑:
import random

# Python 正常随机数列表
a = [[random.random() for _ in range(30)] for _ in range(20)]
b = [[random.random() for _ in range(30)] for _ in range(20)]

# 对应元素相乘  矩阵相乘是 点乘
start_time = time.time()
for i in range(10000):
		# 正常Python逻辑
		c = [[a[k][j]*b[k][j] for j in range(a[0])] for k in range(len(a))]
end_time = time.time()

print(f"Executed time is {end_time-start_time} s.")
可以看到,执行10000次的操作大概需要0.628s,运行时间足足是numpy的100倍:
notion image
 
 
那么,这个加速效果这么优秀,其缺陷或者说不足在哪里?
很难用到我们常用的代码中,虽然机器学习是矩阵密集型计算,但矩阵计算那一部分大多都被开源框架的GPU加速搞定了。真正需要应对加速的,是各种提前需要准备好的一些数据计算,以及一些参数的处理。
 

2. 使用numba库进行加速

这是我最近搞一个参数不是全连接的模型,碰到的麻烦。它需要自己定义梯度和求解梯度,里面会包含大量的非矩阵计算,因此需要进行很多循环(可能也可以将大部分转换成numpy支持的矩阵运算),运行速度很慢。
这个时候我看到了一个第三方库,使用JIT模型进行加速,查一个issue的时候,看到官方对加速的一部分原因做了说明:1. 对需要反复申请的空间进行常态化,并且将类型编译为静态类型;2. 将一些计算分解成矩阵计算,索引过程加快,省去大量查询时间;
我实际使用下来,也是非常的amazing啊,之前跑了4个小时的一个epoch(3000条数据,每一条数据都需要20s,实在是太慢了),优化之后直接用了4s跑完,简直不敢相信。
之前的代码版本
notion image
现在的代码版本
notion image
同样是一个epoch,差距逆天了,这个速度,我觉得都快赶上Java了。
那么,numba怎么使用呢?我查询资料的时候,发现大多数人都语焉不详,要么只贴了常规Python和加速后的Python做对比,要么只说了numbanumpy的数据类型支持的比较好。我也先贴一个使用和加速效果对比,首先是使用,很简单,只需要使用numba提供的装饰器即可。
from numba import njit
import time

@njit
def need_to_accelerate(a, b):
	c = 0
	for i in range(1000000):
		c += (a + b)
	return c

start_time = time.time()
c = need_to_accelerate(1, 2)
end_time = time.time()

print(f"运行时间是 {end_time-start_time} s")
让我们跑一下这个代码看看有jit装饰器和没有的差距,第一个带装饰器:
notion image
不带装饰器的情况下:
notion image
怎么样,迷糊了是吗?以为会有一个巨大的加速?结果却是不加装饰器的时候反而更快了?
是的,这个库的主要作用是不进行频繁的malloc内存,但是这个例子中并没有存在这样的情况,也没有大量索引的元素相乘。所以我的实际应用加速很快,但这里却没有加速到。
(贴出我的加速代码,大家可以运行试一下不加速的运行时间)
@njit
def calculate_second_logit(alpha_h, alpha_i, alpha_eta, x_i, norm_table, f_p):
    s_p = [0.0 for _ in range(len(alpha_h))]
    for i in range(len(alpha_h)):
        for j in range(len(alpha_h[0])):
            temp = sum([m * n for m, n in zip(alpha_i[i], x_i)])
            temp = sum([alpha_h[i][j] * f_p[j] + alpha_eta[i] * k + temp * k for k in norm_table]) / len(norm_table)
            s_p[i] = s_p[i] + temp

    return s_p
啊,写到这里,头有点晕,下次再写。
顺便,使用numba还有两个坑(全文精华):
一个是对 numpy 的数据结构支持地很好,但是 numpy 的函数一个也不行(亲测,就是函数里不要有 numpy 的函数调用);
二是,能循环的尽量循环,不要做一些精简了代码行数,却没有减少循环次数的代码(比如推导式),以及数据结构转成 numpy 的或者 numba 的。
 
总结一下,就是,传入的参数最好是 numba 的数据结构,被加速的函数内部不使用 numpy 的操作。