解决一个问题,想到一个方法,设计个对应的模型,开始训练吧。。。。
是这个思路没错吧(别较真,较真你就输了。)
然后就是慢长的等待,2天,3天。。。
transformer的WMT训练了3.5天,几天的训练是家常便饭了现在。
GEMM
回到一般理论,贾杨青先生曾经得出在用Alex Krizhevsky’s的图像识别方法里的时间占比如下图
在alex net里大量的conv和dense层消耗了几乎所有的计算时间。
conv与dense里的计算都是gemm实现的,拿caffe举例
这里为gemm(GEneral Matrix to Matrix Multiplication)而来。
Python
我们要计算C=AB, 看看有多少升空间。
不说空话,直接来探探吧。
上篇文章里提到numpy[https://www.atatech.org/articles/143184], 但是受好心朋友反馈说不具可比性,虽然我已经开启了多线程,但为免误会(无奈),这里只列numpy代码,结果不列了,也不作评伦。
import numpy as np
import time
t1 = time.time()
a = np.full([1024, 512], 4)
b = np.full([512,1024], 3)
c = np.dot(a, b)
t2 = time.time()
dt = t2 - t1
dt = round(dt * 1000)
print("cost time:{0}ms".format(dt))
虽然没列,但运算结果比tf差不少的,代码可能存在优化空间,但我这篇文章写的是硬算,不是逻辑。
咱们来看看tensorflow吧,算法用tf准没错了
import tensorflow as tf
import time
t1 = time.time()
a = tf.range(0,1024*512)
b = tf.range(0, 512*1024)
a = tf.reshape(a, [1024, 512])
b = tf.reshape(b, [512, 1024])
c = tf.matmul(a, b)
with tf.Session() as sess:
sess.run(c)
t2 = time.time()
dt = t2 - t1
dt = round(dt * 1000)
print("cost time:{0}ms".format(dt))
tf底层是C写的,性能要高一些。
本来还要列caffe的,一个好心的同事说不要跪舔贾先生,我想那还是算了吧。
C
void matrixdot(){
float *a = malloc(1024*512 * sizeof(float));
float *b = malloc(512*1024 * sizeof(float));
float *c = malloc(1024 * 1024 * sizeof(float));
memset(a, 0, 1024*512*sizeof(float));
memset(b, 0, 512*1024*sizeof(float));
int i, j, k;
for( i = 0; i < 1024; i++){
int ci = 1024 * i;
for(j = 0; j < 1024; j++){
int bi = 0;
int cc = 0;
int ai = 512 * i;
for(k=0; k < 512; k++){
cc += a[ai++] * b[bi];
bi += 1024;
}
c[ci++] = cc;
}
}
}
没错,这个是O3的循环,顺便问下算友, 矩阵乘法类似Coppersmith-Winograd复杂度最低能达到多少,有没有实操过计算时间减少了多少。
上面的运算逻辑最基本,也最简单。
线程化提升10倍
增加线程,把资源用到最大化
void matrixdot_single(void *data){
struct thdata* thdata = data;
int i, j, k;
float* a = thdata->a.data;
float* b = thdata->b.data;
float* c = thdata->c;
int atc = thdata->a.column;
int btc = thdata->b.column;
for( i = thdata->cindex; i < thdata->cend; i++){
int ai = atc * i;
int ci = btc * i;
for(j = 0; j < btc; j++){
int bi = 0;
int cc = 0;
float* aip = a + ai;
for(k=0; k < atc; k++){
cc += (*aip++) * b[bi];
bi += btc;
}
c[ci++] = cc;
}
}
*thdata->finish = 1;
}
这个速度已经快了差不多10倍了,还能再快吗。
改变思路再提升10倍
还有别的思路吗, 2个矩阵相乘
A中的每一行都会乘以B中的每一列,那么A中的行之间是同规律的,从数据上来说
A中的第i行仅影响C中的第i行数据。
公式也可以理解为
C[i,j]由A[i,k]产生的增量为A[i, k] * B[k,j]
C[i,j]是由a-column个增量合并而成。
这样我们转而求出所有的增量即可。
比如
从A中的任一行,求出第一个因子给C的所有增量,以此类推到整行
从上图可以看出,在计算上所有的处理都连续了,尤其是B,不用隔行取数了,理论上可以提升不少。
void matrixdot_single_1(void *data){
struct thdata* thdata = data;
int i, j, k;
float* a = thdata->a.data;
float* b = thdata->b.data;
float* c = thdata->c;
int atc = thdata->a.column;
int btc = thdata->b.column;
for( i = thdata->cindex; i < thdata->cend; i++){
int ai = atc * i;
int ci = btc * i;
float* aip = a + ai;
float* cip = c + ci;
float* bjp = b;
float* cjp = cip;
for(j = 0; j < atc; j++){
cjp = cip;
for(k = 0; k < btc; k++){
(*cjp++) += (*aip) * (*bjp++);
}
aip++;
}
}
*thdata->finish = 1;
}
整体性能又提升了10倍,已经秒钉tensorflow的96ms了。
还能再快吗?
SSE与AVX
对于大矩阵的乘法,如果能将多条指令合并,应该会提升不少,于是便 想到了SSE
//sse
void matrixdot_single_2(void *data){
struct thdata* thdata = data;
int i, j, k;
float* a = thdata->a.data;
float* b = thdata->b.data;
float* c = thdata->c;
int atc = thdata->a.column;
int btc = thdata->b.column;
for( i = thdata->cindex; i < thdata->cend; i++){
int ai = atc * i;
int ci = btc * i;
float* aip = a + ai;
float* cip = c + ci;
float* bjp = b;
float* cjp = cip;
int btc2 = btc/4;
for(j = 0; j < atc; j++){
int aipv = *aip;
__m128 ma = _mm_set_ps(aipv, aipv, aipv, aipv);
cjp = cip;
for(k = 0; k < btc2; k++){
__m128 mb = _mm_load_ps(bjp);
__m128 mc = _mm_mul_ps(ma, mb);
_mm_store_ps(cjp, mc);
bjp += 4;
cjp += 4;
}
aip++;
}
}
*thdata->finish = 1;
}
发现效果提升了1ms,再看AVX。
Avx:
//avx
void matrixdot_single_3(void *data){
struct thdata* thdata = data;
int i, j, k;
float* a = thdata->a.data;
float* b = thdata->b.data;
float* c = thdata->c;
int atc = thdata->a.column;
int btc = thdata->b.column;
for( i = thdata->cindex; i < thdata->cend; i++){
int ai = atc * i;
int ci = btc * i;
float* aip = a + ai;
float* cip = c + ci;
float* bjp = b;
float* cjp = cip;
int btc2 = btc/8;
for(j = 0; j < atc; j++){
int aipv = *aip;
__m256 ma = _mm256_set_ps(aipv, aipv, aipv, aipv,aipv, aipv, aipv, aipv);
cjp = cip;
for(k = 0; k < btc2; k++){
__m256 mb = _mm256_load_ps(bjp);
__m256 mc = _mm256_mul_ps(ma, mb);
_mm256_store_ps(cjp, mc);
bjp += 8;
cjp += 8;
}
aip++;
}
}
*thdata->finish = 1;
}
又提升了1ms。
思路
可以看出来,就GEMM乘法言提升空间还是很大的。
同样配置情况下, tf 的96ms是个理想的结果了,但是还是有很大的优化空间,在C上可以优化到13ms,已经提升了将近7倍。
对于这样的结果,你满意吗?不服来辨。