最近在研究Mxnet，準備從底層先嘗試一遍GPU的編程，這樣更容易熟悉mshadow的那部分邏輯。于是用CUDA嘗試了一下寫一個稀疏的Logistic Regression的程序。所謂稀疏的LR，指的是特征非常多，而每個樣本的特征非常少。因為是學習CUDA，所以我構(gòu)造了一個問題，假設特征的個數(shù)是100W，而每個樣本的特征只有100個左右。如果樣本的特征ID中偶數(shù)的多于奇數(shù)的，我們就認為是正樣本，否則是負樣本。如下的代碼用來產(chǎn)生一個樣本：

void mock_sample(const int max_feature_id, vector< pair<int, float> > & out, int * label) {
  int count = rand() % 100 + 10;
  int ret = 0;
  for(int i = 0; i < count; i++) {
    int fid = rand() % max_feature_id;
    if(fid % 2 == 0) ret += 1;
    else ret -= 1;
    if(abs(ret) > 10) break;
    out.push_back(make_pair<int, float>(fid, 1.0));
  }
  *label = (ret > 0) ? 1 : 0;
}

LR的解法有很多種，PRML那本書上就提到了不下5種LR的優(yōu)化方法。因為只是為了學習CUDA，我選擇了最簡單的一種，也就是基于mini-batch的梯度下降法。LR的優(yōu)化核心主要是去學習w矩陣。而所謂稀疏性來自于計算點積<w, x>，這里w是一個100W維的稠密向量，而x是一個100維左右的稀疏向量。因此<w,x>就是一個稠密向量和稀疏向量的點擊。如果我們考慮mini-batch, 假設X是由n個樣本組成的一個batch，那么X就是一個n*100W的稀疏矩陣。而我們需要計算Xw就是一個稀疏矩陣和稠密向量的乘法。

為了解決這個乘法，我們第一個想到的是用 cusparse，這是由CUDA提供的一個稀疏矩陣的計算庫，在level2的函數(shù)中有CSR格式的稀疏矩陣和稠密向量的乘法：

cusparseStatus_t cusparseScsrmv(cusparseHandle_t handle, 
        cusparseOperation_t transA, int m, int n, int nnz, 
        const float *alpha, const cusparseMatDescr_t descrA, 
        const float *csrValA, const int *csrRowPtrA, 
        const int *csrColIndA, const float *x, 
        const float *beta, float *y)

我首先用cusparse實現(xiàn)了一個版本，具體代碼見 dot_cusparse.cu。不過很不幸的是實現(xiàn)出來的速度還不如CPU的版本。所以我也沒繼續(xù)做優(yōu)化，就放棄了CUSPARSE，準備完全自己實現(xiàn)。（按照后來的經(jīng)驗，如果繼續(xù)優(yōu)化，是有可能優(yōu)化的比較快的，所以讀者可以試一試）。

后來我想，其實如果用COO的表示方法，其實可以很簡單的直接實現(xiàn)Xw的乘法，代碼如下：

__global__ void dot(float * val, int *row_ind, int *col_ind, int nnz, float * ret, float * w) {
  const int tid = (blockIdx.x * blockDim.x) + threadIdx.x;
  if (tid < nnz) {
    int r = row_ind[tid];
    int c = col_ind[tid];
    float v = val[tid];
    atomicAdd(&ret[r], v * w[c]);
  }
}

這里，val, row_ind, col_ind 是矩陣X的COO表示：

X[row_ind[i]][col_ind[i]] = val[i]
nnz = len(val) 是X中非0元素的個數(shù)

ret是dot的返回，也就是Xw的計算結(jié)果，是1個n維的向量，下面一步就是計算這個n維向量中每一個元素的sigmoid：

__global__ void vec_sigmoid(float * d, int num) {
  const int tid = (blockIdx.x * blockDim.x) + threadIdx.x;
  if (tid < num) {
    if(d[tid] > 10.0) d[tid] = 1.0;
    else if(d[tid] < -10.0) d[tid] = 0.0;
    else d[tid] = 1.0 / (1.0 + exp(-1.0 * d[tid]));
  }
}

