ちなみに、自分の手法に関してはこちらをご覧ください。

https://www.kaggle.com/competitions/open-problems-multimodal/discussion/366961

Open Problems – Multimodal Single-Cell Integration とは？

コンペのサイトのoverviewに書かれていることをざっくり日本語に訳すと「骨髄幹細胞が血液細胞になるにつれて単一の細胞のDNAとRNA、タンパク質がどのように変化するかを予測するコンペ」ということになります。もう少し専門家の人がわかりやすいように説明すると以下の二つのsingle-cellのデータを使ったモデルを構築するコンペです。

1. Chromatin accessibility(peak countをTF-IDFで変換したデータ)からRNAの発現量 (library-sizeでノーマライズされたcountデータ)の予測
2. RNAの発現量 (library-sizeでノーマライズされたcountデータ)からsurface protein levels (dsb でノーマライズされたデータ)の予測

そもそもなぜ参加したのか？

参加してがんばろうと思ったきっかけをせっかくなので書いておくと

1. Bioチームの同僚がこういうコンペがあるよとslackで教えてくれた (参加したきっかけ)
2. 育休に入る前にやっておくべきことは？ということを会社の先輩パパさんに相談したら、「solo gold medalは取っておけ」と言われた (がんばったきっかけ)

ということがあります。特に2つ目は「確かに！」と思いました。なので、今回が初ソロ参加にして、solo gold medalを取る最後のチャンスということで頑張りました。アドバイスをくれた先輩パパさんには感謝しかありません。

やってよかったこと

さて、ここからやってよかったことについて忘れないように書いておきます。

まずはとにかく簡単な方法でいいのでsubmitする

社内のkaggle強い人に前に言われた気がするので、まず意識したことがこれです。やってみて思ったのですが以下のような効果があることを実感しました。

1. submitして順位やスコアがでるようになるとモチベーションが上がる。
2. 何かベースラインがあると手法開発がしやすい

特に1の効果がすごかった気がします。ちなみに私の場合、KaggleのCodeで公開されてたシンプルな手法をそのままコピペしてsubmitして一番最初のスコアを出しました。シンプルな手法だったため、最初は200位にも入れなかったと記憶していますが、それでもモチベーションはそれまでと比べてすごく上がりました。

毎日決めた数の改良を試す

これはコンペに限らず重要なことだと思いますが、とにかくコンスタントに改良を続けていき、最低限local CV scoreを出すということを意識してました。

今回のコンペでは3，4個の改良を毎日試すことを目標にしてやってました。私の場合、gitの1 branchが1改良になっていて、9/1からがんばり始めて11/15までに最終的には351個branchができていました。なので、通算4から5個くらいの改良を毎日試していたことになります。

このとき改良としてうまくいきそうなものはもちろんですが、思いついたタイミングではあまり筋が良くなさそうだけど他にやることがないというときは、筋が良くないアイディアもダメ元で試すようにしてました。

結果として、やってよかったと感じた理由としては、仮説を立てて試すとうまくいかなかったときの問題への理解度がすごい上がるため、とにかくいろいろ試すことで、そこからいいアイディアをひらめくということが多くあったからです。

今回のコンペで何度か大きくスコアを上げたタイミングがありましたが、総じて何かの失敗から気が付いたアイディアをもとにしたことが多かった印象です。ちなみにどれくらい失敗し続けたのかの体感ですが、1週間スコアが上がらないということがよくあったので、うまくいった改良というのはこれだけ試しても数えるくらいしかなかった印象です。

ensembleはすぐに試さない

これも社内のkaggle強い人に前に言われたことなのですが、ensembleをすぐに試さないというのはやってよかったなと思います。
今回はensembleを試さなくてよかったなと感じた理由としては、ensembleはやれば簡単にスコアが上がるのでいいように感じますが、ensembleは始めるとすごいいろいろ試せることがある一方、おそらくそこまで高いスコア上昇をしていなかったような印象があります。

このためensembleを頑張ることに時間使うよりも1つのモデルのスコアが上がるように頑張るという作戦でいたのですが、結果としてそれが良かった印象です。

最後まで諦めない

最後はこれです。特に今回はsolo gold medalが欲しかったので、一人で出ていたのですが、途中全然順位が上がらず何度ももうやめようかなぁと思うタイミングがありました。結果としてそのときダメもとで試した改良やそこから思いついたアイディアがうまくいってスコアを伸ばすということが何度もあったので諦めなくてよかったです。

課題

次にコンペの参加中、もしくは終わったあと振り返って感じた課題的な部分も列挙しておきます。

モチベーションの維持が大変（特に一人のとき）

「最後まで諦めない」のところに似たようなことを書きましたが、とにかくモチベーション維持が大変でした。今まではチームででていたのと、すごい応援してくれる上司がいてくれたりとこの部分はそこまで問題にならなかった印象でした。

今後一人ででるならこの部分はまだまだどうにかしないと最後まで戦うのは難しい気がしています。

public scoreがどうしても気になってしまう

今回のコンペはデータセットの説明を読んだ段階でpublic scoreとprivate scoreのギャップが激しそうだなぁということを思ってたので、それほどpublic scoreを気にしないほうがいいかも？ということを最初思ってました。ただ、それでも最後はpublic scoreを気にしてモデルの改良をしてしまっていました。

コンペのサイトのデータセットの説明のところに詳しくかかれていますが、今回public scoreは4人中一人の最終日一つ前までのデータを、private scoreは4人全員の最後の日のデータを予測するというものでした。この説明を見ると予測する対象が全然違うことがわかります。実際、public scoreの上位陣が軒並みprivate scoreでは順位を落としていました。なので、後から振り返るとやっぱり最初に思った通りpublic scoreをそれほど気にせずモデルの改良をしていて正解だったと思います。ただ、そうはいっても、public scoreの順位が気になってしまって、結局最後はpublic scoreを気にして最終submittionを決めていました。ただ、ふたを開けたらlocal CVがベストなものがprivate scoreも一番よかったので、public scoreを気にしすぎたなぁと思っていました。

ただ、この部分はふたを開けてみないとわからないところなので、難しいポイントな気がしています。

論文を読んで最新研究をコンペの期間中に試すことが心理的に難しい

期間の前半ならまだましですが、どうしても後半になればなるほど、心理的に追い込まれていきます。なので、試すアイディアがないと思いつつも、自分の全然読んだことない論文の手法を試すということが難しかったです。この結果、特に後半はアイディアが前に試したもののちょっと変更したものばかりになり、結果としてそれほど精度が向上しないということが起きました。ただ、この部分一人でやっていると意識しても難しい部分な気がしています。

