こんにちは，ハカルスでインターンをしている，エッジエンジニアの岸本です．NVIDIAから発売された最近話題の高性能ボードコンピュータ，Jetson nanoを触ってみました．Jetson nanoは4CoreのArm CPUと128基のCUDA Coreを搭載した高性能なEdgeデバイスです．そこで今回はこのCUDA Coreを利用しLasso回帰を行ってみました．

1. Jetson nanoとその開発環境

Jetson nanoはNVIDIAの開発したMaxwell CUDA Coreを128基搭載しており，並列に大量の演算を行うことができます．

CUDAはNVIDIAのGPUの並列計算のプラットフォームで，NVIDIAのGPUを持っていれば手軽に利用することができます．

1.1 開発環境/環境構築について

Jetson nano / Ubuntu 18.04 LTS

Jetson nanoへのOSのインストールはこちらを参照しました．

CUDAのコンパイラであるnvccを使うため，/usr/local/cuda-10.0/binにPATHを通します．.bashrcにexport PATH=/usr/local/cuda-10.0/binと追記します．コードのコンパイルは以下のコマンドで行いました．

nvcc main.cu -O3 –generate-code arch=compute_53,code=sm_53 -lcuda -lcublas_static -lcusolver_static -lcusparse_static

2. ADMMの実装

Lasso のアルゴリズム実装の一つにADMMというものがあり，今回はこのアルゴリズムを使用します．L1ノルムを最適化することで，結果として得られる線形回帰モデルの重みがスパースになります．今回の実装はGitHubで公開しています．実装はハカルスがメンテナンスしているオープンソースのライブラリspm-image を参考にしました．

逆行列計算と行列積などの一般的な線形代数の計算はCUDAのライブラリで既に用意されているため，それを用いました．使ったものは以下になります．

cuBLAS
- cublasSnrm2
  ベクトルのL2ノルムを求める．
- cublasSasum
  ベクトルの合計を求める．
- cublasSgemv
  2次元行列と1次元ベクトルの行列積を求める．
- cublasSgemm
  2次元行列と2次元行列の行列積を求める．

cuSOLVER
- cusolverDnSgetrf
  LU分解を行う関数．
- cusolverDnSgetrs
  LU分解された行列を用いて元の行列の逆行列を求める．

学習のループの部分においては実装するにあたり参考にしたspm-image/admm.pyを見てもらうとわかるのですが，spm-imageではz_kとh_kの演算を別々にしておりメモリ帯域の効率が悪いので節約のため一つのCUDAカーネルにまとめました．

以下のようなコードになります．

__device__ float _soft_threshold(float x, float thresh){
    if(x>thresh){
        return x - thresh;
    }else if(x < -thresh){ return x + thresh; }else{ return 0; } } __global__ void _uzh(float *Dw_t, float *h_k, float *z_k, float *sub_z_h_k, float threshold, int n) { int index = blockIdx.x * blockDim.x + threadIdx.x; if(index>=n) return;
    float dw = Dw_t[index];
    float h = h_k[index];
    float z = _soft_threshold(dw + h, threshold);
    z_k[index] = z;
    h += dw - z;
    h_k[index] = h;
    sub_z_h_k[index] = z - h;
}

void _update_z_h(float *Dw_t, float *h_k, float *z_k, float *sub_z_h_k, float threshold, int n)
{
    _uzh<<>>(Dw_t, h_k, z_k, sub_z_h_k, threshold, n);
}

3. ADMMの実行

3.1 学習に用いたデータセット

ボストンの住宅価格のデータセットを用いました．以下のデータを1~13のデータを入力し14・MEDVの予測をするモデルを学習させます．

CRIM – per capita crime rate by town

ZN – proportion of residential land zoned for lots over 25,000 sq.ft.

INDUS – proportion of non-retail business acres per town.

CHAS – Charles River dummy variable (1 if tract bounds river; 0 otherwise)

NOX – nitric oxides concentration (parts per 10 million)

RM – average number of rooms per dwelling

AGE – proportion of owner-occupied units built prior to 1940

DIS – weighted distances to five Boston employment centres

RAD – index of accessibility to radial highways

TAX – full-value property-tax rate per $10,000

PTRATIO – pupil-teacher ratio by town

B – 1000(Bk – 0.63)^2 where Bk is the proportion of blacks by town

LSTAT – % lower status of the population

MEDV – Median value of owner-occupied homes in $1000’s

3.2 学習の実行

プログラムのコンパイルと実行に使ったコマンドは以下になります．

nvcc lasso.cu -O3 –generate-code arch=compute_53,code=sm_53 -lcuda -lcublas_static -lcusolver_static -lcusparse_static

./a.out

学習器にデータを入れ学習を実行したら学習したスパースベクトルをコンソールに表示します．実行した結果が下の画像になります．表示されているベクトルより，学習したベクトルがスパースになっていることが分かります．

3.3 実行時間の比較

手元のノートパソコン(CPU:i7 6500U)で動かすと16msでJetson nanoで動かすと87msかかりました．かなり遅くなっていることが分かります．学習に使ったデータの規模が小さくGPUの強みを活かせてないようです．次項でもっと大きなデータを扱ってみます．

4. Fused Lassoの実装

Fused Lasso は隣り合う要素同士の差を用いた罰則項を設定することで、隣り合う要素間の変化を抑制することができます．今回はGeneralized LassoからFused LassoやTrend Filteringへ簡単に拡張できるよう，Dに関する演算を構造体にしてadmmに渡すように実装しました．

Dに関する演算はspm-imageの実装においてD.dot(w_k)，D.T.dot(D)，inv_matrix.dot(rho * D.T)の3つあり，Dは2重対角行列でスパースなベクトルのため通常の行列積を用いるのは効率が悪く自前で実装しました．

Generalized Lassoの拡張法に関してはハカルスのスパースモデリングのエバンジェリスト増井さんによる発表資料もありますので、ご興味ありましたらそちらもご覧ください．

4.1. 学習させるデータ

ここでは矩形波にノイズを乗せ，そこからノイズ除去するというタスクを用意しました．

4.2. 学習の実行

プログラムのコンパイルと実行に使ったコマンドは以下になります．

nvcc fused_lasso.cu -O3 –generate-code arch=compute_53,code=sm_53 -lcuda -lcublas_static -lcusolver_static -lcusparse_static

./a.out

学習結果をmatplotlibで表示したものが下図になります．1080次元のノイズを乗せた矩形波ベクトル(青)をFused Lassoを用いてspm-image(橙)と今回の実装(緑)でデノイズしました．左側が少しずれていますがこれはspm-imageがdouble，CUDAがfloatで演算したためではないかと考えます．