为何EDSR模型FP16推理时出现NaN？参数与输入均无异常

我采用FP32训练模型、FP16进行推理。其他网络（如ResNet）用FP16推理正常，但超分辨率模型EDSR在FP16推理时出现问题。对比后发现EDSR与正常网络的差异点：

• 使用带inplace=True的ReLU
• 包含PixelShuffle层
• 无BatchNorm层

我的环境配置：

• CUDA 11.3
• Python 3.8.12
• PyTorch 1.12.1
• cuDNN 8.7.0
• GPU：RTX A6000

推理流程如下：

net_half = net.half()
net_half.eval()
input_half = input.half()

with torch.no_grad():
     output_half = net_half(input_half)

我已经做了以下排查：

1. 通过代码确认模型参数和输入均无NaN：

torch.stack([torch.isnan(p).any() for p in net_half.parameters()]).any()
torch.isnan(input_half).any()

返回结果均为False。
2. 测试EDSR的简化操作模块，未出现NaN：

x = torch.randn(1,4,Ny//2,Nx//2)

test_block1 = nn.Sequential(
    nn.Conv2d(4,64,kernel_size=3,padding=1),
    nn.Conv2d(64,64,kernel_size=3,padding=1,bias=True),
    nn.ReLU(True),
    nn.Conv2d(64,64,kernel_size=3,padding=1,bias=True),
    nn.Conv2d(64,64*4,kernel_size=3,padding=1,bias=True),
    nn.PixelShuffle(2),
    nn.ReLU(True),
    nn.Conv2d(64,4,kernel_size=3,padding=1)
)
x = x.half().to(device)
test_block1 = test_block1.half().to(device)
with torch.no_grad():
    y = test_block1(x)

print(y)

但实际训练中，第1轮推理能得到正常结果，第4轮推理时就出现了NaN值。请问卷积神经网络中是否存在不支持FP16的函数？

内容的提问来源于stack exchange，提问作者SIwoo Lee