このため、日ごろから論文を読んでいろいろ手法を勉強し、できれば実装を動かして感触を確かめておくことが重要かなと思いました。実際、今回のコンペではReactomeのpathwayデータを使ったのですが、これは論文の追試のついでにいろいろ試したことがあったからこそできたことだったと思っています。

ちなみにその時の記事はこちらにあります。

ensembleの準備をせずにいると最後の土壇場で困る

ensembleをすぐに試さないことを意識していた結果、実は今回期間最後のほうでensembleをしようとしたときにいろいろ準備できてないことに気が付いてすごい慌てました。具体的には以下の二つが締め切り1週間前の段階でできていませんでした。

1. 予測結果をどのように集計して最終的なファイルを作るか？
2. ensembleした結果をどのように評価するか？

結局、1はギリギリできたのですが、2の評価方法はとても間に合いそうになかったのでpublic scoreを見てうまくいっているかいっていないか？を判断していました。ただ、これは評価としては微妙なので、次からはそんなに頑張らないにしても中盤くらいにはensembleの準備はしておこうと思います。

最後に

折角なので個人的な振り返りをblog記事にしました。仕事ではない状態でのkaggle参加は初めてだったのですが、思ったよりも疲れたのと、無事子供がコンペ締め切りの前日に生まれて、これから子育てもあるのででしばらくはコンペにでない気がします。ただ、次出たときに今回のコンペで何を思ったのか忘れないようにしておければと思っていたので、記事にできてよかったです。

これが他の人の参考になれば幸いです。

The post Kaggleの「Open Problems – Multimodal Single-Cell Integration」の振り返り first appeared on まったり勉強ノート.

UniProtで検索した配列がほしくなることが度々あるのですが、数が多いとブラウザで検索して配列をコピペするというのが面倒になってきます。このため、UniProtのWeb APIを使ってまとめて取ってくる方法を調べたので、今回はそのまとめの記事になります。

今回はUniProtが提供しているWeb APIの中で、クエリを使ってほしいタンパク質の情報を持ってきてPythonで簡単に使えるようにpandasのDataFrameにするところまでを紹介します。

今回のコードはこちらにあります。

https://github.com/shu65/uniprot_web_api_python_example/blob/main/UniProt_Result_to_Pandas_DataFrame.ipynb

UniProt Web APIについて

UniProtはタンパク質配列や機能に関するデータを集めた有名なデータベースです。このUniProtではブラウザで検索するのはもちろん、Web APIも提供されています。詳細はここに書かれています。

https://www.uniprot.org/help/api

クエリを使った方法はここにまとめられています。詳しくは次で紹介します。

https://www.uniprot.org/help/api_queries

UniProt Web APIの簡単な使い方

説明がわかりやすくなるようにまず具体的な例を示します。ここではヒトのTP53をUniProtKB/Swiss-Protから検索してタブ区切りのフォーマットでidとgeneの名前、アミノ酸配列を受け取ってみます。これを実行するURLとしては以下のようになります。

https://www.uniprot.org/uniprot/?query=reviewed:yes+AND+organism:%22Homo%20sapiens%22+AND+gene_exact:TP53&format=tab&columns=id,genes(PREFERRED),sequence

順番に与えているパラメータを見ていきます。

query

検索するqueryを文字列にして与えます。使えるフィールドは以下に一覧が載っています。

https://www.uniprot.org/help/query-fields

普通は複数のフィールドを使ってand検索かor検索を使いたくなるかと思います。and検索、or検索をする場合は以下のようにします。

• and検索例: human かつ antigenのand検索
  1. https://www.uniprot.org/uniprot/?query=human%20antigen
  2. https://www.uniprot.org/uniprot/?query=human%20AND%20antigen
• or検索例: human または mouseのor検索
  1. https://www.uniprot.org/uniprot/?query=human%20OR%20mouse

それ以外のandやor以外についてもいろいろ使えます。詳細はこちらに書かれています。

https://www.uniprot.org/help/text-search

format

結果のフォーマットを指定するパラメータです。今回示したタブ区切りの場合はtabを指定します。それ以外のフォーマットは以下のページの[Parameter]→[format]に書かれています。

https://www.uniprot.org/help/api_queries

columns

デフォルトではUniProtで取得できる一部のデータしか取得できませんが、columnsを指定するとほしいデータを取得できます。上の例ではidとgeneの名前、アミノ酸配列を指定するために以下のようにしていました。

columns=id,genes(PREFERRED),sequence

カンマでほしいカラム名を区切ることで複数のカラムの情報を得ることができます。他に使えるカラムに関しては以下のページをご覧ください。

https://www.uniprot.org/help/uniprotkb_column_names

このページの[Column names as displayed in URL]側を指定するようにします。

それ以外のパラメータ

それ以外に圧縮をするかどうかなどがパラメータとして指定できます。詳しくはこのページの[Parameter]をご覧ください

https://www.uniprot.org/help/api_queries

UniProtのWeb APIを使って検索した結果をPandasのデータフレームにする

ここまでわかっていれば簡単にできてしまいますが、UniProtのWeb APIを使って上で示した結果をPandasのDataFrameにするところまでのコードを示します。

requestsを使ってUniProtから検索結果をタブ形式で取得します。

import requests

url="https://www.uniprot.org/uniprot/"

params={
    "query":"reviewed:yes AND organism:\"Homo sapiens\" AND gene_exact:TP53",
    "format":"tab",
    "columns":"id,genes(PREFERRED),sequence",
}

response = requests.get(url=url, params=params)
response.raise_for_status()

これを以下のようにpandasで読み込むだけで、pandasのDataFrame形式の結果を取得できます。

import io

import pandas as pd


df = pd.read_csv(io.StringIO(response.text), sep="\t")

できたpandasのDataFrameをprintした結果は以下の通りです。

    Entry Gene names  (primary )  \
0  P04637                   TP53   

                                            Sequence  
0  MEEPQSDPSVEPPLSQETFSDLWKLLPENNVLSPLPSQAMDDLMLS...  

Python以外でweb apiの結果を取得するまでの例はこちらにも書かれています。

https://www.uniprot.org/help/api_idmapping

他の言語でやりたい方は参考にしてみてください。

終わりに

今回はUniProtのWeb APIを使って検索結果をpandasのDataFrameにするところまでの例を紹介しました。当初はほしいアミノ酸配列を取得するのにfastaファイルをがんばってパースしようかな？とおも思ったのですが、調べてみると簡単にWeb APIが使えそうなことがわかったので、試してみました。

機会があれば他の有名データベースのWeb APIもしらべてみようかなと思います。

The post UniProtのWeb APIを使ってほしい遺伝子のタンパク質配列を取ってくる first appeared on まったり勉強ノート.

だいぶ前ですが以下の論文を読み、ReactomeからどのようにしてPathwayデータを取得するのか？というのが気になっていました。

