第二章 计算环境搭建

搭建AI药物发现管线，第一步就是搭环境。这事听着简单，但我在项目里见过太多同事卡在这一步——装完PyTorch跑不了GPU、CUDA版本不对、Python环境冲突……一折腾就是两天。说白了，环境搭不好，后面所有工作都是空中楼阁。

这一章我们就直奔主题，把Linux、Python、Anaconda、GPU驱动和CUDA一次性搞定。我会给出我常用的方案，也会讲讲踩过的坑。

Linux系统 — 选Ubuntu还是CentOS？

我个人习惯用Ubuntu 20.04 LTS或22.04 LTS。为什么？社区活跃、软件源新、NVIDIA驱动兼容性好。CentOS其实也行，但很多生物信息学工具对Ubuntu的支持更及时。

你可能会问：桌面版还是服务器版？如果你有GPU做训练，桌面版更方便调试，但服务器版更稳定。我会建议：

  • 个人工作站：Ubuntu Desktop 22.04 LTS
  • 集群/云端：Ubuntu Server 22.04 LTS（无图形界面，省资源）
小技巧：安装时记得勾选“安装第三方软件”，会自动装好Nouveau开源驱动（后面我们会替换成NVIDIA官方驱动）。

系统装好后，先做三件事：

sudo apt update && sudo apt upgrade -y
sudo apt install build-essential dkms -y
sudo apt install curl wget git vim -y

build-essential是编译工具链，dkms用于动态内核模块支持。嗯，这步必做，否则后面装驱动可能会报错。

Python与Anaconda — 环境管理是命脉

AI药物发现涉及多个项目，不同项目依赖的Python包版本可能完全不同。直接装系统Python？那会死得很惨。我强烈建议用Anaconda来隔离环境。

下载安装Anaconda：

wget https://repo.anaconda.com/archive/Anaconda3-2024.10-1-Linux-x86_64.sh
bash Anaconda3-2024.10-1-Linux-x86_64.sh

安装过程中一路yes就行。注意最后会问是否初始化conda，选yes。然后刷新一下：

source ~/.bashrc

验证：

conda --version

我个人习惯创建一个基础环境叫drugai：

conda create -n drugai python=3.10
conda activate drugai

为什么用Python 3.10？因为PyTorch 2.x、TensorFlow 2.x、RDKit、DeepChem这些主流库，在3.10下兼容性最好。3.11有些库还没适配，3.9又太老。

Python版本推荐/不推荐原因
3.9可用，但偏旧部分新库已停止支持
3.10推荐目前兼容性最广
3.11谨慎TensorFlow、RDKit可能报错
3.12不推荐大量药物发现工具还未适配
注意：不要用系统自带的pip直接装科学计算包，尤其是TensorFlow和PyTorch。我曾经因为用了系统pip导致CUDA版本冲突，重装了三遍。conda会自动处理依赖，省心很多。

GPU驱动与CUDA — 加速的基石

GPU是AI药物发现的命根子。没有GPU，做分子动力学模拟或深度学习训练会慢到怀疑人生。这里分三步走：装NVIDIA驱动、装CUDA Toolkit、验证。

第一步：NVIDIA驱动

先查你的显卡型号：

lspci | grep -i nvidia

我个人建议用官方NVIDIA驱动，不要用系统自带的nouveau开源驱动（性能差且不支持CUDA）。

添加NVIDIA官方PPA并安装驱动（以Ubuntu 22.04为例）：

sudo add-apt-repository ppa:graphics-drivers/ppa
sudo apt update
sudo apt install nvidia-driver-545

545是当前推荐版本，你也可以装更新的560或570，但545最稳。装完重启：

sudo reboot

重启后运行nvidia-smi，如果能看到GPU信息，驱动就装好了。我曾经在一台旧服务器上装驱动，内核版本太新导致dkms编译失败，最后手动降级内核才搞定。所以如果报错，先确认内核版本和驱动版本匹配。

第二步：CUDA Toolkit

驱动装好只完成了30%。CUDA Toolkit才是真正让你能调用GPU做计算的工具包。

去NVIDIA官网下载CUDA 12.1（推荐，与PyTorch兼容性好）。用runfile安装：

wget https://developer.download.nvidia.com/compute/cuda/12.1.1/local_installers/cuda_12.1.1_530.30.02_linux.run
sudo sh cuda_12.1.1_530.30.02_linux.run

安装时注意：

  • 如果已经装了驱动，取消勾选“Driver”选项（只装Toolkit和Samples）
  • 安装路径默认就好

安装完成后，配置环境变量：

echo 'export PATH=/usr/local/cuda-12.1/bin:$PATH' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda-12.1/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc
source ~/.bashrc

验证CUDA版本：

nvcc --version

看到Cuda compilation tools, release 12.1就对了。

关键点：NVIDIA驱动版本与CUDA Toolkit版本有匹配关系。驱动版本必须 ≥ CUDA版本所需最低驱动。例如CUDA 12.1要求驱动≥530.30.02。用nvidia-smi能看到的CUDA Version是驱动支持的最高版本，不是实际安装的Toolkit版本。千万别搞混。

第三步：验证GPU可用性

回到conda环境，装个PyTorch测试：

conda activate drugai
conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia

然后跑个简单脚本：

python -c "import torch; print(torch.cuda.is_available()); print(torch.cuda.get_device_name(0))"

如果输出True和你显卡的型号，恭喜——环境搭建完毕。

环境搭建核心流程

我把整个流程画成了脑图，方便你对照检查：

AI药物发现计算环境搭建流程 选择Linux发行版 安装基础工具 安装Anaconda 安装NVIDIA驱动 安装CUDA Toolkit 验证GPU可用性 提示：每个步骤完成后建议重启终端或source配置文件，避免环境变量不生效 从零到能跑GPU训练，大约需要2-4小时（取决于网络速度）
快速自查清单：
nvidia-smi 能显示GPU → 驱动OK
nvcc --version 显示CUDA版本 → Toolkit OK
python -c "import torch; print(torch.cuda.is_available())" → PyTorch/GPU OK

避坑指南

最后分享几个我踩过的坑：

  • 千万别装错驱动版本：如果装了太新的驱动（比如570），而CUDA 12.1需要530～545，可能不兼容。我曾经为了尝鲜装570，结果DCGAN训练直接崩，换回545才正常。
  • conda环境名尽量用英文：有次我用中文名药物计算，结果conda activate死活找不到，最后发现是编码问题。老老实实用drugai。
  • CUDA安装路径别改：默认/usr/local/cuda-12.1就行。有同事改成自定义路径，结果PyTorch编译时死活找不到库，折腾了半天。
  • 多版本CUDA共存：如果你需要同时用不同CUDA版本（比如11.8和12.1），可以用软连接管理。我一般把/usr/local/cuda链接到当前使用的版本，切换时更新链接就行。
← 上一章 📖 返回目录 下一章 →