PyTorch中如何强制启用Tensor Core以提升CNN模型推理吞吐量？

关于多任务并行吞吐量未达预期的原因 #

你同时运行全精度与AMP推理时，吞吐量无法达到两者单独运行的总和，并非PyTorch未调用Tensor Core，核心原因是GPU的计算资源（CUDA核心、Tensor Core）与显存是共享的，两个任务会互相抢占资源，导致整体效率无法叠加。单独运行AMP推理时如果吞吐量高于全精度，说明Tensor Core已经在正常工作。

确保Tensor Core启用的关键步骤 #

1. 验证GPU架构兼容性
   Tensor Core仅支持Volta（sm_70）及以上架构的GPU，可通过以下代码确认：

   print(torch.cuda.get_device_capability())
   

   输出如(7, 0)及以上则支持。

2. 正确配置AMP推理流程
   推理时无需使用GradScaler，只需用autocast上下文管理器包裹推理逻辑，同时确保输入张量为半精度（torch.float16）或让AMP自动转换：

   model = model.cuda()
   model.eval()
   input_data = input_data.cuda()
   
   with torch.no_grad(), torch.cuda.amp.autocast(dtype=torch.float16):
       output = model(input_data)
   

3. 满足Tensor Core的操作形状要求
   Tensor Core对张量形状有特定要求：

   • 卷积层：输入/输出通道数建议为8的倍数（Volta/Turing）或4的倍数（Ampere+）
   • 矩阵乘法：维度建议为16的倍数（Volta/Turing）或8的倍数（Ampere+）
     若模型通道数不符合，可微调模型（如修改最后一层卷积通道数为8的倍数）或对输入做padding调整。

4. 启用TF32加速（Ampere及以上GPU）
   Ampere及以上架构支持TF32格式，该格式由Tensor Core加速，精度接近FP32，PyTorch默认启用，可手动确认并强制开启：

   torch.backends.cuda.matmul.allow_tf32 = True
   torch.backends.cudnn.allow_tf32 = True
   

严格用Tensor Core运行全精度预训练模型的方法 #

全精度模型直接跑无法触发Tensor Core，需通过以下方式强制转换：

• 方法一：半精度转换+AMP
  将预训练模型的参数与输入转换为torch.float16，配合autocast：

  model = model.cuda().half()
   model.eval()
   input_data = input_data.cuda().half()
  
   with torch.no_grad():
       output = model(input_data)
  

  此方法会让所有支持的层（卷积、矩阵乘法）使用Tensor Core。

• 方法二：TF32强制加速
  针对Ampere及以上GPU，开启TF32后，全精度模型的矩阵乘法和卷积会自动用Tensor Core加速，无需修改模型精度：

  torch.backends.cuda.matmul.allow_tf32 = True
  torch.backends.cudnn.allow_tf32 = True
  
  model = model.cuda().eval()
  input_data = input_data.cuda()
  
  with torch.no_grad():
      output = model(input_data)
  

内容的提问来源于stack exchange，提问作者MT 16