Elmarakeby, H.A., Hwang, J., Arafeh, R. et al. Biologically informed deep neural network for prostate cancer discovery. Nature 598, 348–352 (2021). https://doi.org/10.1038/s41586-021-03922-4

こちらの論文自体はReactomeというPathwayデータベースのデータを組み合わせてP-NETという深層学習モデルを構築して前立腺がんの患者のデータに対して適用した研究になります。

この深層学習のモデルを構築する際、ReactomeからPathwayに関するデータを取得して利用しているのですが、ReactomeからどうPathwayのデータを取得して、それをどう加工すればよさそうか？が論文を読んだだけでは分からず、コードを読んで調べたので本日の記事はそのまとめになります。

今回の記事で利用したコードはこちらに置いてあります。

https://github.com/shu65/reactome-example/blob/main/Reactome_gene_pathway_hierarchy_relationship.ipynb

P-Netで利用されているReactomeのデータ

P-NetではPathwayの階層情報に基づいて深層学習のモデルを構築しています。この際、利用されているデータは以下のものになります。

1. Pathway毎に関連するGeneの集合
2. Pathwayの階層構造
3. Pathwayの名前と種

これらをReactomeから取得する方法を紹介していきます。わかりやすいようにReactomeのリンクやページのスクリーンショットも合わせのせています。これらは2022/05/06時点のものになります。アクセスする時期によってはこれらが変わっている可能性もあるので注意してください。

Reactomeから大本のデータを取得する

Reactomeは各種データを様々なフォーマットでダウンロードできるようになっています。それらはここにまとまっています。

https://reactome.org/download-data

今回利用するデータはそれぞれ以下のファイルに書かれた情報から抽出することができます。

Pathway毎に関連するGeneの集合

Pathway毎に関連するGeneの集合のデータは[Specialized data formats]→[Reactome Pathways Gene Set.]から落とすことができます。

こちらはGMTフォーマットで書かれたファイルになります。具体的には行ごとに1つのPathwayについて書かれており、カラムはPathwayの名前、Reactome Stable identifiers (ST_ID)が続き、それ以降は登場するGeneが並んでいます。

Pathwayの階層構造

Pathwayの階層構造のデータは[Pathways] → [Pathways hierarchy relationship]から落とすことができます。

[Read more] のリンクに詳しいフォーマットが書かれていますが、タブ区切りで最初のカラムが親のPathwayのST_ID、2つ目が子のST_IDになります。

Pathwayの名前と種

Pathwayの名前と種に関するデータはPathwayの階層構造のデータは[Pathways] → [Complete List of Pathways]から落とすことができます。

[Read more] のリンクに詳しいフォーマットが書かれていますが、タブ区切りで最初のカラムがPathwayのReactome Stable identifiers (ST_ID)、2つ目がPathwayの名前、3つ目が種になります。

Reactomeから落としたデータをPythonで読み込み

ここまででほしいデータがどこから落とすことができるか説明しました。ここからは実際に使う際に利用しやすいようにPythonで読み込むパーサーを書いたのでその紹介をします。

Pathway毎に関連するGeneの集合のデータのパース

Reactomeから落としてきたGMTファイルをパースするスクリプトは以下の通りです。

import pandas as pd


def read_reactome_gmt(file_path):
  data_dict_list = []
  with open(file_path) as f:
      for i, line in enumerate(f):
          values = line.strip().split("\t")
          st_id = values[1]
          genes = values[3:]
          for gene in genes:
              data_dict_list.append({'st_id': st_id, 'gene': gene})
  df = pd.DataFrame(data_dict_list)
  return df

以下のようにようにファイルのパスを渡すとpandasのデータフレームでファイルから読み込んだ結果を返すようにしてあります。

pathway_gene_df = read_reactome_gmt("ReactomePathways.gmt")
print(pathway_gene_df)

出力は以下の通りです。

               st_id   gene
0       R-HSA-164843  HMGA1
1       R-HSA-164843   LIG4
2       R-HSA-164843  PSIP1
3       R-HSA-164843  XRCC4
4       R-HSA-164843  XRCC5
...              ...    ...
121252  R-HSA-192905     NP
121253  R-HSA-192905     NS
121254  R-HSA-192905     PA
121255  R-HSA-192905    PB1
121256  R-HSA-192905    PB2

[121257 rows x 2 columns]

Pathwayの階層構造の読み込み

Reactomeから落としたPathwayの階層構造のデータを読み込むパーサーは以下の通りです。

import pandas as pd


def read_reactome_pathway_hierarchy_relationship(file_path):
  df = pd.read_csv(file_path, sep='\t')
  df.columns = ['parent_st_id', 'child_st_id']
  return df

Reactomeから落としたファイルパスを渡すと以下のようなpandasのデータフレームを返します。

pathway_hierarchy_df = read_reactome_pathway_hierarchy_relationship("ReactomePathwaysRelation.txt")
print(pathway_hierarchy_df)

出力は以下の通りです。

       parent_st_id    child_st_id
0      R-BTA-109581  R-BTA-5357769
1      R-BTA-109581    R-BTA-75153
2      R-BTA-109582   R-BTA-140877
3      R-BTA-109582   R-BTA-202733
4      R-BTA-109582   R-BTA-418346
...             ...            ...
21521  R-XTR-983705  R-XTR-5690714
21522  R-XTR-983705   R-XTR-983695
21523  R-XTR-983712  R-XTR-2672351
21524  R-XTR-983712   R-XTR-936837
21525  R-XTR-991365   R-XTR-997272

[21526 rows x 2 columns]

Pathwayの名前と種の読み込み

ReactomeのPathwayの名前と種のパーサーは以下の通りです。

import pandas as pd


def read_reactome_complete_list(file_path):
  df = pd.read_csv(file_path, sep='\t')
  df.columns = ['st_id', 'pathway_name', 'species']
  return df

こちらも以下のようにしてファイルパスを指定するとpandasのデータフレームを返します。

complete_list = read_reactome_complete_list("ReactomePathways.txt")
print(complete_list)

complete_list = read_reactome_complete_list("ReactomePathways.txt")
print(complete_list)
               st_id                                       pathway_name  \
0      R-BTA-1971475  A tetrasaccharide linker sequence is required ...   
1      R-BTA-1369062              ABC transporters in lipid homeostasis   
2       R-BTA-382556             ABC-family proteins mediated transport   
3      R-BTA-9033807                       ABO blood group biosynthesis   
4       R-BTA-418592          ADP signalling through P2Y purinoceptor 1   
...              ...                                                ...   
21417   R-XTR-193639                           p75NTR signals via NF-kB   
21418   R-XTR-111995                               phospho-PLA2 pathway   
21419   R-XTR-191859                                     snRNP Assembly   
21420   R-XTR-379724                                tRNA Aminoacylation   
21421   R-XTR-199992                trans-Golgi Network Vesicle Budding   

                  species  
