SuperAPI AI 工具集

fair-esm2

通用与其它

查看项目

更新于 10-15 22:15

技能介绍

用途:使用 Meta AI 的 ESM-2(fair-esm 包,import esm)为蛋白质序列提取嵌入,供下游机器学习使用。注意与 Biohub 的 esm 分支区分,后者提供 esm.models.esmfold2,属于 esmfold2 技能。

主要能力:逐残基或逐序列嵌入;掩码语言模型打分(可对突变位点做 logit 比较);接触预测(return_contacts=True)。

环境要求:Python 3.8+(推荐 3.11),CUDA 11.7+(推荐 12.x),显存 8GB(8M 模型)、16GB(650M)、24GB+(650M/3B)。

用法:安装 pip install fair-esm,加载 esm.pretrained.esm2_t33_650M_UR50D(),用 alphabet.get_batch_converter() 编码序列,调用 model(toks.cuda(), repr_layers=[33]) 获取 out["representations"][33],形状为 (B, L+2, D),切片 [:, 1:-1, :] 去掉 BOS/EOS;逐序列可取均值。接触预测需 return_contacts=True,输出 (B, L, L)。

模型:esm2_t6_8M_UR50D(6 层,320 维,8M)、esm2_t33_650M_UR50D(33 层,1280 维,650M,默认)、esm2_t36_3B_UR50D(36 层,2560 维,3B,需 24GB+)。

远程计算:650M 模型需 ≥16GB 显存,并需预缓存 .pt 权重或可访问 dl.fbaipublicfiles.com;可提交 GPU 作业,设置 TORCH_HOME 指向缓存挂载,作业完成后保存产物。

限制:首次调用可能因 torch.hub 下载权重而变慢;若报 ModuleNotFoundError: No module named 'esm.models',说明误用了 Biohub 分支。结构预测请使用 esmfold2 技能。