Python GPU加速深度学习需打通CUDA→cuDNN→框架→TensorRT链路，关键在版本对齐、算子编译、内存布局与图优化；先用nvidia-smi和nvcc -V确认环境，再验证torch.cuda.is_available()或tf.config.list_physical_devices('GPU')。

Python 中用 GPU 加速深度学习，核心是打通 CUDA → cuDNN → 框架（PyTorch/TensorFlow）→ TensorRT 推理优化 这条链路。光装驱动或只跑个 .py 脚本不等于真正加速——关键在版本对齐、算子编译、内存布局和推理阶段的图优化。

确认 GPU 与 CUDA 环境是否真正就绪

很多“GPU 不加速”问题其实卡在底层环境没通。先运行：

nvidia-smi  # 看驱动版本和 GPU 状态
nvcc -V     # 看 CUDA 编译器版本（不是驱动版本！）

然后在 Python 里验证 PyTorch 或 TensorFlow 是否识别到 GPU：

• PyTorch：torch.cuda.is_available() 返回 True，且 torch.cuda.device_count() > 0
• TensorFlow：tf.config.list_physical_devices('GPU') 应返回非空列表

⚠️ 注意：CUDA Toolkit 版本必须与你安装的 PyTorch/TensorFlow 预编译包要求严格匹配。比如 PyTorch 2.1 官方 wheel 通常只支持 CUDA 11.8 或 12.1 —— 装错版本会导致 cuda runtime error 或静默退化到 CPU。

训练阶段提速：不只是 batch_size 加大

训练快 ≠ 光靠调大 batch_size。真正起效的是以下组合：

• 混合精度训练（AMP）：用 torch.cuda.amp 自动管理 float16/float32 切换，显存减半、速度提升 1.5–3 倍，对大部分模型无损精度
• 梯度检查点（Gradient Checkpointing）：用 torch.utils.checkpoint 换取显存，适合大模型微调
• 数据加载优化：DataLoader(num_workers>0, pin_memory=True, persistent_workers=True) 减少 CPU→GPU 数据搬运瓶颈
• 避免 CPU-GPU 频繁同步：删掉训练循环里的 .item()、.numpy()、print(loss) —— 改用 torch.no_grad() 下统计或日志异步上报

导出模型为 ONNX：跨框架部署的第一步

TensorRT 不直接读 PyTorch 或 TF 模型，需先转 ONNX。关键点：

• 导出前确保模型处于 eval() 模式，并关闭 dropout/batchnorm 更新
• 输入 tensor 必须是固定 shape（如 torch.randn(1,3,224,224)），动态轴要显式标注：dynamic_axes={'input': {0: 'batch'}, 'output': {0: 'batch'}}
• 用 opset_version=17（PyTorch 2.0+ 推荐），避免旧 opset 在 TensorRT 中不支持
• 导出后用 onnx.checker.check_model() 和 onnxsim.simplify() 简化图结构，减少冗余节点

TensorRT 加速推理：不止是“load + run”

TensorRT 的性能优势来自编译时优化，不是简单加载模型。典型流程：

• 用 trt.Builder 创建 builder，设置 max_workspace_size（建议 ≥ 1GB）和 fp16_mode=True（开启半精度）
• 解析 ONNX 模型时启用 strict_type_constraints=True 避免类型推断错误
• 构建 engine 后序列化保存（engine.serialize()），下次直接反序列化加载，跳过耗时编译
• 推理时使用 context.execute_v2(bindings)，bindings 是 GPU 内存地址数组，务必提前 cuda.mem_alloc() 分配并绑定

常见提速点：对重复调用的小 batch（如 B=1），开启 BuilderConfig.set_flag(trt.BuilderFlag.FP16) 和 BuilderFlag.OBEY_PRECISION_CONSTRAINTS；对大 batch，可尝试 INT8 量化（需校准数据集）。

不复杂但容易忽略：每个环节的版本兼容表得查官方文档，比如 TensorRT 8.6 支持最高 CUDA 11.8，而 CUDA 12.x 需用 TensorRT 8.7+。跑通比写模型更花时间，但一旦串起来，端到端推理延迟能压到毫秒级。