2018年3月18日に作成したメモですが、「Tensor Comprehensions」の活動がほぼ止まっているので、公開しときます。
ソースコードを調べているのですが、なかなか手ごわいです。
とりあえず、ソースコードを調査結果をまとめる前に、
また、Twitterにツイートしたものをメモとしてまとめておきます。
Tensor Comprehensions の特徴。 カスタムレイヤーへの入力テンソルのサイズは、パラメータ化するのではなく、 決まったサイズとして扱うことで、最適化を図るというもの。 Poolingも、2x2 や 3x3 とかもそれぞれ別のレイヤーとしている。 基本的には、カスタムレイヤーを作るための道具。。 なので、PyTorchやCaffe2だけでなく、 他のMLフレームワークでも利用可能だと。 現時点では、CUDA 専用って感じだけど。
Tensor Comprehensionsのレイヤー例 Max Poolingは、 def maxpool(float(B, C, H, W) input) -> (output) {{ output(b, c, h, w) max= input(b, c, h * {sH} + kh, w * {sW} + kw) where kh in 0:{kH}, kw in 0:{kW} }} のように定義するけど、使うときに、kHとkWが決まって、最適化する
matmul = tc.define(lang, name="matmul")
を定義し、
mat1, mat2 = torch.randn(3, 4).cuda(), torch.randn(4, 5).cuda()
にて、CUDA用のTensor (入力)を宣言
out = matmul(mat1, mat2)
にて、out (出力)を計算する。
lang = """
def matmul(float(M,N) A, float(N,K) B) -> (output) {
output(i, j) +=! A(i, kk) * B(kk, j)
}
"""
から Halide IR に変換し、Range Inference and Specialization を行う。
Tensor Comprehensionsの目的は、 カスタムレイヤーをかける人が限られているので、 それを普通の人でもある程度の性能が出るようにするための道具。 だと思う。
TensorFlow XLAやNNVMのようなグラフ最適化(複数レイヤー間の最適化)をするようなものではない。
Autotuner は、
matmul = tc.define(lang, name="matmul")
mat1, mat2 = torch.randn(100, 400).cuda(), torch.randn(400, 500).cuda()
matmul.autotune(mat1, mat2, **tc.autotuner_settings)
out = matmul(mat1, mat2)
のように、関数を実行する前に、autotune すればいい。
out1 = matmul(mat1, mat2, cache=cache_file)
# the second time we run the kernel, we skip the compilation since it was
# already compiled earlier
out2 = matmul(mat1, mat2)
ということで、cache引数に、キャッシュファイル名を指定できる
matmul = tc.define(lang, name="matmul") matmul.autotune((3, 4), (4, 5), cache=True, **tc.small_sizes_autotuner_settings) matmul.autotune((100, 400), (400, 500), cache=True, **tc.autotuner_settings) の入力テンソルのサイズを分けて、autotune すればいい。
matmul = tc.define(lang, name="matmul") matmul.autotune((3, 4), (4, 5), cache=cache, **tc.small_sizes_autotuner_settings) tc.decode(cache + ".options") キャッシュしてファイルは、tc.decode にて、デコードできる。
tc.GlobalDebugInit(["--dump_cuda=true"]) matmul = tc.define(tc.database['matmul']['lang'], name='matmul') mat1, mat2 = torch.randn(3, 4).cuda(), torch.randn(4, 5).cuda() out = matmul(mat1, mat2) -dump_cuda=true で生成されたCUDAコードが表示される
Caffe2との統合は、推論のみ。 The integration with Caffe2 is very basic at the moment. We do not provide autotuner support for Caffe2 and welcome contributions from community. ということで、Caffe2との統合はまだまだだと。
今気づいたんだけど、Tensor Comprehensionsの Andre Adams さん https://andrew.adams.pub/ GoogleからFacebookに移っていたんですね。。