最后一步就是更新w向量，這里我們不考慮正則化，函數(shù)如下：

__global__ void grad(float * val, int * row_ind, int *col_ind, float * mat_err,
             int nnz, float *act, float *label, 
             float *w, float learning_rate) {
  const int tid = (blockIdx.x * blockDim.x) + threadIdx.x;
  if (tid < nnz) {
    int r = row_ind[tid];
    int c = col_ind[tid];
    float v = val[tid];
    //mat_err是為了后面打印訓練集的誤差而紀錄的
    mat_err[tid] = abs(label[r] - act[r]);
    float err = v * (label[r] - act[r]);
    atomicAdd(&w[c], learning_rate * err);
  }
}

整個LR的函數(shù)如下：

void lr(const vector< vector< pair<int, float> > > & data, 
    const vector<float> & label,
    CooMatrixHost * coo_mat_host, 
    CooMatrix * coo_mat,
    float * w, int ncol, int batch) {
  vec2coo(data, coo_mat_host, coo_mat);
  CUDA_CALL(cudaMemcpyAsync(coo_mat->label, label.data(), sizeof(float) * label.size(), cudaMemcpyHostToDevice, stream));
  CUDA_CALL(cudaMemset(coo_mat->act, 0, sizeof(float) * data.size()));

  int shared_memory_usage = 1;
  int num_blocks = ((coo_mat->nnz + (NUM_THREADS - 1)) / NUM_THREADS);
  dot<<<num_blocks, NUM_THREADS, shared_memory_usage, stream>>>(coo_mat->val,
                                coo_mat->row_ind,
                                coo_mat->col_ind,
                                coo_mat->nnz, 
                                coo_mat->act, w);
  
  num_blocks = ((data.size() + (NUM_THREADS - 1)) / NUM_THREADS);
  vec_sigmoid<<<num_blocks, NUM_THREADS, shared_memory_usage, stream>>>(coo_mat->act, data.size());
  
  num_blocks = ((coo_mat->nnz + (NUM_THREADS - 1)) / NUM_THREADS);
  grad<<<num_blocks, NUM_THREADS, shared_memory_usage, stream>>>(coo_mat->val,
                                 coo_mat->row_ind,
                                 coo_mat->col_ind,
                                 coo_mat->err,
                                 coo_mat->nnz, 
                                 coo_mat->act,
                                 coo_mat->label, 
                                 w, 0.01);
  if (batch % 10000 == 0){
    float * err = (float*) malloc(sizeof(float) * coo_mat->nnz);
    CUDA_CALL(cudaMemcpyAsync(err, coo_mat->err, sizeof(float) * coo_mat->nnz, cudaMemcpyDeviceToHost, stream));
    float total = 0.;
    for(int i = 0; i < coo_mat->nnz; i++) total += err[i];
    cout << total / (float) coo_mat->nnz << endl;
  }
}

整個程序并不復雜，所有代碼見這里。這里也沒有考慮多線程和多卡。其中優(yōu)化的點有幾個：

1. 樣本是在CPU內(nèi)存中生成的，而計算是在GPU中完成的。所以就涉及到CPU向GPU內(nèi)存拷貝的問題。我們是在vec2coo函數(shù)中完成這一步的。
2. 因為不考慮多線程和多卡，因此CPU的內(nèi)存可以預先分配好（zeroCooMatrixHost)。GPU的內(nèi)存也可以事先分配好(zeroCooMatrix)。否則malloc和cudaMalloc將是最耗時的函數(shù)。
3. 使用stream，cudaMemoryAsync。

這里還有一些沒有考慮到的優(yōu)化點，后面需要再試試。

1. 多線程
2. 多個stream
3. 多個卡

因為剛開始寫CUDA的程序，讀者發(fā)現(xiàn)這個代碼有任何問題請發(fā)issue告訴我。