Docker容器部署FastAPI+DINOv2遇memory_efficient_attention_forward错误

核心原因 #

DINOv2默认启用的memory_efficient_attention是为CUDA环境优化的算子，纯CPU环境下不支持float32数据类型，导致兼容性报错。

解决方案 #

1. 强制禁用内存高效注意力，回退原生实现 #

在加载DINOv2模型的代码中添加参数，关闭内存高效注意力：

import torch
from dinov2.models.vision_transformer import dinov2_vitl14

# 加载预训练模型时禁用内存高效注意力
model = dinov2_vitl14(pretrained=True, use_memory_efficient_attention=False)

如果是自定义构建模型，同样添加该参数：

from torch.nn import LayerNorm
from functools import partial
from dinov2.models.vision_transformer import vit_large

model = vit_large(
    patch_size=16,
    embed_dim=1024,
    depth=24,
    num_heads=16,
    mlp_ratio=4,
    qkv_bias=True,
    norm_layer=partial(LayerNorm, eps=1e-6),
    use_memory_efficient_attention=False  # 关键配置
)

2. 优化Docker镜像的CPU适配（可选） #

若要提升CPU运行性能，使用支持AVX2等高级CPU指令集的基础镜像，修改Dockerfile：

# 选用官方CPU优化版PyTorch镜像
FROM pytorch/pytorch:2.1.0-cpu

WORKDIR /app

COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

COPY API.py .

CMD ["uvicorn", "API:app", "--host", "0.0.0.0", "--port", "8000"]

3. 对齐依赖版本 #

确保容器内的PyTorch、torchvision版本和本地环境一致，避免算子兼容性问题，示例requirements.txt：

fastapi>=0.100.0
uvicorn>=0.23.2
torch>=2.0.0
torchvision>=0.15.0
dinov2>=0.1.0

验证步骤 #

1. 重新构建镜像：docker build -t dinov2-api .
2. 启动容器：docker run -p 8000:8000 dinov2-api
3. 调用API接口测试，确认算子错误消失

内容的提问来源于stack exchange，提问作者MOUAD EL HAOUS