HOOOS

显存不够也能玩转AI制药：本地低配环境搭建 RFdiffusion + ProteinMPNN 工作流指南

0 5 药研极客 蛋白质设计
Apple

作为蛋白质 de novo 设计领域的“黄金搭档”，RFdiffusion（负责骨架生成）和 ProteinMPNN（负责序列设计）几乎是目前计算生物学研究的标配。然而，官方文档中动辄要求 A100 或 24G 显存显卡的配置，让许多只有单张消费级游戏卡（如 8G 显存的 RTX 3060/4060）甚至只有 CPU 的开发者望而却步。

实际上，通过合理的参数调整、依赖精简以及合理的软硬件协同，我们完全可以在**本地低配环境（甚至无 GPU 环境）**下顺畅运行这套工作流。本文将手把手带你完成环境搭建与性能压榨。

核心痛点与低配优化思路

在低配环境下运行该工作流，主要面临两个瓶颈：

  1. RFdiffusion 的显存占用（VRAM OOM）：RFdiffusion 依赖 PyTorch 和 SE3-Transformer，在处理长蛋白（>150 aa）时极易发生显存溢出。
  2. 环境依赖冲突：DGL、PyTorch 与 CUDA 版本的匹配极其繁琐，稍有不慎就会导致编译失败。

我们的低配优化策略：

  • 精简环境：避开复杂的 Docker，使用轻量化的 Conda 环境。
  • 分阶段运行：RFdiffusion 进行严格的长度控制（建议首试 100 aa 以内）并启用半精度/CPU 垫底；ProteinMPNN 运算量极小，直接强制运行在 CPU 上，把珍贵的显存全部留给 RFdiffusion。
  • 批处理限制：将生成数量（num_designs）设为 1，采用循环串行生成，而非并行生成。

第一阶段：搭建极简 Conda 运行环境

建议在 Linux（如 Ubuntu 20.04/22.04）或 Windows WSL2 下操作。以下以 WSL2/Linux 为例。

1. 创建干净的 Python 环境

建议使用 Python 3.9，这是目前对 PyTorch 1.12+ 和相关几何深度学习库兼容性最好的版本。

conda create -n protein-design python=3.9 -y
conda activate protein-design

2. 安装 PyTorch 和 CUDA Toolkit

针对低配显卡，我们无需追求最新的 CUDA 版本，稳定且兼容性广的 CUDA 11.7/11.8 是首选。

# 推荐安装 PyTorch 1.13.1 + CUDA 11.7
conda install pytorch==1.13.1 torchvision==0.14.1 torchaudio==0.13.1 pytorch-cuda=11.7 -c pytorch -c nvidia -y

3. 安装 RFdiffusion 核心依赖 (SE3-Transformer)

这是最容易报错的一步。为了避免本地编译失败，我们直接使用已编译好的 wheel 包，或者通过特定分支安装：

# 安装必要的科学计算库
pip install decorator scipy dgl==1.0.2+cu117 -f https://data.dgl.ai/wheels/cu117/repo.html
pip install icecream Opti-Oversampling

# 克隆并安装 SE3-Transformer
git clone https://github.com/vgsatorras/se3_transformer.git
cd se3_transformer
pip install --no-cache-dir -r requirements.txt
python setup.py install
cd ..

第二阶段：部署与配置 RFdiffusion

1. 下载源码与权重文件

git clone https://github.com/RosettaCommons/RFdiffusion.git
cd RFdiffusion
pip install -e .

接着下载模型权重。低配环境下，不要尝试运行复杂的 Base 节点模型，建议先下载基础的 ActiveSite 和 Monomer 权重：

mkdir models && cd models
# 下载单体生成模型 (约 300MB)
wget http://files.ipd.uw.edu/pub/RFdiffusion/ActiveSite_v1.pt
wget http://files.ipd.uw.edu/pub/RFdiffusion/Base_v1.pt
cd ..

2. 低配运行参数微调（保姆级命令行）

要在 8G 甚至 6G 显存上强行运行 RFdiffusion，必须在启动脚本中加入限制。

低显存（8G VRAM）测试命令：

./scripts/run_inference.py \
  --config-name monomer \
  inference.output_prefix=outputs/test_monomer \
  inference.num_designs=1 \
  'contigmap.contigs=[100-100]' \
  inference.model_directory_path=models \
  inference.checkpoint_path=models/Base_v1.pt
  • 优化解析：
    • 'contigmap.contigs=[100-100]'：将生成长度限制在 100 aa。
    • inference.num_designs=1：单次只生成一个，防止显存积压。

