引言

某一天心血来潮准备拿laptop跑一个刚剪枝完的demo模型，装完pytorch之后想测试一下1660ti的性能，便跑了一下yolov5作为benchmark，没想到推理得到的矩阵都是NaN，着实令人费解，下面说说踩坑的经历以及结果。

🐛 错误以及重现

在带有 Cuda 11.1 的 GTX 1660 Ti 上运行时，半精度推理返回多个模型的 NaN

测试代码及环境

环境：

Collecting environment information...
PyTorch version: 1.8.1+cu111
Is debug build: False
CUDA used to build PyTorch: 11.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.2 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: 6.0.0-1ubuntu2 (tags/RELEASE_600/final)
CMake version: version 3.20.2

Python version: 3.7 (64-bit runtime)
Is CUDA available: True
CUDA runtime version: 11.1.105
GPU models and configuration: GPU 0: NVIDIA GeForce GTX 1660
Nvidia driver version: 465.19.01
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.0.5
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.20.3
[pip3] torch==1.8.1+cu111
[pip3] torchvision==0.9.1+cu111
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               11.1.74              h6bb024c_0    nvidia
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2021.2.0           h06a4308_296  
[conda] mkl-service               2.3.0            py39h27cfd23_1  
[conda] mkl_fft                   1.3.0            py39h42c9631_2  
[conda] mkl_random                1.2.1            py39ha9443f7_2  
[conda] numpy                     1.20.1           py39h93e21f0_0  
[conda] numpy-base                1.20.1           py39h7d8b39e_0  
[conda] pytorch                   1.8.1           py3.9_cuda11.1_cudnn8.0.5_0    pytorch
[conda] torchaudio                0.8.1                      py39    pytorch
[conda] torchvision               0.9.1                py39_cu111    pytorch

测试用例：

import torch
import urllib
from PIL import Image
from torchvision import transforms

model = torch.hub.load('pytorch/vision:v0.9.0', 'mobilenet_v2', pretrained=True)
url, filename = ("https://github.com/pytorch/hub/raw/master/images/dog.jpg", "dog.jpg")
input_image = Image.open(filename)
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]),
])
input_tensor = preprocess(input_image)
input_batch = input_tensor.unsqueeze(0)

model = model.cuda().float()
input_batch = input_batch.cuda().float()

with torch.no_grad():
    output = model(input_batch)
print("FP32 output:", output)

model = model.cuda().half()
input_batch = input_batch.cuda().half()

with torch.no_grad():
    output = model(input_batch)
print("FP16 output:", output)

预期输出

我认为 FP16 的输出与 FP32 的输出大致相同，但上面的运行结果都为nan（为清楚起见，以下为截断后的输出）：

FP32 output: tensor([[-1.8283e+00, -1.4972e+00, -1.1716e+00, ..., -5.8914e-01,  9.7267e-01,  1.9510e+00]],
       device='cuda:0')

FP16 output: tensor([[nan, nan, nan, ..., nan, nan, nan, nan, nan]],
       device='cuda:0', dtype=torch.float16)

用其他网络进行测试，仍然产生相同的错误：

import torch

model = torch.hub.load('pytorch/vision:v0.9.0', 'resnet18', pretrained=True)
input_tensor = torch.rand((1,3,512,512))

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

print("FP16 output:", output)

model = model.cuda().float()
input_tensor = input_tensor.cuda().float()
with torch.no_grad():
    output = model(input_tensor)

print("FP32 output:", output)

事实上，在其他 GPU（例如 1080ti、2080ti）上运行相同的测试均可以得到有效的fp16输出，而在另一台1660Ti的机器上进行测试也产生了上述相同的错误。

令人费解的是，在带有 torch 1.8.1 / cuda 10.2 的 1660 上进行相同的测试无法重现该问题，只要cuda版本高于11则问题复现。

错误原因以及解决方案

经过不懈的努力，我从Nvidia官网上的公告发现了蛛丝马迹：这其实是Nvidia的锅，GTX1600系列中的FP16核心对于 Turing Minor 来说是全新的，在过去的任何 NVIDIA GPU 架构中都没有出现过（TU116专用阉割核心）。GTX 1600 系列有专用的 FP16 CUDA 内核，但其 1408 个内核中只有 128 个是这些专用的 FP16 内核。它们的目的在功能上与通过 Turing Major 上的张量核心运行 FP16 操作相同：允许 NVIDIA 在每个 SM 分区内同时执行 FP16 操作和 FP32 或 INT32 操作。而且因为它们只是 FP16 核心，所以它们非常小。NVIDIA 没有给出具体细节，但仅从吞吐量来看，它们应该只是它们所取代的张量核心大小的一小部分。

NVIDIA声称这个问题已经在8.2.2版本的cudnn上解决了，但非常遗憾，问题依旧存在。

用人话说就是，Nvidia在GTX 1600系列的显卡中使用了特殊的fp16运算单元，导致GPU在计算fp16时会发生错误。

解决的办法就是：在训练或者推理时使用全精度（fp32）即可解决问题，这意味着你需要关闭AMP（自动混合精度(Automatic Mixed Precision)）：amp=False

--precision full --no-half
#在训练或者推理时加上这一行参数即可

总结

虽说推理/训练是正常了，但VRAM的占用变大了100%，失去fp16的加持，运算速度也慢了不少，对于1660这区区6G的显存，至于为什么cuda10.2上没有这个问题，原因是cuda10.2压根就不支持混合精度运算… 不支持混合精度运算意味着没调用fp16的计算单元（全程fp32计算）。