微调MedSAM模型时遇RuntimeError：无法找到计算引擎求助

问题背景 #

在RedHat Linux高性能计算集群上微调MedSAM（专为医学图像优化的Segment Anything分支）时，触发以下错误：

Traceback (most recent call last):
File "/mnt/beegfs/khans24/medsam_finetuning/minimal.py", line 30, in
embedding = sam_model.image_encoder(input_image)
...
RuntimeError: GET was unable to find an engine to execute this computation

环境信息：

• PyTorch 2.0/1.13
• Runtime CUDA 11.0，Runtime cuDNN 8.2.1
• torch.backends.cudnn.version()输出为8500，CUDA可用性检查结果：True, 4, True, 8500
• 无root权限，通过modules加载CUDA工具包

最小复现代码：

import torch
import numpy as np
from skimage import io, transform
from segment_anything import SamPredictor, sam_model_registry
from segment_anything.utils.transforms import ResizeLongestSide

# Set up the model and device
model_type = 'vit_b'
checkpoint = 'load/medsam_20230423_vit_b_0.0.1.pth'
device = 'cuda:0'

sam_model = sam_model_registry[model_type](checkpoint=checkpoint).to(device)

# Generate a random image
image_size = 256
random_image = np.random.randint(0, 256, (image_size, image_size, 3), dtype=np.uint8)

# Resize the random image
sam_transform = ResizeLongestSide(sam_model.image_encoder.img_size)
resized_image = sam_transform.apply_image(random_image)

# Convert the resized image to a PyTorch tensor
resized_image_tensor = torch.as_tensor(resized_image.transpose(2, 0, 1)).to(device)

# Preprocess the image tensor
input_image = sam_model.preprocess(resized_image_tensor[None, :, :, :])

# Compute the image embedding using the sam_model
with torch.no_grad():
    embedding = sam_model.image_encoder(input_image)
    print(embedding.shape)

解决方案 #

1. 禁用cuDNN自动调优 #

该错误常与cuDNN的自动调优机制和集群硬件不兼容有关，在代码开头添加以下配置：

import torch
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
# 若上述无效，可尝试临时禁用cuDNN（会降低性能，仅作测试）
# torch.backends.cudnn.enabled = False

2. 对齐PyTorch、CUDA、cuDNN版本匹配 #

当前Runtime CUDA 11.0与PyTorch 2.0/1.13版本不匹配（PyTorch2.0最低要求CUDA11.7，1.13对应CUDA11.6），通过modules加载匹配的工具包：

# 例如加载适配PyTorch2.0的CUDA和cuDNN
module load cuda/11.7 cudnn/8.5.0
# 加载后重新激活conda环境
conda activate sam

3. 确保输入张量与模型的dtype一致 #

检查模型参数的 dtype，强制输入张量与其匹配：

# 检查模型dtype
print(sam_model.dtype)
# 预处理后强制转换输入张量dtype
input_image = sam_model.preprocess(resized_image_tensor[None, :, :, :])
input_image = input_image.to(sam_model.dtype).to(device)

4. 禁用PyTorch 2.0的编译优化 #

若使用PyTorch2.0，torch.compile可能与集群CUDA环境冲突，强制使用eager模式：

sam_model = sam_model_registry[model_type](checkpoint=checkpoint).to(device)
# 禁用编译优化
sam_model = torch.compile(sam_model, backend="eager")

5. 强制指定单GPU运行 #

集群多GPU环境可能引发设备分配冲突，在代码开头添加：

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

或提交作业时指定单GPU资源（如通过SLURM的--gres=gpu:1参数）

内容的提问来源于stack exchange，提问作者Sameed Khan