概述
AI 开发环境搭建,是项目成功的一半。没有恰当的环境,实验的可复现性、团队协作、快速迭代都无从谈起。本指南涵盖 AI/ML 开发者需要掌握的全部环境配置内容 — 从 GPU 服务器的初始设置,到日常开发工作流。
内容包括:在 Ubuntu 上安装 CUDA、管理 Python 环境、像专家一样使用 JupyterLab 和 VS Code,以及用 Docker 容器搭建可复现的实验环境,逐一分步说明。
1. AI 开发环境要求
1.1 硬件推荐
GPU(最优先)
在 AI 研究中,GPU 不是可选项,而是必需品。按用途推荐如下:
- 入门/个人研究:NVIDIA RTX 4080/4090(16-24GB 显存)
- 团队共享服务器:NVIDIA A100 40GB 或 80GB
- 大规模 LLM 训练:NVIDIA H100 或 H200(80GB+ 显存)
- 云端:AWS p3/p4d/p5、GCP A100/H100、Lambda Labs
CPU
GPU 训练时,CPU 主要负责数据预处理和加载。
- 最低:8 核 16 线程(Intel Core i9 或 AMD Ryzen 9)
- 推荐:16 核以上(AMD Threadripper、Intel Xeon)
- DataLoader worker 数量 = CPU 核心数的一半较为合适
内存
- 最低:32GB
- 推荐:64GB(显存的 4 倍以上)
- 大规模 NLP:128GB 以上(分词、数据加载)
存储
/home/user/ → NVMe SSD(操作系统、代码、环境)
/data/ → NVMe SSD 或高速 HDD(训练数据)
/models/ → HDD 或 NAS(模型检查点)
- 操作系统及软件包:NVMe SSD 500GB 以上
- 数据集:NVMe SSD(I/O 密集型训练)或高性能 HDD
- 检查点:HDD 4TB 以上(性价比高的大容量存储)
1.2 操作系统选择
强烈推荐 Ubuntu 22.04 LTS
# 查看 Ubuntu 版本
lsb_release -a
# 示例输出:
# Ubuntu 22.04.3 LTS (Jammy Jellyfish)
选择 Ubuntu 的理由如下:
- NVIDIA 官方支持(第一时间提供最新驱动、CUDA、cuDNN 版本)
- 丰富的 AI/ML 社区文档与软件包
- 与 Docker、Kubernetes 兼容性最佳
- 5 年 LTS 支持,服务器运行稳定
macOS(M1/M2/M3)
Apple Silicon 通过 Metal Performance Shaders 支持 GPU 加速。
# 确认 MPS 加速可用性(PyTorch)
python -c "import torch; print(torch.backends.mps.is_available())"
1.3 云端 vs 本地开发
| 项目 | 本地服务器 | 云端 |
|---|---|---|
| 初始成本 | 高(硬件) | 低 |
| 运营成本 | 低(电费) | 高(按小时计费) |
| 等待时间 | 无 | 实例启动时间 |
| 扩展性 | 有限 | 无限 |
| 数据安全 | 高 | 取决于策略 |
| 推荐用途 | 持续性研究 | 一次性大规模实验 |
2. NVIDIA GPU 驱动与 CUDA 安装
2.1 安装 nvidia-driver(Ubuntu)
# 1. 移除已有的 NVIDIA 软件包
sudo apt-get purge nvidia*
sudo apt autoremove
# 2. 查看可用驱动
ubuntu-drivers devices
# 3. 自动安装推荐驱动
sudo ubuntu-drivers autoinstall
# 或安装指定版本
sudo apt install nvidia-driver-535
# 4. 重启
sudo reboot
# 5. 确认安装
nvidia-smi
nvidia-smi 输出示例:
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 535.154.05 Driver Version: 535.154.05 CUDA Version: 12.2 |
|-------------------------------+----------------------+----------------------+
| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|===============================+======================+======================|
| 0 NVIDIA A100-SXM4-40GB Off| 00000000:00:04.0 Off | 0 |
| N/A 34C P0 56W / 400W | 1024MiB / 40960MiB | 0% Default |
+-----------------------------------------------------------------------------+
2.2 安装 CUDA Toolkit
在 NVIDIA 官网生成适合你环境的安装命令。
# 安装 CUDA 12.2(以 Ubuntu 22.04 为例)
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb
sudo dpkg -i cuda-keyring_1.1-1_all.deb
sudo apt-get update
sudo apt-get -y install cuda-toolkit-12-2
# 设置环境变量(追加到 ~/.bashrc 或 ~/.zshrc)
echo 'export PATH=/usr/local/cuda/bin:$PATH' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc
source ~/.bashrc
# 确认 CUDA 版本
nvcc --version
2.3 安装 cuDNN
# 需要 NVIDIA Developer 账号
# 下载 cuDNN 后:
# 以 Ubuntu 22.04、CUDA 12.x 为例
sudo dpkg -i cudnn-local-repo-ubuntu2204-8.9.7.29_1.0-1_amd64.deb
sudo cp /var/cudnn-local-repo-ubuntu2204-8.9.7.29/cudnn-local-*-keyring.gpg /usr/share/keyrings/
sudo apt-get update
sudo apt-get install libcudnn8 libcudnn8-dev libcudnn8-samples
# 确认 cuDNN 版本
cat /usr/include/cudnn_version.h | grep CUDNN_MAJOR -A 2
2.4 验证安装
# verify_gpu.py
import subprocess
import sys
def check_nvidia_smi():
result = subprocess.run(['nvidia-smi'], capture_output=True, text=True)
if result.returncode == 0:
print("nvidia-smi 运行正常")
# 提取 GPU 信息
lines = result.stdout.split('\n')
for line in lines:
if 'NVIDIA' in line and 'Driver' in line:
print(f" {line.strip()}")
else:
print("nvidia-smi 出错:", result.stderr)
def check_pytorch_cuda():
try:
import torch
print(f"\nPyTorch 版本:{torch.__version__}")
print(f"CUDA 可用:{torch.cuda.is_available()}")
if torch.cuda.is_available():
print(f"CUDA 版本:{torch.version.cuda}")
print(f"cuDNN 版本:{torch.backends.cudnn.version()}")
print(f"GPU 数量:{torch.cuda.device_count()}")
for i in range(torch.cuda.device_count()):
props = torch.cuda.get_device_properties(i)
print(f" GPU {i}: {props.name} ({props.total_memory / 1024**3:.1f}GB)")
# 简单的张量运算测试
x = torch.randn(1000, 1000).cuda()
y = torch.randn(1000, 1000).cuda()
z = x @ y
print(f"GPU 张量运算测试:成功(shape: {z.shape})")
except ImportError:
print("未安装 PyTorch")
def check_tensorflow_gpu():
try:
import tensorflow as tf
print(f"\nTensorFlow 版本:{tf.__version__}")
gpus = tf.config.list_physical_devices('GPU')
print(f"检测到的 GPU:{len(gpus)}个")
for gpu in gpus:
print(f" {gpu}")
except ImportError:
print("未安装 TensorFlow")
if __name__ == "__main__":
check_nvidia_smi()
check_pytorch_cuda()
check_tensorflow_gpu()
python verify_gpu.py
3. Python 环境管理
3.1 用 pyenv 管理 Python 版本
# 安装 pyenv
curl https://pyenv.run | bash
# 追加到 ~/.bashrc 或 ~/.zshrc
echo 'export PYENV_ROOT="$HOME/.pyenv"' >> ~/.bashrc
echo 'command -v pyenv >/dev/null || export PATH="$PYENV_ROOT/bin:$PATH"' >> ~/.bashrc
echo 'eval "$(pyenv init -)"' >> ~/.bashrc
source ~/.bashrc
# 查看可用的 Python 版本列表
pyenv install --list | grep -E "^\s+3\.(10|11|12)"
# 安装 Python
pyenv install 3.11.7
pyenv install 3.10.13
# 设置全局版本
pyenv global 3.11.7
# 设置项目专用版本(在该目录下执行)
cd my-project
pyenv local 3.10.13
# 确认当前 Python 版本
python --version
pyenv versions
3.2 conda 环境(GPU 库)
conda 在安装 RAPIDS、cuML 等 GPU 加速库时特别有用。
# 安装 Miniconda
wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
bash Miniconda3-latest-Linux-x86_64.sh -b
eval "$($HOME/miniconda3/bin/conda shell.bash hook)"
conda init
# 配置渠道
conda config --add channels conda-forge
conda config --add channels nvidia
conda config --set channel_priority strict
# 创建 AI/ML 环境
conda create -n aiml python=3.11 -y
conda activate aiml
# 安装带 CUDA 的 PyTorch
conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia
# RAPIDS(GPU 加速数据科学)
conda install -c rapidsai -c conda-forge -c nvidia rapids=23.10 cuda-version=12.0
# 导出与导入环境
conda env export > environment.yml
conda env create -f environment.yml
# 环境列表
conda env list
environment.yml 示例:
name: aiml
channels:
- pytorch
- nvidia
- conda-forge
- defaults
dependencies:
- python=3.11
- pytorch>=2.1
- torchvision
- cudatoolkit=12.1
- numpy>=1.24
- pandas>=2.0
- scikit-learn>=1.3
- pip:
- transformers>=4.35
- wandb>=0.16
- pydantic>=2.0
3.3 用 poetry 管理依赖
# 安装 poetry
curl -sSL https://install.python-poetry.org | python3 -
# 创建项目
poetry new ml-research
cd ml-research
# 添加依赖
poetry add torch torchvision numpy pandas transformers
poetry add --group dev pytest black ruff mypy jupyter
# 指定版本或包含 extra
poetry add "torch[cuda]>=2.1"
poetry add "transformers[torch]>=4.35"
# 安装
poetry install
# 查看虚拟环境信息
poetry env info
poetry env list
3.4 uv(超高速包管理器)
# 安装 uv
curl -LsSf https://astral.sh/uv/install.sh | sh
# 比 pip 快 10-100 倍的包安装
uv pip install torch torchvision numpy pandas
# 创建虚拟环境
uv venv .venv --python 3.11
source .venv/bin/activate
# 从 requirements.txt 安装(速度极快)
uv pip install -r requirements.txt
# 管理 Python 版本
uv python install 3.11 3.12
uv python list
4. JupyterLab 高级配置
4.1 安装 JupyterLab 及扩展
# 安装 JupyterLab
pip install jupyterlab
# 安装常用扩展
pip install jupyterlab-git # Git 集成
pip install jupyterlab-lsp # 语言服务器协议(自动补全)
pip install python-lsp-server # Python LSP
pip install jupyterlab-code-formatter # 代码格式化工具
pip install black isort # 格式化工具
pip install jupyterlab-vim # Vim 键位绑定
pip install ipywidgets # 交互式控件
# 查看扩展列表
jupyter labextension list
# 启动 JupyterLab
jupyter lab --ip=0.0.0.0 --port=8888 --no-browser
4.2 内核管理
# 查看内核列表
jupyter kernelspec list
# 将当前 conda/venv 环境注册为内核
pip install ipykernel
python -m ipykernel install --user --name myenv --display-name "Python (myenv)"
# 将指定 conda 环境注册为内核
conda activate aiml
conda install ipykernel
python -m ipykernel install --user --name aiml --display-name "Python (aiml GPU)"
# 删除内核
jupyter kernelspec remove old-kernel
# 修改自定义内核规格
# ~/.local/share/jupyter/kernels/aiml/kernel.json
kernel.json 示例:
{
"argv": [
"/home/user/miniconda3/envs/aiml/bin/python",
"-m",
"ipykernel_launcher",
"-f",
"{connection_file}"
],
"display_name": "Python (aiml GPU)",
"language": "python",
"env": {
"CUDA_VISIBLE_DEVICES": "0",
"PYTHONPATH": "/home/user/projects"
}
}
4.3 远程 Jupyter(SSH 隧道)
# 在服务器上启动 Jupyter(不使用 token)
jupyter lab --no-browser --port=8888 --ip=127.0.0.1
# 在本地创建 SSH 隧道
ssh -N -L 8888:localhost:8888 user@your-server.com
# 浏览器中访问
# http://localhost:8888
# 设置密码(更安全)
jupyter lab password
# 之后可在 localhost:8888 用密码登录
用于自动启动的 systemd 服务配置:
# /etc/systemd/system/jupyter.service
cat > /tmp/jupyter.service << 'EOF'
[Unit]
Description=Jupyter Lab
After=network.target
[Service]
Type=simple
User=your-username
WorkingDirectory=/home/your-username
ExecStart=/home/your-username/miniconda3/envs/aiml/bin/jupyter lab --no-browser --port=8888
Restart=on-failure
[Install]
WantedBy=multi-user.target
EOF
sudo mv /tmp/jupyter.service /etc/systemd/system/
sudo systemctl daemon-reload
sudo systemctl enable jupyter
sudo systemctl start jupyter
4.4 Jupyter Magic 命令
# 计时
%time result = expensive_function()
%timeit -n 100 result = fast_function() # 重复 100 次
# 逐行性能分析
%load_ext line_profiler
%lprun -f my_function my_function(data)
# 内存性能分析
%load_ext memory_profiler
%memit result = memory_heavy_function()
# 将单元格内容保存为文件
%%writefile my_script.py
import numpy as np
# ...
# 将文件内容加载到单元格
%load my_script.py
# 执行 shell 命令
!nvidia-smi
!pip list | grep torch
files = !ls -la *.py
# 环境变量
%env CUDA_VISIBLE_DEVICES=0
# 自动重新加载(代码修改后立即生效)
%load_ext autoreload
%autoreload 2
# matplotlib 内嵌显示
%matplotlib inline
# 当前变量列表
%who
%whos
# 查看之前的输出结果
print(_) # 最后一次结果
print(__) # 倒数第二次结果
4.5 nbconvert(Notebook 转换)
# 将 notebook 转换为脚本
jupyter nbconvert --to script notebook.ipynb
# 转换为 HTML(便于分享)
jupyter nbconvert --to html notebook.ipynb
# 转换为 PDF(需要 LaTeX)
jupyter nbconvert --to pdf notebook.ipynb
# 转换的同时执行 notebook,输出 HTML
jupyter nbconvert --to html --execute notebook.ipynb
# 在命令行中执行 notebook
jupyter nbconvert --to notebook --execute --inplace notebook.ipynb
# 参数化执行(papermill)
pip install papermill
papermill input.ipynb output.ipynb -p lr 0.001 -p epochs 100
5. AI 场景下的 VS Code
5.1 安装必备扩展
# 在 VS Code 命令行中安装扩展
code --install-extension ms-python.python # Python
code --install-extension ms-toolsai.jupyter # Jupyter
code --install-extension ms-python.pylance # Pylance(LSP)
code --install-extension charliermarsh.ruff # Ruff 检查工具
code --install-extension github.copilot # GitHub Copilot
code --install-extension github.copilot-chat # Copilot Chat
code --install-extension ms-vscode-remote.remote-ssh # Remote SSH
code --install-extension ms-vscode-remote.remote-containers # Dev Containers
code --install-extension eamodio.gitlens # GitLens
code --install-extension njpwerner.autodocstring # 自动生成 docstring
code --install-extension tabnine.tabnine-vscode # Tabnine AI
5.2 VS Code 设置(settings.json)
{
"python.defaultInterpreterPath": "${workspaceFolder}/.venv/bin/python",
"python.formatting.provider": "none",
"[python]": {
"editor.defaultFormatter": "charliermarsh.ruff",
"editor.formatOnSave": true,
"editor.codeActionsOnSave": {
"source.fixAll.ruff": true,
"source.organizeImports.ruff": true
}
},
"pylance.typeCheckingMode": "basic",
"python.analysis.autoImportCompletions": true,
"python.analysis.indexing": true,
"python.analysis.packageIndexDepths": [
{ "name": "torch", "depth": 5 },
{ "name": "transformers", "depth": 5 }
],
"jupyter.askForKernelRestart": false,
"jupyter.interactiveWindow.creationMode": "perFile",
"editor.inlineSuggest.enabled": true,
"github.copilot.enable": {
"*": true,
"yaml": true,
"plaintext": false
},
"files.exclude": {
"**/__pycache__": true,
"**/*.pyc": true,
".mypy_cache": true,
".ruff_cache": true
},
"terminal.integrated.env.linux": {
"CUDA_VISIBLE_DEVICES": "0"
}
}
5.3 Remote SSH 配置
# 本地 ~/.ssh/config
Host ml-server
HostName 192.168.1.100
User your-username
IdentityFile ~/.ssh/id_rsa
ForwardAgent yes
ServerAliveInterval 60
ServerAliveCountMax 3
Host gpu-cloud
HostName gpu-server.example.com
User ubuntu
IdentityFile ~/.ssh/cloud-key.pem
Port 22
在 VS Code 中使用 Remote SSH:
- Ctrl+Shift+P(Mac 上为 Cmd+Shift+P)
- 选择 "Remote-SSH: Connect to Host"
- 选择已配置的主机名
远程服务器上自动选择 Python 环境:
{
"remote.SSH.defaultExtensions": ["ms-python.python", "ms-toolsai.jupyter", "charliermarsh.ruff"]
}
5.4 Dev Containers
.devcontainer/devcontainer.json 示例:
{
"name": "AI Development",
"build": {
"dockerfile": "Dockerfile",
"args": {
"CUDA_VERSION": "12.1.1",
"PYTHON_VERSION": "3.11"
}
},
"runArgs": ["--gpus", "all", "--shm-size", "8g"],
"mounts": [
"source=/data,target=/data,type=bind",
"source=${localEnv:HOME}/.cache,target=/root/.cache,type=bind"
],
"customizations": {
"vscode": {
"extensions": [
"ms-python.python",
"ms-toolsai.jupyter",
"charliermarsh.ruff",
"github.copilot"
],
"settings": {
"python.defaultInterpreterPath": "/usr/local/bin/python"
}
}
},
"postCreateCommand": "pip install -e '.[dev]'",
"remoteUser": "root"
}
6. GPU Docker 容器
6.1 安装 NVIDIA Container Toolkit
# 安装 NVIDIA Container Toolkit
distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey \
| sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
curl -s -L https://nvidia.github.io/libnvidia-container/$distribution/libnvidia-container.list \
| sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' \
| sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
sudo apt-get update
sudo apt-get install -y nvidia-container-toolkit
sudo nvidia-ctk runtime configure --runtime=docker
sudo systemctl restart docker
# 测试 GPU 访问
docker run --rm --gpus all nvidia/cuda:12.2.0-base-ubuntu22.04 nvidia-smi
6.2 AI 项目用 Dockerfile
# Dockerfile
FROM nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04
# 环境变量
ENV DEBIAN_FRONTEND=noninteractive
ENV PYTHONUNBUFFERED=1
ENV PYTHONDONTWRITEBYTECODE=1
ENV PIP_NO_CACHE_DIR=1
# 系统软件包
RUN apt-get update && apt-get install -y \
python3.11 \
python3.11-dev \
python3.11-venv \
python3-pip \
git \
wget \
curl \
vim \
htop \
tmux \
&& rm -rf /var/lib/apt/lists/*
# 设置默认 Python
RUN update-alternatives --install /usr/bin/python python /usr/bin/python3.11 1
RUN update-alternatives --install /usr/bin/pip pip /usr/bin/pip3 1
# 升级 pip
RUN pip install --upgrade pip setuptools wheel
# 工作目录
WORKDIR /workspace
# Python 依赖(优化分层缓存)
COPY requirements.txt .
RUN pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
RUN pip install -r requirements.txt
# 复制源码
COPY . .
# 安装包
RUN pip install -e ".[dev]"
# 创建非 root 用户(安全考虑)
RUN useradd -m -u 1000 researcher
RUN chown -R researcher:researcher /workspace
USER researcher
# 暴露端口
EXPOSE 8888 6006
CMD ["jupyter", "lab", "--ip=0.0.0.0", "--port=8888", "--no-browser"]
6.3 用 Docker Compose 搭建服务栈
# docker-compose.yml
version: '3.8'
services:
# 训练容器
trainer:
build:
context: .
dockerfile: Dockerfile
runtime: nvidia
environment:
- NVIDIA_VISIBLE_DEVICES=0
- CUDA_VISIBLE_DEVICES=0
- WANDB_API_KEY=${WANDB_API_KEY}
volumes:
- ./src:/workspace/src
- ./configs:/workspace/configs
- /data:/data:ro
- model-checkpoints:/workspace/checkpoints
- ~/.cache/huggingface:/root/.cache/huggingface
shm_size: '8g'
command: python train.py
# Jupyter notebook 服务器
notebook:
build:
context: .
dockerfile: Dockerfile
runtime: nvidia
environment:
- NVIDIA_VISIBLE_DEVICES=1
ports:
- '8888:8888'
volumes:
- ./notebooks:/workspace/notebooks
- ./src:/workspace/src
- /data:/data:ro
command: jupyter lab --ip=0.0.0.0 --no-browser --allow-root
# TensorBoard
tensorboard:
image: tensorflow/tensorflow:latest
ports:
- '6006:6006'
volumes:
- model-checkpoints:/workspace/checkpoints:ro
command: tensorboard --logdir=/workspace/checkpoints --host=0.0.0.0
# MLflow 追踪服务器
mlflow:
image: ghcr.io/mlflow/mlflow:v2.8.0
ports:
- '5000:5000'
volumes:
- mlflow-data:/mlflow
command: mlflow server --host=0.0.0.0 --port=5000 --default-artifact-root=/mlflow/artifacts
volumes:
model-checkpoints:
mlflow-data:
# 启动完整服务栈
docker-compose up -d
# 查看日志
docker-compose logs -f trainer
# 重启指定服务
docker-compose restart notebook
# 查看 GPU 使用情况(容器内)
docker-compose exec trainer nvidia-smi
# 清理
docker-compose down --volumes
6.4 GPU 共享策略
# 分配指定 GPU
docker run --gpus '"device=0,1"' myimage # 使用 GPU 0、1
docker run --gpus '"device=2"' myimage # 仅使用 GPU 2
# MIG(Multi-Instance GPU)配置(A100/H100)
sudo nvidia-smi mig -i 0 --create-gpu-instance 3g.20gb
sudo nvidia-smi mig -i 0 --create-compute-instance 3g.20gb
# 为容器分配 MIG 实例
docker run --gpus '"MIG-GPU-xxxxxxxx-xx-xx-xxxx-xxxxxxxxxxxx/x/x"' myimage
# GPU 内存分数分配(time-slicing)
# 在 /etc/nvidia-container-runtime/config.toml 中设置
7. 远程开发环境
7.1 基于 SSH 的远程开发
# 生成 SSH 密钥(本地)
ssh-keygen -t ed25519 -C "your-email@example.com"
# 将公钥复制到服务器
ssh-copy-id -i ~/.ssh/id_ed25519.pub user@server.com
# 或手动复制
cat ~/.ssh/id_ed25519.pub | ssh user@server.com "mkdir -p ~/.ssh && cat >> ~/.ssh/authorized_keys"
# SSH 连接优化(~/.ssh/config)
Host ml-server
HostName server.example.com
User researcher
IdentityFile ~/.ssh/id_ed25519
ControlMaster auto
ControlPath ~/.ssh/cm_%r@%h:%p
ControlPersist 10m
ServerAliveInterval 30
Compression yes
7.2 tmux(保持会话)
# 安装 tmux
sudo apt install tmux
# 启动新会话
tmux new-session -s training
# 重新连接会话
tmux attach-session -t training
# 会话列表
tmux list-sessions
# 常用快捷键(Prefix:Ctrl+B)
# Ctrl+B d :分离会话(SSH 断开后依然运行)
# Ctrl+B c :新建窗口
# Ctrl+B n/p :下一个/上一个窗口
# Ctrl+B % :垂直分屏
# Ctrl+B " :水平分屏
# Ctrl+B z :最大化/还原面板
~/.tmux.conf 配置:
# ~/.tmux.conf
# 启用鼠标支持
set -g mouse on
# 增加历史缓冲区
set -g history-limit 50000
# 窗口编号从 1 开始
set -g base-index 1
setw -g pane-base-index 1
# 颜色支持
set -g default-terminal "screen-256color"
set -ga terminal-overrides ",xterm-256color:Tc"
# 自定义状态栏
set -g status-bg colour235
set -g status-fg colour136
set -g status-right '#[fg=colour166]%d %b #[fg=colour136]%H:%M '
# 修改 Prefix 键(Ctrl+A,screen 风格)
set -g prefix C-a
unbind C-b
bind C-a send-prefix
# 快速切换面板
bind -n M-Left select-pane -L
bind -n M-Right select-pane -R
bind -n M-Up select-pane -U
bind -n M-Down select-pane -D
7.3 文件同步(rsync、sshfs)
# 用 rsync 同步代码
# 本地 -> 服务器
rsync -avz --exclude '__pycache__' --exclude '.git' \
./my-project/ user@server:/home/user/my-project/
# 服务器 -> 本地(下载结果)
rsync -avz user@server:/home/user/my-project/outputs/ ./outputs/
# 自动同步脚本
cat > sync.sh << 'EOF'
#!/bin/bash
rsync -avz --exclude '__pycache__' \
--exclude '.git' \
--exclude '.venv' \
--exclude '*.pyc' \
--delete \
./ user@ml-server:/home/user/project/
echo "同步完成:$(date)"
EOF
chmod +x sync.sh
# 用 sshfs 挂载远程文件系统
sudo apt install sshfs
mkdir -p ~/remote-server
sshfs user@server:/home/user ~/remote-server -o reconnect,ServerAliveInterval=15
# 卸载挂载
fusermount -u ~/remote-server
8. 监控工具
8.1 nvidia-smi watch
# 每 1 秒刷新一次 GPU 状态
watch -n 1 nvidia-smi
# 仅查看 GPU 内存使用量
nvidia-smi --query-gpu=index,name,memory.used,memory.total,utilization.gpu \
--format=csv -l 1
# 记录 CSV 日志
nvidia-smi --query-gpu=timestamp,index,utilization.gpu,memory.used,temperature.gpu \
--format=csv -l 1 > gpu_log.csv
8.2 nvtop
# 安装 nvtop(htop 的 GPU 版本)
sudo apt install nvtop
# 运行
nvtop
8.3 用 Python 监控 GPU
# gpu_monitor.py
import subprocess
import time
import json
from dataclasses import dataclass
from typing import List
@dataclass
class GPUStats:
index: int
name: str
temperature: float
utilization: float
memory_used: int
memory_total: int
power_draw: float
def get_gpu_stats() -> List[GPUStats]:
"""通过 nvidia-smi 查询 GPU 统计信息"""
cmd = [
'nvidia-smi',
'--query-gpu=index,name,temperature.gpu,utilization.gpu,memory.used,memory.total,power.draw',
'--format=csv,noheader,nounits'
]
result = subprocess.run(cmd, capture_output=True, text=True)
gpus = []
for line in result.stdout.strip().split('\n'):
parts = [p.strip() for p in line.split(',')]
gpus.append(GPUStats(
index=int(parts[0]),
name=parts[1],
temperature=float(parts[2]),
utilization=float(parts[3]),
memory_used=int(parts[4]),
memory_total=int(parts[5]),
power_draw=float(parts[6]) if parts[6] != 'N/A' else 0.0
))
return gpus
def monitor_training(interval: int = 10, duration: int = 3600):
"""训练过程中监控 GPU"""
print(f"开始 GPU 监控(间隔:{interval}秒)")
history = []
start_time = time.time()
while time.time() - start_time < duration:
stats = get_gpu_stats()
timestamp = time.time() - start_time
for gpu in stats:
mem_pct = gpu.memory_used / gpu.memory_total * 100
print(
f"[{timestamp:.0f}s] GPU {gpu.index}: "
f"利用率={gpu.utilization:.0f}% "
f"内存={gpu.memory_used}/{gpu.memory_total}MB ({mem_pct:.0f}%) "
f"温度={gpu.temperature:.0f}C "
f"功率={gpu.power_draw:.0f}W"
)
history.append({
'timestamp': timestamp,
'gpu_index': gpu.index,
'utilization': gpu.utilization,
'memory_used': gpu.memory_used,
'temperature': gpu.temperature
})
time.sleep(interval)
return history
if __name__ == "__main__":
history = monitor_training(interval=5, duration=60)
print(f"\n共收集 {len(history)} 个测量值")
8.4 系统监控
# 安装并使用 htop
sudo apt install htop
htop
# 用 iotop 监控 I/O
sudo apt install iotop
sudo iotop -o # 仅显示有 I/O 活动的进程
# 查看磁盘 I/O
iostat -x 1 5
# 内存使用量
free -h
vmstat 1 10
# 网络监控
sudo apt install nethogs
sudo nethogs
# 综合监控面板(glances)
pip install glances
glances
9. 代码质量工具
9.1 black + ruff 配置
# 安装
pip install black ruff pre-commit
# 运行 black
black .
black --line-length 88 src/
# 运行 ruff
ruff check .
ruff check --fix . # 自动修复
ruff format . # 格式化(可替代 black)
pyproject.toml 配置:
[tool.black]
line-length = 88
target-version = ['py310', 'py311']
include = '\.pyi?$'
[tool.ruff]
line-length = 88
target-version = "py310"
[tool.ruff.lint]
select = [
"E", # pycodestyle errors
"W", # pycodestyle warnings
"F", # pyflakes
"I", # isort
"N", # pep8-naming
"UP", # pyupgrade
"B", # flake8-bugbear
]
ignore = ["E501", "B008"]
[tool.ruff.lint.isort]
known-first-party = ["my_package"]
9.2 pre-commit hooks
# .pre-commit-config.yaml
repos:
- repo: https://github.com/pre-commit/pre-commit-hooks
rev: v4.5.0
hooks:
- id: trailing-whitespace
- id: end-of-file-fixer
- id: check-yaml
- id: check-json
- id: check-merge-conflict
- id: detect-private-key
- id: check-added-large-files
args: ['--maxkb=10000']
- repo: https://github.com/psf/black
rev: 23.11.0
hooks:
- id: black
- repo: https://github.com/astral-sh/ruff-pre-commit
rev: v0.1.6
hooks:
- id: ruff
args: [--fix]
- repo: https://github.com/pre-commit/mirrors-mypy
rev: v1.7.1
hooks:
- id: mypy
additional_dependencies:
- types-requests
- pydantic
# 安装并启用 pre-commit
pre-commit install
# 对全部文件执行
pre-commit run --all-files
# 执行指定 hook
pre-commit run black --all-files
9.3 带覆盖率的 pytest
# 安装
pip install pytest pytest-cov pytest-xdist
# 基本执行
pytest tests/ -v
# 包含覆盖率
pytest tests/ --cov=src --cov-report=html --cov-report=term-missing
# 并行执行
pytest tests/ -n auto # 按 CPU 数量并行
# 显示耗时最长的测试
pytest tests/ --durations=10
# 仅执行指定标记
pytest tests/ -m "not slow"
pytest.ini 或 pyproject.toml:
[tool.pytest.ini_options]
testpaths = ["tests"]
python_files = ["test_*.py", "*_test.py"]
python_classes = ["Test*"]
python_functions = ["test_*"]
markers = [
"slow: marks tests as slow",
"gpu: marks tests requiring GPU",
"integration: integration tests"
]
addopts = [
"-v",
"--tb=short",
"--cov=src",
"--cov-report=html",
"--cov-fail-under=80"
]
10. Weights & Biases 配置
10.1 安装与初始化 W&B
# 安装
pip install wandb
# 登录(设置 API key)
wandb login
# 或通过环境变量设置
export WANDB_API_KEY=your-api-key-here
10.2 W&B 基本用法
import wandb
import numpy as np
import torch
# 初始化实验
run = wandb.init(
project="bert-sentiment-analysis",
name="run-001-baseline",
config={
"learning_rate": 2e-5,
"batch_size": 32,
"epochs": 10,
"model": "bert-base-uncased",
"optimizer": "adamw"
},
tags=["baseline", "bert", "nlp"],
notes="基础 BERT 微调实验"
)
# 访问配置
config = wandb.config
print(f"学习率:{config.learning_rate}")
# 记录指标
for epoch in range(config.epochs):
train_loss = 2.0 - epoch * 0.15 + np.random.randn() * 0.1
val_loss = 2.2 - epoch * 0.12 + np.random.randn() * 0.1
val_acc = 0.5 + epoch * 0.04 + np.random.randn() * 0.01
wandb.log({
"epoch": epoch,
"train/loss": train_loss,
"val/loss": val_loss,
"val/accuracy": val_acc,
"learning_rate": config.learning_rate * (0.95 ** epoch)
})
# 梯度直方图(PyTorch)
# wandb.log({"gradients": wandb.Histogram(model.fc.weight.grad)})
# 记录可视化图表
import matplotlib.pyplot as plt
fig, ax = plt.subplots()
x = np.linspace(0, 10, 100)
ax.plot(x, np.sin(x), label='预测')
ax.plot(x, np.cos(x), label='实际')
ax.legend()
wandb.log({"prediction_plot": wandb.Image(fig)})
plt.close()
# 结束实验
wandb.finish()
10.3 Artifact 管理
import wandb
import os
# 保存 artifact(模型检查点)
def save_model_artifact(model_path: str, run, metadata: dict = None):
artifact = wandb.Artifact(
name="model-checkpoint",
type="model",
metadata=metadata or {}
)
artifact.add_file(model_path)
run.log_artifact(artifact)
print(f"模型 artifact 已保存:{model_path}")
# 保存 artifact(数据集)
def save_dataset_artifact(data_dir: str, run):
artifact = wandb.Artifact(
name="training-dataset",
type="dataset",
description="预处理后的训练数据集"
)
artifact.add_dir(data_dir)
run.log_artifact(artifact)
# 加载 artifact
def load_model_artifact(artifact_name: str, version: str = "latest"):
run = wandb.init(project="my-project", job_type="inference")
artifact = run.use_artifact(f"{artifact_name}:{version}")
artifact_dir = artifact.download()
print(f"artifact 下载位置:{artifact_dir}")
return artifact_dir
# 使用示例
run = wandb.init(project="bert-training")
# 模拟保存临时模型文件
with open("/tmp/model.pt", "w") as f:
f.write("model_weights")
save_model_artifact(
"/tmp/model.pt",
run,
metadata={"accuracy": 0.94, "epoch": 10, "val_loss": 0.28}
)
wandb.finish()
10.4 用于超参数优化的 Sweep
import wandb
import numpy as np
# Sweep 配置
sweep_config = {
"method": "bayes", # random, grid, bayes
"metric": {
"name": "val/accuracy",
"goal": "maximize"
},
"parameters": {
"learning_rate": {
"distribution": "log_uniform_values",
"min": 1e-5,
"max": 1e-3
},
"batch_size": {
"values": [16, 32, 64, 128]
},
"hidden_size": {
"values": [256, 512, 768, 1024]
},
"dropout": {
"distribution": "uniform",
"min": 0.1,
"max": 0.5
},
"num_layers": {
"values": [2, 4, 6, 8]
}
},
"early_terminate": {
"type": "hyperband",
"min_iter": 3
}
}
def train_sweep():
"""Sweep 调用的训练函数"""
run = wandb.init()
config = wandb.config
# 按配置模拟训练模型
best_val_acc = 0.0
for epoch in range(10):
# 实际场景中会真正训练模型
train_loss = 1.0 / (config.learning_rate * 1000) * np.random.uniform(0.8, 1.2)
val_acc = min(0.99, config.hidden_size / 2000 + np.random.uniform(0, 0.1))
if val_acc > best_val_acc:
best_val_acc = val_acc
wandb.log({
"epoch": epoch,
"train/loss": train_loss,
"val/accuracy": val_acc
})
wandb.finish()
# 创建并运行 Sweep
sweep_id = wandb.sweep(sweep_config, project="hyperparameter-search")
print(f"Sweep ID:{sweep_id}")
# 运行 agent(尝试 N 次)
wandb.agent(sweep_id, function=train_sweep, count=20)
如需在多台服务器上分布式运行 Sweep,可使用如下方式:
# 服务器 1
wandb agent username/project/sweep-id
# 服务器 2(同时运行)
wandb agent username/project/sweep-id
结语
AI 开发环境搭建,前期投入的时间越多,后期的生产力提升就越明显。本指南涉及的核心要点总结如下:
必备优先级:
- CUDA 环境搭建:正确安装 GPU 驱动和 CUDA Toolkit 是一切的基础。
- Python 环境隔离:用 pyenv + conda/poetry 组合,按项目分离环境。
- JupyterLab + VS Code:探索性分析用 Jupyter,正式开发用 VS Code。
- Docker 容器化:保证可复现的实验环境,简化团队协作。
- W&B 实验追踪:追踪全部实验,获得可复现性与洞察。
值得长期投入的方向:
- 用 pre-commit hooks 和 CI/CD 自动化代码质量管控。
- 熟练掌握远程开发环境(Remote SSH + tmux),随时随地高效工作。
- 通过 GPU 监控,及早发现训练过程中的问题。
参考资料
현재 단락 (1/810)
AI 开发环境搭建,是项目成功的一半。没有恰当的环境,实验的可复现性、团队协作、快速迭代都无从谈起。本指南涵盖 AI/ML 开发者需要掌握的全部环境配置内容 — 从 GPU 服务器的初始设置,到日常开...