最近の型付きPythonについて

はじめに

株式会社みらい翻訳でカスタマーサクセスを担当している @shhshn です。
日常的に使用するものの業務とは特に関係のないPythonの話をします。

製品開発にPythonを使用する上で、型の明示が重要になっています。
「Pythonで型を極める」話もありますが、主な理由は以下2つです。

• 静的型チェックを使いたい
• ついでにPythonを高速化したい

前者は「Dropboxの取り組み」が有名ですが、後者をあまり見かけません。
この記事ではNumbaやTorchScriptなど3種類の高速化方法を紹介します。

Numba

condaで有名なAnaconda主導で開発されているJITコンパイラです。

from numba import jit
import random

@jit(nopython=True)
def monte_carlo_pi(nsamples):
    acc = 0
    for i in range(nsamples):
        x = random.random()
        y = random.random()
        if (x ** 2 + y ** 2) < 1.0:
            acc += 1
    return 4.0 * acc / nsamples

例えば、上記のコードは特に型を明示していませんが、型推論によりLLVMでコンパイルされます。その結果、C++で書いたコードと同等の速度を発揮します。C++からの呼び出しも可能です。

Numbaの強み

• 既存コードに @jit とアノテーションするだけで高速化できます
• numpy との互換性が高く、ほぼそのままJIT化できます
• @jit(nogil=True) でGILを外すことでCPythonの限界を突破できます

Numbaの弱み

• Python言語仕様の一部しかサポートしていません
  • 最近まで dict すらありませんでした
• 型推論でエラーになると原因究明が困難です

TorchScript

FacebookがPyTorchの一部として開発しているコンパイル機能です。

@torch.jit.script
def scripted_fn(x : torch.Tensor):
    for i in range(12):
        x = x + x
    return x

def fn(x):
    x = torch.neg(x)
    import pdb; pdb.set_trace()
    return scripted_fn(x)

traced_fn = torch.jit.trace(fn, (torch.rand(4, 5),))
traced_fn(torch.rand(3, 4))

この例では2通りの方法で型付きPythonとして処理しています。
torch.jit.trace はNumbaのように関数を丸ごと型推論します。
一方 @torch.jit.script はmypy形式で明示された型を使用します。

TorchScriptの強み

• Python言語仕様のほとんどをサポートしています
  • comprehensionなど着々と対応されています
• 型推論エラーが分かりやすく親切です

TorchScriptの弱み

• 公式ドキュメント以外の情報源が薄いです
• C++からの呼び出しにはコード全体のアノテーションが必要です
• もはやPythonではない別の言語になります

別の言語を使うという選択肢

Pythonより高速性に主眼を置くスクリプト言語も多いです。例えばJuliaとはよく比較されますが、JuliaはJIT対応に一日の長があります。正直、Pythonの文法・エコシステムのような些細な違いを除けば、他の言語は十分選択肢として考えられます。

また始めから終わりまでC++で書いていくストロングスタイルもあります。保守運用（特に人的資源）に問題がなければ、C++で書くのが一番です。

例: PythonからC++へのコンパイル

https://github.com/shhshn/sonish （概念実証コードです）

最近はC++にも型推論が備わっているため、型を明示せずともPythonからC++に変換できます。

def main():
    print("Hello and, again!")
    l = [3, 2]
    l.append(1)
    for x in l:
        print(x)
    d = {4: "HUGE", 5: "SUCCESS"}
    print(d[5])

if __name__ == "__main__":
    main()

例えば、上記Pythonを全く同じ出力のC++へ変換することができます。

#include <iostream>
#include <vector>
#include <unordered_map>
#include <iomanip>

int main() {
std::cout << std::setprecision(17) << "Hello and, again!" << std::endl;
auto l = std::vector<int>{3, 2};
l.push_back(1);
for (auto x : l) {
std::cout << std::setprecision(17) << x << std::endl;
}
auto d = std::unordered_map<int, const char *>{{4, "HUGE"}, {5, "SUCCESS"}};
std::cout << std::setprecision(17) << d[5] << std::endl;return 0;
}

おわりに

Pythonでのラピッド・プロトタイピングとC++での製品化のギャップを埋める話をしました。

この記事はみらい翻訳 Advent Calendar 2020に参加しています。
次はAWSマスターの @kobarasukimaro が再度登場するようです
その前に @reonyanarticle が「CLI作成モジュールclickについて」話してくれるそうです！