テンソルの各要素にPytorchの学習済みネットワークを適用したい(GPU上で)

やりたいこと

  • 例えば, 以下のような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などになってくると実装がめんどくさくなると思われる. もっといい方法をご存知でしたら教えて下さい.

Leave a Reply

CAPTCHA