昇思MindSpore再突破！蛋白质结构预测训练推理全流程开源，助力生物医药发展
昇思MindSpore再突破！蛋白质结构预测训练推理全流程开源，助力生物医药发展
近日，昇思MindSpore团队与昌平实验室、北京大学生物医学前沿创新中心（BIOPIC）和化学与分子工程学院、深圳湾实验室高毅勤教授课题组及鹏城实验室陈杰团队****基于全场景AI框架昇思MindSpore实现AlphaFold2蛋白质结构训练。
这是继2021年11月发布推理工具之后又一次突破，本次训练意味着国产AI框架具备了强大的AI for Science底层软件能力，同时也为相关科研工作者提供新的选择。
该联合工作依托鹏城云脑II、昇腾AI Atlas集群进行，单步迭代性能提升40%，TM-score达85分（国际权威评测数据集CASP14）。相关训练代码将在昇思MindSpore社区进行开源并定期扩展与维护。
蛋白质结构预测是获得蛋白质功能结构和构象的过程，近半个世纪以来，这一问题一直被誉为“21世纪的生物物理学”最重要的课题之一。
在过去，因蛋白质构象数量巨大，计算过程复杂，通过AI来对蛋白质结构进行预测一直未能取得实质性突破，获取蛋白质空间结构的方法仍然以冷冻电镜、X-ray等实验技术为主，单个蛋白质的观测成本高达数月及数百万人民币。
直至AlphaFold2的出现，使得这一问题迎来新的曙光。AlphaFold2凭借其接近实验精度的成绩取得CASP14蛋白质空间结构预测比赛的榜首，这一成就也被Nature誉为“前所未有的进步”。
2021年7月DeepMind团队宣布对AlphaFold2的推理源代码进行开源，昇思团队与高毅勤课题组第一时间对其进行了复现及优化，并于同年11月开源了基于昇思MindSpore的推理工具，效率同比提升2-3倍。
由于开源范围仅限推理，相关从业者无法基于此进行优化，因此许多团队积极地投入训练过程的复现（Pytorch、JAX等）。AlphaFold2模型本身存在内存需求大，数据处理繁琐，控制编译复杂等特点，对基础AI框架存在着巨大挑战。
近期，昇思MindSpore研发团队联合高毅勤课题组（刘思睿，张骏等）、鹏城实验室全面打通AlphaFold2的训练。相同配置下采用混合精度，单步迭代时间由20秒（V100）优化为12秒（昇腾910），性能提升40%。依托昇思MindSpore内存复用能力, 训练序列长度由384提升至512。
为了尽可能客观地评估训练结果，我们选取了AlphaFold2论文附录中提到的87条验证集进行验证。如图2所示，平均TM-score达到85分，基本持平AlphaFold2。
表1. T1052-D1推理精度
图.1 T1052-D1 预测结构图
图2. CASP14 87 targets TM-score 对比
昇思MindSpore对蛋白质结构预测训练推理的支持填补了国产AI软硬件的空白。
在训练精度接近AlphaFold2的基础上，研发团队将在算法、规模和软硬件支持等方向上持续改进，并计划开放共享训练数据集供同仁使用。
同时也期望与更多学术界和工业界伙伴合作，进一步提升模型精度、扩展应用场景。
代码开源路径
https://gitee.com/mindspore/mindscience/tree/dev/MindSPONGE/mindsponge/fold
MindSpore官方资料
官方QQ群 : 486831414
论坛：https://bbs.huaweicloud.com/forum/forum-1076-1.html
Gitee : https : //gitee.com/mindspore/mindspore