0              Bos taurus  
1              Bos taurus  
2              Bos taurus  
3              Bos taurus  
4              Bos taurus  
...                   ...  
21417  Xenopus tropicalis  
21418  Xenopus tropicalis  
21419  Xenopus tropicalis  
21420  Xenopus tropicalis  
21421  Xenopus tropicalis  

[21422 rows x 3 columns]

Reactomeのデータを組み合わせて使う

これらのデータを組み合わせて利用する例として、ヒトのPathwayの中で一番上の親のPathwayとその一つ下のPathway、2つの階層のPathwayの子のリストを取得し、その子の一つのPathwayの名前とGeneのリストを取得するコードを示します。

まずはヒトのPathwayを抽出してPythonのNetworkXを利用して有向グラフを作り、一番上の親をrootというノードにつなげたグラフを作ります。

import networkx as nx

human_pathway_ids = complete_list[complete_list["species"] == 'Homo sapiens']["st_id"]
human_pathway_hierarchy_df = pathway_hierarchy_df[pathway_hierarchy_df["parent_st_id"].isin(human_pathway_ids) & pathway_hierarchy_df["child_st_id"].isin(human_pathway_ids)]
human_pathway_graph = nx.from_pandas_edgelist(human_pathway_hierarchy_df, source="parent_st_id", target="child_st_id", create_using=nx.DiGraph())
root_pathways = [n for n, d in human_pathway_graph.in_degree() if d==0] 
root_edges = [("root", n) for n in root_pathways] 
human_pathway_graph.add_edges_from(root_edges)

これでヒトのPathwayの階層構造を示した有向グラフができました。これを使えばrootからの最小距離が2以下のnodeを列挙することで、一番上の親のPathwayとその一つ下のPathway、2階層のPathwayが取得できます。

selected_ids = nx.single_source_shortest_path_length(human_pathway_graph, source="root", cutoff=2)

selected_idsをprintするとこのような形になります。

