说到矩阵运算,最简单的粗暴的就是三重循环直接遍历:
def matrix_multiplication_loop(A,B):
m = A.shape[0]
n = A.shape[1]
l = B.shape[1]
C = np.zeros([m,l])
for i in xrange(m):
for j in xrange(l):
for k in xrange(n):
C += A[i][k]*B[k][j]
return C
A = np.random.random([300,12])
B = np.random.random([12,256])
%timeit C = matrix_multiplication_loop(A,B)
1 loop, best of 3: 2.22 s per loop
简直龟速了,可不可再快一点?当然,上numpy
%timeit C = np.dot(A,B)
10000 loops, best of 3: 105 μs per loop
numpy还是牛牛哒,一下子快了2万倍~
可不可再快一点?当然,JIT听过吗?just in time-即时编译。我第一次听到这个词是在工业工程的精益制造里,它的含义是生产线上即时生产,需要什么马上预定什么,没有库存。numba就是just in time的一个编译器,让我们来试试:
import numba
@numba.autojit
def matrix_multiplication_numba(A,B):
return np.dot(A,B)
%timeit C = matrix_multiplication_numba(D,E)
10000 loops, best of 3: 55 μs per loop
又快了将近一倍~
可不可再快一点?当然,只是今天没时间了,未完待续。
numpy本身是非常优秀的,把速度优化就极佳了,要打败它并不容易,我们需要借助上古的力量C语言和blas库。cython是python里实现C语言的一座桥梁,下面是用cython实现的矩阵乘法:
%load_ext Cython
%%cython
#!python
#cython: boundscheck=False, wraparound=False, nonecheck=False
#cython: cdivision=True
from scipy.linalg.cython_blas cimport dgemm
cpdef void cython_blas_MatrixMul(double[::1,:] a, double[::1,:] b, double[::1,:] out, char* TransA, char* TransB) nogil:
cdef:
char* Trans='T'
char* No_Trans='N'
int m, n, k, lda, ldb, ldc
int col_a, col_b
double alpha, beta
#dimensions of input arrays
lda = a.shape[0]
col_a = a.shape[1]
ldb = b.shape[0]
col_b = b.shape[1]
ldc = m
alpha = 1.0
beta = 0.0
dgemm(TransA, TransB, &m, &n, &k, &alpha, &a[0,0], &lda, &b[0,0], &ldb, &beta, &out[0,0], &ldc)
%timeit cython_blas_MatrixMul(A,B,C,b"T",b"T")
100000 loops, best of 3: 9.34 μs per loop
厉害吧!又快了五倍,比最开始的实现方法已经快了20万倍!这性能也已经逼近C语言了。
可不可以再快一点?嘿嘿,当然!现在已经接近CPU的极限了,要更快我们就要买入GPU的世界了~
你们感兴趣,超过一百赞,我就写怎么使用python做GPU计算,让计算速度快破天际
谢谢大家捧场,这么快就过100赞了。来来来,让我们继续飙车~
GPU相比CPU并非在所有情况下都更快,小矩阵时,矩阵可以直接存储在CPU的cache里,CPU可以快速访问,这个时候CPU会比GPU快。但是当遇到大矩阵时,GPU的威力就显示出来了。让我们先把矩阵扩大一千倍来看看:
A = np.random.random([3000,1280])
B = np.random.random([1280,2560])
C = np.zeros([3000,2560])
先用numpy做baseline:
%timeit C = np.dot(A,B)
1 loop, best of 3: 582 ms per loop
可怕,一下子慢了5000倍。来试试,cython:
%timeit cython_blas_MatrixMul(A,B,C,b"T",b"T")
1 loop, best of 3: 280 ms per loop
快了一倍,可是还要280ms。让我们来试试GPU吧。先用pyculib走一波,pyculib是cuda在Python里的一个开源库,集成了cudablas一系列算法,非常好用:
from pyculib import blas
%timeit Cres = blas.gemm('N', 'N', alpha, A, B)
1 loop, best of 3: 140 ms per loop
哇塞,一下快了一倍,GPU果然厉害~
可不可以再快一点?那是必须的。tensorflow是Google开源的深度学习框架,矩阵方面内部优化很多:
import tensorflow as tf
A = tf.random_normal([3000,1280])
B = tf.random_normal([1280,2560])
C = tf.matmul(A,B)
with tf.Session() as sess:
%timeit result = sess.run(C)
100 loops, best of 3: 4.83 ms per loop
哇咔咔,比numpy快了100倍!tensorflow果然是Google的技术名不虚传!
这就是终点了吗?还能更快吗?答案是肯定的,我听NVIDIA的工程师说,如果你用C语言编写的cuDNN直接操作GPU指针还能比tensorflow快3倍~但那就脱离python的范畴了。看了这么多,有木有觉得计算机真是博大精深!勇敢的少年们,快来拥抱CS吧~