A100 GPU上PyTorch ViT推理速度异常缓慢的排查求助

1. 环境变量拼写错误（关键问题） #

你运行命令中的环境变量名写错了：CUDA_VISIBLE_DEVICE 应该是 CUDA_VISIBLE_DEVICES（末尾多一个s）。拼写错误会导致该环境变量不生效，程序可能没有绑定到你指定的MIG实例，而是默认使用了资源受限的实例，直接拉低推理速度。

• 修正后的运行命令：

  CUDA_VISIBLE_DEVICES=$UUID python myfile.py
  

• 代码内设置方式（需替换为具体UUID值，不能用$变量）：

  import os
  os.environ["CUDA_VISIBLE_DEVICES"] = "你的MIG实例UUID"
  

2. MIG实例资源分配不足 #

A100开启MIG后会被分割为多个独立小实例，若你分配的实例规格过小（比如1g.5gb），其计算能力可能不如GTX 1650 Super。

• 查看当前MIG实例的规格与资源：

  nvidia-smi mig -l
  nvidia-smi mig -i $UUID --query-gpu=mig.mode.current,mig.instance.profile.name,memory.total --format=csv
  

• 确保分配的是3g.20gb及以上规格的MIG实例，这类实例的计算核心与显存足够支撑ViT批量推理。

3. 数据加载瓶颈 #

你的ImageNet数据集放在共享存储/home/shared/Imagenet_data/val，共享存储IO速度远慢于本地SSD，且默认DataLoader未启用多进程加载，会导致GPU长时间等待数据，表现为推理速度慢。

• 优化DataLoader配置：

  val_loader = DataLoader(
      imagenet_val,
      batch_size=256,
      shuffle=False,
      num_workers=8,  # 根据服务器CPU核心数调整，比如16、32
      pin_memory=True  # 加速数据从CPU到GPU的传输
  )
  

• 临时解决方案：将数据集拷贝到服务器本地SSD目录，再修改代码中的路径。

4. 模型推理未做性能优化 #

仅开启model.eval()和torch.no_grad()是基础优化，还可通过以下方式提升GPU利用率：

• 开启cuDNN基准测试（自动选择最优卷积算法）：

  torch.backends.cudnn.benchmark = True
  

• 使用TorchScript编译模型：

  model = timm.create_model('vit_base_patch16_224', pretrained=True)
  model.eval()
  model = torch.jit.trace(model, torch.randn(1, 3, 224, 224).to(device))
  model = model.to(device)
  

• 启用自动混合精度推理：

  def evaluate_accuracy(model, data_loader):
      model.eval()
      correct = 0
      total = 0
      
      with torch.no_grad(), torch.cuda.amp.autocast():
          for images, labels in tqdm(data_loader):
              images, labels = images.to(device), labels.to(device)
              outputs = model(images)
              _, predicted = outputs.max(1)
              correct += (predicted == labels).sum().item()
              total += labels.size(0)
  
      accuracy = 100 * correct / total
      return accuracy
  

5. Conda环境依赖不兼容 #

PyTorch、CUDA驱动与MIG的兼容性问题可能导致GPU无法满负载运行：

• 检查PyTorch的CUDA版本与服务器驱动版本是否匹配：

  import torch
  print(torch.version.cuda)  # 输出PyTorch内置的CUDA版本
  

  运行nvidia-smi查看驱动支持的CUDA版本，确保PyTorch的CUDA版本不超过驱动支持的版本。
• 更新timm到最新版本：

  pip install --upgrade timm
  

6. 检查MIG实例是否被占用 #

使用nvidia-smi查看当前MIG实例的进程和资源占用情况，确认没有其他进程共享该实例的计算资源，导致你的推理任务被抢占资源。

内容的提问来源于stack exchange，提问作者Hyeongjun Cho