length 1 {'R-HSA-109581': 2,
 'R-HSA-109582': 1,
 'R-HSA-112307': 2,
 'R-HSA-112315': 2,
 'R-HSA-112316': 1,
 'R-HSA-1181150': 2,
 'R-HSA-1187000': 2,
 'R-HSA-1266738': 1,
...

ここからsuccessors()を使って各Pathwayの子のPathwayを取得します。

for pathway_id, shortest_path_length in selected_ids.items():
  if shortest_path_length > 0:
    children = list(human_pathway_graph.successors(pathway_id))
    print("length", shortest_path_length, "parent", pathway_id, "children", children)

ちなみに、shortest_path_length=0はrootだけなのでskipしています。出力としては以下のようになります。

length 1 parent R-HSA-1852241 children ['R-HSA-1592230', 'R-HSA-5617833']
length 1 parent R-HSA-5357801 children ['R-HSA-109581', 'R-HSA-5218859']
length 1 parent R-HSA-1266738 children ['R-HSA-1181150', 'R-HSA-186712', 'R-HSA-381340', 'R-HSA-452723', 'R-HSA-525793', 'R-HSA-5619507', 'R-HSA-5682910', 'R-HSA-6805567', 'R-HSA-9616222', 'R-HSA-9675108', 'R-HSA-9690406']
length 1 parent R-HSA-4839726 children ['R-HSA-3247509']
length 1 parent R-HSA-9709957 children ['R-HSA-2187338', 'R-HSA-381753', 'R-HSA-9659379', 'R-HSA-9717189']
length 1 parent R-HSA-1474244 children ['R-HSA-1474228', 'R-HSA-1474290', 'R-HSA-1566948', 'R-HSA-1566977', 'R-HSA-216083', 'R-HSA-3000157', 'R-HSA-3000171', 'R-HSA-3000178', 'R-HSA-8941237']
length 1 parent R-HSA-9612973 children ['R-HSA-1632852', 'R-HSA-9613829', 'R-HSA-9615710']
length 1 parent R-HSA-397014 children ['R-HSA-390522', 'R-HSA-445355', 'R-HSA-5576891']
...

ここで試しにR-HSA-1592230の名前を出してみます。コードとしては以下のようになります。

print(complete_list[complete_list["st_id"] == "R-HSA-1592230"])

出力は以下の通りです。

               st_id              pathway_name       species
11395  R-HSA-1592230  Mitochondrial biogenesis  Homo sapiens

また、R-HSA-1592230の関連するgeneは以下のようにして出力できます。

print(pathway_gene_df[pathway_gene_df["st_id"] == "R-HSA-1592230"])

出力は以下の通りです。

               st_id    gene
65977  R-HSA-1592230   ALAS1
65978  R-HSA-1592230    APOO
65979  R-HSA-1592230   APOOL
65980  R-HSA-1592230    ATF2
65981  R-HSA-1592230   ATP5B
...              ...     ...
66066  R-HSA-1592230    TFAM
66067  R-HSA-1592230   TFB1M
66068  R-HSA-1592230   TFB2M
66069  R-HSA-1592230    TGS1
66070  R-HSA-1592230  TMEM11

[94 rows x 2 columns]

終わりに

Reactomeで公開されているデータを読み込んでPathwayの階層構造やPathwayに関連するGeneのデータを読み込む方法を紹介しました。最初はGraph Databaseから頑張って抽出しないといけないのかと思っていましたが、調べてみると簡単なことがわかりました。ただ、何も知らない状態ではどうしていいのかわからないことが多かったので記事にまとめてみました。同じように悩んでいる方の参考になれば幸いです。

The post ReactomeからPathwayの階層構造とPathwayに関連するGeneのデータを取得する first appeared on まったり勉強ノート.

論文は以下のものです。

[1] Petti, S., Bhattacharya, N., Rao, R., Dauparas, J., Thomas, N., Zhou, J., … Ovchinnikov, S. (2021). End-to-end learning of multiple sequence alignments with differentiable Smith-Waterman. BioRxiv, 2021.10.23.465204. https://doi.org/10.1101/2021.10.23.465204

また、著者の実装はこちらに公開されています。

https://github.com/spetti/SMURF

今回は主に私がJAXの勉強をしたかったということもあり、いくつか実装を作ってパフォーマンスを測定して、「JAXって速いの？」という疑問にある程度答えられればと思い、記事を書いています。今回の実装はすべてこちらにありますので参考にしてみてください。

https://github.com/shu65/blog-jax-notebook/blob/main/JAX_Smooth_Smith_Waterman.ipynb

また、計算時間測定はすべてGoogle Cloab上のCPUで行っています。

論文概要

この論文では教師なし学習によるコンタクト予測において、前処理で使われるSmith Watermanアルゴリズムを微分可能なものに置き換えて、Smith Watermanアルゴリズムの中で使われるパラメータ（置換スコア）も含めて学習する手法 SMURFを提案した論文です。論文自体にはコンタクト予測の精度なども書かれていますが、微分可能なSmith Watermanの紹介をメインにしたいため、今回は割愛します。

微分可能な Smith Waterman アルゴリズム「Smooth Smith Waterman」とは？

Smith Watermanアルゴリズムを微分可能にするためには、微分可能ではない関数を微分可能なものに置き換えて、近似することで実現します。まずは大本のSmith Watermanアルゴリズムの説明をしたあと、微分可能なものに変更する方法を紹介していきます。

Smith Watermanアルゴリズムとは

Smith Watermanアルゴリズム は2つのDNAやタンパク質の配列の類似度、特にローカルアライメントのスコアと呼ばれる類似度を計算するアルゴリズムです。ローカルアラインメントとは2配列間の類似度の高い部分的な文字列を発見するときに使われます。これは以下のように行列の要素を計算する動的計画法 (Dynamic Programming, DP) により計算します。

$$ H_{i0} = H_{0j} = 0 \\ H_{ij} = \max\begin{cases} H_{i-1,j-1} + s(a_i,b_j), \\ H_{i-k,j} + g, \\ H_{i,j-l} +g, \\
0 \\
\end{cases} $$

\( s(a_i,b_j) \) は 配列Aのi番目の文字と配列Bのj番目の文字の置換スコアと呼ばれるもので、同じ文字、もしくは類似度の高い文字のペアはプラス、類似度の低い文字のペアはマイナスにするのが一般的です。また、 \( g \) はギャップペナルティと呼ばれるもので、1文字飛ばしのペナルティを表します。

Smith Watermanアルゴリズムを微分可能にする

先ほど説明したとおり、Smith Watermanアルゴリズムではmax関数があります。この部分が微分可能ではないため、SmithWatermanアルゴリズムは微分可能ではありません。このため、このmax関数を微分可能な何等かの関数で置き換える必要があります。この論文ではmax関数を「logsumexp」で置き換えることで微分可能にします。

logsumexpはmax関数を滑らかに近似するための関数として使われる関数で、微分可能な関数です。このためmax関数を logsumexp に置き換えればSmith Watermanアルゴリズムの計算全体が微分可能になります。論文中ではこの微分可能なSmtth Watermanアルゴリズムを「Smooth Smith Waterman」と呼んでいます。
なぜlogsumexpがmax関数の近似になるかを詳しく知りたい方は、こちらのブログ記事がわかりやすかったのでお勧めです。

numpyによるシンプルな Smooth Smith Waterman

後ほどJAXの実装を示しますが、高速化したあとのJAXのコードは初見では分かりづらいため、先にシンプルなnumpyの実装を示します。この実装は著者の実装にあわせつつ、numpyとのパフォーマンス実装をするために以下のようにしています。

1. 配列Aと配列Bの全文字ペアの置換スコアの行列 score_matrix（置換スコアの行列のサイズは|A|×|B|）と2つの配列の長さlengths、その他のパラメータを入力とする
2. この記事では勾配を計算できないnumpyとの比較のために、著者実装ではscore_matrixの勾配を返すのに対して、今回の記事では2配列の最大スコアを返す。

Smith Watermanアルゴリズムをご存じの方は戸惑うかもしれませんが、Smooth Smith Watermanアルゴリズムではscore_matrixの勾配を出力として返す関数になっています。このため、あらかじめ配列Aと配列Bの全文字ペアの置換スコアの行列を用意して入力にします。
このため、配列Aと配列Bは入力に出てきませんし、PAMやBLOSUMなどの置換スコアもでてきません。

このSmooth Smith Watermanアルゴリズムをシンプルにnumpyで実装すると以下の通りになります。

def sw_np(NINF=-1e30):
    
    def _logsumexp(y, axis):
        y = np.maximum(y,NINF)
        return y.max(axis) + np.log(np.sum(np.exp(y - y.max(axis, keepdims=True)), axis=axis))

    def _soft_maximum(x, temp, axis=None):
        return temp*_logsumexp(x/temp, axis)

    def _sw(score_matrix, lengths, gap=0, temp=1.0):
        real_a, real_b = lengths
        hij = np.full((real_a + 1, real_b + 1), fill_value=NINF, dtype=np.float32)
        for i in range(real_a):
            for j in range(real_b):
                s = score_matrix[i, j]
                m = hij[i, j] + s
                g0 = hij[i + 1, j] + gap
                g1 = hij[i, j + 1] + gap

                h = np.stack([m, g0, g1, s], -1)
                hij[i + 1, j + 1] = _soft_maximum(h, temp=temp, axis=-1)
        hij = hij[1:, 1:]
        score = _soft_maximum(hij, temp=temp)
        return score
    return _sw

こちらの実装で通常のSmith Watermanアルゴリズムと違う点は以下の2点です

1. DPの行列の要素更新のところでmax関数をlogsumexpで実装した_soft_maximum()という関数に置き換えている。
2. DPの行列の各要素を入れるところで最大値を取るところで0以下にならないようにmax関数の入力の一つとして0を入れるところを、置換スコア(s)を入れている。

1が微分可能とするための改良した箇所です。一方、2に関しては私が読み飛ばしてしまった可能性がありますが、特に論文中に説明が見当たらなかった変更点です。なんとなくSmooth Smith Watermanアルゴリズムを使って深層学習のモデル更新をするときにうまく勾配が置換スコアに流れるようにするためでは？と思っているのですが、未確認な状態です。何かご存じの方がいれば教えていただければと思っています。

この実装を実行したときの計算時間を%timeで測定すると以下の通りです。

CPU times: user 735 ms, sys: 4 ms, total: 739 ms
Wall time: 743 ms

JAXを使ったSmooth Smith アルゴリズム

ここからnumpyの部分をJAXに置き換えてSmooth Smith Watermanアルゴリズムを実装し、徐々に改良していくという順番で説明していきます。まずはJAXをご存じない方のためにJAXを簡単に説明します。

JAXってなに？

JAXはPythonやnumpyの関数を微分可能なものにし、XLAというコンパイラを使ってGPUやTPUで実行で実行できるようにしたライブラリです。
JAXでは勾配が計算できることと、jitをはじめとした様々な高速化する仕組みが用意されているため、最近論文で利用しているケースが増えてきた印象です。特に今回紹介した論文のような、従来では微分可能でなかった計算を微分可能なものに置き換え、深層学習のモデル学習の中で利用するという手法の実装にJAXが使われるのをよく目にします。今回紹介したものの他には BRAXがあります。

https://github.com/google/brax

単純なJAX実装

JAXはnumpyの関数と同じAPIの関数があるので、まずはそれをそのまま利用してみます。

def sw_v0(NINF=-1e30):
    
    def _logsumexp(y, axis):
        y = jnp.maximum(y,NINF)
        return jax.nn.logsumexp(y, axis=axis)

    def _soft_maximum(x, temp, axis=None):
        return temp*_logsumexp(x/temp, axis)

    def _sw(score_matrix, lengths, gap=0, temp=1.0):
        real_a, real_b = lengths
        hij = jnp.full((real_a + 1, real_b + 1), fill_value=NINF, dtype=jnp.float32)
        for i in range(real_a):
            for j in range(real_b):
                s = score_matrix[i, j]
                m = hij[i, j] + s
                g0 = hij[i + 1, j] + gap
                g1 = hij[i, j + 1] + gap
                h = jnp.stack([m, g0, g1, s], -1)
                hij = hij.at[i + 1, j + 1].set(_soft_maximum(h, -1))
        hij = hij[1:, 1:]
        score = _soft_maximum(hij)
        return score
    return _sw

これも動くには動くのですが、あまりにも遅いため、まったく使い物になりません。このためJAXを使う際はもう少し真面目に高速に動くアルゴリズムで実装する必要があります。

Striped Smith-Watermanベースの実装

論文でも紹介されているStriped Smith-Watermanベースで実装してみます。 Smith-Waterman アルゴリズムをSIMDなどで並列化する方法として、依存関係のないDP行列の斜めのセルを同時に埋めていくという方法がしばしば取られます。詳しくはこちらをご覧ください。

Farrar, M. (2007). Striped Smith-Waterman speeds database searches six times over other SIMD implementations. Bioinformatics (Oxford, England), 23(2), 156–161. https://doi.org/10.1093/bioinformatics/btl582

これをJAXで実装するにあたり、著者はDP行列を回転させ、依存関係のない斜めに並んだセルを横1列に並べて計算するようにしています。

こうすることで内側のforループをJAXのベクトルの計算で実行できるようにしています。個人的にはここがこの論文の最大の貢献な気がしています。具体的にJAXで実装すると以下の通りです。

def sw_v1(unroll=2, NINF=-1e30):
        
    def _rotate(score_matrix):
        a,b = score_matrix.shape
        n,m = (a+b-1),(a+b)//2
        ar,br = jnp.arange(a)[::-1,None], jnp.arange(b)[None,:]
        i,j = (br-ar)+(a-1),(ar+br)//2
        rotated_score_matrix = jnp.full([n,m],NINF).at[i,j].set(score_matrix)
        reverse_idx = (i, j)
        return rotated_score_matrix, reverse_idx
    
    def _rotate_in_reverse(rotated_dp_matrix, reverse_idx):
        return rotated_dp_matrix[reverse_idx]

    def _logsumexp(y, axis):
        y = jnp.maximum(y,NINF)
        return jax.nn.logsumexp(y, axis=axis)

    def _soft_maximum(x, temp, axis=None):
        return temp*_logsumexp(x/temp, axis)

    def _step(prev, gap_cell_condition, rotated_score_matrix, gap, temp):
        h2,h1 = prev   # previous two rows of scoring (hij) mtx
        h1_T = jax.lax.cond(
            gap_cell_condition,
            lambda x: jnp.pad(x[:-1], [1,0], constant_values=(NINF,NINF)),
            lambda x: jnp.pad(x[1:], [0,1], constant_values=(NINF,NINF)),
            h1,
        )

        a = h2 + rotated_score_matrix
        g0 = h1 + gap
        g1 = h1_T + gap
        s = rotated_score_matrix

        h0 = jnp.stack([a, g0, g1, s], -1)
        h0 = _soft_maximum(h0, temp, -1)
        return (h1,h0), h0

    def _sw(score_matrix, lengths, gap=0, temp=1.0):
        rotated_score_matrix, reverse_idx = _rotate(score_matrix)
        
        a,b = score_matrix.shape
        n,m = rotated_score_matrix.shape
        
        gap_cell_condition = (jnp.arange(n)+a%2)%2
        prev = (jnp.full(m, NINF), jnp.full(m, NINF))
        rotated_hij = []
        for i in range(n):
            prev, h = _step(prev, gap_cell_condition[i], rotated_score_matrix[i], gap, temp)
            rotated_hij.append(h)
        rotated_hij = jnp.stack(rotated_hij)
        hij = _rotate_in_reverse(rotated_hij, reverse_idx)
        score = _soft_maximum(hij, temp=temp)
        return score
    return _sw

この実装では置換行列score_matrixを回転させて、DP行列のセルを埋めていき、そのあとDP行列元の方向に戻すということをしています。
回転させたときの注意点として、DP行列の列番号が偶数か奇数かでギャップペナルティのスコアを加算するセルの相対座標が変わります。このため、jax.lax.cond()を利用して使うセルを分岐しています。

この実装をそのまま実行したときとjitを利用したときの計算時間は以下の通りです。

jax default first call
CPU times: user 17.7 s, sys: 177 ms, total: 17.8 s
Wall time: 17.8 s
jax default second call
CPU times: user 17.6 s, sys: 153 ms, total: 17.7 s
Wall time: 17.7 s


jax jit first call
CPU times: user 2min 20s, sys: 715 ms, total: 2min 21s
Wall time: 2min 20s
jax jit second call
CPU times: user 1.98 ms, sys: 0 ns, total: 1.98 ms
Wall time: 1.81 ms

jitなしでそのまま実行するのはnumpyよりもかなり遅い印象です。またjitを使う場合も最初の呼び出しはコンパイルが走ることもあり、jitなしに比べるとさらに遅くなっています。さすがに1回目とはいえ、ここまで時間がかかると使いづらいと思われます。このため、まだ工夫する必要があります。

外側のforループをjax.lax.scan()に置き換える

1つ前の実装で遅い原因がどこか？というとforループです。これを速くする方法としてJAXのforループと類似する処理を実行するための関数を利用します。今回はforループ部分を jax.lax.scan() に置き換えます。

実装は以下の通りです。

def sw_v2(unroll=2, NINF=-1e30):

    def _rotate(score_matrix):
        a,b = score_matrix.shape
        n,m = (a+b-1),(a+b)//2
        ar,br = jnp.arange(a)[::-1,None], jnp.arange(b)[None,:]
        i,j = (br-ar)+(a-1),(ar+br)//2
        rotated_score_matrix = jnp.full([n,m],NINF).at[i,j].set(score_matrix)
        reverse_idx = (i, j)
        return rotated_score_matrix, reverse_idx

    def _prepare_scan_inputs(score_matrix, rotated_score_matrix, gap, temp):
        def scan_f(prev, scan_xs):
            h2, h1 = prev
            h1_T = jax.lax.cond(
                scan_xs["gap_cell_condition"],
                lambda x: jnp.pad(x[:-1], [1,0], constant_values=(NINF,NINF)),
                lambda x: jnp.pad(x[1:], [0,1], constant_values=(NINF,NINF)),
                h1,
            )
            a = h2 + scan_xs["rotated_score_matrix"]
            g0 = h1 + gap
            g1 = h1_T + gap
            s = scan_xs["rotated_score_matrix"]

            h0 = jnp.stack([a, g0, g1, s], -1)
            h0 = _soft_maximum(h0, temp, -1)
            return (h1,h0), h0
        
        a,b = score_matrix.shape
        n,m = rotated_score_matrix.shape

        scan_xs = {
            "rotated_score_matrix": rotated_score_matrix,
            "gap_cell_condition": (jnp.arange(n)+a%2)%2
        }
        scan_init = (jnp.full(m, NINF), jnp.full(m, NINF))
        return scan_f, scan_xs, scan_init

    def _rotate_in_reverse(rotated_dp_matrix, reverse_idx):
        return rotated_dp_matrix[reverse_idx]

    def _logsumexp(y, axis):
        y = jnp.maximum(y,NINF)
        return jax.nn.logsumexp(y, axis=axis)

    def _soft_maximum(x, temp, axis=None):
        return temp*_logsumexp(x/temp, axis)
    
    def _sw(score_matrix, lengths, gap=0, temp=1.0):
        rotated_score_matrix, reverse_idx = _rotate(score_matrix)
        scan_f, scan_xs, scan_init = _prepare_scan_inputs(score_matrix, rotated_score_matrix, gap, temp)
        rotated_hij = jax.lax.scan(scan_f, scan_init, scan_xs, unroll=unroll)[-1]
        hij = _rotate_in_reverse(rotated_hij, reverse_idx)
        score = _soft_maximum(hij, temp, axis=None)
        return score
    return _sw

この実装でforループがなくなりました。実行した結果は以下の通りです。

jax default first call
CPU times: user 739 ms, sys: 18 ms, total: 757 ms
Wall time: 758 ms
jax default second call
CPU times: user 666 ms, sys: 1.98 ms, total: 668 ms
Wall time: 671 ms

jax jit first call
CPU times: user 1 s, sys: 5.01 ms, total: 1.01 s
Wall time: 1.01 s
jax jit second call
CPU times: user 339 µs, sys: 989 µs, total: 1.33 ms
Wall time: 1.14 ms

先ほどに比べるとjitなしでも速くなりましたが、jitありの1回目の実行もかなり速くなった印象です。これなら十分使えるのではないか？と思っています。

jax.lax.condの置き換え

著者の実装では jax.lax.cond()を使わずに加算と乗算だけで実装されています。試しに同様の実装にしたバージョンも示します。具体的な実装は以下の通りです。

def sw_v3(unroll=2, NINF=-1e30):

    def _rotate(score_matrix):
        a,b = score_matrix.shape
        n,m = (a+b-1),(a+b)//2
        ar,br = jnp.arange(a)[::-1,None], jnp.arange(b)[None,:]
        i,j = (br-ar)+(a-1),(ar+br)//2
        rotated_score_matrix = jnp.full([n,m],NINF).at[i,j].set(score_matrix)
        reverse_idx = (i, j)
        return rotated_score_matrix, reverse_idx

    def _prepare_scan_inputs(score_matrix, rotated_score_matrix, gap, temp):
        def scan_f(prev, scan_xs):
            h2, h1 = prev
            h1_T = _get_prev_gap_cell_score(
                scan_xs["gap_cell_condition"],
                jnp.pad(h1[:-1], [1,0], constant_values=(NINF,NINF)),
                jnp.pad(h1[1:], [0,1], constant_values=(NINF,NINF)),
            )
            a = h2 + scan_xs["rotated_score_matrix"]
            g0 = h1 + gap
            g1 = h1_T + gap
            s = scan_xs["rotated_score_matrix"]

            h0 = jnp.stack([a, g0, g1, s], -1)
            h0 = _soft_maximum(h0, temp, -1)
            return (h1,h0), h0
        
        a,b = score_matrix.shape
        n,m = rotated_score_matrix.shape

        scan_xs = {
            "rotated_score_matrix": rotated_score_matrix,
            "gap_cell_condition": (jnp.arange(n)+a%2)%2
        }
        scan_init = (jnp.full(m, NINF), jnp.full(m, NINF))
        return scan_f, scan_xs, scan_init

    def _rotate_in_reverse(rotated_dp_matrix, reverse_idx):
        return rotated_dp_matrix[reverse_idx]

    def _logsumexp(y, axis):
        y = jnp.maximum(y,NINF)
        return jax.nn.logsumexp(y, axis=axis)

    def _soft_maximum(x, temp, axis=None):
        return temp*_logsumexp(x/temp, axis)

    def _get_prev_gap_cell_score(cond, true, false): 
        return cond*true + (1-cond)*false
    
    def _sw(score_matrix, lengths, gap=0, temp=1.0):
        rotated_score_matrix, reverse_idx = _rotate(score_matrix)
        scan_f, scan_xs, scan_init = _prepare_scan_inputs(score_matrix, rotated_score_matrix, gap, temp)
        rotated_hij = jax.lax.scan(scan_f, scan_init, scan_xs, unroll=unroll)[-1]
        hij = _rotate_in_reverse(rotated_hij, reverse_idx)
        score = _soft_maximum(hij, temp, axis=None)
        return score
    return _sw

この時のパフォーマンスは以下の通りです。

jax defaujax default first call
CPU times: user 599 ms, sys: 1.99 ms, total: 601 ms
Wall time: 608 ms
jax default second call
CPU times: user 599 ms, sys: 3.02 ms, total: 602 ms
Wall time: 607 ms

jax jit first call
CPU times: user 940 ms, sys: 2.01 ms, total: 942 ms
Wall time: 947 ms
jax jit second call
CPU times: user 4.9 ms, sys: 0 ns, total: 4.9 ms
Wall time: 3.41 ms

jitなしの時は速くなっている印象ですが、jitありのときは少し遅くなっています。ただ、何度か実行してみると逆転することもあるようなので、誤差の範囲かもしれません。また、JAX特有のパフォーマンス測定のお作法をし忘れている可能性もあります。もしご存じの方があればコメントいただければと思います。

Batch実行用の実装

著者のSmooth Smith Watermanは2つの配列のペアを1つだけ実行するのではなく、複数のペアをまとめて実行することを想定されて実装してあります。ここでも同様に複数のペアをまとめて計算するのもやってみようと思います。

簡単な実装

複数のペアをまとめて実装する際、ペア毎に配列の長さが違っても動作するようにします。このため、置換スコアのうち必要な部分だけmaskするようにします。

実装は以下の通りです。

def sw_v4(unroll=2, NINF=-1e30):
    
    def _make_mask(score_matrix, lengths):
        a,b = score_matrix.shape
        real_a, real_b = lengths
        mask = (jnp.arange(a) < real_a)[:,None] * (jnp.arange(b) < real_b)[None,:]
        return mask

    def _rotate(score_matrix):
        a,b = score_matrix.shape
        n,m = (a+b-1),(a+b)//2
        ar,br = jnp.arange(a)[::-1,None], jnp.arange(b)[None,:]
        i,j = (br-ar)+(a-1),(ar+br)//2
        rotated_score_matrix = jnp.full([n,m],NINF).at[i,j].set(score_matrix)
        reverse_idx = (i, j)
        return rotated_score_matrix, reverse_idx

    def _prepare_scan_inputs(score_matrix, rotated_score_matrix, gap, temp):
        def scan_f(prev, scan_xs):
            h2, h1 = prev
            h1_T = _get_prev_gap_cell_score(
                scan_xs["gap_cell_condition"],
                jnp.pad(h1[:-1], [1,0], constant_values=(NINF,NINF)),
                jnp.pad(h1[1:], [0,1], constant_values=(NINF,NINF)),
            )
            a = h2 + scan_xs["rotated_score_matrix"]
            g0 = h1 + gap
            g1 = h1_T + gap
            s = scan_xs["rotated_score_matrix"]

            h0 = jnp.stack([a, g0, g1, s], -1)
            h0 = _soft_maximum(h0, temp, -1)
            return (h1,h0), h0
        
        a,b = score_matrix.shape
        n,m = rotated_score_matrix.shape

        scan_xs = {
            "rotated_score_matrix": rotated_score_matrix,
            "gap_cell_condition": (jnp.arange(n)+a%2)%2
        }
        scan_init = (jnp.full(m, NINF), jnp.full(m, NINF))
        return scan_f, scan_xs, scan_init

    def _rotate_in_reverse(rotated_dp_matrix, reverse_idx):
        return rotated_dp_matrix[reverse_idx]

    def _logsumexp(y, axis):
        y = jnp.maximum(y,NINF)
        return jax.nn.logsumexp(y, axis=axis)

    def _logsumexp_with_mask(y, axis, mask):
        y = jnp.maximum(y,NINF)
        return y.max(axis) + jnp.log(jnp.sum(mask * jnp.exp(y - y.max(axis, keepdims=True)), axis=axis))

    def _soft_maximum(x, temp, axis=None):
        return temp*_logsumexp(x/temp, axis)

    def _soft_maximum_with_mask(x, temp, mask, axis=None):
        return temp*_logsumexp_with_mask(x/temp, axis, mask)

    def _get_prev_gap_cell_score(cond, true, false): 
        return cond*true + (1-cond)*false
    
    def _sw(score_matrix, lengths, gap=0, temp=1.0):
        mask = _make_mask(score_matrix, lengths)
        masked_score_matrix = score_matrix + NINF * (1 - mask)
        rotated_score_matrix, reverse_idx = _rotate(masked_score_matrix)
        scan_f, scan_xs, scan_init = _prepare_scan_inputs(score_matrix, rotated_score_matrix, gap, temp)
        rotated_hij = jax.lax.scan(scan_f, scan_init, scan_xs, unroll=unroll)[-1]
        hij = _rotate_in_reverse(rotated_hij, reverse_idx)
        score = _soft_maximum_with_mask(hij, temp, mask=mask, axis=None)
        return score
    return _sw

この実装をペアの数分、forループで計算していくようにします。

def batch_sw_v0(NINF=-1e30):
    def _batch_sw(batch_score_matrix, batch_lengths, gap=0, temp=1.0):
        n_batches = batch_score_matrix.shape[0]
        sw_func = jax.jit(sw_v4())
        ret = [sw_func(batch_score_matrix[i], batch_lengths[i], gap, temp) 
               for i in range(n_batches)]
        return jnp.array(ret)
    return _batch_sw

これを実行すると計算時間は以下の通りでした。

batch jax default first call
CPU times: user 1.31 s, sys: 13 ms, total: 1.33 s
Wall time: 1.32 s
batch jax default second call
CPU times: user 1.3 s, sys: 5.02 ms, total: 1.3 s
Wall time: 1.3 s

batch jax default first call
CPU times: user 10min 43s, sys: 2.99 s, total: 10min 46s
Wall time: 10min 45s
batch jax default second call
CPU times: user 279 ms, sys: 2 ms, total: 281 ms
Wall time: 281 ms

forループでそのまま実装すると、jitありのときはやはり1度目の実行に非常に時間がかかるようです。このため、この部分を速くします。

forループをjax.vmap()で置き換える

ここではforループを jax.vmap() で置き換えます。

def batch_sw_v1(unroll=2, NINF=-1e30):
    sw_func = sw_v4(unroll=unroll, NINF=NINF)
    batch_sw_func = jax.vmap(sw_func, (0, 0, None, None))
    return batch_sw_func

この時の計算時間は以下の通りです。

batch jax default first call
CPU times: user 1.04 s, sys: 11 ms, total: 1.05 s
Wall time: 1.03 s
batch jax default second call
CPU times: user 1.04 s, sys: 7.97 ms, total: 1.04 s
Wall time: 1.01 s

batch jax default first call
CPU times: user 1.51 s, sys: 10 ms, total: 1.52 s
Wall time: 1.5 s
batch jax default second call
CPU times: user 120 ms, sys: 9 µs, total: 120 ms
Wall time: 97 ms

先ほどと比べるとかなり高速化できました。ちなみにこれがほぼ著者の実装と同じものになります。

結果まとめ

ここまでの計算時間の結果をまとめると以下の通りです。

各実装を比較するとforループのありなしでかなり実行時間やコンパイル時間が変化していることがわかります。このためnumpyの実装をそのままJAXにすればそれだけで速くなることはまずなさそうです。また、何も考えずに実装してjitを使うと、コンパイル時間が長すぎて使い物にならないというケースが多そうな気がしています。このため、JAXを使いこなすにはどのような計算は遅いかを理解して使うことが重要そうな印象です。

おわりに

今回、初のJAX使用だったため、パフォーマンス測定や高速化にはもっとやり方があるかもしれないと思っています。もしお気づきの点がありましたら気兼ねなくコメントいただければと思っています。

次はできればJAXとPyTorchのjitとどちらが速いのか試せればと思っています。

The post JAXによる微分可能Smith Watermanアルゴリズムのパフォーマンス測定 first appeared on まったり勉強ノート.