近日，昇思MindSpore與昌平實驗室、北京大學生物醫(yī)學前沿創(chuàng)新中心（BIOPIC）和化學與分子工程學院、深圳灣實驗室高毅勤教授課題組及鵬城實驗室陳杰團隊基于全場景AI框架昇思MindSpore實現(xiàn)AlphaFold2蛋白質(zhì)結(jié)構(gòu)訓練。繼2021年11月發(fā)布推理工具后，本次訓練意味著國產(chǎn)AI框架具備了強大的AI for Science底層軟件能力，同時也為相關(guān)科研工作者提供新的選擇。該聯(lián)合工作依托鵬城云腦II 昇騰AI 集群進行，單步迭代性能提升超過60%，TM-score達85分（國際權(quán)威評測數(shù)據(jù)集CASP14）。相關(guān)訓練代碼已在昇思MindSpore社區(qū)開源，后續(xù)也會在Openl啟智社區(qū)進行開源并定期擴展與維護。

T1052-D1 預測結(jié)構(gòu)圖（左）CASP14 87 targets TM-score 對比（右）

蛋白質(zhì)結(jié)構(gòu)預測是獲得蛋白質(zhì)功能結(jié)構(gòu)和構(gòu)象的過程，近半個世紀以來，這一問題一直被譽為“21世紀的生物物理學”最重要的課題之一。在過去，因蛋白質(zhì)構(gòu)象數(shù)量巨大，計算過程復雜，通過AI來對蛋白質(zhì)結(jié)構(gòu)進行預測一直未能取得實質(zhì)性突破，獲取蛋白質(zhì)空間結(jié)構(gòu)的方法仍然以冷凍電鏡、X-ray等實驗技術(shù)為主，單個蛋白質(zhì)的觀測成本高達數(shù)月及數(shù)百萬人民幣。直至AlphaFold2的出現(xiàn)，使得這一問題迎來新的曙光。AlphaFold2憑借其接近實驗精度的成績?nèi)〉?span>CASP14蛋白質(zhì)空間結(jié)構(gòu)預測比賽的榜首，這一成就也被Nature譽為“前所未有的進步”。

2021年7月DeepMind宣布對AlphaFold2的推理代碼進行開源，昇思與高毅勤課題組第一時間對其進行了復現(xiàn)及優(yōu)化，并于同年11月開源了基于昇思MindSpore的推理工具，效率同比提升2-3倍。由于開源范圍僅限推理，相關(guān)從業(yè)者無法基于此進行優(yōu)化，因此許多團隊積極地投入訓練過程的復現(xiàn)。AlphaFold2模型本身存在內(nèi)存需求大，數(shù)據(jù)處理繁瑣，控制編譯復雜等特點，對基礎AI框架存在著巨大挑戰(zhàn)。

近期，昇思MindSpore聯(lián)合高毅勤課題組、鵬城實驗室陳杰團隊全面打通AlphaFold2的訓練。采用昇騰基礎軟硬件平臺后，在混合精度下，單步迭代時間由20秒縮短到12秒，性能提升超過60%。依托昇思MindSpore內(nèi)存復用能力, 訓練序列長度由384提升至512。

為了盡可能客觀地評估訓練結(jié)果，昇思MindSpore選取了AlphaFold2論文附錄中提到的87條驗證集進行驗證，平均TM-score達到85分，基本持平AlphaFold2。

昇思MindSpore對蛋白質(zhì)結(jié)構(gòu)預測訓練推理的支持填補了國產(chǎn)AI軟硬件的空白。在訓練精度接近AlphaFold2的基礎上，昇思MindSpore將在算法、規(guī)模和軟硬件支持等方向上持續(xù)改進，并計劃開放共享訓練數(shù)據(jù)集供同仁使用。昇思MindSpore也期望與更多學術(shù)界和工業(yè)界伙伴合作，進一步提升模型精度、擴展應用場景。

代碼開源路徑：

https://gitee.com/mindspore/mindscience/tree/dev/MindSPONGE/mindsponge/fold

昇思MindSpore：

Gitee：https://gitee.com/mindspore/mindspore

Github：https://github.com/mindspore-ai/mindspore