Python 调用 GPU 进行图像识别：环境搭建到推理执行

1. 环境搭建

(1) 硬件准备

• 确认 NVIDIA GPU 支持 CUDA 计算
• 安装最新 GPU 驱动（NVIDIA 官网下载）

(2) 软件依赖

# 安装 CUDA 工具包 (以 CUDA 11.8 为例)
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-ubuntu2204.pin
sudo mv cuda-ubuntu2204.pin /etc/apt/preferences.d/cuda-repository-pin-600
sudo apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/3bf863cc.pub
sudo add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/ /"
sudo apt-get install cuda-11-8

# 安装 cuDNN (需注册 NVIDIA 开发者账号)
# 下载对应版本后执行:
sudo dpkg -i cudnn-local-repo-ubuntu2204-8.9.4.25_1.0-1_amd64.deb

(3) Python 环境

# 创建虚拟环境
python -m venv gpu_env
source gpu_env/bin/activate

# 安装深度学习框架 (PyTorch 示例)
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

2. 验证 GPU 可用性

import torch
print(f"GPU 可用: {torch.cuda.is_available()}")
print(f"设备数量: {torch.cuda.device_count()}")
print(f"当前设备: {torch.cuda.get_device_name(0)}")

3. 图像识别模型部署

(1) 加载预训练模型

from torchvision import models
import torch.nn as nn

# 加载 ResNet50 模型
model = models.resnet50(weights='IMAGENET1K_V2')
model.fc = nn.Linear(2048, 1000)  # 适配 ImageNet 类别数

# 转移到 GPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = model.to(device)
model.eval()  # 设置为评估模式

(2) 图像预处理

from torchvision import transforms
from PIL import Image

preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    )
])

# 加载测试图像
img = Image.open("test.jpg")
input_tensor = preprocess(img).unsqueeze(0).to(device)  # 添加 batch 维度

4. GPU 推理执行

import time

# 启动 GPU 计算
with torch.no_grad():
    start_time = time.time()
    output = model(input_tensor)
    inference_time = time.time() - start_time

# 获取预测结果
_, predicted_idx = output.max(1)
print(f"推理耗时: {inference_time:.4f} 秒")
print(f"预测类别: {predicted_idx.item()}")

5. 性能优化技巧

• 批处理加速：同时处理多张图像

  batch_tensor = torch.cat([input_tensor]*16)  # 模拟 16 张图像批次
  

• 混合精度训练：减少显存占用

  from torch.cuda import amp
  with amp.autocast():
      output = model(input_tensor)
  

• TensorRT 加速：模型转换优化

  pip install nvidia-tensorrt
  # 转换模型为 TensorRT 格式
  

6. 常见问题排查

问题现象	解决方案
CUDA out of memory	减小批次大小，启用混合精度
Driver/library mismatch	重装匹配版本的 CUDA/cuDNN
推理速度慢	使用 torch.backends.cudnn.benchmark=True

关键提示：通过 nvidia-smi 命令实时监控 GPU 使用情况，确保计算负载正确分配到 GPU。

执行效果示例

GPU 可用: True
设备数量: 1
当前设备: NVIDIA GeForce RTX 4090
推理耗时: 0.0152 秒
预测类别: 285 (埃及猫)

此流程完整覆盖从环境搭建到推理执行的各个环节，可根据实际需求替换模型架构（如 YOLO 用于目标检测，UNet 用于图像分割）。