PyTorch : Tutorial 上級 : カスタム C++ と CUDA エクステンション (翻訳/解説)
翻訳 : (株)クラスキャット セールスインフォメーション
作成日時 : 05/22/2018 (0.4.0)
* 本ページは、PyTorch Intermidiate Tutorials の – Custom C++ and CUDA Extensions を
動作確認・翻訳した上で適宜、補足説明したものです：
* サンプルコードの動作確認はしておりますが、適宜、追加改変している場合もあります。
* ご自由にリンクを張って頂いてかまいませんが、sales-info@classcat.com までご一報いただけると嬉しいです。
PyTorch はニューラルネットワーク、任意の tensor 代数、データ論争そして他の目的に関連する多すぎるほどの演算子を提供します。けれども、貴方は依然として更なるカスタマイズされた演算子を必要とすることを見い出すかもしれません。例えば、ペーパーで見つけた新しい活性化関数を使用したり、研究の一部で貴方が開発した演算子を実装することを望むかもしれません。
そのようなカスタム演算を PyTorch で統合する最も簡単な方法は ここ で概説されているように関数とモジュールを拡張することによって Python でそれを書くことです。これは Python の通常の表現能力に加えて自動微分のフルパワー (導関数を書くことから救われます) を与えます。けれども貴方の演算が C++ で実装される方が良いことがあるかもしれません。例えば、貴方のコードが本当に速い必要があるかもしれません、何故ならばそれがモデルで非常に頻繁に呼び出されるか数回の呼び出しでさえも非常に高くつくためです。もう一つのもっともらしい理由としてそれが他の C や C++ ライブラリに依拠するか相互作用することです。そのようなケースに対処するために、PyTorch はカスタム C++ 拡張を書く非常に簡単な方法を提供しています。
C++ エクステンションは、ソース外で定義された i.e. PyTorch バックエンドから分離した PyTorch 演算子をユーザに作成することを可能にするために私達が開発したメカニズムです。このアプローチは native PyTorch 演算が実装されている方法とは異なります。C++ エクステンションは演算を PyTorch のバックエンドと統合することに関連するボイラープレートの大半から解放されることが意図されていますが、その一方で貴方の PyTorch ベースのプロジェクトのための高度な柔軟性も提供します。それにもかかわらず、ひとたび貴方の演算子を C++ エクステンションとして定義したのであれば、それを native PyTorch 関数に変えることは大半はコード体系化の問題であり、演算をアップストリームに寄与することを決めた場合にこれは事後に取り組むことが可能です。
動機とサンプル
このノートの残りは C++ (そして CUDA) エクステンションを書いて使用する実際的なサンプルをウォークスルーします。もし貴方が追い立てられているか、一日の終わりまでにその op をやり終えない場合に誰かが貴方をお払い箱にするのであれば、このセクションをスキップして次のセクションの実装詳細に直行することができます。
先端技術に比較して優れた特質を持つことを見いだせる新しい種類のリカレント・ユニットを貴方が考えついたと仮定しましょう。このリカレント・ユニットは LSTM に類似していますが、それは忘却ゲートが欠落していてその内部活性化関数として Exponential Linear Unit (ELU) を使用する点で異なります。このユニットは決して忘却しないので、それを LLTM, あるいは Long-Long-Term-Memory ユニットと呼称します。
LLTM が vanilla LSTM と異なる2 つの点は私達の目的のために PyTorch の LSTMCell を構成することができないほどに本質的ですので、カスタム・セルを作成しなければならないでしょう。このための最初のそして最も簡単な – そして総てのケースで良い最初のステップになりがちな – アプローチは Python で plain PyTorch 内で望まれる機能を実装することです。このためには、torch.nn.Module をサブクラス化して LLTM の forward パスを実装する必要があります。これはこのように見えるでしょう :
class LLTM(torch.nn.Module):
def __init__(self, input_features, state_size):
super(LLTM, self).__init__()
self.input_features = input_features
self.state_size = state_size
# 3 * state_size for input gate, output gate and candidate cell gate.
# input_features + state_size because we will multiply with [input, h].
self.weights = torch.nn.Parameter(
torch.Tensor(3 * state_size, input_features + state_size))
self.bias = torch.nn.Parameter(torch.Tensor(3 * state_size))
self.reset_parameters()
def reset_parameters(self):
stdv = 1.0 / math.sqrt(self.state_size)
for weight in self.parameters():
weight.data.uniform_(-stdv, +stdv)
def forward(self, input, state):
old_h, old_cell = state
X = torch.cat([old_h, input], dim=1)
# Compute the input, output and candidate cell gates with one MM.
gate_weights = F.linear(X, self.weights, self.bias)
# Split the combined gate weight matrix into its components.
gates = gate_weights.chunk(3, dim=1)
input_gate = F.sigmoid(gates[0])
output_gate = F.sigmoid(gates[1])
# Here we use an ELU instead of the usual tanh.
candidate_cell = F.elu(gates[2])
# Compute the new cell state.
new_cell = old_cell + candidate_cell * input_gate
# Compute the new hidden state and output.
new_h = F.tanh(new_cell) * output_gate
return new_h, new_cell
そしてそれは期待されるように使用できるでしょう :
import torch X = torch.randn(batch_size, input_features) h = torch.randn(batch_size, state_size) C = torch.randn(batch_size, state_size) rnn = LLTM(input_features, state_size) new_h, new_C = rnn(X, (h, C))
当然、もし可能でありもっともらしい場合には、PyTorch を拡張するためにこのアプローチを使用するべきです。NVIDIA cuDNN, Intel MKL あるいは NNPACK のようなライブラリで駆動される、CPU と GPU のための演算の高度に最適化された実装を PyTorch は持ちますから、上のような PyTorch コードはしばしば十分に速いです。けれども、特定の条件下では、何故更なるパフォーマンス改良のための余地があるかを見ることもできます。最も明白な理由は PyTorch は貴方が実装しているアルゴリズムの知識を持たないことです。それは貴方のアルゴリズムを構成するために使用している個々の演算のみを知ります。そのようなものとして、PyTorch は貴方の演算を個々に順々に実行しなければなりません。CUDA カーネルの launch を伴うかもしれない、演算の実装 (or カーネル) への各々の個々の呼び出しはオーバーヘッドのある量を持ちますので、このオーバーヘッドは多くの関数呼び出しに渡り本質的になるかもしれません。更に、私達のコードを実行する Python インタープリタはそれ自身がプログラムをスローダウンするかもしれません。
高速化する明確な方法は従ってパーツを C++ (or CUDA) で書き換えて演算の特定のグループを融合することです。融合は多くの関数の実装を単一の関数に結合することを意味しています、これはより少ないカーネル launch に加えてデータの大域フローの増大した可視性を伴って遂行可能な他の最適化から利益を得ます。
LLTM の融合バージョンを実装するために C++ エクステンションをどのように使用できるかを見てみましょう。PyTorch のバックエンドの大半を駆動している ATen ライブラリを使用してそれを plain C++ で書くことから始めて、そしてそれが私達の Python コードをどれほど簡単に私達に翻訳させるかを見ます。それからGPU が提供する大規模な並列処理から利益を得るためにモデルのパーツを CUDA カーネルに移すことにより更に高速化します。
C++ エクステンションを書く
C++ エクステンションは 2 つのフレーバーがあります : それらは setuptools で “ahead of time” に、または torch.utils.cpp_extension.load() を通して “just in time” にビルドできます。最初のアプローチから始めて後者は後で議論します。
setuptools でビルドする
“ahead of time” フレーバーのためには、私達の C++ コードをコンパイルするために setuptools を使用する setup.py スクリプトを書くことにより C++ エクステンションをビルドします。LLTM のためには、それはこのように単純に見えます :
from setuptools import setup
from torch.utils.cpp_extension import CppExtension, BuildExtension
setup(name='lltm',
ext_modules=[CppExtension('lltm', ['lltm.cpp'])]
cmdclass={'build_ext': BuildExtension})
このコードで、:class:CppExtension は setuptools.Extension まわりの便利なラッパーで正しい include パスを渡してエクステンションの言語を C++ に設定します。同値な vanilla setuptools コードは単純に :
setuptools.Extension( name='lltm', sources=['lltm.cpp'], include_dirs=torch.utils.cpp_extension.include_paths(), language='c++')
BuildExtension は多くの必要な configuration ステップを遂行して混合 C++/CUDA エクステンションの場合に混合コンパイル (= mixed compilation) をチェックして更に管理もします。そしてそれが C++ エクステンションについて当面実際に知る必要があることの総てです！今私達の C++ エクステンションの実装を見てみましょう、これは lltm.cpp に入ります。
C++ Op を書く
C++ で LLTM の実装を始めましょう！backward パスのために必要な 1 つの関数は sigmoid の導関数です。C++ エクステンションを書くときに私達に利用可能な環境全体を議論するために、これは十分に小さいコード・ピースです :
#include <torch/torch.h>
#include <iostream>
at::Tensor d_sigmoid(at::Tensor z) {
auto s = at::sigmoid(z);
return (1 - s) * s;
}
<torch/torch.h> は C++ エクステンションを書くために必要な PyTorch の細々としたもの総てを含むワンストップ・ヘッダです。それは以下を含みます :
- ATen ライブラリ、これは tensor 計算のための主要な API です。
- pybind11, これは私達の C++ コードのための Python バインディングをどのように作成するかです。
- ヘッダは ATen と pybind11 の間の相互作用の詳細を管理します。
d_sigmoid() の実装は ATen API をどのように使用するかを示します。PyTorch の tensor と variable インターフェイスは ATen ライブラリから自動的に生成されますので、私達の Python 実装を多かれ少なかれ 1:1 で C++ に翻訳できます。総ての計算のための基本 datatype は at::Tensor になります。その完全な API は ここ で調べることができます。<iostream> や任意の他の C または C++ ヘッダをインクルードすることができることにも注意してください – C++11 のフルパワーを自由に持ちます。
Forward パス
次に forward パス全体を C++ にポートできます :
#include <vector>
std::vector<at::Tensor> lltm_forward(
at::Tensor input,
at::Tensor weights,
at::Tensor bias,
at::Tensor old_h,
at::Tensor old_cell) {
auto X = at::cat({old_h, input}, /*dim=*/1);
auto gate_weights = at::addmm(bias, X, weights.transpose(0, 1));
auto gates = gate_weights.chunk(3, /*dim=*/1);
auto input_gate = at::sigmoid(gates[0]);
auto output_gate = at::sigmoid(gates[1]);
auto candidate_cell = at::elu(gates[2], /*alpha=*/1.0);
auto new_cell = old_cell + candidate_cell * input_gate;
auto new_h = at::tanh(new_cell) * output_gate;
return {new_h,
new_cell,
input_gate,
output_gate,
candidate_cell,
X,
gate_weights};
}
Backward パス
現時点で、PyTorch の C++ インターフェイスは自動微分をサポートません。これは PyTorch チームが作業しているものですが、まだそれは利用可能ではありません。そういうものとして、forward パスの各入力に関する損失の導関数を計算する、LLTM の backward パスもまた実装しなければなりません、最終的には、素晴らしい Python バインディングを作成するために forward と backward 関数の両者を torch.nn.Function にどんと置きます (= plop)。backward 関数は少しより難解ですので、コードにより深くは入りません (もし興味があれば、Alex Graves の学位論文 はこれについての更なる情報のために良い読み物です) :
// tanh'(z) = 1 - tanh^2(z)
at::Tensor d_tanh(at::Tensor z) {
return 1 - z.tanh().pow(2);
}
// elu'(z) = relu'(z) + { alpha * exp(z) if (alpha * (exp(z) - 1)) < 0, else 0}
at::Tensor d_elu(at::Tensor z, at::Scalar alpha = 1.0) {
auto e = z.exp();
auto mask = (alpha * (e - 1)) < 0;
return (z > 0).type_as(z) + mask.type_as(z) * (alpha * e);
}
std::vector lltm_backward(
at::Tensor grad_h,
at::Tensor grad_cell,
at::Tensor new_cell,
at::Tensor input_gate,
at::Tensor output_gate,
at::Tensor candidate_cell,
at::Tensor X,
at::Tensor gate_weights,
at::Tensor weights) {
auto d_output_gate = at::tanh(new_cell) * grad_h;
auto d_tanh_new_cell = output_gate * grad_h;
auto d_new_cell = d_tanh(new_cell) * d_tanh_new_cell + grad_cell;
auto d_old_cell = d_new_cell;
auto d_candidate_cell = input_gate * d_new_cell;
auto d_input_gate = candidate_cell * d_new_cell;
auto gates = gate_weights.chunk(3, /*dim=*/1);
d_input_gate *= d_sigmoid(gates[0]);
d_output_gate *= d_sigmoid(gates[1]);
d_candidate_cell *= d_elu(gates[2]);
auto d_gates =
at::cat({d_input_gate, d_output_gate, d_candidate_cell}, /*dim=*/1);
auto d_weights = d_gates.t().mm(X);
auto d_bias = d_gates.sum(/*dim=*/0, /*keepdim=*/true);
auto d_X = d_gates.mm(weights);
const auto state_size = grad_h.size(1);
auto d_old_h = d_X.slice(/*dim=*/1, 0, state_size);
auto d_input = d_X.slice(/*dim=*/1, state_size);
return {d_old_h, d_input, d_weights, d_bias, d_old_cell};
}
Python へのバインディング
ひとたび貴方の演算が C++ と Aten で書かれたのであれば、貴方の C++ 関数かクラスを非常に単純な方法で Python にバインドするために pybind11 が使用できます。PyTorch C++ エクステンションのこの部分について貴方が持つ疑問や問題は pybind11 ドキュメント により広く対処されます。
私達のエクステンションのためには、必要なバインディング・コードは 4 行だけにわたります :
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("forward", &lltm_forward, "LLTM forward");
m.def("backward", &lltm_backward, "LLTM backward");
}
ここで注意すべき一つのことはマクロ TORCH_EXTENSION_NAME です。torch エクステンション・ビルドはそれを貴方が setup.py スクリプトでエクステンションとして与える名前として定義します。このケースでは、TORCH_EXTENSION_NAME の値は “lltm” でしょう。これはエクステンションの名前を 2 つの場所 (ビルド・スクリプトと貴方の C++ コード) で維持しなければならないことを回避するためで、何故ならば 2 つの間のミスマッチは扱いにくいものとなり問題を追跡することが困難になります。
貴方のエクステンションを使用する
今 PyTorch でエクステンションをインポートするために設定を行ないます。この時点で、貴方のディレクトリ構造はこのようなものに見えるでしょう :
pytorch/
lltm-extension/
lltm.cpp
setup.py
さて、貴方のエクステンションをビルドしてインストールするために python setup.py install を実行します。これはこのようなものに見えるはずです :
running install running bdist_egg running egg_info writing lltm.egg-info/PKG-INFO writing dependency_links to lltm.egg-info/dependency_links.txt writing top-level names to lltm.egg-info/top_level.txt reading manifest file 'lltm.egg-info/SOURCES.txt' writing manifest file 'lltm.egg-info/SOURCES.txt' installing library code to build/bdist.linux-x86_64/egg running install_lib running build_ext building 'lltm' extension gcc -Wsign-compare -DNDEBUG -g -fwrapv -O3 -Wall -Wstrict-prototypes -fPIC -I~/local/miniconda/lib/python3.6/site-packages/torch/lib/include -I~/local/miniconda/lib/python3.6/site-packages/torch/lib/include/TH -I~/local/miniconda/lib/python3.6/site-packages/torch/lib/include/THC -I~/local/miniconda/include/python3.6m -c lltm.cpp -o build/temp.linux-x86_64-3.6/lltm.o -DTORCH_EXTENSION_NAME=lltm -std=c++11 cc1plus: warning: command line option ‘-Wstrict-prototypes’ is valid for C/ObjC but not for C++ g++ -pthread -shared -B ~/local/miniconda/compiler_compat -L~/local/miniconda/lib -Wl,-rpath=~/local/miniconda/lib -Wl,--no-as-needed -Wl,--sysroot=/ build/temp.linux-x86_64-3.6/lltm.o -o build/lib.linux-x86_64-3.6/lltm.cpython-36m-x86_64-linux-gnu.so creating build/bdist.linux-x86_64/egg copying build/lib.linux-x86_64-3.6/lltm_cuda.cpython-36m-x86_64-linux-gnu.so -> build/bdist.linux-x86_64/egg copying build/lib.linux-x86_64-3.6/lltm.cpython-36m-x86_64-linux-gnu.so -> build/bdist.linux-x86_64/egg creating stub loader for lltm.cpython-36m-x86_64-linux-gnu.so byte-compiling build/bdist.linux-x86_64/egg/lltm.py to lltm.cpython-36.pyc creating build/bdist.linux-x86_64/egg/EGG-INFO copying lltm.egg-info/PKG-INFO -> build/bdist.linux-x86_64/egg/EGG-INFO copying lltm.egg-info/SOURCES.txt -> build/bdist.linux-x86_64/egg/EGG-INFO copying lltm.egg-info/dependency_links.txt -> build/bdist.linux-x86_64/egg/EGG-INFO copying lltm.egg-info/top_level.txt -> build/bdist.linux-x86_64/egg/EGG-INFO writing build/bdist.linux-x86_64/egg/EGG-INFO/native_libs.txt zip_safe flag not set; analyzing archive contents... __pycache__.lltm.cpython-36: module references __file__ creating 'dist/lltm-0.0.0-py3.6-linux-x86_64.egg' and adding 'build/bdist.linux-x86_64/egg' to it removing 'build/bdist.linux-x86_64/egg' (and everything under it) Processing lltm-0.0.0-py3.6-linux-x86_64.egg removing '~/local/miniconda/lib/python3.6/site-packages/lltm-0.0.0-py3.6-linux-x86_64.egg' (and everything under it) creating ~/local/miniconda/lib/python3.6/site-packages/lltm-0.0.0-py3.6-linux-x86_64.egg Extracting lltm-0.0.0-py3.6-linux-x86_64.egg to ~/local/miniconda/lib/python3.6/site-packages lltm 0.0.0 is already the active version in easy-install.pth Installed ~/local/miniconda/lib/python3.6/site-packages/lltm-0.0.0-py3.6-linux-x86_64.egg Processing dependencies for lltm==0.0.0 Finished processing dependencies for lltm==0.0.0
コンパイラの小さいノート: ABI バージョン問題のため、貴方の C++ エクステンションをビルドするために使用するコンパイラは PyTorch がビルドされたコンパイラと ABI-互換でなければなりません。実際には、これは GCC version 4.9 とそれ以上を使用しなければならないことを意味します。Ubuntu 16.04 と他のより最近の Linux ディストリビューションのためには、これは既にデフォルトのコンパイラであるはずです。MacOS 上では、GCC をダウンロードしなければならないでしょう (e.g. これを書いている時点で brew install gcc が GCC 7 を与えるでしょう)。最悪の場合には、PyTorch を貴方のコンパイラでソースからビルドしてから同じコンパイラでエクステンションをビルドできます。
ひとたび貴方のエクステンションがビルドされれば、setup.py スクリプトで指定した名前を使用してそれを Python で単純にインポートすることたできます。
単に torch を最初にインポートすることは確実にしてください、何故ならばこれはダイナミック・リンカが見なければならない幾つかのシンボルを解決するからです :
In [1]: import torch In [2]: import lltm In [3]: lltm.forward Out[3]: <function lltm.PyCapsule.forward>
もし関数かモジュール上で help() を呼びだせば、そのシグネチャが私達の C++ コードにマッチすることを見て取れます :
In[4] help(lltm.forward)
forward(...) method of builtins.PyCapsule instance
forward(arg0: at::Tensor, arg1: at::Tensor, arg2: at::Tensor, arg3: at::Tensor, arg4: at::Tensor) -> List[at::Tensor]
LLTM forward
今は私達の C++ 関数を Python から呼び出すことができるので、それらを PyTorch の第一級オブジェクト (= first class citizens) とするためにそれらを torch.nn.Function と torch.nn.Module でラップすることができます :
import math
import torch
# Our module!
import lltm
class LLTMFunction(torch.nn.Function):
@staticmethod
def forward(ctx, input, weights, bias, old_h, old_cell):
outputs = lltm.forward(input, weights, bias, old_h, old_cell)
new_h, new_cell = outputs[:2]
variables = outputs[1:] + [weights, old_cell]
ctx.save_for_backward(*variables)
return new_h, new_cell
@staticmethod
def backward(ctx, grad_h, grad_cell):
outputs = lltm.backward(
grad_h.contiguous(), grad_cell.contiguous(), *ctx.saved_variables)
d_old_h, d_input, d_weights, d_bias, d_old_cell, d_gates = outputs
return d_input, d_weights, d_bias, d_old_h, d_old_cell
class LLTM(torch.nn.Module):
def __init__(self, input_features, state_size):
super(LLTM, self).__init__()
self.input_features = input_features
self.state_size = state_size
self.weights = torch.nn.Parameter(
torch.Tensor(3 * state_size, input_features + state_size))
self.bias = torch.nn.Parameter(torch.Tensor(3 * state_size))
self.reset_parameters()
def reset_parameters(self):
stdv = 1.0 / math.sqrt(self.state_size)
for weight in self.parameters():
weight.data.uniform_(-stdv, +stdv)
def forward(self, input, state):
return LLTMFunction.apply(input, self.weights, self.bias, *state)
パフォーマンス比較
私達の C++ コードを PyTorch から使用して呼び出せるようになった今、私達の op を C++ で書き換えることからどのくらいのパフォーマンスを得たかを見るために小さいベンチマークを実行することができます。LLTM forward と backward を数回実行して時間を測定します :
import torch
batch_size = 16
input_features = 32
state_size = 128
X = torch.randn(batch_size, input_features)
h = torch.randn(batch_size, state_size)
C = torch.randn(batch_size, state_size)
rnn = LLTM(input_features, state_size)
forward = 0
backward = 0
for _ in range(100000):
start = time.time()
new_h, new_C = rnn(X, (h, C))
forward += time.time() - start
start = time.time()
(new_h.sum() + new_C.sum()).backward()
backward += time.time() - start
print('Forward: {:.3f} us | Backward {:.3f} us'.format(forward * 1e6/1e5, backward * 1e6/1e5))
この投稿の最初に pure Python で書いた元の LLTM でこのコードを実行する場合、(私のマシン上で) 次の数字を得ます :
Forward: 506.480 us | Backward 444.694 us
and with our new C++ version:
Forward: 349.335 us | Backward 443.523 us
forward 関数について本質的なスピードアップを既に見ることができます (30% 以上)。backward 関数についてスピードアップを見れます、大きなものではありませんが。上で書いた backward パスは特に最適化されておらず確実に改善できるでしょう。また、PyTorch の自動微分エンジンは自動的に計算グラフを並列化でき、演算全体のより効率的なフローを使用するかもしれません、そしてまた C++ で実装されていますので、高速であることが期待されます。それでも、これは良いスタートです。
GPU デバイス上のパフォーマンス
PyTorch の ATen バックエンドについての素晴らしい事実はそれが貴方がその上で実行している計算デバイスを抽象化することです。これは CPU のために書いた同じコードが GPU 上でも動作できることを意味し、個々の演算は対応する GPU-最適化実装へディスパッチされるでしょう。(mm や admm のような) 行列乗算のような確かな演算について、これは大きな成功です。CUDA tensor で私達の C++ コードを実行することからどれだけのパフォーマンスを得られるかを見てみましょう。私達の実装への変更は必要ありません、Python から .cuda() で単に tensor を GPU メモリに移動する必要があるだけです :
mport torch
assert torch.cuda.is_available()
batch_size = 16
input_features = 32
state_size = 128
# Note the .cuda() calls here
X = torch.randn(batch_size, input_features).cuda()
h = torch.randn(batch_size, state_size).cuda()
C = torch.randn(batch_size, state_size).cuda()
rnn = LLTM(input_features, state_size).cuda()
forward = 0
backward = 0
for _ in range(100000):
start = time.time()
new_h, new_C = rnn(X, (h, C))
forward += time.time() - start
start = time.time()
(new_h.sum() + new_C.sum()).backward()
backward += time.time() - start
print('Forward: {:.3f} us | Backward {:.3f} us'.format(forward * 1e6/1e5, backward * 1e6/1e5))
もう一度 plain PyTorch コードを私達の C++ バージョンと比較し、今は両者を CUDA デバイス上で実行すると、再びパフォーマンス向上を見ます。Python/PyTorch について :
Forward: 187.719 us | Backward 410.815 us
そして C++/ATen:
Forward: 149.802 us | Backward 393.458 us
それは non-CUDA コードに比較して全体として大きなスピードアップです。けれども、カスタム CUDA カーネルを書くことにより私達の C++ コードから更なるパフォーマンスさえ引き出すことができます、間もなくこれに飛び込みます。その前に、貴方の C++ エクステンションをビルドするもう一つの方法を議論しましょう。
JIT コンパイル・エクステンション
前に、C++ エクステンションをビルドする 2 つの方法があることに言及しました : setuptools か just in time (JIT) を使用します。前者はカバーしましたので、後者を詳述しましょう。JIT コンパイル・メカニズムは torch.utils.cpp_extension.load() と呼ばれる PyTorch の API の単純な関数を呼び出すことにより貴方のエクステンションを on the fly でコンパイルしてロードする方法を提供します。LLTM のためには、これはこのように単純に見えるでしょう :
from torch.utils.cpp_extension import load lltm = load(name="lltm", sources=["lltm.cpp"])
ここで、この関数に setuptools のためのものと同じ情報を提供します。background では、これは以下を行ないます :
- 一時ディレクトリ /tmp/torch_extensions/lltm を作成する、
- Ninja ビルドファイルをその一時ディレクトリに吐く、
- 貴方のソースファイルを共有ライブラリにコンパイルする、
- この共有ライブラリを Python モジュールとしてインポートする。
実際に、verbose=True を cpp_extension.load() に渡す場合、そのプロセスについて知らされるでしょう :
Using /tmp/torch_extensions as PyTorch extensions root... Creating extension directory /tmp/torch_extensions/lltm... Emitting ninja build file /tmp/torch_extensions/lltm/build.ninja... Building extension module lltm... Loading extension module lltm...
結果としての Python モジュールは setuptools により生成されたものと正確に同じですが、分離した setup.py ビルドファイルを維持しなければならないという要件は除去します。もし貴方のセットアップがより複雑で setuptools のフルパワーを必要とするのであれば、貴方自身の setup.py を書くことができます – しかし多くの場合この JIT テクニックで十分です。最初にこの行を通して実行すると、それはある程度時間がかかります、何故ならばエクステンションがバックラウンドでコンパイルされているからです。貴方のソースをビルドするために Ninja ビルドシステムを使用していますので、再コンパイルはインクリメンタルでエクステンションのソースファイルを変更しなかった場合にはエクステンションの再ロードは Python モジュールを 2 度目に実行するときは高速で低いオーバーヘッドを持ちます。
混合 C++/CUDA エクステンションを書く
実際に私達の実装を次のレベルにするために、カスタム CUDA カーネルで forward と backward パスの一部を手書きすることができます。LLTM については、これは特に効果的であるという見通しを持ちます、何故ならば巨大な数のシーケンスの pointwise な演算があるからで、これらは総て単一の CUDA カーネルで融合して並列化できます。エクステンション・メカニズムを使用してどのようにそのような CUDA カーネルを書いてそれを PyTorch に統合するかを見ましょう。
CUDA エクステンションを書くための一般的なストラテジーは最初に C++ ファイルを書くことです、これは Python から呼び出される関数を定義してそれらの関数を pybind11 で Python にバインドします。更に、このファイルはまた CUDA (.cu) ファイルで定義される関数を宣言します。それから C++ 関数は幾つかのチェックを行なって最終的にはその呼び出しを CUDA 関数に forward します。CUDA ファイルで、実際の CUDA カーネルを書きます。それから cpp_extension パッケージは gcc のような C++ コンパイラで C++ ソースをそして NVIDIA の nvcc コンパイラで CUDA ソースをコンパイルすることをケアします。これは各コンパイラがそれがコンパイルするのに最善と知るファイルをケアすることを確かなものにします。最終的に、それらはPython コードから利用可能な一つの共有ライブラリにリンクされます。
C++ ファイルから始めます、これを lltm_cuda.cpp と呼びましょう、例えば :
#include <torch/torch.h>
#include <vector>
// CUDA forward declarations
std::vector<at::Tensor> lltm_cuda_forward(
at::Tensor input,
at::Tensor weights,
at::Tensor bias,
at::Tensor old_h,
at::Tensor old_cell);
std::vector<at::Tensor> lltm_cuda_backward(
at::Tensor grad_h,
at::Tensor grad_cell,
at::Tensor new_cell,
at::Tensor input_gate,
at::Tensor output_gate,
at::Tensor candidate_cell,
at::Tensor X,
at::Tensor gate_weights,
at::Tensor weights);
// C++ interface
#define CHECK_CUDA(x) AT_ASSERT(x.type().is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) AT_ASSERT(x.is_contiguous(), #x " must be contiguous")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
std::vector<at::Tensor> lltm_forward(
at::Tensor input,
at::Tensor weights,
at::Tensor bias,
at::Tensor old_h,
at::Tensor old_cell) {
CHECK_INPUT(input);
CHECK_INPUT(weights);
CHECK_INPUT(bias);
CHECK_INPUT(old_h);
CHECK_INPUT(old_cell);
return lltm_cuda_forward(input, weights, bias, old_h, old_cell);
}
std::vector<at::Tensor> lltm_backward(
at::Tensor grad_h,
at::Tensor grad_cell,
at::Tensor new_cell,
at::Tensor input_gate,
at::Tensor output_gate,
at::Tensor candidate_cell,
at::Tensor X,
at::Tensor gate_weights,
at::Tensor weights) {
CHECK_INPUT(grad_h);
CHECK_INPUT(grad_cell);
CHECK_INPUT(input_gate);
CHECK_INPUT(output_gate);
CHECK_INPUT(candidate_cell);
CHECK_INPUT(X);
CHECK_INPUT(gate_weights);
CHECK_INPUT(weights);
return lltm_cuda_backward(
grad_h,
grad_cell,
new_cell,
input_gate,
output_gate,
candidate_cell,
X,
gate_weights,
weights);
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("forward", &lltm_forward, "LLTM forward (CUDA)");
m.def("backward", &lltm_backward, "LLTM backward (CUDA)");
}
見て取れるように、それは大部分はボイラープレートで、チェックと CUDA ファイルで定義する関数へのフォワードです。このファイルを lltm_cuda_kernel.cu (.cu エクステンションに注意！) と名前付けます。NVCC は合理的に C++11 をコンパイルできて、依然として ATen と C++ 標準ライブラリが利用可能です (しかし not torch.h)。setuptools は同じ名前でしかし異なるエクステンションのファイルを処理できませんので、JIT メソッドの代わりに setup.py メソッドを使用する場合、CUDA ファイルに C++ ファイルとは異なる名前を与えなければなりません (JIT メソッドのためには、lltm.cpp と lltm.cu は上手く動作します)。
このファイルがどのように見えるか覗いてみましょう :
#include <ATen/ATen.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <vector>
template <typename scalar_t>
__device__ __forceinline__ scalar_t sigmoid(scalar_t z) {
return 1.0 / (1.0 + exp(-z));
}
ここで丁度説明したヘッダを見ます、そして __device__ と __forceinline__ のような CUDA-特有の宣言と exp のような関数を使用しています。必要となる幾つかの更なるヘルパー関数を続けましょう :
template <typename scalar_t>
__device__ __forceinline__ scalar_t d_sigmoid(scalar_t z) {
const auto s = sigmoid(z);
return (1.0 - s) * s;
}
template <typename scalar_t>
__device__ __forceinline__ scalar_t d_tanh(scalar_t z) {
const auto t = tanh(z);
return 1 - (t * t);
}
template <typename scalar_t>
__device__ __forceinline__ scalar_t elu(scalar_t z, scalar_t alpha = 1.0) {
return fmax(0.0, z) + fmin(0.0, alpha * (exp(z) - 1.0));
}
template <typename scalar_t>
__device__ __forceinline__ scalar_t d_elu(scalar_t z, scalar_t alpha = 1.0) {
const auto e = exp(z);
const auto d_relu = z < 0.0 ? 0.0 : 1.0;
return d_relu + (((alpha * (e - 1.0)) < 0.0) ? (alpha * e) : 0.0);
}
今実際に関数を実装するためには、再び 2 つのことが必要です : 明示的に手書きを望まない演算と CUDA カーネルへの呼び出しを遂行する 1 つの関数、そしてそれからスピードアップを望む部分のための実際の CUDA カーネルです。forward パスについて、最初の関数はこのように見えるはずです :
std::vector<at::Tensor> lltm_cuda_forward(
at::Tensor input,
at::Tensor weights,
at::Tensor bias,
at::Tensor old_h,
at::Tensor old_cell) {
auto X = at::cat({old_h, input}, /*dim=*/1);
auto gates = at::addmm(bias, X, weights.transpose(0, 1));
const auto batch_size = old_cell.size(0);
const auto state_size = old_cell.size(1);
auto new_h = at::zeros_like(old_cell);
auto new_cell = at::zeros_like(old_cell);
auto input_gate = at::zeros_like(old_cell);
auto output_gate = at::zeros_like(old_cell);
auto candidate_cell = at::zeros_like(old_cell);
const int threads = 1024;
const dim3 blocks((state_size + threads - 1) / threads, batch_size);
AT_DISPATCH_FLOATING_TYPES(gates.type(), "lltm_forward_cuda", ([&] {
lltm_cuda_forward_kernel<scalar_t><<<blocks, threads>>>(
gates.data<scalar_t>(),
old_cell.data<scalar_t>(),
new_h.data<scalar_t>(),
new_cell.data<scalar_t>(),
input_gate.data<scalar_t>(),
output_gate.data<scalar_t>(),
candidate_cell.data<scalar_t>(),
state_size);
}));
return {new_h, new_cell, input_gate, output_gate, candidate_cell, X, gates};
}
ここでの興味の主要な点は AT_DISPATCH_FLOATING_TYPES マクロとカーネル launch (<<<...>>> で示されます) です。ATen がデバイスと扱う tensor のデータ型から離れて抽象化する一方で、実行時、tensor は依然として具体的なデバイス上の具体的な型のメモリにより支援されます。そのようなものとして、実行時に tensor がどのような型であるかを決定する方法が必要でそしてそれから対応する正しい型シグネチャで関数を選択的に呼び出します。手動で成されたとき、これはこのようなものとして (概念的に) 見えます :
switch (tensor.type().scalarType()) {
case at::ScalarType::Double:
return function<double>(tensor.data<double>());
case at::ScalarType::Float:
return function<float>(tensor.data<float>());
...
}
AT_DISPATCH_FLOATING_TYPES の目的はこのディスパッチを私達のためにケアすることです。それは型 (私達のケースでは gates.type())、名前 (エラーメッセージのため) そして lambda 関数を取ります。lambda 関数の内部では、型エイリアス scalar_t が利用可能でそれは tensor がそのコンテキストで実行時に実際に在る型として定義されます。そのようなものとして、(CUDA カーネルがそうである) テンプレート関数を持つ場合、それをこの scalar_t エイリアスでインスタンス化できて、正しい関数が呼び出されます。この場合、tensor のデータ・ポインタをその scalar_t 型のポインタとして取得することもまた望みます。もし貴方が浮動小数点型 (Float と Double) だけでなく総ての型に渡りディスパッチすることを望んだ場合には、AT_DISPATCH_ALL_TYPES を使用できます。
幾つかの演算を plain ATen で遂行していることに注意してください。これらの演算は依然として GPU 上で動作しますが、ATen のデフォルト実装を使用しています。これは意味があります、何故ならば ATen は私達自身で実装して改良するのが非常により困難な行列乗算 (e.g. addmm) や畳み込みのようなものに対して高度に最適化されたルーチンを使用するからです。
カーネル launch 自身に関しては、ここでは各 CUDA ブロックが 1024 スレッドを持ち、そして GPU グリッド全体は私達の行列を成分毎に 1 スレッドで満たすのに必要とされるだけの 1 x 1024 スレッドのブロックに分割されるように指定します。例えば、私達のステート・サイズが 2048 でバッチサイズが 4 であったならば、それぞれ 1024 スレッドを持つ総計 4 x 2 = 8 ブロックを launch したでしょう。もし貴方が CUDA 「ブロック」や「グリッド」を前に聞いたことがないのであれば、introductory read about CUDA が役に立つかもしれません。
実際の CUDA カーネルは非常に単純です (以前に GPU をプログラムしたことさえあれば) :
template <typename scalar_t>
__global__ void lltm_cuda_forward_kernel(
const scalar_t* __restrict__ gates,
const scalar_t* __restrict__ old_cell,
scalar_t* __restrict__ new_h,
scalar_t* __restrict__ new_cell,
scalar_t* __restrict__ input_gate,
scalar_t* __restrict__ output_gate,
scalar_t* __restrict__ candidate_cell,
size_t state_size) {
const int column = blockIdx.x * blockDim.x + threadIdx.x;
const int index = blockIdx.y * state_size + column;
const int gates_row = blockIdx.y * (state_size * 3);
if (column < state_size) {
input_gate[index] = sigmoid(gates[gates_row + column]);
output_gate[index] = sigmoid(gates[gates_row + state_size + column]);
candidate_cell[index] = elu(gates[gates_row + 2 * state_size + column]);
new_cell[index] =
old_cell[index] + candidate_cell[index] * input_gate[index];
new_h[index] = tanh(new_cell[index]) * output_gate[index];
}
}
ここで第一に興味深いものは gate 行列の各々の個別の成分のためのこれらの pointwise な演算の総てを全体的に並列で計算することができることです。百万の成分に渡りシリアルに巨大な for ループでこれを行わなければならないことを想像すれば、何故これが遥かに速いかが分かるでしょう。
backward パスは同じパターンを追いますのでそれを更に詳述はしません :
template <typename scalar_t>
__global__ void lltm_cuda_backward_kernel(
scalar_t* __restrict__ d_old_cell,
scalar_t* __restrict__ d_gates,
const scalar_t* __restrict__ grad_h,
const scalar_t* __restrict__ grad_cell,
const scalar_t* __restrict__ new_cell,
const scalar_t* __restrict__ input_gate,
const scalar_t* __restrict__ output_gate,
const scalar_t* __restrict__ candidate_cell,
const scalar_t* __restrict__ gate_weights,
size_t state_size) {
const int column = blockIdx.x * blockDim.x + threadIdx.x;
const int index = blockIdx.y * state_size + column;
const int gates_row = blockIdx.y * (state_size * 3);
if (column < state_size) {
const auto d_output_gate = tanh(new_cell[index]) * grad_h[index];
const auto d_tanh_new_cell = output_gate[index] * grad_h[index];
const auto d_new_cell =
d_tanh(new_cell[index]) * d_tanh_new_cell + grad_cell[index];
d_old_cell[index] = d_new_cell;
const auto d_candidate_cell = input_gate[index] * d_new_cell;
const auto d_input_gate = candidate_cell[index] * d_new_cell;
const auto input_gate_index = gates_row + column;
const auto output_gate_index = gates_row + state_size + column;
const auto candidate_cell_index = gates_row + 2 * state_size + column;
d_gates[input_gate_index] =
d_input_gate * d_sigmoid(gate_weights[input_gate_index]);
d_gates[output_gate_index] =
d_output_gate * d_sigmoid(gate_weights[output_gate_index]);
d_gates[candidate_cell_index] =
d_candidate_cell * d_elu(gate_weights[candidate_cell_index]);
}
}
std::vector<at::Tensor> lltm_cuda_backward(
at::Tensor grad_h,
at::Tensor grad_cell,
at::Tensor new_cell,
at::Tensor input_gate,
at::Tensor output_gate,
at::Tensor candidate_cell,
at::Tensor X,
at::Tensor gate_weights,
at::Tensor weights) {
auto d_old_cell = at::zeros_like(new_cell);
auto d_gates = at::zeros_like(gate_weights);
const auto batch_size = new_cell.size(0);
const auto state_size = new_cell.size(1);
const int threads = 1024;
const dim3 blocks((state_size + threads - 1) / threads, batch_size);
AT_DISPATCH_FLOATING_TYPES(X.type(), "lltm_forward_cuda", ([&] {
lltm_cuda_backward_kernel<scalar_t><<<blocks, threads>>>(
d_old_cell.data<scalar_t>(),
d_gates.data<scalar_t>(),
grad_h.contiguous().data<scalar_t>(),
grad_cell.contiguous().data<scalar_t>(),
new_cell.contiguous().data<scalar_t>(),
input_gate.contiguous().data<scalar_t>(),
output_gate.contiguous().data<scalar_t>(),
candidate_cell.contiguous().data<scalar_t>(),
gate_weights.contiguous().data<scalar_t>(),
state_size);
}));
auto d_weights = d_gates.t().mm(X);
auto d_bias = d_gates.sum(/*dim=*/0, /*keepdim=*/true);
auto d_X = d_gates.mm(weights);
auto d_old_h = d_X.slice(/*dim=*/1, 0, state_size);
auto d_input = d_X.slice(/*dim=*/1, state_size);
return {d_old_h, d_input, d_weights, d_bias, d_old_cell, d_gates};
C++/CUDA 演算を PyTorch と統合する
CUDA-enabled op の PyTorch との統合は再度非常に率直です。setup.py スクリプトを書くことを望むのであれば、それはこのように見えるでしょう :
from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
setup(
name='lltm',
ext_modules=[
CUDAExtension('lltm_cuda', [
'lltm_cuda.cpp',
'lltm_cuda_kernel.cu',
])
],
cmdclass={
'build_ext': BuildExtension
})
CppExtension() の代わりに、今は CUDAExtension() を使用します。.cu ファイルを .cpp ファイルと一緒に単に指定できます - ライブラリはこれが伴う総ての面倒を貴方のためにケアします。JIT メカニズムはより単純でさえあります :
from torch.utils.cpp_extension import load lltm = load(name='lltm', sources=['lltm_cuda.cpp', 'lltm_cuda_kernel.cu'])
パフォーマンス比較
私達の希望は、CUDA による私達のコードの pointwise 演算の並列化と融合が私達の LLTM のパフォーマンスを改良することした。それが事実であることを見てみましょう。ベンチマークの実行のために先にリストしたコードを実行することができます。私達の前の最速のバージョンは CUDA-based C++ コードでした :
Forward: 149.802 us | Backward 393.458 us
そして今私達のカスタム CUDA カーネルでは :
And now with our custom CUDA kernel:
Forward: 129.431 us | Backward 304.641 us
更にパフォーマンスが上がります！
結論
貴方は今 PyTorch の C++ エクステンション・メカニズムとそれらを使用するための動機の良い概要を持ちあわせているはずです。このノートで示されたコードサンプルは ここ で見つけることができます。
以上