千萬張臉的特徵向量，計算相似度提速？

上篇日誌人臉特徵向量用整數存儲精度損失多少？裡面寫到把浮點特徵向量轉成整型，可以提速幾倍，今天做了一番實驗。

摘要：

如果你對性能一無所知，把計算相似度的代碼寫成下面這樣

float dot_float(int len, const float* v1, const float* v2){ float sumf = 0; for (int i = 0; i < len; i++) { sumf += v1[i] * v2[i]; } return sumf;}

然後在外層這樣調用

for (int pp = 0; pp < person_num; pp++) scoresf[pp] = dot_float(N, valf1, valf2 + pp*N);

與這種最粗糙的寫法相比，本文介紹的寫法最終可以提速12.6倍左右。

基本思路：

1.float轉short， 相似度損失小於萬分之五（見最上面的鏈接），位寬縮小一半，定點比浮點要快。

2.利用SSE加速。

3.批量處理減少函數調用次數。

4. 展開循環

第一部分：數據準備

注意不用SSE時，兩short相乘會溢出，因此直接用int存儲。如果用short存儲，相乘時才轉成int，應該會更慢。

float scale = (int)0x7fff;void init(){ person_num = total_elements / N; printf("N = %d, person_num = %d
", N, person_num); valf1 = new float[N]; valf2 = new float[N*person_num]; vali1 = new int[N]; vali2 = new int[N*person_num]; valshort1 = new short[N]; valshort2 = new short[N*person_num]; scoresf = new float[person_num]; scoresi = new int[person_num]; float len1 = 0; for (int i = 0; i < N; i++) { valf1[i] = i; len1 += valf1[i] * valf1[i]; } len1 = sqrt(len1); for (int i = 0; i < N; i++) { valf1[i] /= len1; vali1[i] = __min(scale, valf1[i] * scale + 0.5f); valshort1[i] = vali1[i]; } for (int pp = 0; pp < person_num; pp++) { double len2 = 0; for (int i = 0; i < N; i++) { int tmp = rand() % N; valf2[pp*N + i] = tmp; len2 += tmp*tmp; } len2 = sqrt(len2); for (int i = 0; i < N; i++) { valf2[pp*N + i] /= len2; vali2[pp*N + i] = __min(scale, valf2[pp*N + i] * scale + 0.5f); valshort2[pp*N + i] = vali2[pp*N + i]; } }}

第二部分： 不用SSE時浮點點積和整型點積

float dot_float(int len, const float* v1, const float* v2){ float sumf = 0; for (int i = 0; i < len; i++) { sumf += v1[i] * v2[i]; } return sumf;}int dot_int(int len, const int* v1, const int* v2){ int sumi = 0; for (int i = 0; i < len; i++) { sumi += v1[i] * v2[i]; } return sumi;}

len =128， 256, 512, 1024, 2048時，整型比浮點大概快2.3-2.9倍

第三部分：使用SSE （位寬256，頭文件在#include <immintrin.h>），浮點與short

注意地址對齊，否則會慢很多，包括前面申請數組內存的時候也要，我是因為測試的都是整16倍的，所以沒有加對齊代碼，對齊方式參照float tmp[8 + 8], *q = (float*)(((long long)tmp+8)>>5<<5)，請讀者自己琢磨。

浮點改SSE加速更明顯，浮點改SSE後加速2倍左右，整型改SSE之後加速1.5倍左右。

#ifdef use_aligned#define my_mm256_load_ps _mm256_load_ps#define my_mm256_store_ps _mm256_store_ps#define my_mm256_load_si256 _mm256_load_si256#define my_mm256_store_si256 _mm256_store_si256#else#define my_mm256_load_ps _mm256_loadu_ps#define my_mm256_store_ps _mm256_storeu_ps#define my_mm256_load_si256 _mm256_loadu_si256#define my_mm256_store_si256 _mm256_storeu_si256#endiffloat dot_float_SSE(int len, const float* v1, const float* v2){ __m256 result = _mm256_setzero_ps(); __m256 val1, val2; float sumf;#ifdef use_aligned float tmp[8 + 8], *q = (float*)(((long long)tmp+8)>>5<<5);#else float q[8];#endif int i = 0; for (; i < len; i += 8) { val1 = my_mm256_load_ps(v1 + i); val2 = my_mm256_load_ps(v2 + i); val1 = _mm256_mul_ps(val1, val2); result = _mm256_add_ps(result, val1); } my_mm256_store_ps(q, result); sumf = q[0] + q[1] + q[2] + q[3] + q[4] + q[5] + q[6] + q[7]; for (; i < len; i++) { sumf += v1[i] * v2[i]; } return sumf;}int dot_short_SSE(int len, const short* v1, const short* v2){ int sumi;#ifdef use_aligned int tmp[8 + 8], *q = (int*)(((long long)tmp + 8) >> 5 << 5);#else int q[8];#endif __m256i result = _mm256_setzero_si256(); __m256i val1, val2; int i = 0; const __m256i* vv1 = (const __m256i*)v1; const __m256i* vv2 = (const __m256i*)v2; for (; i < len; i += 16) { val1 = my_mm256_load_si256(vv1++); val2 = my_mm256_load_si256(vv2++); val1 = _mm256_madd_epi16(val1, val2); result = _mm256_add_epi32(result, val1); } my_mm256_store_si256((__m256i*)q, result); sumi = q[0] + q[1] + q[2] + q[3] + q[4] + q[5] + q[6] + q[7]; for (; i < len; i++) { sumi += (int)v1[i] * (int)v2[i]; } return sumi;}

第四部分：SSE批量計算與SSE不批量

看下面這個函數意思意思一下，就是盡量少調用函數。

void dot_float_batch(int len, const float* v1, const float* v2, int num, float* scores){ for (int pp = 0; pp < num; pp++) { float sumf = 0; for (int i = 0; i < len; i++) { sumf += v1[i] * v2[pp*N+i]; } scores[pp] = sumf; }}

SSE批量比不批量提速較為明顯，維度越低越明顯。

第五部分：展開循環

只展開了256維。為什麼選256維，不防看看這裡SpherefaceNet-04, SpherefaceNet-06 Release · Issue #81 · wy1iu/sphereface

void dot_short_SSE_batch_dim256(const short* v1, const short* v2, int num, int* scores){ int sumi;#ifdef use_aligned int tmp[8 + 8], *q = (int*)(((long long)tmp + 8) >> 5 << 5);#else int q[8];#endif for (int nn = 0; nn < num; nn++) { __m256i result = _mm256_setzero_si256(); __m256i val1, val2; int i = 0; const __m256i* vv1 = (const __m256i*)v1; const __m256i* vv2 = (const __m256i*)v2; val1 = my_mm256_load_si256(vv1++); val2 = my_mm256_load_si256(vv2++); val1 = _mm256_madd_epi16(val1, val2); result = _mm256_add_epi32(result, val1); val1 = my_mm256_load_si256(vv1++); val2 = my_mm256_load_si256(vv2++); val1 = _mm256_madd_epi16(val1, val2); result = _mm256_add_epi32(result, val1); val1 = my_mm256_load_si256(vv1++); val2 = my_mm256_load_si256(vv2++); val1 = _mm256_madd_epi16(val1, val2); result = _mm256_add_epi32(result, val1); val1 = my_mm256_load_si256(vv1++); val2 = my_mm256_load_si256(vv2++); val1 = _mm256_madd_epi16(val1, val2); result = _mm256_add_epi32(result, val1); val1 = my_mm256_load_si256(vv1++); val2 = my_mm256_load_si256(vv2++); val1 = _mm256_madd_epi16(val1, val2); result = _mm256_add_epi32(result, val1); val1 = my_mm256_load_si256(vv1++); val2 = my_mm256_load_si256(vv2++); val1 = _mm256_madd_epi16(val1, val2); result = _mm256_add_epi32(result, val1); val1 = my_mm256_load_si256(vv1++); val2 = my_mm256_load_si256(vv2++); val1 = _mm256_madd_epi16(val1, val2); result = _mm256_add_epi32(result, val1); val1 = my_mm256_load_si256(vv1++); val2 = my_mm256_load_si256(vv2++); val1 = _mm256_madd_epi16(val1, val2); result = _mm256_add_epi32(result, val1); val1 = my_mm256_load_si256(vv1++); val2 = my_mm256_load_si256(vv2++); val1 = _mm256_madd_epi16(val1, val2); result = _mm256_add_epi32(result, val1); val1 = my_mm256_load_si256(vv1++); val2 = my_mm256_load_si256(vv2++); val1 = _mm256_madd_epi16(val1, val2); result = _mm256_add_epi32(result, val1); val1 = my_mm256_load_si256(vv1++); val2 = my_mm256_load_si256(vv2++); val1 = _mm256_madd_epi16(val1, val2); result = _mm256_add_epi32(result, val1); val1 = my_mm256_load_si256(vv1++); val2 = my_mm256_load_si256(vv2++); val1 = _mm256_madd_epi16(val1, val2); result = _mm256_add_epi32(result, val1); val1 = my_mm256_load_si256(vv1++); val2 = my_mm256_load_si256(vv2++); val1 = _mm256_madd_epi16(val1, val2); result = _mm256_add_epi32(result, val1); val1 = my_mm256_load_si256(vv1++); val2 = my_mm256_load_si256(vv2++); val1 = _mm256_madd_epi16(val1, val2); result = _mm256_add_epi32(result, val1); val1 = my_mm256_load_si256(vv1++); val2 = my_mm256_load_si256(vv2++); val1 = _mm256_madd_epi16(val1, val2); result = _mm256_add_epi32(result, val1); val1 = my_mm256_load_si256(vv1++); val2 = my_mm256_load_si256(vv2++); val1 = _mm256_madd_epi16(val1, val2); result = _mm256_add_epi32(result, val1); my_mm256_store_si256((__m256i*)q, result); sumi = q[0] + q[1] + q[2] + q[3] + q[4] + q[5] + q[6] + q[7]; scores[nn] = sumi; }}

展開之後性能得到意想不到的提升。

實驗：維度*人臉數=25.6百萬，外層循環10次，每次隨機生成向量，依次統計各種方法時間，每次執行100次，相當於維度*人臉數=2560百萬。

符號說明：

浮點：不使用SSE的float點積

浮點*：批量float點積

浮點SSE：SSE float點積

浮點SSE*：SSE批量float點積

浮點SSE**：展開循環的SSE批量float點積

整型：不使用SSE的int點積

整型*：批量int點積

整型SSE：SSE short點積

整型SSE*： SSE批量short點積

整型SSE**：展開循環的SSE批量short點積

時間表

維度 人數 浮點 浮點* 浮點SSE 浮點SSE* 浮點SSE** 整型 整型* 整型SSE 整型SSE* 整型SSE**128 20M 1.927 1.900 1.057 0.752 - 0.824 0.808 0.547 0.346 - 256 10M 2.166 2.082 1.006 0.786 0.231 0.784 0.802 0.507 0.330 0.172 512 5M 2.205 2.174 1.027 0.889 - 0.845 0.786 0.495 0.313 - 1024 2.5M 2.276 2.225 1.085 0.933 - 0.792 0.814 0.483 0.365 - 2048 1.25M 2.272 2.307 1.096 1.049 - 0.815 0.827 0.471 0.391 -

加速表

維度 人數 浮點 浮點* 浮點SSE 浮點SSE* 浮點SSE** 整型 整型* 整型SSE 整型SSE* 整型SSE**128 20M 1.000 1.014 1.824 2.563 - 2.337 2.385 3.522 5.574 -256 10M 1.000 1.040 2.153 2.757 9.375 2.762 2.701 4.273 6.563 12.605 512 5M 1.000 1.014 2.146 2.481 - 2.609 2.805 4.455 7.045 -1024 2.5M 1.000 1.023 2.098 2.440 - 2.875 2.797 4.717 6.237 -2048 1.25M 1.000 0.985 2.073 2.165 - 2.788 2.746 4.825 5.813 -

可以看到，與最粗糙的寫法相比，不展開的時候也能得到5.5-7.0倍加速，展開的時候達到12.6倍加速

結論：

將浮點特徵向量轉為short型存儲，並使用SSE加速，對特定維數展開循環，可以獲得一個量級的性能提升。同時，存儲空間也減半。其代價只是不到萬分之五的精度損失。