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

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

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

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

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

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

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

代碼開源路徑：

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

昇思MindSpore：

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

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