BUG:pytorch测试模型时结果为nan或全0

代码片段

import torch
import clip
from PIL import Image

# device = "cpu" 运行正常
device = "cuda" if torch.cuda.is_available() else "cpu"
model, preprocess = clip.load("ViT-B/32", device=device)

image = preprocess(Image.open("CLIP.png")).unsqueeze(0).to(device)
text = clip.tokenize(["a diagram", "a dog", "a cat"]).to(device)

with torch.no_grad():
    image_features = model.encode_image(image)
    text_features = model.encode_text(text)
    
    logits_per_image, logits_per_text = model(image, text)
    probs = logits_per_image.softmax(dim=-1).cpu().numpy()

print("Label probs:", probs)  # prints: [[0.9927937  0.00421068 0.00299572]]

DEBUG过程

在运行CLIP实例代码时，发现image_feature为全0.

调试后发现：输入数据没问题，模型参数加载也没问题，但是一旦使用cnv2D卷积操作feature map就变为全0。但我随便初始化了一个卷积核发现是可以运行的。

将数据改到cpu设备上运行，结果正常，推测可能是cuda的问题。

解决方法

重新安装pytorch以及对应的cudatoolkit，注意用nvidia-smi查看的是CUDA最大支持版本，而pytorch安装应该与nvcc -V中的版本对应。


注意，可以使用conda给不同的虚拟环境安装对应的cudatoolkit（cudatoolkit是cuda的一个子集，足够pytorch使用），pytorch运行时优先使用对应环境下的cudatoolkit，这样方便修改cuda版本。

最后我是用的命令为：

conda install pytorch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1 cudatoolkit=11.3 -c pytorch