共计 3264 个字符,预计需要花费 9 分钟才能阅读完成。
认识你的 RTX 4090
作为 NVIDIA Ampere 架构的旗舰显卡,RTX 4090 拥有一些令人印象深刻的技术指标:

- CUDA 核心:16,384 个,相比上一代 RTX 3090 增加约 50%
- Tensor Core:第 3 代 Tensor Core 支持 FP16、TF32 和 INT8 精度
- 显存:24GB GDDR6X,带宽达到 1TB/s
- Boost 频率:可达 2.52GHz
这些硬件特性使得 4090 在 AI 训练和推理任务中都能提供卓越的性能。但要想充分发挥它的潜力,正确的环境配置是关键。
环境配置
驱动安装
在 Ubuntu 20.04/22.04 或 CentOS 7/ 8 上安装驱动:
- 首先禁用 Secure Boot(否则可能导致驱动加载失败)
- 添加官方驱动仓库:
sudo add-apt-repository ppa:graphics-drivers/ppa sudo apt update - 安装推荐版本驱动:
sudo apt install nvidia-driver-525 - 重启后验证:
nvidia-smi
CUDA Toolkit 与 cuDNN
版本匹配非常重要,以下是经过测试的稳定组合:
| CUDA 版本 | cuDNN 版本 | PyTorch 版本 |
|---|---|---|
| 11.7 | 8.5.0 | 1.13.1 |
| 11.8 | 8.6.0 | 2.0.0 |
安装示例(CUDA 11.8):
wget https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_520.61.05_linux.run
sudo sh cuda_11.8.0_520.61.05_linux.run
Docker 环境
配置 NVIDIA Container Toolkit 让 Docker 支持 GPU:
- 安装工具包:
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 -fsSL 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 install -y nvidia-container-toolkit sudo systemctl restart docker - 测试 GPU 访问:
docker run --gpus all nvidia/cuda:11.8.0-base-ubuntu20.04 nvidia-smi
实战:ResNet-18 图像分类
基础 PyTorch 实现
创建一个简单的图像分类流程:
import torch
import torchvision
# 初始化模型
model = torchvision.models.resnet18(pretrained=True).cuda()
model.eval()
# 模拟输入数据
input_tensor = torch.rand(1, 3, 224, 224).cuda()
# 推理
with torch.no_grad():
output = model(input_tensor)
TensorRT 加速
安装 Torch-TensorRT:
pip install torch-tensorrt
转换并运行模型:
import torch_tensorrt
# 编译模型
trt_model = torch_tensorrt.compile(model,
inputs=[torch_tensorrt.Input((1, 3, 224, 224))],
enabled_precisions={torch.float32, torch.float16}
)
# 测试速度
import time
start = time.time()
for _ in range(100):
trt_model(input_tensor)
print(f"TensorRT 平均推理时间: {(time.time()-start)/100*1000:.2f}ms")
在我的测试中,TensorRT 可以将 ResNet-18 的推理延迟从 3.2ms 降低到 1.8ms,提升约 44%。
显存监控
进阶使用 nvidia-smi:
# 持续监控显存变化(每秒刷新)nvidia-smi -l 1
# 查看进程详情的显存使用
nvidia-smi -q -i 0 -d MEMORY
# 输出 CSV 格式便于分析
nvidia-smi --query-gpu=timestamp,name,utilization.gpu,memory.used --format=csv -l 1
避坑指南
多进程显存泄漏
当使用 DataLoader 时,设置正确的 num_workers:
# 错误示范(可能导致显存累积)train_loader = DataLoader(dataset, batch_size=32, num_workers=8)
# 正确做法(限制子进程显存)torch.multiprocessing.set_sharing_strategy('file_system')
train_loader = DataLoader(dataset, batch_size=32, num_workers=4)
FP16 精度问题
混合精度训练的最佳实践:
- 使用梯度缩放
- 检查模型中有不支持 FP16 的操作
- 逐步降低精度
scaler = torch.cuda.amp.GradScaler()
with torch.cuda.amp.autocast():
output = model(input)
loss = criterion(output, target)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
PCIe 带宽瓶颈
检测方法:
# 查看 PCIe 链接速度
nvidia-smi -q | grep "Link Width"
# 使用带宽测试工具
sudo apt install pciutils
lspci -vvv | grep -i nvidia
如果显示 x8 而不是 x16,可能需要:
1. 检查主板插槽是否支持 x16
2. 更新 BIOS 设置
3. 避免使用 PCIe 延长线
完整 Dockerfile 示例
FROM nvidia/cuda:11.8.0-base-ubuntu20.04
# 安装基础工具
RUN apt update && apt install -y python3 python3-pip
# 设置工作目录
WORKDIR /app
# 安装依赖
COPY requirements.txt .
RUN pip install -r requirements.txt
# 复制代码
COPY . .
# 启动命令
CMD ["python3", "inference.py"]
requirements.txt 内容:
torch==2.0.0
torchvision==0.15.1
torch-tensorrt==1.4.0
numpy==1.24.3
思考题
如何设计 benchmark 测试 4090 在不同 batch size 下的吞吐量瓶颈?这里有几个方向:
- 逐步增加 batch size 直到显存耗尽
- 监控 GPU 利用率与功率限制的关系
- 分析 PCIe 总线传输数据量
- 比较不同精度 (FP32/FP16/INT8) 下的最佳 batch size
通过系统化的测试,你可以找到适合你特定模型的最佳 batch size 配置,充分发挥 4090 的强大性能。
正文完
发表至: 未分类
近三天内
