scikit-learnの学習済みモデルをONNX形式に変換して配布する
だいぶ昔の話ですが、日本語テキストをネガ/ポジ分類するソフトウェアとして、scikit-learnを用いて『asari』を作り、Pythonパッケージとして公開したことがあります。作った自分でも存在をほぼ忘れていたのですが、ときどき使うことを試みる方がいて、Issueを上げてくれることがありました。そこで、重い腰を上げて、Issueをすべて解決し、v0.2.0をリリースしました。
asariの学習済みモデルはJoblibを用いて永続化していたのですが、scikit-learnのドキュメントにもあるように、保守性とセキュリティ面で課題がある状態でした。保守性の観点からいうと、この手法では、あるscikit-learnのバージョンで永続化したモデルが別のバージョンで動くとは限らず、動いたとしても結果が変わることもありえます。また、セキュリティ面では、読み込み時に悪意のあるコードが実行される可能性があります。
実際、Issueでは動かないという報告が上げられていました。そこで、再現性を担保したり異なる環境でも動くようにするため、学習済みモデルをONNX形式に変換して提供することにしました。ONNXは、機械学習モデルを表現するために構築されたオープンフォーマットのことを指します。異なる機械学習フレームワーク間でのモデルの変換を容易にしたり、モデルの移植性を向上させることを目的としています。今回は、学習環境とは異なる環境での予測に役立つと考えて採用しました。
scikit-learnのモデルをONNXに変換するためには、sklearn-onnxを使います。モデルの学習から推論までの大まかな手順は次のとおりです。
- scikit-learnでモデルを学習
- sklearn-onnxを用いて、モデルをONNX形式へ変換
- ONNX Runtimeを用いて推論
モデルの学習
まずはモデルを学習します。この辺はとくに難しいところはないでしょう。ただし、1つ注意点があります。asariの場合はトークン化にJanomeを使っているのですが、現状ではカスタムトークナイザーをONNX形式へ変換することはサポートされていません。そのため、TfidfVectorizerにトークナイザーを指定する代わりに、パイプラインの外側でトークン化をして、空白文字で結合しています。
from sklearn.calibration import CalibratedClassifierCV from sklearn.feature_extraction.text import TfidfVectorizer from sklearn.model_selection import train_test_split from sklearn.pipeline import Pipeline from sklearn.svm import LinearSVC from asari.preprocess import tokenize X, y = load_jsonl(...) X = [tokenize(x) for x in X] x_train, x_test, y_train, y_test = train_test_split( X, y, test_size=0.1, random_state=42 ) pipe = Pipeline( [ ("vectorizer", TfidfVectorizer(ngram_range=(1, 2))), ("classifier", CalibratedClassifierCV(LinearSVC(), ensemble=False)), ] ) pipe.fit(x_train, y_train)
ONNX形式への変換
次に、学習したモデルをONNX形式に変換します。scikit-learnのモデルを変換するためには、次の2つの関数のうち、どちらかを使います。
convert_sklearnを使う場合、次のように、modelとinitial_typesの2つを必ず渡す必要があります。modelにはscikit-learnのモデルかパイプラインを渡します。initial_typesには入力の名前と型の情報を渡します。今回はテキストを入力するので、[('text', StringTensorType(None))]のようにしています。textは入力の名前であり、2番目の値は型と形状を表しています。1次元目は行数ですが、こちらは事前に予測できないためNoneにします。
from skl2onnx import convert_sklearn from skl2onnx.common.data_types import StringTensorType initial_type = [('text', StringTensorType(None))] onx = convert_sklearn(pipe, initial_types=initial_type) with open("pipeline.onnx", "wb") as f: f.write(onx.SerializeToString())
to_onnxを使う場合、学習データセットのうちの1件を元に適切な型を推論します。convert_sklearnでは入力の名前と型を指定していたところを、次のように学習データセットの1行を渡します。asariでは、現時点ではこちらを使っています。
from skl2onnx import to_onnx onx = to_onnx(pipe, np.array(x_train)[1:]) with open("pipeline.onnx", "wb") as f: f.write(onx.SerializeToString())
残念ながら、TfidfVectorizerのONNX版はscikit-learn版と全く同じ結果を出力するわけではありません。現在のところ、tokenexpとseparatorsというパラメータを変換時に渡すことができます。実はこの辺の値を設定しておかないと、正規表現が自動的に置き換わって、日本語が上手く分割されなくなり、推論がデタラメになりました。今回は空白文字で区切った入力を与えているので、separatorsに空白だけを指定しています。
seps = {
TfidfVectorizer: {
"separators": [
" ",
],
}
}
onx = to_onnx(pipe, np.array(x_train)[1:], options=seps)
ONNX形式に変換できたら、Netronを使って可視化してみましょう。Netronはブラウザやデスクトップアプリで機械学習モデルを可視化できるツールです。変換したモデルをアップロードするだけで、次の画像のように、モデルの構造を可視化したりプロパティを確認できます。ONNX Optimizerなどでモデルを最適化した際に、変更箇所の確認やデバッグに使うと便利です。
ONNX Runtimeによる推論
モデルをONNX形式に変換したら、ONNX Runtimeを使って推論します。カスタムトークナイザーはONNX形式に変換できなかったので、推論前に自分でトークン化しています。パイプラインに組み込めるのであれば、学習時と推論時で前処理（今回の場合はトークン化だけ）が一致するので、組み込んだほうがよいと思います。
import onnxruntime as rt sess = rt.InferenceSession("pipeline.onnx") input_name = sess.get_inputs()[0].name label_name = sess.get_outputs()[0].name tokenized = tokenize(text) pred = sess.run([label_name], {input_name: [tokenized]})
確率値を得たい場合は次のようにします。
prob_name = sess.get_outputs()[1].name
proba = sess.run([prob_name], {input_name: [tokenized]})
参考資料
部分的にアノテーションされたデータからの固有表現認識器の学習
固有表現認識は、テキスト中に含まれる人名や地名、組織名といった固有表現を自動的に認識する技術です。たとえば、「太郎は10時30分に東京駅に着いた」というテキストからは、人名として「太郎」、時間として「10時30分」、地名として「東京駅」を認識できます。固有表現認識は、情報抽出や情報検索、質問応答などの自然言語処理のアプリケーションで広く利用されています。
最近では、固有表現認識でも教師あり学習を用いた手法が主流となっていますが、ラベル付きデータは人手で用意するのが一般的なので、十分な量を用意するためには相応のコストがかかります。とくに固有表現認識の場合、分類タスクと比較すると、単語や文字ごとにラベル付けするため、用意するのは大変です。
アノテーションコストを減らすため、部分的にアノテーションされたデータの活用が考えられます。部分的にアノテーションされたデータとは、データセットの一部分だけにアノテーションされたデータのことです。たとえば、冒頭の例で言えば、「太郎」と「東京駅」のみにアノテーションされていれば、それは部分的にアノテーションされたデータと言えます。このようなデータは固有表現辞書やルールによって作成できます。
辞書やルールで部分的にアノテーションするだけでもよいのですが、こうして作成したデータからモデルを学習する方法も考えられます（下図）。一般的に、辞書を用いたアノテーションでは、辞書に格納されている固有表現とテキストをマッチングします。この方法は、精度は高いのですが、辞書サイズが小さかったり表記ゆれがある場合、再現率が低くなるのが欠点です。そのため、直接的に使うのではなく、いったんモデルに学習させます。
部分的にアノテーションされたデータをモデルに学習させる場合、本来は固有表現があるべき場所にないこと（タグの欠落）が問題になりえます。通常の教師あり学習では、このような状態を考慮していないため、学習が上手く進まず、性能が低くなると考えられます。辞書を用いてアノテーションしたデータの場合、とくに再現率が低くなると考えられます。
この問題に対処するために、2021年のTACLでは期待固有表現比率（Expected Entity Ratio：EER）と呼ばれる損失が提案されました[2]。この手法では、データセット全体での固有表現の出現割合が一定の範囲内になるように、モデルを学習します。この損失を組み込むことで、タグが欠落している場合にもなるべく固有表現の予測を試みてくれます。
では、実際にEERを組み込んだモデルを試してみましょう。
実験
検証内容
今回は、辞書データを用いてアノテーションしたデータに対して、EERを組み込んだモデルの性能を検証します。実際の状況では、辞書サイズが限られている場合もあると考えられるため、辞書サイズを25~100%まで25%ずつ変化させたときの性能を検証します。実装としては、EERのspaCy実装であるspacy-partial-taggerを利用します。
固有表現認識用のデータセットとしては、BC5CDRを使います。このデータセットは医療系のデータセットであり、固有表現としては「Chemical」と「Disease」の2つが付いています。また、辞書としては先行研究で公開されているものを使用します。この辞書に登録されている固有表現数は2482件です。
事前学習済みモデルとしては、Microsoftが公開しているPubMedBERTを使用します。このモデルはPubMedのテキストで学習されているため、今回のデータセットに適していると考えられます。
以下の4つのモデルについての性能を見てみましょう。
- 完全な教師データで学習したPubMedBERTモデル（Fully supervised）
- 部分的にアノテーションされたデータで学習したPubMedBERTモデル（PubMedBERT）
- 部分的にアノテーションされたデータで学習したPubMedBERT + EERモデル（+EER）
- 辞書マッチ（Dictionary Match）
検証結果
検証結果は次のとおりです。今回の場合は、辞書マッチの結果をそのまま使うよりも、いったんモデルに学習させたほうが性能が高いことがわかります。また、EERを使うことで、とくに辞書サイズが小さな場合に性能が向上していることがわかります。これは主に再現率が向上しているためです。また、Fully supervisedと比べると、7ポイントほど性能が低いという結果になりました。
終わりに
本記事では、部分的にアノテーションされたデータから固有表現認識モデルを学習する方法としてEERを検証しました。実験の結果、今回のデータセットと辞書の組み合わせでは、とくに辞書サイズが小さな場合でのEERの有効性を確認できました。辞書からはランダムにサンプリングしており、どんな場合でも有効とは限りませんが、固有表現辞書とアノテーション対象のテキストが用意できる場合には利用することを検討してみます。
参考資料
M1チップ上でのspaCyの高速化
導入されたのはだいぶ前ですが、spaCy v3.2からM1チップ上での学習と予測が最大で8倍高速化できるということで試してみました。以前は行列積の演算にBLISを使っていたところを、Appleのネイティブのライブラリに切り替えることで実現しています。その中核となっているのがthinc-apple-opsであり、こちらでライブラリを置き換えています。
公式でも、spaCyを使った予測と学習に対して、高速化前と高速化後のベンチマークを出しています。予測では、de_core_news_lgを使った1秒あたりの処理単語数を比較しています。結果として、Intel Macだと速度にほとんど変化がない一方、M1 Macだとおおよそ4.3倍になっています。学習では1秒あたりのイテレーション数を比較しており、3倍程度高速化できています。
ちょうど手元にM1 Proを積んだMacがあったので試してみました。通常のspaCyのインストール（pip install spacy）とM1に対応した版（pip install spacy[apple]）で固有表現認識モデルの学習をして、実行時間を比較してみます。設定としてはほぼ以下のノートブックのとおりですが、結果を比較するために、最大ステップ数が同じになるように調整しました。
結果は次のようになりました。値は3回実行したときの平均をとっています。結果を見ると、今回の設定の場合、M1版のほうがおおよそ4.3倍程度高速という結果になっています。
- 通常のspaCy：680秒
- M1版のspaCy：160秒
重たいモデルをローカルで学習することはほとんどないとは思いますが、インストールを除いて、ユーザー側でコードを変更する必要がないので、M1 Macであればお手軽にできる高速化といえます。今回は試していませんが、予測も高速化されるので、ちょっとした分析をローカルでするときに使ってみることにします。