镜像名称

esmfold-rtx5090-cu128-py310

镜像简介

本镜像用于在 NVIDIA RTX 5090 / Blackwell 架构 GPU 上运行 ESMFold 蛋白质结构预测任务，同时兼容 Facebook Research ESM 相关示例代码。

• 功能： 该镜像主要用于基于蛋白质氨基酸序列进行结构预测，支持通过 HuggingFace Transformers 加载 facebook/esmfold_v1 模型，也可安装并运行 facebookresearch/esm 源码中的相关工具。

• 特点：

  • 预装 Python 3.10、PyTorch CUDA 12.8 环境；
  • 支持 RTX 5090 的 sm_120 GPU 架构；
  • 预装 ESMFold 推理所需依赖；
  • 支持 Jupyter Notebook / IPython 交互式运行；
  • 可用于蛋白质序列结构预测、PDB 文件生成、模型验证等任务。

环境与依赖

本镜像构建和运行所需的基础环境如下：

• 操作系统： Linux x86_64

• Python 版本： Python 3.10

• 框架及版本：

  • PyTorch 2.11.0
  • Transformers
  • fair-esm

• CUDA 版本： CUDA 12.8

• GPU 支持：

  • NVIDIA RTX 5090
  • CUDA Compute Capability sm_120

• 其他依赖：

  • numpy<2
  • scipy
  • pandas
  • matplotlib
  • tqdm
  • biopython
  • biotite==0.39.0
  • safetensors
  • accelerate
  • ipykernel

如果需要运行 ESM inverse folding / ESM-IF1 示例，还需要：

• torch_geometric
• pyg_lib
• torch_scatter
• torch_sparse
• torch_cluster

环境介绍

本镜像已安装好emsfold环境，测试可运行examples下的contact_prediction.ipynb、esm_structural_dataset.ipynb、sup_variant_prediction.ipynb

环境验证代码

1. 验证 PyTorch 与 GPU

python - <<'PY'
import torch

print("torch:", torch.__version__)
print("cuda version:", torch.version.cuda)
print("cuda available:", torch.cuda.is_available())

if torch.cuda.is_available():
    print("device:", torch.cuda.get_device_name(0))
    print("capability:", torch.cuda.get_device_capability(0))
    print("arch list:", torch.cuda.get_arch_list())

    x = torch.randn(1024, 1024, device="cuda")
    y = x @ x
    print("matmul ok:", y.mean().item())
PY

预期结果中应包含：

cuda available: True
device: NVIDIA GeForce RTX 5090
capability: (12, 0)

2. 验证 ESM 包

python - <<'PY'
import esm
print("esm path:", esm.__file__)
PY