やりたいこと

• 例えば, 以下のような3次元の物体を表示するpytorchテンソルaがあったとする
  a.shape = (x方向の画素数, y方向の画素数, z方向の画素数, n個の特徴量)
• 個々の座標(x’, y’, z’)に対してn個の特徴量をinputとするpytorchの学習済みネットワーク
  f(n個の特徴量) = 出力
  を適応したい。なおfはtorch.nn.Linearを積み重ねたシンプルなネットワークとする
• numpyでは便利なapply_along_axisという関数が用意されており, 以下の要領で任意の関数fを個々の要素に適用することができる。しかしPytorchに類似の機能をもつ関数は実装されてない

np.apply_along_axis(f, 3, a) Code language: CSS (css)

• numpyのGPUバージョンcupyにはcupy.apply_along_axisがある. 一見これを使えばやりたいことができるように思える.
• しかしcupyはあくまでcupy tensorしかインプットとアウトプットにとることができない
  • つまりPytorchの学習済みネットワークをcp.apply_along_axisに渡す場合, cp.asnumpy(torch.tensor)といった変換処理をかます必要がある
  • しかし要素ごとにこの変換処理が入るので激烈に遅い。なんならCPU上でnp.apply_along_axisを使ったほうが早いくらい

• => なんとかGPU上で素早く処理できないのか？

結論：cupy.ElementwiseKernelを使って自分でcuda関数を自作する

• cupy.elementwisekernelはcupyで任意の自作コードをいい感じにcuda対応してくれるいい感じの関数である（雑）
• ただし例のごとくcupy.tensorしかインプットとアウトプットにとることができない
• 従って自分でtorchの学習済みモデルから係数をぶっこ抜いてきてcupy.elementwisekernelに渡す.
• 今回は以下の関数を自作する.
  • 学習済みネットワークの係数とある3次元画像をインプットとして
  • 学習済みネットワークの出力をアウトプットする
• cupy.elementwisekernelの使い方はCuPyのElementwiseKernelで楽にGPUの恩恵を受けるを参考にした.

1. 学習済みのPytorchモデルからlayer層のbiasとweightを抽出してcupy.tensorにする

先にも述べたようにまずは従って自分でtorchの学習済みモデルから係数をぶっこ抜いてきてcupy.tensor形式に変換してやる必要がある.

単純化するために今回は学習済みモデルは2層のレイヤーから成るシンプルなものだとする. learned_modelは学習済みのモデルを指す. モデル内において2層のレイヤーはlinear1およびlinear2という名前で定義されている.

この場合, 以下のコードでbiasとweightを抽出できる. 勾配は要らないのでdetachする.

layer1_weight = cp.asarray(learned_model.linear_layer.linear1.weight.detach())
layer1_bias = cp.asarray(learned_model.linear_layer.linear1.bias.detach())
layer2_weight = cp.asarray(learned_model.linear_layer.linear2.weight.detach())
layer2_bias = cp.asarray(learned_model.linear_layer.linear2.bias.detach())

2. cp.ElementwiseKernelを使って自作のcuda関数を定義する

次にcp.ElementwiseKernelを使って以下のような関数を自作する.

• 学習済みネットワークの係数と3次元画像をインプットとして
• 学習済みネットワークの出力をアウトプットする

変数の意味を以下に示しておく.

• input_matrix: 以下のような任意の次元数のcupy.tensor. 最後は特徴量でければならない.
  • (x方向の画素数, y方向の画素数, z方向の画素数, n個の特徴量)
  • (x方向の画素数, y方向の画素数, n個の特徴量)

• n_of_current_layer: 現在の層のニューロン数
• n_of_next_layer: 次の層のニューロン数
• weight: 現在の層のweight
• bias: 現在の層のbias
• output_matrix: 学習済みモデルにinput_marixを渡したときの出力. 形式はもちろんcupy.tensor

elementwise_linear_layer = \
    cp.ElementwiseKernel(
            in_params='raw float32 input_matrix, int16 n_of_current_layer, int16 n_of_next_layer, raw float32 weight, raw float32 bias',
            out_params='raw float32 output_matrix',
            operation= \
                '''
                float elementwise_product;
                for (int r = 0; r < n_of_next_layer; r++){
                    elementwise_product = 0;
                    for (int c = 0; c < n_of_current_layer; c++){
                        elementwise_product += weight[r * n_of_current_layer + c] * input_matrix[i * n_of_current_layer + c];
                    }
                    output_matrix[i * n_of_next_layer + r] = elementwise_product + bias[r];
                }
                ''',
            name='elementwise_linear_layer')Code language: PHP (php)

3. 自作関数に1.で抽出した要素をぶっこむ

今回は2層レイヤーから成るので, 一層目のアウトプットに相当するlayer1_resultを適宜アウトプットを次のレイヤーに渡している. layer2_resultが最終出力である.

layer1_result = cp.zeros(x方向の画素数, y方向の画素数, z方向の画素数, layer1_weight.shape[0])).astype(cp.float32)
layer2_result = cp.zeros(x方向の画素数, y方向の画素数, z方向の画素数, layer2_weight.shape[0])).astype(cp.float32)


elementwise_linear_layer(
    input_matrix,
    layer1_weight.shape[1],
    layer1_weight.shape[0],
    layer1_weight,
    layer1_bias,
    layer1_result,
    size = (x方向の画素数, y方向の画素数, z方向の画素数)
)

elementwise_linear_layer(
    layer1_result,
    layer2_weight.shape[1],
    layer2_weight.shape[0],
    layer2_weight,
    layer2_bias,
    layer2_result,
    size = (x方向の画素数, y方向の画素数, z方向の画素数)
)

4. あるテンソルに対する実行速度のおおまかな例

• np.apply_along_axis (CPU)
  10秒
• cupy.apply_along_axis (GPU)
  60秒 (毎回変換処理が入るのでくっそおそい)
• 自作関数:
  0.2秒
  => 感涙レベルで早い

今回はモデルが単純だったのでcupy.ElementwiseKernelで自作できたが, これがCNNなどになってくると実装がめんどくさくなると思われる. もっといい方法をご存知でしたら教えて下さい.