真·无 GPU 环境（纯 CPU 运行）：
如果你的电脑没有 NVIDIA 显卡，只需在命令最后加上 inference.device=cpu。

  • 注意：CPU 运行生成一个结构大约需要 10-30 分钟，适合验证工作流是否打通。
./scripts/run_inference.py \
  --config-name monomer \
  inference.output_prefix=outputs/cpu_test \
  inference.num_designs=1 \
  'contigmap.contigs=[80-80]' \
  inference.device=cpu

第三阶段：部署轻量级 ProteinMPNN

相比 RFdiffusion，ProteinMPNN 的计算效率极高，完全可以在 CPU 上瞬间秒杀。

1. 克隆与环境配置

直接进入你的工作目录，克隆 ProteinMPNN 仓库：

git clone https://github.com/dauparas/ProteinMPNN.git
cd ProteinMPNN

由于我们已经在前面的 protein-design 环境中安装了 PyTorch，这里几乎不需要额外配置，直接可用。

2. 极简运行命令（强制 CPU 运行）

将 RFdiffusion 输出的 .pdb 文件复制到当前目录下。执行以下命令，强制使用 CPU 并在极短时间内生成序列：

python protein_mpnn_run.py \
        --pdb_path "../RFdiffusion/outputs/test_monomer_0.pdb" \
        --pdb_path_chains "A" \
        --out_folder "mpnn_outputs" \
        --num_seq_per_target 5 \
        --sampling_temp "0.1" \
        --seed 37 \
        --batch_size 1 \
        --device cpu
  • --num_seq_per_target 5：针对这一个骨架生成 5 条候选序列。
  • --device cpu：强制使用 CPU，释放显卡压力。

第四阶段：一键式本地极简工作流脚本

为了避免手动在两个文件夹之间拷贝 PDB 文件，我们可以写一个简单的 Bash 脚本 pipeline.sh，实现一键式“骨架生成 -> 序列设计”：

#!/bin/bash

# 配置路径
RF_PATH="/path/to/RFdiffusion"
MPNN_PATH="/path/to/ProteinMPNN"
OUT_DIR="./workspace"
mkdir -p $OUT_DIR

echo "[Step 1] 正在通过 RFdiffusion 生成蛋白质主链骨架..."
python $RF_PATH/scripts/run_inference.py \
  --config-name monomer \
  inference.output_prefix=$OUT_DIR/rf_output \
  inference.num_designs=1 \
  'contigmap.contigs=[90-90]' \
  inference.model_directory_path=$RF_PATH/models \
  inference.checkpoint_path=$RF_PATH/models/Base_v1.pt

echo "[Step 2] 骨架生成完毕，正在通过 ProteinMPNN 设计氨基酸序列..."
python $MPNN_PATH/protein_mpnn_run.py \
        --pdb_path "$OUT_DIR/rf_output_0.pdb" \
        --pdb_path_chains "A" \
        --out_folder "$OUT_DIR/mpnn_results" \
        --num_seq_per_target 10 \
        --sampling_temp "0.1" \
        --batch_size 1 \
        --device cpu

echo "[Finished] 结果已保存在 $OUT_DIR/mpnn_results/seqs 目录下！"

避坑指南与常见报错解决

1. 报错：RuntimeError: CUDA out of memory.

  • 对策：
    1. 缩短蛋白长度（尝试 [60-60]）。
    2. 检查后台是否有其他占用显存的进程（使用 nvidia-smi 查看）。
    3. 如果是多核 CPU，可以通过设置环境变量限制 PyTorch 线程数防止内存溢出：export OMP_NUM_THREADS=4。

2. 报错：ModuleNotFoundError: No module named 'dgl'

  • 对策：由于 DGL 库和 CUDA 版本的绑定非常严格，务必根据你的 CUDA 版本去 DGL 官网寻找对应的 wheel 链接进行下载，不要直接 pip install dgl。

3. 本地实在带不动怎么办？

如果你连 8G 显存都没有（例如集显轻薄本），但仍想做本地化探索，推荐使用 LocalColabFold + Colab 远程挂载 的折中方案。利用 Google Colab 的免费 T4 显卡跑完 RFdiffusion 步骤并导出 PDB，再在本地用自己的 CPU 跑 ProteinMPNN 进行精细化多序列调整。

点评评价

captcha
健康